fix: google gen ai format error fix (#9147)
* 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 <noreply@letta.com> --------- Co-authored-by: Letta <noreply@letta.com>
This commit is contained in:
Kian Jones
@@ -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
|
||||
|
||||
178
tests/test_google_schema_refs.py
Normal file
178
tests/test_google_schema_refs.py
Normal file
@@ -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"]
|
||||
Reference in New Issue
Block a user