From 36ccb2e98cc7734f20da766657102574a16b3516 Mon Sep 17 00:00:00 2001 From: Haroon <106879583+haroon0x@users.noreply.github.com> Date: Sun, 26 Oct 2025 12:30:02 +0530 Subject: [PATCH 1/8] add longprobs support --- .../langchain_google_genai/chat_models.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index 4e39cb425..6713067b1 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -932,12 +932,13 @@ def _parse_response_candidate( # Add function call signatures to content only if there's already other content # This preserves backward compatibility where content is "" for # function-only responses - if function_call_signatures and content is not None: - for sig_block in function_call_signatures: - content = _append_to_content(content, sig_block) - - if content is None: - content = "" + if function_call_signatures and content is not None: + for sig_block in function_call_signatures: + content = _append_to_content(content, sig_block) + if hasattr(response_candidate, "logprobs_result"): + response_metadata["logprobs"] = proto.Message.to_dict(response_candidate.logprobs_result) + if content is None: + content = "" if isinstance(content, list) and any( isinstance(item, dict) and "executable_code" in item for item in content ): @@ -1825,6 +1826,9 @@ class Joke(BaseModel): stop: Optional[List[str]] = None """Stop sequences for the model.""" + logprobs: Optional[int] = None + """The number of logprobs to return.""" + streaming: Optional[bool] = None """Whether to stream responses from the model.""" @@ -2037,6 +2041,7 @@ def _prepare_params( "max_output_tokens": self.max_output_tokens, "top_k": self.top_k, "top_p": self.top_p, + "logprobs": getattr(self, "logprobs", None), "response_modalities": self.response_modalities, "thinking_config": ( ( @@ -2058,6 +2063,8 @@ def _prepare_params( }.items() if v is not None } + if getattr(self, "logprobs", None) is not None: + gen_config["response_logprobs"] = True if generation_config: gen_config = {**gen_config, **generation_config} From 639f8efe62d370d8c14440c3550bb71b89f55852 Mon Sep 17 00:00:00 2001 From: Haroon <106879583+haroon0x@users.noreply.github.com> Date: Sun, 26 Oct 2025 12:47:34 +0530 Subject: [PATCH 2/8] add test --- .../langchain_google_genai/chat_models.py | 14 +++--- .../tests/unit_tests/test_chat_models.py | 49 +++++++++++++++++++ 2 files changed, 57 insertions(+), 6 deletions(-) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index 6713067b1..a9acc2f69 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -929,14 +929,16 @@ def _parse_response_candidate( } function_call_signatures.append(sig_block) - # Add function call signatures to content only if there's already other content - # This preserves backward compatibility where content is "" for - # function-only responses + # Add function call signatures to content only if there's already other content + # This preserves backward compatibility where content is "" for + # function-only responses if function_call_signatures and content is not None: for sig_block in function_call_signatures: content = _append_to_content(content, sig_block) if hasattr(response_candidate, "logprobs_result"): - response_metadata["logprobs"] = proto.Message.to_dict(response_candidate.logprobs_result) + response_metadata["logprobs"] = proto.Message.to_dict( + response_candidate.logprobs_result + ) if content is None: content = "" if isinstance(content, list) and any( @@ -2063,8 +2065,8 @@ def _prepare_params( }.items() if v is not None } - if getattr(self, "logprobs", None) is not None: - gen_config["response_logprobs"] = True + if getattr(self, "logprobs", None) is not None: + gen_config["response_logprobs"] = True if generation_config: gen_config = {**gen_config, **generation_config} diff --git a/libs/genai/tests/unit_tests/test_chat_models.py b/libs/genai/tests/unit_tests/test_chat_models.py index c729cad09..cf4283c67 100644 --- a/libs/genai/tests/unit_tests/test_chat_models.py +++ b/libs/genai/tests/unit_tests/test_chat_models.py @@ -137,6 +137,55 @@ def test_initialization_inside_threadpool() -> None: ).result() +def test_logprobs() -> None: + """Test that logprobs parameter is set correctly and is in the response.""" + llm = ChatGoogleGenerativeAI( + model=MODEL_NAME, + google_api_key=SecretStr("secret-api-key"), + logprobs=10, + ) + assert llm.logprobs == 10 + + raw_response = { + "candidates": [ + { + "content": {"parts": [{"text": "Test response"}]}, + "logprobs_result": { + "top_candidates": [ + { + "candidates": [ + {"token": "Test", "log_probability": -0.1}, + ] + } + ] + }, + } + ], + } + response = GenerateContentResponse(raw_response) + + with patch( + "langchain_google_genai.chat_models._chat_with_retry" + ) as mock_chat_with_retry: + mock_chat_with_retry.return_value = response + llm = ChatGoogleGenerativeAI( + model=MODEL_NAME, + google_api_key="test-key", + logprobs=1, + ) + result = llm.invoke("test") + assert "logprobs" in result.response_metadata + assert result.response_metadata["logprobs"] == { + "top_candidates": [ + { + "candidates": [ + {"token": "Test", "log_probability": -0.1}, + ] + } + ] + } + + def test_client_transport() -> None: """Test client transport configuration.""" model = ChatGoogleGenerativeAI(model=MODEL_NAME, google_api_key="fake-key") From 4bb12102a0f7bc9ab5a5aaa605e5d5570af72535 Mon Sep 17 00:00:00 2001 From: Haroon <106879583+haroon0x@users.noreply.github.com> Date: Sun, 26 Oct 2025 17:20:17 +0530 Subject: [PATCH 3/8] update test function --- .../langchain_google_genai/chat_models.py | 8 +++--- .../tests/unit_tests/test_chat_models.py | 27 +++++++++++++++++-- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index a9acc2f69..5b2adc607 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -935,9 +935,10 @@ def _parse_response_candidate( if function_call_signatures and content is not None: for sig_block in function_call_signatures: content = _append_to_content(content, sig_block) - if hasattr(response_candidate, "logprobs_result"): - response_metadata["logprobs"] = proto.Message.to_dict( - response_candidate.logprobs_result + if hasattr(response_candidate, "logprobs_result") and response_candidate.logprobs_result: + response_metadata["logprobs"] = MessageToDict( + response_candidate.logprobs_result._pb, + preserving_proto_field_name=True, ) if content is None: content = "" @@ -1969,6 +1970,7 @@ def _identifying_params(self) -> Dict[str, Any]: "media_resolution": self.media_resolution, "thinking_budget": self.thinking_budget, "include_thoughts": self.include_thoughts, + "logprobs": self.logprobs, } def invoke( diff --git a/libs/genai/tests/unit_tests/test_chat_models.py b/libs/genai/tests/unit_tests/test_chat_models.py index cf4283c67..aef4a3cd5 100644 --- a/libs/genai/tests/unit_tests/test_chat_models.py +++ b/libs/genai/tests/unit_tests/test_chat_models.py @@ -146,10 +146,13 @@ def test_logprobs() -> None: ) assert llm.logprobs == 10 + # Create proper mock response with logprobs_result raw_response = { "candidates": [ { "content": {"parts": [{"text": "Test response"}]}, + "finish_reason": 1, + "safety_ratings": [], "logprobs_result": { "top_candidates": [ { @@ -161,6 +164,12 @@ def test_logprobs() -> None: }, } ], + "prompt_feedback": {"block_reason": 0, "safety_ratings": []}, + "usage_metadata": { + "prompt_token_count": 5, + "candidates_token_count": 2, + "total_token_count": 7, + }, } response = GenerateContentResponse(raw_response) @@ -185,25 +194,39 @@ def test_logprobs() -> None: ] } + mock_chat_with_retry.assert_called_once() + request = mock_chat_with_retry.call_args.kwargs["request"] + assert request.generation_config.logprobs == 1 + assert request.generation_config.response_logprobs is True -def test_client_transport() -> None: +@patch("langchain_google_genai._genai_extension.v1betaGenerativeServiceAsyncClient") +@patch("langchain_google_genai._genai_extension.v1betaGenerativeServiceClient") +def test_client_transport(mock_client: Mock, mock_async_client: Mock) -> None: """Test client transport configuration.""" + mock_client.return_value.transport = Mock() + mock_client.return_value.transport.kind = "grpc" model = ChatGoogleGenerativeAI(model=MODEL_NAME, google_api_key="fake-key") assert model.client.transport.kind == "grpc" + mock_client.return_value.transport.kind = "rest" model = ChatGoogleGenerativeAI( model=MODEL_NAME, google_api_key="fake-key", transport="rest" ) assert model.client.transport.kind == "rest" async def check_async_client() -> None: + mock_async_client.return_value.transport = Mock() + mock_async_client.return_value.transport.kind = "grpc_asyncio" model = ChatGoogleGenerativeAI(model=MODEL_NAME, google_api_key="fake-key") + _ = model.async_client assert model.async_client.transport.kind == "grpc_asyncio" # Test auto conversion of transport to "grpc_asyncio" from "rest" model = ChatGoogleGenerativeAI( model=MODEL_NAME, google_api_key="fake-key", transport="rest" ) + model.async_client_running = None + _ = model.async_client assert model.async_client.transport.kind == "grpc_asyncio" asyncio.run(check_async_client()) @@ -1907,4 +1930,4 @@ def test_chat_google_genai_invoke_with_audio_mocked() -> None: audio_block = audio_blocks[0] assert audio_block["type"] == "audio" assert "base64" in audio_block - assert audio_block["base64"] == base64.b64encode(wav_bytes).decode() + assert audio_block["base64"] == base64.b64encode(wav_bytes).decode() \ No newline at end of file From 46d932dc6a794c4d65312b9d67524eab3e8f3765 Mon Sep 17 00:00:00 2001 From: Haroon <106879583+haroon0x@users.noreply.github.com> Date: Sun, 26 Oct 2025 23:15:05 +0530 Subject: [PATCH 4/8] fix tests --- libs/genai/langchain_google_genai/chat_models.py | 5 ++++- libs/genai/tests/unit_tests/test_chat_models.py | 7 ++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index 5b2adc607..acd675737 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -935,7 +935,10 @@ def _parse_response_candidate( if function_call_signatures and content is not None: for sig_block in function_call_signatures: content = _append_to_content(content, sig_block) - if hasattr(response_candidate, "logprobs_result") and response_candidate.logprobs_result: + if ( + hasattr(response_candidate, "logprobs_result") + and response_candidate.logprobs_result + ): response_metadata["logprobs"] = MessageToDict( response_candidate.logprobs_result._pb, preserving_proto_field_name=True, diff --git a/libs/genai/tests/unit_tests/test_chat_models.py b/libs/genai/tests/unit_tests/test_chat_models.py index aef4a3cd5..314a0b763 100644 --- a/libs/genai/tests/unit_tests/test_chat_models.py +++ b/libs/genai/tests/unit_tests/test_chat_models.py @@ -199,6 +199,8 @@ def test_logprobs() -> None: assert request.generation_config.logprobs == 1 assert request.generation_config.response_logprobs is True + +@pytest.mark.enable_socket @patch("langchain_google_genai._genai_extension.v1betaGenerativeServiceAsyncClient") @patch("langchain_google_genai._genai_extension.v1betaGenerativeServiceClient") def test_client_transport(mock_client: Mock, mock_async_client: Mock) -> None: @@ -240,6 +242,7 @@ def test_initalization_without_async() -> None: assert chat.async_client is None +@pytest.mark.enable_socket def test_initialization_with_async() -> None: async def initialize_chat_with_async_client() -> ChatGoogleGenerativeAI: model = ChatGoogleGenerativeAI( @@ -1356,6 +1359,7 @@ def test_grounding_metadata_multiple_parts() -> None: assert grounding["grounding_supports"][0]["segment"]["part_index"] == 1 +@pytest.mark.enable_socket @pytest.mark.parametrize( "is_async,mock_target,method_name", [ @@ -1482,6 +1486,7 @@ def mock_stream() -> Iterator[GenerateContentResponse]: assert "timeout" not in call_kwargs +@pytest.mark.enable_socket @pytest.mark.parametrize( "is_async,mock_target,method_name", [ @@ -1930,4 +1935,4 @@ def test_chat_google_genai_invoke_with_audio_mocked() -> None: audio_block = audio_blocks[0] assert audio_block["type"] == "audio" assert "base64" in audio_block - assert audio_block["base64"] == base64.b64encode(wav_bytes).decode() \ No newline at end of file + assert audio_block["base64"] == base64.b64encode(wav_bytes).decode() From 2f932f5f42ca3df58faa244ea452291cdc7d34d3 Mon Sep 17 00:00:00 2001 From: Haroon <106879583+haroon0x@users.noreply.github.com> Date: Tue, 28 Oct 2025 23:13:12 +0530 Subject: [PATCH 5/8] Update chat_models.py --- .../langchain_google_genai/chat_models.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index acd675737..ba52d28f0 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -788,6 +788,8 @@ def _parse_response_candidate( except (AttributeError, TypeError): thought_sig = None + has_function_call = hasattr(part, "function_call") and part.function_call + if hasattr(part, "thought") and part.thought: thinking_message = { "type": "thinking", @@ -797,7 +799,7 @@ def _parse_response_candidate( if thought_sig: thinking_message["signature"] = thought_sig content = _append_to_content(content, thinking_message) - elif text is not None and text: + elif text is not None and text.strip() and not has_function_call: # Check if this text Part has a signature attached if thought_sig: # Text with signature needs structured block to preserve signature @@ -935,16 +937,16 @@ def _parse_response_candidate( if function_call_signatures and content is not None: for sig_block in function_call_signatures: content = _append_to_content(content, sig_block) - if ( - hasattr(response_candidate, "logprobs_result") - and response_candidate.logprobs_result - ): - response_metadata["logprobs"] = MessageToDict( - response_candidate.logprobs_result._pb, - preserving_proto_field_name=True, - ) - if content is None: - content = "" + + if content is None: + content = "" + + if hasattr(response_candidate, "logprobs_result") and response_candidate.logprobs_result: + response_metadata["logprobs"] = MessageToDict( + response_candidate.logprobs_result._pb, + preserving_proto_field_name=True, + ) + if isinstance(content, list) and any( isinstance(item, dict) and "executable_code" in item for item in content ): From 44ed74cbefa57aafdd08eee0089cc0c8eeb27dde Mon Sep 17 00:00:00 2001 From: Haroon <106879583+haroon0x@users.noreply.github.com> Date: Tue, 28 Oct 2025 23:19:58 +0530 Subject: [PATCH 6/8] run lint and format check --- libs/genai/langchain_google_genai/chat_models.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index ba52d28f0..6a4f3df7b 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -789,7 +789,7 @@ def _parse_response_candidate( thought_sig = None has_function_call = hasattr(part, "function_call") and part.function_call - + if hasattr(part, "thought") and part.thought: thinking_message = { "type": "thinking", @@ -937,16 +937,19 @@ def _parse_response_candidate( if function_call_signatures and content is not None: for sig_block in function_call_signatures: content = _append_to_content(content, sig_block) - + if content is None: content = "" - - if hasattr(response_candidate, "logprobs_result") and response_candidate.logprobs_result: + + if ( + hasattr(response_candidate, "logprobs_result") + and response_candidate.logprobs_result + ): response_metadata["logprobs"] = MessageToDict( response_candidate.logprobs_result._pb, preserving_proto_field_name=True, ) - + if isinstance(content, list) and any( isinstance(item, dict) and "executable_code" in item for item in content ): From 8767476391921980a18bf21a41ac0fa56240a24e Mon Sep 17 00:00:00 2001 From: Mason Daugherty Date: Mon, 10 Nov 2025 16:24:15 -0500 Subject: [PATCH 7/8] . --- libs/genai/langchain_google_genai/chat_models.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index 1f459d407..e563e9512 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -933,6 +933,8 @@ def _parse_response_candidate( hasattr(response_candidate, "logprobs_result") and response_candidate.logprobs_result ): + # Note: logprobs is flaky, sometimes available, sometimes not + # https://discuss.ai.google.dev/t/logprobs-is-not-enabled-for-gemini-models/107989/15 response_metadata["logprobs"] = MessageToDict( response_candidate.logprobs_result._pb, preserving_proto_field_name=True, From dd4086161ddeff4dbe5827c17aadc4795dee715c Mon Sep 17 00:00:00 2001 From: Haroon <106879583+haroon0x@users.noreply.github.com> Date: Fri, 14 Nov 2025 21:40:28 +0530 Subject: [PATCH 8/8] Update chat_models.py --- libs/genai/langchain_google_genai/chat_models.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index 6a4f3df7b..016c527ea 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -2053,7 +2053,6 @@ def _prepare_params( "max_output_tokens": self.max_output_tokens, "top_k": self.top_k, "top_p": self.top_p, - "logprobs": getattr(self, "logprobs", None), "response_modalities": self.response_modalities, "thinking_config": ( ( @@ -2075,7 +2074,9 @@ def _prepare_params( }.items() if v is not None } - if getattr(self, "logprobs", None) is not None: + logprobs = getattr(self, "logprobs", None) + if logprobs: + gen_config["logprobs"] = logprobs gen_config["response_logprobs"] = True if generation_config: gen_config = {**gen_config, **generation_config}