From 424a1ada6401d5eef99ce724ba6538e4cd210a34 Mon Sep 17 00:00:00 2001 From: Kian Jones <11655409+kianjones9@users.noreply.github.com> Date: Tue, 10 Feb 2026 20:32:45 -0800 Subject: [PATCH] fix: google gen ai format error fix (#9147) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * google gen ai format error fix * fix(core): add $ref safety net, warning log, and unit tests for Google schema resolution - Add `$ref` to unsupported_keys in `_clean_google_ai_schema_properties` so unresolvable refs (e.g. `#/properties/...` style) are stripped as a safety net instead of crashing the Google SDK - Add warning log when `_resolve_json_schema_refs` encounters a ref it cannot resolve - Deduplicate the `#/$defs/` and `#/definitions/` resolution branches - Add 11 unit tests covering: single/multiple $defs, nested refs, refs in anyOf/allOf, array items, definitions key, unresolvable refs, and the full resolve+clean pipeline 🤖 Generated with [Letta Code](https://letta.com) Co-Authored-By: Letta --------- Co-authored-by: Letta --- letta/llm_api/google_vertex_client.py | 47 ++++++- tests/test_google_schema_refs.py | 178 ++++++++++++++++++++++++++ 2 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 tests/test_google_schema_refs.py diff --git a/letta/llm_api/google_vertex_client.py b/letta/llm_api/google_vertex_client.py index 18a1a1b1..9b41e918 100644 --- a/letta/llm_api/google_vertex_client.py +++ b/letta/llm_api/google_vertex_client.py @@ -247,7 +247,7 @@ class GoogleVertexClient(LLMClientBase): # Per https://ai.google.dev/gemini-api/docs/function-calling?example=meeting#notes_and_limitations # * Only a subset of the OpenAPI schema is supported. # * Supported parameter types in Python are limited. - unsupported_keys = ["default", "exclusiveMaximum", "exclusiveMinimum", "additionalProperties", "$schema", "const"] + unsupported_keys = ["default", "exclusiveMaximum", "exclusiveMinimum", "additionalProperties", "$schema", "const", "$ref"] keys_to_remove_at_this_level = [key for key in unsupported_keys if key in schema_part] for key_to_remove in keys_to_remove_at_this_level: logger.debug(f"Removing unsupported keyword '{key_to_remove}' from schema part.") @@ -274,6 +274,49 @@ class GoogleVertexClient(LLMClientBase): for item_schema in schema_part[key]: self._clean_google_ai_schema_properties(item_schema) + def _resolve_json_schema_refs(self, schema: dict, defs: dict = None) -> dict: + """ + Recursively resolve $ref in JSON schema by inlining definitions. + Google GenAI SDK does not support $ref. + """ + if defs is None: + # Look for definitions at the top level + defs = schema.get("$defs") or schema.get("definitions") or {} + + if not isinstance(schema, dict): + return schema + + # If this is a ref, resolve it + if "$ref" in schema: + ref = schema["$ref"] + if isinstance(ref, str): + for prefix in ("#/$defs/", "#/definitions/"): + if ref.startswith(prefix): + ref_name = ref.split("/")[-1] + if ref_name in defs: + resolved = defs[ref_name].copy() + return self._resolve_json_schema_refs(resolved, defs) + break + + logger.warning(f"Could not resolve $ref '{ref}' in schema — will be stripped by schema cleaner") + + # Recursively process children + new_schema = schema.copy() + + # We need to remove $defs/definitions from the output schema as Google doesn't support them + if "$defs" in new_schema: + del new_schema["$defs"] + if "definitions" in new_schema: + del new_schema["definitions"] + + for k, v in new_schema.items(): + if isinstance(v, dict): + new_schema[k] = self._resolve_json_schema_refs(v, defs) + elif isinstance(v, list): + new_schema[k] = [self._resolve_json_schema_refs(i, defs) if isinstance(i, dict) else i for i in v] + + return new_schema + def convert_tools_to_google_ai_format(self, tools: List[Tool], llm_config: LLMConfig) -> List[dict]: """ OpenAI style: @@ -336,6 +379,8 @@ class GoogleVertexClient(LLMClientBase): # Google AI API only supports a subset of OpenAPI 3.0, so unsupported params must be cleaned if "parameters" in func and isinstance(func["parameters"], dict): + # Resolve $ref in schema because Google AI SDK doesn't support them + func["parameters"] = self._resolve_json_schema_refs(func["parameters"]) self._clean_google_ai_schema_properties(func["parameters"]) # Add inner thoughts diff --git a/tests/test_google_schema_refs.py b/tests/test_google_schema_refs.py new file mode 100644 index 00000000..ff10e1cc --- /dev/null +++ b/tests/test_google_schema_refs.py @@ -0,0 +1,178 @@ +"""Unit tests for GoogleVertexClient._resolve_json_schema_refs and $ref safety net.""" + +import pytest + +from letta.llm_api.google_vertex_client import GoogleVertexClient + + +@pytest.fixture +def client(): + return GoogleVertexClient() + + +class TestResolveJsonSchemaRefs: + def test_single_def_with_ref(self, client): + schema = { + "type": "object", + "properties": { + "status": {"$ref": "#/$defs/StatusEnum"}, + }, + "$defs": { + "StatusEnum": {"type": "string", "enum": ["active", "inactive"]}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert "$defs" not in result + assert result["properties"]["status"] == {"type": "string", "enum": ["active", "inactive"]} + + def test_multiple_defs(self, client): + schema = { + "type": "object", + "properties": { + "ticket": {"$ref": "#/$defs/TicketStatus"}, + "report": {"$ref": "#/$defs/ReportType"}, + }, + "$defs": { + "TicketStatus": {"type": "string", "enum": ["open", "closed"]}, + "ReportType": {"type": "string", "enum": ["summary", "detailed"]}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert "$defs" not in result + assert result["properties"]["ticket"] == {"type": "string", "enum": ["open", "closed"]} + assert result["properties"]["report"] == {"type": "string", "enum": ["summary", "detailed"]} + + def test_nested_ref_in_def(self, client): + schema = { + "type": "object", + "properties": { + "order": {"$ref": "#/$defs/Order"}, + }, + "$defs": { + "Order": { + "type": "object", + "properties": { + "status": {"$ref": "#/$defs/OrderStatus"}, + }, + }, + "OrderStatus": {"type": "string", "enum": ["pending", "shipped"]}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert "$defs" not in result + assert result["properties"]["order"]["properties"]["status"] == {"type": "string", "enum": ["pending", "shipped"]} + + def test_ref_inside_anyof(self, client): + schema = { + "type": "object", + "properties": { + "value": { + "anyOf": [ + {"$ref": "#/$defs/StringVal"}, + {"type": "null"}, + ] + }, + }, + "$defs": { + "StringVal": {"type": "string", "maxLength": 100}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert "$defs" not in result + assert result["properties"]["value"]["anyOf"][0] == {"type": "string", "maxLength": 100} + assert result["properties"]["value"]["anyOf"][1] == {"type": "null"} + + def test_ref_inside_allof(self, client): + schema = { + "type": "object", + "properties": { + "item": {"allOf": [{"$ref": "#/$defs/Base"}, {"type": "object", "properties": {"extra": {"type": "string"}}}]}, + }, + "$defs": { + "Base": {"type": "object", "properties": {"name": {"type": "string"}}}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert result["properties"]["item"]["allOf"][0] == {"type": "object", "properties": {"name": {"type": "string"}}} + + def test_no_defs_is_noop(self, client): + schema = { + "type": "object", + "properties": { + "name": {"type": "string"}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert result == schema + + def test_definitions_key(self, client): + schema = { + "type": "object", + "properties": { + "role": {"$ref": "#/definitions/Role"}, + }, + "definitions": { + "Role": {"type": "string", "enum": ["admin", "user"]}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert "definitions" not in result + assert result["properties"]["role"] == {"type": "string", "enum": ["admin", "user"]} + + def test_unresolvable_ref_logged(self, client): + schema = { + "type": "object", + "properties": { + "thing": {"$ref": "#/properties/other/nested"}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert "$ref" in result["properties"]["thing"] + + def test_ref_in_array_items(self, client): + schema = { + "type": "object", + "properties": { + "tags": { + "type": "array", + "items": {"$ref": "#/$defs/Tag"}, + }, + }, + "$defs": { + "Tag": {"type": "string", "enum": ["urgent", "low"]}, + }, + } + result = client._resolve_json_schema_refs(schema) + assert "$defs" not in result + assert result["properties"]["tags"]["items"] == {"type": "string", "enum": ["urgent", "low"]} + + +class TestCleanSchemaStripsUnresolvedRefs: + def test_ref_stripped_by_cleaner(self, client): + schema = { + "type": "object", + "properties": { + "thing": {"$ref": "#/properties/other/nested", "type": "string"}, + }, + } + client._clean_google_ai_schema_properties(schema) + assert "$ref" not in schema["properties"]["thing"] + assert schema["properties"]["thing"]["type"] == "string" + + def test_full_pipeline_resolves_then_cleans(self, client): + schema = { + "type": "object", + "properties": { + "status": {"$ref": "#/$defs/Status"}, + "weird": {"$ref": "#/properties/foo/bar", "type": "string"}, + }, + "$defs": { + "Status": {"type": "string", "enum": ["a", "b"], "default": "a"}, + }, + } + resolved = client._resolve_json_schema_refs(schema) + client._clean_google_ai_schema_properties(resolved) + assert "$defs" not in resolved + assert "$ref" not in resolved["properties"]["weird"] + assert resolved["properties"]["status"]["enum"] == ["a", "b"] + assert "default" not in resolved["properties"]["status"]