fix(tools): sanitize anyOf schemas for Vertex AI function declarations

src/google/adk/tools/_function_tool_declarations.py

@@ -40,6 +40,9 @@
 from pydantic import create_model
 from pydantic import fields as pydantic_fields
 
+from ..utils.variant_utils import get_google_llm_variant
+from ..utils.variant_utils import GoogleLLMVariant
+
 
 def _get_function_fields(
     func: Callable[..., Any],
@@ -97,6 +100,64 @@ def _get_function_fields(
   return fields
 
 
+def _flatten_optional_any_of(schema: dict[str, Any]) -> dict[str, Any]:
+  """Flattens `Optional[X]`-style `anyOf` schemas for Vertex AI.
+
+  Pydantic serializes `Optional[X]` fields as
+  `{"anyOf": [<X schema>, {"type": "null"}], ...}`. Vertex AI rejects such
+  schemas because the wrapping schema itself has no top-level `type` field
+  (see https://github.com/googleapis/python-genai/issues/1807), even though
+  the Gemini Developer API (AI Studio) accepts them as-is. This merges the
+  non-null branch into the parent schema and marks it `nullable` instead,
+  which Vertex AI accepts.
+
+  True unions with more than one non-null variant (e.g. `Union[int, str]`)
+  can't be losslessly flattened this way and are left untouched.
+  """
+  any_of = schema.get('anyOf')
+  if not isinstance(any_of, list) or len(any_of) != 2:
+    return schema
+
+  null_variants = [
+      variant
+      for variant in any_of
+      if isinstance(variant, dict) and variant.get('type') == 'null'
+  ]
+  non_null_variants = [
+      variant for variant in any_of if variant not in null_variants
+  ]
+  if len(null_variants) != 1 or len(non_null_variants) != 1:
+    return schema
+
+  flattened = dict(non_null_variants[0])
+  for key, value in schema.items():
+    if key != 'anyOf':
+      flattened.setdefault(key, value)
+  flattened['nullable'] = True
+  return flattened
+
+
+def _sanitize_json_schema_for_vertex(schema: Any) -> Any:
+  """Recursively rewrites `Optional[X]` `anyOf` schemas for Vertex AI.
+
+  Vertex AI's schema validator requires every (sub)schema to declare a
+  top-level `type`, which Pydantic's `anyOf`-based representation of
+  `Optional`/`Union` fields does not provide. This is only applied for the
+  Vertex AI backend since the Gemini Developer API (AI Studio) already
+  accepts the unmodified Pydantic schema.
+  """
+  if isinstance(schema, list):
+    return [_sanitize_json_schema_for_vertex(item) for item in schema]
+  if not isinstance(schema, dict):
+    return schema
+
+  sanitized = {
+      key: _sanitize_json_schema_for_vertex(value)
+      for key, value in schema.items()
+  }
+  return _flatten_optional_any_of(sanitized)
+
+
 def _build_parameters_json_schema(
     func: Callable[..., Any],
     ignore_params: Optional[list[str]] = None,
@@ -228,9 +289,13 @@ def build_function_declaration_with_json_schema(
     >>> decl.name
     'paint_room'
   """
+  is_vertex_ai = get_google_llm_variant() == GoogleLLMVariant.VERTEX_AI
+
   # Handle Pydantic BaseModel classes
   if isinstance(func, type) and issubclass(func, pydantic.BaseModel):
     schema = func.model_json_schema()
+    if is_vertex_ai:
+      schema = _sanitize_json_schema_for_vertex(schema)
     description = inspect.cleandoc(func.__doc__) if func.__doc__ else None
     return types.FunctionDeclaration(
         name=func.__name__,
@@ -248,10 +313,14 @@ def build_function_declaration_with_json_schema(
 
   parameters_schema = _build_parameters_json_schema(func, ignore_params)
   if parameters_schema:
+    if is_vertex_ai:
+      parameters_schema = _sanitize_json_schema_for_vertex(parameters_schema)
     declaration.parameters_json_schema = parameters_schema
 
   response_schema = _build_response_json_schema(func)
   if response_schema:
+    if is_vertex_ai:
+      response_schema = _sanitize_json_schema_for_vertex(response_schema)
     declaration.response_json_schema = response_schema
 
   return declaration

tests/unittests/tools/test_function_tool_declarations.py

@@ -35,6 +35,7 @@
 from pydantic import BaseModel
 from pydantic import Field
 from pydantic.dataclasses import dataclass as pyd_dataclass
+import pytest
 
 
 class Color(Enum):
@@ -837,6 +838,83 @@ def complex_fn(
     )
 
 
+class TestVertexAiAnyOfSanitization(parameterized.TestCase):
+  """Tests that `Optional`/`Union` `anyOf` schemas are made Vertex-safe.
+
+  Vertex AI rejects schemas where an `anyOf` wrapper has no top-level `type`
+  (see https://github.com/googleapis/python-genai/issues/1807), while AI
+  Studio accepts them unmodified. This is only exercised when the
+  enterprise/Vertex variant is active.
+  """
+
+  @pytest.fixture(autouse=True)
+  def _use_vertex_variant(self, monkeypatch):
+    monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "true")
+
+  def test_optional_field_flattened_to_nullable(self):
+    """Optional[list[str]] should become a plain array schema + nullable."""
+
+    def search(query: str, sources: Optional[list[str]] = None) -> str:
+      """Search using optional sources."""
+      return query
+
+    decl = build_function_declaration_with_json_schema(search)
+    schema = decl.parameters_json_schema
+
+    sources_schema = schema["properties"]["sources"]
+    self.assertNotIn("anyOf", sources_schema)
+    self.assertEqual(sources_schema["type"], "array")
+    self.assertEqual(sources_schema["items"]["type"], "string")
+    self.assertTrue(sources_schema["nullable"])
+    self.assertIsNone(sources_schema["default"])
+
+  def test_optional_pydantic_model_field_flattened(self):
+    """Optional[BaseModel] fields should also be flattened, not just primitives."""
+
+    def save_address(address: Optional[Address] = None) -> str:
+      """Save an optional address."""
+      return "ok"
+
+    decl = build_function_declaration_with_json_schema(save_address)
+    schema = decl.parameters_json_schema
+
+    address_schema = schema["properties"]["address"]
+    self.assertNotIn("anyOf", address_schema)
+    self.assertIn("$ref", address_schema)
+    self.assertTrue(address_schema["nullable"])
+
+  def test_output_schema_pydantic_model_flattened(self):
+    """Optional fields on a BaseModel passed directly should be flattened."""
+
+    class CoordinatorResponse(BaseModel):
+      """A coordinator response."""
+
+      answer: str
+      sources: Optional[list[str]] = None
+
+    decl = build_function_declaration_with_json_schema(CoordinatorResponse)
+    schema = decl.parameters_json_schema
+
+    sources_schema = schema["properties"]["sources"]
+    self.assertNotIn("anyOf", sources_schema)
+    self.assertEqual(sources_schema["type"], "array")
+    self.assertTrue(sources_schema["nullable"])
+
+  def test_true_union_left_untouched(self):
+    """A real multi-variant union (not Optional) is not flattened."""
+
+    def process(value: int | str) -> str:
+      """Process a value."""
+      return str(value)
+
+    decl = build_function_declaration_with_json_schema(process)
+    schema = decl.parameters_json_schema
+
+    value_schema = schema["properties"]["value"]
+    self.assertIn("anyOf", value_schema)
+    self.assertLen(value_schema["anyOf"], 2)
+
+
 class TestPydanticModelAsFunction(parameterized.TestCase):
   """Tests for using Pydantic BaseModel directly."""