fix: Convert float values with no decimals to integers in FunctionCall.

PiperOrigin-RevId: 695551016

Author: vertex-sdk-bot

tests/unit/vertexai/test_generative_models.py

@@ -18,7 +18,7 @@
 # pylint: disable=protected-access,bad-continuation
 import io
 import pytest
-from typing import Iterable, MutableSequence, Optional
+from typing import Dict, Iterable, List, MutableSequence, Optional
 from unittest import mock
 
 import vertexai
@@ -300,12 +300,31 @@ def mock_generate_content(
  if has_function_declarations and not had_any_function_calls:
  # response_part_struct = _RESPONSE_FUNCTION_CALL_PART_STRUCT
  # Workaround for the proto library bug
- response_part_struct = dict(
- function_call=gapic_tool_types.FunctionCall(
- name="get_current_weather",
- args={"location": "Boston"},
+ first_tool_with_function_declarations = next(
+ tool for tool in request.tools if tool.function_declarations
+ )
+ if (
+ first_tool_with_function_declarations.function_declarations[0].name
+ == update_weather_data.__name__
+ ):
+ response_part_struct = dict(
+ function_call=gapic_tool_types.FunctionCall(
+ name=update_weather_data.__name__,
+ args={
+ "location": "Boston",
+ "temperature": 60,
+ "forecasts": [61, 62],
+ "extra_info": {"humidity": 50},
+ },
+ )
+ )
+ else:
+ response_part_struct = dict(
+ function_call=gapic_tool_types.FunctionCall(
+ name="get_current_weather",
+ args={"location": "Boston"},
+ )
  )
- )
  elif has_function_declarations and latest_user_message_function_responses:
  function_response = latest_user_message_function_responses[0]
  function_response_dict = type(function_response).to_dict(function_response)
@@ -537,6 +556,19 @@ def get_current_weather(location: str, unit: Optional[str] = "centigrade"):
  )
 
 
+def update_weather_data(
+ location: str, temperature: int, forecasts: List[int], extra_info: Dict[str, int]
+):
+ """Updates the weather data in the specified location."""
+ return dict(
+ location=location,
+ temperature=temperature,
+ forecasts=forecasts,
+ extra_info=extra_info,
+ result="Success",
+ )
+
+
 @pytest.mark.usefixtures("google_auth_mock")
 @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
 class TestGenerativeModels:
@@ -839,6 +871,46 @@ def test_generate_content_streaming(self, generative_models: generative_models):
  for chunk in stream:
  assert chunk.text
 
+ @patch_genai_services
+ @pytest.mark.parametrize(
+ "generative_models",
+ [generative_models, preview_generative_models],
+ )
+ def test_generate_content_with_function_calling_integer_args(
+ self, generative_models: generative_models
+ ):
+ model = generative_models.GenerativeModel("gemini-pro")
+ weather_tool = generative_models.Tool(
+ function_declarations=[
+ generative_models.FunctionDeclaration.from_func(update_weather_data)
+ ],
+ )
+ response = model.generate_content(
+ "Please update the weather data in Boston, with temperature 60, "
+ "forecasts are 61 and 62, and humidity is 50",
+ tools=[weather_tool],
+ )
+ assert (
+ response.candidates[0].content.parts[0].function_call.name
+ == "update_weather_data"
+ )
+ assert [
+ function_call.name
+ for function_call in response.candidates[0].function_calls
+ ] == ["update_weather_data"]
+ for args in (
+ response.candidates[0].function_calls[0].args,
+ response.candidates[0].function_calls[0].to_dict()["args"],
+ ):
+ assert args["location"] == "Boston"
+ assert args["temperature"] == 60
+ assert isinstance(args["temperature"], int)
+ assert 61 in args["forecasts"] and 62 in args["forecasts"]
+ assert all(isinstance(forecast, int) for forecast in args["forecasts"])
+ assert args["extra_info"]["humidity"] == 50
+ assert isinstance(args["extra_info"]["humidity"], int)
+ assert args["location"] == "Boston"
+
  @patch_genai_services
  @pytest.mark.parametrize(
  "generative_models",
@@ -1052,6 +1124,52 @@ def test_chat_function_calling(self, generative_models: generative_models):
  assert "nice" in response2.text
  assert not response2.candidates[0].function_calls
 
+ @patch_genai_services
+ @pytest.mark.parametrize(
+ "generative_models",
+ [generative_models, preview_generative_models],
+ )
+ def test_chat_function_calling_with_integer_args(
+ self, generative_models: generative_models
+ ):
+ get_current_weather_func = generative_models.FunctionDeclaration.from_func(
+ update_weather_data
+ )
+ weather_tool = generative_models.Tool(
+ function_declarations=[get_current_weather_func],
+ )
+ model = generative_models.GenerativeModel(
+ "gemini-pro",
+ # Specifying the tools once to avoid specifying them in every request
+ tools=[weather_tool],
+ )
+ chat = model.start_chat()
+
+ response1 = chat.send_message(
+ "Please update the weather data in Boston, with temperature 60, "
+ "forecasts are 61 and 62, and humidity is 50"
+ )
+ assert (
+ response1.candidates[0].content.parts[0].function_call.name
+ == "update_weather_data"
+ )
+ assert [
+ function_call.name
+ for function_call in response1.candidates[0].function_calls
+ ] == ["update_weather_data"]
+ for args in (
+ response1.candidates[0].function_calls[0].args,
+ response1.candidates[0].function_calls[0].to_dict()["args"],
+ ):
+ assert args["location"] == "Boston"
+ assert args["temperature"] == 60
+ assert isinstance(args["temperature"], int)
+ assert 61 in args["forecasts"] and 62 in args["forecasts"]
+ assert all(isinstance(forecast, int) for forecast in args["forecasts"])
+ assert args["extra_info"]["humidity"] == 50
+ assert isinstance(args["extra_info"]["humidity"], int)
+ assert args["location"] == "Boston"
+
  @patch_genai_services
  @pytest.mark.parametrize(
  "generative_models",

vertexai/generative_models/_generative_models.py

@@ -2631,7 +2631,11 @@ def _from_gapic(cls, raw_message: aiplatform_types.FunctionCall) -> "FunctionCal
  return response
 
  def to_dict(self) -> Dict[str, Any]:
- return _proto_to_dict(self._raw_message)
+ function_call_dict = _proto_to_dict(self._raw_message)
+ function_call_dict["args"] = self._convert_number_values(
+ function_call_dict["args"]
+ )
+ return function_call_dict
 
  def __repr__(self) -> str:
  return self._raw_message.__repr__()
@@ -2644,7 +2648,21 @@ def name(self) -> str:
  def args(self) -> Dict[str, Any]:
  # We cannot use `type(self.args).to_dict(self.args)`
  # due to: AttributeError: type object 'MapComposite' has no attribute 'to_dict'
- return self.to_dict().get("args")
+ args_dict = self.to_dict().get("args")
+ return self._convert_number_values(args_dict)
+
+ def _convert_number_values(self, data: Any) -> Any:
+ """Converts float values with no decimal part to integers."""
+ if isinstance(data, float) and data.is_integer():
+ return int(data)
+ elif isinstance(data, dict):
+ return {
+ key: self._convert_number_values(value) for key, value in data.items()
+ }
+ elif isinstance(data, list):
+ return [self._convert_number_values(item) for item in data]
+ else:
+ return data
 
 
 class SafetySetting: