fix: convert Union[Pydantic, Pydantic] tool args at runtime

FunctionTool._preprocess_args only converted dict args to a Pydantic
model for single-model and Optional[Model] annotations. A
Union[ModelA, ModelB] parameter was left as a raw dict, so
isinstance checks inside the tool failed with "Unexpected entity
type: <class 'dict'>"

Use pydantic.TypeAdapter to validate against the full Union so
pydantic picks the matching member. None and instances of any
declared union member pass through unchanged; instances of
unrelated BaseModels fall back to the existing graceful-failure
warning path.

Close #5799

Change-Id: Ie69f8efc8395162eac375a0eaad0c77ed2097cec

Author: GWeale

contributing/samples/tools/pydantic_argument/tests/test_create_company.json

@@ -53,66 +53,23 @@
               "id": "fc-1",
               "name": "create_entity_profile",
               "response": {
-                "message": "Unexpected entity type: <class 'dict'>",
-                "status": "error"
-              }
-            }
-          }
-        ],
-        "role": "user"
-      },
-      "id": "e-3",
-      "invocationId": "i-1",
-      "nodeInfo": {
-        "path": "profile_agent@1"
-      }
-    },
-    {
-      "author": "profile_agent",
-      "content": {
-        "parts": [
-          {
-            "functionCall": {
-              "args": {
-                "entity": {
+                "entity_type": "company",
+                "message": "Company profile created for Acme Corp!",
+                "profile": {
                   "company_name": "Acme Corp",
                   "employee_count": 50,
-                  "industry": "tech"
-                }
-              },
-              "id": "fc-2",
-              "name": "create_entity_profile"
-            }
-          }
-        ],
-        "role": "model"
-      },
-      "finishReason": "STOP",
-      "id": "e-4",
-      "invocationId": "i-1",
-      "longRunningToolIds": [],
-      "nodeInfo": {
-        "path": "profile_agent@1"
-      }
-    },
-    {
-      "author": "profile_agent",
-      "content": {
-        "parts": [
-          {
-            "functionResponse": {
-              "id": "fc-2",
-              "name": "create_entity_profile",
-              "response": {
-                "message": "Unexpected entity type: <class 'dict'>",
-                "status": "error"
+                  "industry": "tech",
+                  "model_type": "CompanyProfile",
+                  "website": "Not provided"
+                },
+                "status": "company_profile_created"
               }
             }
           }
         ],
         "role": "user"
       },
-      "id": "e-5",
+      "id": "e-3",
       "invocationId": "i-1",
       "nodeInfo": {
         "path": "profile_agent@1"
@@ -123,13 +80,13 @@
       "content": {
         "parts": [
           {
-            "text": "I apologize for the repeated errors. It seems I am having trouble with the tool's expected input format. I will try again to create the company profile for Acme Corp. I believe the issue was in how I was structuring the data for the `create_entity_profile` function. I will now use the correct dataclass constructor to create the profile. Please bear with me."
+            "text": "I have created a profile for Acme Corp in the tech industry with 50 employees."
           }
         ],
         "role": "model"
       },
       "finishReason": "STOP",
-      "id": "e-6",
+      "id": "e-4",
       "invocationId": "i-1",
       "nodeInfo": {
         "path": "profile_agent@1"

src/google/adk/tools/function_tool.py

@@ -148,6 +148,24 @@ def _preprocess_args(self, args: dict[str, Any]) -> dict[str, Any]:
             ]
             if len(non_none_types) == 1:
               target_type = non_none_types[0]
+            elif len(non_none_types) > 1 and all(
+                inspect.isclass(t) and issubclass(t, pydantic.BaseModel)
+                for t in non_none_types
+            ):
+              if args[param_name] is None or isinstance(
+                  args[param_name], tuple(non_none_types)
+              ):
+                continue
+              try:
+                converted_args[param_name] = pydantic.TypeAdapter(
+                    param.annotation
+                ).validate_python(args[param_name])
+              except Exception as e:
+                logger.warning(
+                    f"Failed to convert argument '{param_name}' to"
+                    f' {param.annotation}: {e}'
+                )
+              continue
 
           # Check if the target type is a Pydantic model
           if inspect.isclass(target_type) and issubclass(

tests/unittests/tools/test_function_tool_pydantic.py

@@ -15,6 +15,7 @@
 # Pydantic model conversion tests
 
 from typing import Optional
+from typing import Union
 from unittest.mock import MagicMock
 
 from google.adk.agents.invocation_context import InvocationContext
@@ -40,6 +41,14 @@ class PreferencesModel(pydantic.BaseModel):
   notifications: bool = True
 
 
+class CompanyModel(pydantic.BaseModel):
+  """Test Pydantic model for company data."""
+
+  company_name: str
+  industry: str
+  employee_count: int
+
+
 def sync_function_with_pydantic_model(user: UserModel) -> dict:
   """Sync function that takes a Pydantic model."""
   return {
@@ -370,3 +379,145 @@ def place_order(orders: list[UserModel]) -> int:
   result = await tool.run_async(args=args, tool_context=tool_context_mock)
 
   assert result == 50
+
+
+def _function_with_union_of_basemodels(
+    entity: Union[UserModel, CompanyModel],
+) -> str:
+  return type(entity).__name__
+
+
+def test_preprocess_args_with_union_of_basemodels_picks_user():
+  """Dict matching UserModel is converted to UserModel."""
+  tool = FunctionTool(_function_with_union_of_basemodels)
+
+  processed_args = tool._preprocess_args(
+      {"entity": {"name": "Diana", "age": 32, "email": "d@example.com"}}
+  )
+
+  assert isinstance(processed_args["entity"], UserModel)
+  assert processed_args["entity"].name == "Diana"
+
+
+def test_preprocess_args_with_union_of_basemodels_picks_company():
+  """Dict matching CompanyModel is converted to CompanyModel."""
+  tool = FunctionTool(_function_with_union_of_basemodels)
+
+  processed_args = tool._preprocess_args({
+      "entity": {
+          "company_name": "Acme Corp",
+          "industry": "tech",
+          "employee_count": 50,
+      }
+  })
+
+  assert isinstance(processed_args["entity"], CompanyModel)
+  assert processed_args["entity"].company_name == "Acme Corp"
+
+
+def test_preprocess_args_with_union_of_basemodels_existing_instance_unchanged():
+  """Existing instance of any union member is left unchanged."""
+  tool = FunctionTool(_function_with_union_of_basemodels)
+
+  user = UserModel(name="Bob", age=25)
+  assert tool._preprocess_args({"entity": user})["entity"] is user
+
+  company = CompanyModel(
+      company_name="Acme", industry="tech", employee_count=10
+  )
+  assert tool._preprocess_args({"entity": company})["entity"] is company
+
+
+def test_preprocess_args_with_union_of_basemodels_unrelated_instance_passthrough():
+  """A BaseModel instance not in the union is not silently accepted."""
+  tool = FunctionTool(_function_with_union_of_basemodels)
+
+  class UnrelatedModel(pydantic.BaseModel):
+    name: str
+    age: int
+
+  unrelated = UnrelatedModel(name="Carol", age=20)
+  processed_args = tool._preprocess_args({"entity": unrelated})
+
+  # Conversion fails (UnrelatedModel is not in the union); value is left
+  # alone so the function receives it and raises a clear error itself.
+  assert processed_args["entity"] is unrelated
+
+
+def test_preprocess_args_with_optional_union_of_basemodels_none():
+  """Optional[Union[A, B]] passes None through unchanged."""
+
+  def fn(entity: Optional[Union[UserModel, CompanyModel]] = None) -> str:
+    return type(entity).__name__
+
+  tool = FunctionTool(fn)
+
+  processed_args = tool._preprocess_args({"entity": None})
+
+  assert processed_args["entity"] is None
+
+
+def test_preprocess_args_with_optional_union_of_basemodels_dict():
+  """Optional[Union[A, B]] converts a dict to the matching model."""
+
+  def fn(entity: Optional[Union[UserModel, CompanyModel]] = None) -> str:
+    return type(entity).__name__
+
+  tool = FunctionTool(fn)
+
+  processed_args = tool._preprocess_args({"entity": {"name": "Eve", "age": 40}})
+
+  assert isinstance(processed_args["entity"], UserModel)
+  assert processed_args["entity"].name == "Eve"
+
+
+def test_preprocess_args_with_union_of_basemodels_invalid_data():
+  """Invalid data for Union[BaseModel, BaseModel] is kept unchanged."""
+  tool = FunctionTool(_function_with_union_of_basemodels)
+
+  # Dict matches neither model.
+  processed_args = tool._preprocess_args(
+      {"entity": {"unrelated_field": "value"}}
+  )
+
+  assert processed_args["entity"] == {"unrelated_field": "value"}
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_union_of_basemodels():
+  """run_async end-to-end converts dict to the matching union member."""
+
+  def create_entity_profile(
+      entity: Union[UserModel, CompanyModel],
+  ) -> dict:
+    if isinstance(entity, UserModel):
+      return {"entity_type": "user", "name": entity.name}
+    if isinstance(entity, CompanyModel):
+      return {"entity_type": "company", "name": entity.company_name}
+    return {"entity_type": "unknown"}
+
+  tool = FunctionTool(create_entity_profile)
+
+  tool_context_mock = MagicMock(spec=ToolContext)
+  invocation_context_mock = MagicMock(spec=InvocationContext)
+  session_mock = MagicMock(spec=Session)
+  invocation_context_mock.session = session_mock
+  tool_context_mock.invocation_context = invocation_context_mock
+
+  user_result = await tool.run_async(
+      args={"entity": {"name": "Diana", "age": 32}},
+      tool_context=tool_context_mock,
+  )
+  assert user_result == {"entity_type": "user", "name": "Diana"}
+
+  company_result = await tool.run_async(
+      args={
+          "entity": {
+              "company_name": "Acme Corp",
+              "industry": "tech",
+              "employee_count": 50,
+          }
+      },
+      tool_context=tool_context_mock,
+  )
+  assert company_result == {"entity_type": "company", "name": "Acme Corp"}