Skip to content

Commit c745d73

Browse files
committed
refactor(antigravity): 🔨 inline JSON schema references before sanitization
Introduced `_inline_schema_refs()` function to resolve local $ref definitions in JSON schemas before applying Claude-specific sanitization. This ensures that schema references are properly expanded and circular references are handled gracefully. - Added new helper function to recursively resolve $ref pointers from $defs and definitions - Circular reference detection prevents infinite loops by tracking seen references - Applied inlining step before `_clean_claude_schema()` in both tool transformation flows - Updated docstring to reflect the new two-step sanitization process (inline, then clean) - Removed extraneous `$schema` pop that was no longer needed This change improves schema compatibility with Antigravity/Gemini's Proto-based API by ensuring all references are resolved before unsupported JSON Schema keywords are stripped.
1 parent e4bf852 commit c745d73

File tree

1 file changed

+34
-8
lines changed

1 file changed

+34
-8
lines changed

‎src/rotator_library/providers/antigravity_provider.py‎

Lines changed: 34 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,33 @@ def _recursively_parse_json_strings(obj: Any) -> Any:
340340
return obj
341341

342342

343+
def _inline_schema_refs(schema: Dict[str, Any]) -> Dict[str, Any]:
344+
"""Inline local $ref definitions before sanitization."""
345+
if not isinstance(schema, dict):
346+
return schema
347+
348+
defs = schema.get("$defs", schema.get("definitions", {}))
349+
if not defs:
350+
return schema
351+
352+
def resolve(node, seen=()):
353+
if not isinstance(node, dict):
354+
return [resolve(x, seen) for x in node] if isinstance(node, list) else node
355+
if "$ref" in node:
356+
ref = node["$ref"]
357+
if ref in seen: # Circular - drop it
358+
return {k: resolve(v, seen) for k, v in node.items() if k != "$ref"}
359+
for prefix in ("#/$defs/", "#/definitions/"):
360+
if isinstance(ref, str) and ref.startswith(prefix):
361+
name = ref[len(prefix) :]
362+
if name in defs:
363+
return resolve(copy.deepcopy(defs[name]), seen + (ref,))
364+
return {k: resolve(v, seen) for k, v in node.items() if k != "$ref"}
365+
return {k: resolve(v, seen) for k, v in node.items()}
366+
367+
return resolve(schema)
368+
369+
343370
def _clean_claude_schema(schema: Any) -> Any:
344371
"""
345372
Recursively clean JSON Schema for Antigravity/Google's Proto-based API.
@@ -397,7 +424,6 @@ def _clean_claude_schema(schema: Any) -> Any:
397424
return first_option
398425

399426
cleaned = {}
400-
401427
# Handle 'const' by converting to 'enum' with single value
402428
if "const" in schema:
403429
const_value = schema["const"]
@@ -2507,8 +2533,10 @@ def _build_tools_payload(
25072533

25082534
if params and isinstance(params, dict):
25092535
schema = dict(params)
2510-
schema.pop("$schema", None)
25112536
schema.pop("strict", None)
2537+
# Inline $ref definitions, then strip unsupported keywords
2538+
schema = _inline_schema_refs(schema)
2539+
schema = _clean_claude_schema(schema)
25122540
schema = _normalize_type_arrays(schema)
25132541

25142542
# Workaround: Antigravity/Gemini fails to emit functionCall
@@ -2667,18 +2695,16 @@ def _apply_claude_tool_transform(self, payload: Dict[str, Any]) -> None:
26672695
"""Apply Claude-specific tool schema transformations.
26682696
26692697
Converts parametersJsonSchema to parameters and applies Claude-specific
2670-
schema cleaning (removes unsupported JSON Schema fields).
2698+
schema sanitization (inlines $ref, removes unsupported JSON Schema fields).
26712699
"""
26722700
tools = payload["request"].get("tools", [])
26732701
for tool in tools:
26742702
for func_decl in tool.get("functionDeclarations", []):
26752703
if "parametersJsonSchema" in func_decl:
26762704
params = func_decl["parametersJsonSchema"]
2677-
params = (
2678-
_clean_claude_schema(params)
2679-
if isinstance(params, dict)
2680-
else params
2681-
)
2705+
if isinstance(params, dict):
2706+
params = _inline_schema_refs(params)
2707+
params = _clean_claude_schema(params)
26822708
func_decl["parameters"] = params
26832709
del func_decl["parametersJsonSchema"]
26842710

0 commit comments

Comments
 (0)