ADK changes

PiperOrigin-RevId: 813321782

Author: xuanyang15

src/google/adk/tools/_gemini_schema_util.py

@@ -99,6 +99,44 @@ def _sanitize_schema_type(schema: dict[str, Any]) -> dict[str, Any]:
   return schema
 
 
+def _dereference_schema(schema: dict[str, Any]) -> dict[str, Any]:
+  """Resolves $ref pointers in a JSON schema."""
+
+  defs = schema.get("$defs", {})
+
+  def _resolve_refs(sub_schema: Any) -> Any:
+    if isinstance(sub_schema, dict):
+      if "$ref" in sub_schema:
+        ref_key = sub_schema["$ref"].split("/")[-1]
+        if ref_key in defs:
+          # Found the reference, replace it with the definition.
+          resolved = defs[ref_key].copy()
+          # Merge properties from the reference, allowing overrides.
+          sub_schema_copy = sub_schema.copy()
+          del sub_schema_copy["$ref"]
+          resolved.update(sub_schema_copy)
+          # Recursively resolve refs in the newly inserted part.
+          return _resolve_refs(resolved)
+        else:
+          # Reference not found, return as is.
+          return sub_schema
+      else:
+        # No $ref, so traverse deeper into the dictionary.
+        return {key: _resolve_refs(value) for key, value in sub_schema.items()}
+    elif isinstance(sub_schema, list):
+      # Traverse into lists.
+      return [_resolve_refs(item) for item in sub_schema]
+    else:
+      # Not a dict or list, return as is.
+      return sub_schema
+
+  dereferenced_schema = _resolve_refs(schema)
+  # Remove the definitions block after resolving.
+  if "$defs" in dereferenced_schema:
+    del dereferenced_schema["$defs"]
+  return dereferenced_schema
+
+
 def _sanitize_schema_formats_for_gemini(
     schema: dict[str, Any],
 ) -> dict[str, Any]:
@@ -109,7 +147,10 @@ def _sanitize_schema_formats_for_gemini(
       "any_of",  # 'one_of', 'all_of', 'not' to come
   }
   snake_case_schema = {}
-  dict_schema_field_names: tuple[str] = ("properties",)  # 'defs' to come
+  dict_schema_field_names: tuple[str, ...] = (
+      "properties",
+      "defs",
+  )
   for field_name, field_value in schema.items():
     field_name = _to_snake_case(field_name)
     if field_name in schema_field_names:
@@ -151,8 +192,9 @@ def _to_gemini_schema(openapi_schema: dict[str, Any]) -> Schema:
   if not isinstance(openapi_schema, dict):
     raise TypeError("openapi_schema must be a dictionary")
 
-  openapi_schema = _sanitize_schema_formats_for_gemini(openapi_schema)
+  dereferenced_schema = _dereference_schema(openapi_schema)
+  sanitized_schema = _sanitize_schema_formats_for_gemini(dereferenced_schema)
   return Schema.from_json_schema(
-      json_schema=_ExtendedJSONSchema.model_validate(openapi_schema),
+      json_schema=_ExtendedJSONSchema.model_validate(sanitized_schema),
       api_option=get_google_llm_variant(),
   )

tests/unittests/tools/test_gemini_schema_util.py

@@ -224,6 +224,64 @@ def test_to_gemini_schema_remove_unrecognized_fields(self):
     assert gemini_schema.type == Type.STRING
     assert not gemini_schema.format
 
+  def test_to_gemini_schema_nested_dict_with_defs_and_ref(self):
+    """Test that nested dict with $defs and $refs is converted correctly."""
+    openapi_schema = {
+        "$defs": {
+            "DeviceEnum": {
+                "enum": ["GLOBAL", "desktop", "mobile"],
+                "title": "DeviceEnum",
+                "type": "string",
+            },
+            "DomainPayload": {
+                "properties": {
+                    "adDomain": {
+                        "description": "List of one or many domains.",
+                        "items": {"type": "string"},
+                        "title": "Addomain",
+                        "type": "array",
+                    },
+                    "device": {
+                        "$ref": "#/$defs/DeviceEnum",
+                        "default": "GLOBAL",
+                        "description": (
+                            "Filter by device. All devices are returned by"
+                            " default."
+                        ),
+                    },
+                },
+                "required": ["adDomain"],
+                "title": "DomainPayload",
+                "type": "object",
+            },
+        },
+        "properties": {"payload": {"$ref": "#/$defs/DomainPayload"}},
+        "required": ["payload"],
+        "title": "query_domainsArguments",
+        "type": "object",
+    }
+    gemini_schema = _to_gemini_schema(openapi_schema)
+    assert gemini_schema.type == Type.OBJECT
+    assert gemini_schema.properties["payload"].type == Type.OBJECT
+    assert (
+        gemini_schema.properties["payload"].properties["adDomain"].type
+        == Type.ARRAY
+    )
+    assert (
+        gemini_schema.properties["payload"].properties["adDomain"].items.type
+        == Type.STRING
+    )
+    assert (
+        gemini_schema.properties["payload"].properties["device"].type
+        == Type.STRING
+    )
+    assert gemini_schema.properties["payload"].properties["device"].enum == [
+        "GLOBAL",
+        "desktop",
+        "mobile",
+    ]
+    assert gemini_schema.properties["payload"].required == ["adDomain"]
+
   def test_sanitize_integer_formats(self):
     """Test that int32 and int64 formats are preserved for integer types"""
     openapi_schema = {