fix: sanitize anyOf schemas for Vertex AI function declarations

Merge #6381

Flatten the common Optional[X] -> anyOf: [X, {type: null}] pattern into X + nullable: true for parameters_json_schema and response_json_schema when targeting Vertex AI, gated on GoogleLLMVariant.VERTEX_AI.

Closes: #6373
PiperOrigin-RevId: 971577962

Author: ShauryaaSharma

src/google/adk/cli/conformance/_conformance_test_google_llm.py

@@ -61,7 +61,14 @@ def _resolve_refs(data: Any, defs: dict[str, Any]) -> Any:
       if ref_path.startswith('#/$defs/'):
         def_name = ref_path.split('/')[-1]
         if def_name in defs:
-          return _resolve_refs(defs[def_name], defs)
+          resolved = _resolve_refs(defs[def_name], defs)
+          if isinstance(resolved, dict):
+            ref_copy = data.copy()
+            del ref_copy['$ref']
+            resolved_copy = resolved.copy()
+            resolved_copy.update(ref_copy)
+            return resolved_copy
+          return resolved
     return {k: _resolve_refs(v, defs) for k, v in data.items()}
   elif isinstance(data, list):
     return [_resolve_refs(x, defs) for x in data]

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],
@@ -114,6 +117,77 @@ def get_callable_name(func: Callable[..., Any]) -> str:
   return getattr(func, '__name__', None) or func.__class__.__name__
 
 
+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):
+    return schema
+
+  # We need at least one null variant and at least one non-null variant
+  # to flatten an optional/nullable schema.
+  if len(any_of) < 2:
+    return schema
+
+  null_variants = [
+      variant
+      for variant in any_of
+      if isinstance(variant, dict) and variant.get('type') == 'null'
+  ]
+  # We only flatten if there is exactly one null variant (i.e. it is optional).
+  if len(null_variants) != 1:
+    return schema
+
+  non_null_variants = [v for v in any_of if v not in null_variants]
+  if len(non_null_variants) == 1:
+    # Optional[X] -> X + nullable: true
+    flattened = dict(non_null_variants[0])
+    for key, value in schema.items():
+      if key != 'anyOf':
+        flattened.setdefault(key, value)
+    flattened['nullable'] = True
+    return flattened
+  else:
+    # Optional[Union[A, B, ...]] -> Union[A, B, ...] + nullable: true
+    # We keep the anyOf but remove the null variant.
+    flattened = dict(schema)
+    flattened['anyOf'] = non_null_variants
+    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,
@@ -245,9 +319,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__,
@@ -265,10 +343,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

@@ -23,18 +23,21 @@
 from collections.abc import Sequence
 import dataclasses
 from enum import Enum
+import os
 from typing import Any
 from typing import AsyncGenerator
 from typing import Generator
 from typing import Literal
 from typing import Optional
+from unittest import mock
 
 from absl.testing import parameterized
 from google.adk.tools._function_tool_declarations import build_function_declaration_with_json_schema
 from google.adk.tools.tool_context import ToolContext
 from pydantic import BaseModel
 from pydantic import Field
 from pydantic.dataclasses import dataclass as pyd_dataclass
+import pytest
 
 
 class Color(Enum):
@@ -837,6 +840,97 @@ def complex_fn(
     )
 
 
+@mock.patch.dict(os.environ, {"GOOGLE_GENAI_USE_ENTERPRISE": "true"})
+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.
+  """
+
+  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)
+
+  def test_optional_union_flattened(self):
+    """Optional[Union[int, str]] should flatten to Union[int, str] + nullable."""
+
+    def process(value: Optional[int | str] = None) -> str:
+      """Process an optional union 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)
+    types = {v.get("type") for v in value_schema["anyOf"]}
+    self.assertEqual(types, {"integer", "string"})
+    self.assertTrue(value_schema["nullable"])
+
+
 class TestPydanticModelAsFunction(parameterized.TestCase):
   """Tests for using Pydantic BaseModel directly."""