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 = {