4040from pydantic import create_model
4141from pydantic import fields as pydantic_fields
4242
43+ from ..utils .variant_utils import get_google_llm_variant
44+ from ..utils .variant_utils import GoogleLLMVariant
45+
4346
4447def _get_function_fields (
4548 func : Callable [..., Any ],
@@ -114,6 +117,77 @@ def get_callable_name(func: Callable[..., Any]) -> str:
114117 return getattr (func , '__name__' , None ) or func .__class__ .__name__
115118
116119
120+ def _flatten_optional_any_of (schema : dict [str , Any ]) -> dict [str , Any ]:
121+ """Flattens `Optional[X]`-style `anyOf` schemas for Vertex AI.
122+
123+ Pydantic serializes `Optional[X]` fields as
124+ `{"anyOf": [<X schema>, {"type": "null"}], ...}`. Vertex AI rejects such
125+ schemas because the wrapping schema itself has no top-level `type` field
126+ (see https://github.com/googleapis/python-genai/issues/1807), even though
127+ the Gemini Developer API (AI Studio) accepts them as-is. This merges the
128+ non-null branch into the parent schema and marks it `nullable` instead,
129+ which Vertex AI accepts.
130+
131+ True unions with more than one non-null variant (e.g. `Union[int, str]`)
132+ can't be losslessly flattened this way and are left untouched.
133+ """
134+ any_of = schema .get ('anyOf' )
135+ if not isinstance (any_of , list ):
136+ return schema
137+
138+ # We need at least one null variant and at least one non-null variant
139+ # to flatten an optional/nullable schema.
140+ if len (any_of ) < 2 :
141+ return schema
142+
143+ null_variants = [
144+ variant
145+ for variant in any_of
146+ if isinstance (variant , dict ) and variant .get ('type' ) == 'null'
147+ ]
148+ # We only flatten if there is exactly one null variant (i.e. it is optional).
149+ if len (null_variants ) != 1 :
150+ return schema
151+
152+ non_null_variants = [v for v in any_of if v not in null_variants ]
153+ if len (non_null_variants ) == 1 :
154+ # Optional[X] -> X + nullable: true
155+ flattened = dict (non_null_variants [0 ])
156+ for key , value in schema .items ():
157+ if key != 'anyOf' :
158+ flattened .setdefault (key , value )
159+ flattened ['nullable' ] = True
160+ return flattened
161+ else :
162+ # Optional[Union[A, B, ...]] -> Union[A, B, ...] + nullable: true
163+ # We keep the anyOf but remove the null variant.
164+ flattened = dict (schema )
165+ flattened ['anyOf' ] = non_null_variants
166+ flattened ['nullable' ] = True
167+ return flattened
168+
169+
170+ def _sanitize_json_schema_for_vertex (schema : Any ) -> Any :
171+ """Recursively rewrites `Optional[X]` `anyOf` schemas for Vertex AI.
172+
173+ Vertex AI's schema validator requires every (sub)schema to declare a
174+ top-level `type`, which Pydantic's `anyOf`-based representation of
175+ `Optional`/`Union` fields does not provide. This is only applied for the
176+ Vertex AI backend since the Gemini Developer API (AI Studio) already
177+ accepts the unmodified Pydantic schema.
178+ """
179+ if isinstance (schema , list ):
180+ return [_sanitize_json_schema_for_vertex (item ) for item in schema ]
181+ if not isinstance (schema , dict ):
182+ return schema
183+
184+ sanitized = {
185+ key : _sanitize_json_schema_for_vertex (value )
186+ for key , value in schema .items ()
187+ }
188+ return _flatten_optional_any_of (sanitized )
189+
190+
117191def _build_parameters_json_schema (
118192 func : Callable [..., Any ],
119193 ignore_params : Optional [list [str ]] = None ,
@@ -245,9 +319,13 @@ def build_function_declaration_with_json_schema(
245319 >>> decl.name
246320 'paint_room'
247321 """
322+ is_vertex_ai = get_google_llm_variant () == GoogleLLMVariant .VERTEX_AI
323+
248324 # Handle Pydantic BaseModel classes
249325 if isinstance (func , type ) and issubclass (func , pydantic .BaseModel ):
250326 schema = func .model_json_schema ()
327+ if is_vertex_ai :
328+ schema = _sanitize_json_schema_for_vertex (schema )
251329 description = inspect .cleandoc (func .__doc__ ) if func .__doc__ else None
252330 return types .FunctionDeclaration (
253331 name = func .__name__ ,
@@ -265,10 +343,14 @@ def build_function_declaration_with_json_schema(
265343
266344 parameters_schema = _build_parameters_json_schema (func , ignore_params )
267345 if parameters_schema :
346+ if is_vertex_ai :
347+ parameters_schema = _sanitize_json_schema_for_vertex (parameters_schema )
268348 declaration .parameters_json_schema = parameters_schema
269349
270350 response_schema = _build_response_json_schema (func )
271351 if response_schema :
352+ if is_vertex_ai :
353+ response_schema = _sanitize_json_schema_for_vertex (response_schema )
272354 declaration .response_json_schema = response_schema
273355
274356 return declaration
0 commit comments