fix(genai): Correct int/float type conversion in tool call args (#1110)

Author: codeonbush

libs/genai/langchain_google_genai/chat_models.py

@@ -92,13 +92,7 @@
 )
 from langchain_core.utils.pydantic import is_basemodel_subclass
 from langchain_core.utils.utils import _build_model_kwargs
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-    SecretStr,
-    model_validator,
-)
+from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
 from pydantic.v1 import BaseModel as BaseModelV1
 from tenacity import (
     before_sleep_log,
@@ -664,9 +658,15 @@ def _parse_response_candidate(
             function_call = {"name": part.function_call.name}
             # dump to match other function calling llm for now
             function_call_args_dict = proto.Message.to_dict(part.function_call)["args"]
-            function_call["arguments"] = json.dumps(
-                {k: function_call_args_dict[k] for k in function_call_args_dict}
-            )
+
+            # Fix: Correct integer-like floats from protobuf conversion
+            # The protobuf library sometimes converts integers to floats
+            corrected_args = {
+                k: int(v) if isinstance(v, float) and v.is_integer() else v
+                for k, v in function_call_args_dict.items()
+            }
+
+            function_call["arguments"] = json.dumps(corrected_args)
             additional_kwargs["function_call"] = function_call
 
             if streaming:
@@ -1541,18 +1541,21 @@ def _prepare_params(
                 "response_modalities": self.response_modalities,
                 "thinking_config": (
                     (
-                        {"thinking_budget": self.thinking_budget}
-                        if self.thinking_budget is not None
-                        else {}
-                    )
-                    | (
-                        {"include_thoughts": self.include_thoughts}
-                        if self.include_thoughts is not None
-                        else {}
+                        (
+                            {"thinking_budget": self.thinking_budget}
+                            if self.thinking_budget is not None
+                            else {}
+                        )
+                        | (
+                            {"include_thoughts": self.include_thoughts}
+                            if self.include_thoughts is not None
+                            else {}
+                        )
                     )
-                )
-                if self.thinking_budget is not None or self.include_thoughts is not None
-                else None,
+                    if self.thinking_budget is not None
+                    or self.include_thoughts is not None
+                    else None
+                ),
             }.items()
             if v is not None
         }

libs/genai/tests/unit_tests/test_chat_models_protobuf_fix.py

@@ -0,0 +1,132 @@
+"""Test for protobuf integer/float conversion fix in chat models."""
+
+import json
+
+from google.ai.generativelanguage_v1beta.types import (
+    Candidate,
+    Content,
+    FunctionCall,
+    Part,
+)
+from google.protobuf.struct_pb2 import Struct
+
+from langchain_google_genai.chat_models import _parse_response_candidate
+
+
+def test_parse_response_candidate_corrects_integer_like_floats() -> None:
+    """
+    Test that _parse_response_candidate correctly handles integer-like floats
+    in tool call arguments from the Gemini API response.
+
+    This test addresses a bug where proto.Message.to_dict() converts integers
+    to floats, causing downstream type casting errors.
+    """
+    # Create a mock Protobuf Struct for the arguments with problematic float values
+    args_struct = Struct()
+    args_struct.update(
+        {
+            "entity_type": "table",
+            "upstream_depth": 3.0,  # The problematic float value that should be int
+            "downstream_depth": 5.0,  # Another problematic float value
+            "fqn": "test.table.name",
+            "valid_float": 3.14,  # This should remain as float
+            "string_param": "test_string",  # This should remain as string
+            "bool_param": True,  # This should remain as boolean
+        }
+    )
+
+    # Create the mock API response candidate
+    candidate = Candidate(
+        content=Content(
+            parts=[
+                Part(
+                    function_call=FunctionCall(
+                        name="get_entity_lineage",
+                        args=args_struct,
+                    )
+                )
+            ]
+        )
+    )
+
+    # Call the function we are testing
+    result_message = _parse_response_candidate(candidate)
+
+    # Assert that the parsed tool_calls have the correct integer types
+    assert len(result_message.tool_calls) == 1
+    tool_call = result_message.tool_calls[0]
+    assert tool_call["name"] == "get_entity_lineage"
+    assert tool_call["args"]["upstream_depth"] == 3
+    assert tool_call["args"]["downstream_depth"] == 5
+    assert isinstance(tool_call["args"]["upstream_depth"], int)
+    assert isinstance(tool_call["args"]["downstream_depth"], int)
+
+    # Assert that non-integer values are preserved correctly
+    assert tool_call["args"]["valid_float"] == 3.14
+    assert isinstance(tool_call["args"]["valid_float"], float)
+    assert tool_call["args"]["string_param"] == "test_string"
+    assert isinstance(tool_call["args"]["string_param"], str)
+    assert tool_call["args"]["bool_param"] is True
+    assert isinstance(tool_call["args"]["bool_param"], bool)
+
+    # Assert that the additional_kwargs also contains corrected JSON
+    function_call_args = json.loads(
+        result_message.additional_kwargs["function_call"]["arguments"]
+    )
+    assert function_call_args["upstream_depth"] == 3
+    assert function_call_args["downstream_depth"] == 5
+    assert isinstance(function_call_args["upstream_depth"], int)
+    assert isinstance(function_call_args["downstream_depth"], int)
+
+    # Assert that non-integer values are preserved in additional_kwargs too
+    assert function_call_args["valid_float"] == 3.14
+    assert isinstance(function_call_args["valid_float"], float)
+
+
+def test_parse_response_candidate_handles_no_function_call() -> None:
+    """Test that the function works correctly when there's no function call."""
+    candidate = Candidate(
+        content=Content(
+            parts=[Part(text="This is a regular text response without function calls")]
+        )
+    )
+
+    result_message = _parse_response_candidate(candidate)
+
+    assert (
+        result_message.content
+        == "This is a regular text response without function calls"
+    )
+    assert len(result_message.tool_calls) == 0
+    assert "function_call" not in result_message.additional_kwargs
+
+
+def test_parse_response_candidate_handles_empty_args() -> None:
+    """Test that the function works correctly with empty function call arguments."""
+    args_struct = Struct()
+    # Empty struct - no arguments
+
+    candidate = Candidate(
+        content=Content(
+            parts=[
+                Part(
+                    function_call=FunctionCall(
+                        name="no_args_function",
+                        args=args_struct,
+                    )
+                )
+            ]
+        )
+    )
+
+    result_message = _parse_response_candidate(candidate)
+
+    assert len(result_message.tool_calls) == 1
+    tool_call = result_message.tool_calls[0]
+    assert tool_call["name"] == "no_args_function"
+    assert tool_call["args"] == {}
+
+    function_call_args = json.loads(
+        result_message.additional_kwargs["function_call"]["arguments"]
+    )
+    assert function_call_args == {}