Mirrowel commited on
Commit
c745d73
·
1 Parent(s): e4bf852

refactor(antigravity): 🔨 inline JSON schema references before sanitization

Browse files

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.

src/rotator_library/providers/antigravity_provider.py CHANGED
@@ -340,6 +340,33 @@ def _recursively_parse_json_strings(obj: Any) -> Any:
340
  return obj
341
 
342
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
343
  def _clean_claude_schema(schema: Any) -> Any:
344
  """
345
  Recursively clean JSON Schema for Antigravity/Google's Proto-based API.
@@ -397,7 +424,6 @@ def _clean_claude_schema(schema: Any) -> Any:
397
  return first_option
398
 
399
  cleaned = {}
400
-
401
  # Handle 'const' by converting to 'enum' with single value
402
  if "const" in schema:
403
  const_value = schema["const"]
@@ -2507,8 +2533,10 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface):
2507
 
2508
  if params and isinstance(params, dict):
2509
  schema = dict(params)
2510
- schema.pop("$schema", None)
2511
  schema.pop("strict", None)
 
 
 
2512
  schema = _normalize_type_arrays(schema)
2513
 
2514
  # Workaround: Antigravity/Gemini fails to emit functionCall
@@ -2667,18 +2695,16 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface):
2667
  """Apply Claude-specific tool schema transformations.
2668
 
2669
  Converts parametersJsonSchema to parameters and applies Claude-specific
2670
- schema cleaning (removes unsupported JSON Schema fields).
2671
  """
2672
  tools = payload["request"].get("tools", [])
2673
  for tool in tools:
2674
  for func_decl in tool.get("functionDeclarations", []):
2675
  if "parametersJsonSchema" in func_decl:
2676
  params = func_decl["parametersJsonSchema"]
2677
- params = (
2678
- _clean_claude_schema(params)
2679
- if isinstance(params, dict)
2680
- else params
2681
- )
2682
  func_decl["parameters"] = params
2683
  del func_decl["parametersJsonSchema"]
2684
 
 
340
  return obj
341
 
342
 
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
+
370
  def _clean_claude_schema(schema: Any) -> Any:
371
  """
372
  Recursively clean JSON Schema for Antigravity/Google's Proto-based API.
 
424
  return first_option
425
 
426
  cleaned = {}
 
427
  # Handle 'const' by converting to 'enum' with single value
428
  if "const" in schema:
429
  const_value = schema["const"]
 
2533
 
2534
  if params and isinstance(params, dict):
2535
  schema = dict(params)
 
2536
  schema.pop("strict", None)
2537
+ # Inline $ref definitions, then strip unsupported keywords
2538
+ schema = _inline_schema_refs(schema)
2539
+ schema = _clean_claude_schema(schema)
2540
  schema = _normalize_type_arrays(schema)
2541
 
2542
  # Workaround: Antigravity/Gemini fails to emit functionCall
 
2695
  """Apply Claude-specific tool schema transformations.
2696
 
2697
  Converts parametersJsonSchema to parameters and applies Claude-specific
2698
+ schema sanitization (inlines $ref, removes unsupported JSON Schema fields).
2699
  """
2700
  tools = payload["request"].get("tools", [])
2701
  for tool in tools:
2702
  for func_decl in tool.get("functionDeclarations", []):
2703
  if "parametersJsonSchema" in func_decl:
2704
  params = func_decl["parametersJsonSchema"]
2705
+ if isinstance(params, dict):
2706
+ params = _inline_schema_refs(params)
2707
+ params = _clean_claude_schema(params)
 
 
2708
  func_decl["parameters"] = params
2709
  del func_decl["parametersJsonSchema"]
2710