ADK changes

PiperOrigin-RevId: 749973427
This commit is contained in:
Che Liu 2025-04-21 17:28:20 -07:00 committed by Copybara-Service
parent a5f191650b
commit 21d2047ddc
4 changed files with 127 additions and 36 deletions

View File

@ -14,20 +14,12 @@
import inspect
from textwrap import dedent
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Union
from typing import Any, Dict, List, Optional, Union
from fastapi.encoders import jsonable_encoder
from fastapi.openapi.models import Operation
from fastapi.openapi.models import Parameter
from fastapi.openapi.models import Schema
from fastapi.openapi.models import Operation, Parameter, Schema
from ..common.common import ApiParameter
from ..common.common import PydocHelper
from ..common.common import to_snake_case
from ..common.common import ApiParameter, PydocHelper, to_snake_case
class OperationParser:
@ -110,7 +102,8 @@ class OperationParser:
description = request_body.description or ''
if schema and schema.type == 'object':
for prop_name, prop_details in schema.properties.items():
properties = schema.properties or {}
for prop_name, prop_details in properties.items():
self.params.append(
ApiParameter(
original_name=prop_name,

View File

@ -17,6 +17,7 @@ from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import Union
@ -59,6 +60,40 @@ def snake_to_lower_camel(snake_case_string: str):
])
# TODO: Switch to Gemini `from_json_schema` util when it is released
# in Gemini SDK.
def normalize_json_schema_type(
json_schema_type: Optional[Union[str, Sequence[str]]],
) -> tuple[Optional[str], bool]:
"""Converts a JSON Schema Type into Gemini Schema type.
Adopted and modified from Gemini SDK. This gets the first available schema
type from JSON Schema, and use it to mark Gemini schema type. If JSON Schema
contains a list of types, the first non null type is used.
Remove this after switching to Gemini `from_json_schema`.
"""
if json_schema_type is None:
return None, False
if isinstance(json_schema_type, str):
if json_schema_type == "null":
return None, True
return json_schema_type, False
non_null_types = []
nullable = False
# If json schema type is an array, pick the first non null type.
for type_value in json_schema_type:
if type_value == "null":
nullable = True
else:
non_null_types.append(type_value)
non_null_type = non_null_types[0] if non_null_types else None
return non_null_type, nullable
# TODO: Switch to Gemini `from_json_schema` util when it is released
# in Gemini SDK.
def to_gemini_schema(openapi_schema: Optional[Dict[str, Any]] = None) -> Schema:
"""Converts an OpenAPI schema dictionary to a Gemini Schema object.
@ -82,13 +117,6 @@ def to_gemini_schema(openapi_schema: Optional[Dict[str, Any]] = None) -> Schema:
if not openapi_schema.get("type"):
openapi_schema["type"] = "object"
# Adding this to avoid "properties: should be non-empty for OBJECT type" error
# See b/385165182
if openapi_schema.get("type", "") == "object" and not openapi_schema.get(
"properties"
):
openapi_schema["properties"] = {"dummy_DO_NOT_GENERATE": {"type": "string"}}
for key, value in openapi_schema.items():
snake_case_key = to_snake_case(key)
# Check if the snake_case_key exists in the Schema model's fields.
@ -99,7 +127,17 @@ def to_gemini_schema(openapi_schema: Optional[Dict[str, Any]] = None) -> Schema:
# Format: properties[expiration].format: only 'enum' and 'date-time' are
# supported for STRING type
continue
if snake_case_key == "properties" and isinstance(value, dict):
elif snake_case_key == "type":
schema_type, nullable = normalize_json_schema_type(
openapi_schema.get("type", None)
)
# Adding this to force adding a type to an empty dict
# This avoid "... one_of or any_of must specify a type" error
pydantic_schema_data["type"] = schema_type if schema_type else "object"
pydantic_schema_data["type"] = pydantic_schema_data["type"].upper()
if nullable:
pydantic_schema_data["nullable"] = True
elif snake_case_key == "properties" and isinstance(value, dict):
pydantic_schema_data[snake_case_key] = {
k: to_gemini_schema(v) for k, v in value.items()
}

View File

@ -164,6 +164,18 @@ def test_process_request_body_no_name():
assert parser.params[0].param_location == 'body'
def test_process_request_body_empty_object():
"""Test _process_request_body with a schema that is of type object but with no properties."""
operation = Operation(
requestBody=RequestBody(
content={'application/json': MediaType(schema=Schema(type='object'))}
)
)
parser = OperationParser(operation, should_parse=False)
parser._process_request_body()
assert len(parser.params) == 0
def test_dedupe_param_names(sample_operation):
"""Test _dedupe_param_names method."""
parser = OperationParser(sample_operation, should_parse=False)

View File

@ -14,11 +14,9 @@
import json
from unittest.mock import MagicMock
from unittest.mock import patch
from unittest.mock import MagicMock, patch
from fastapi.openapi.models import MediaType
from fastapi.openapi.models import Operation
from fastapi.openapi.models import MediaType, Operation
from fastapi.openapi.models import Parameter as OpenAPIParameter
from fastapi.openapi.models import RequestBody
from fastapi.openapi.models import Schema as OpenAPISchema
@ -27,13 +25,13 @@ from google.adk.tools.openapi_tool.auth.auth_helpers import token_to_scheme_cred
from google.adk.tools.openapi_tool.common.common import ApiParameter
from google.adk.tools.openapi_tool.openapi_spec_parser.openapi_spec_parser import OperationEndpoint
from google.adk.tools.openapi_tool.openapi_spec_parser.operation_parser import OperationParser
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import snake_to_lower_camel
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import (
RestApiTool,
snake_to_lower_camel,
to_gemini_schema,
)
from google.adk.tools.tool_context import ToolContext
from google.genai.types import FunctionDeclaration
from google.genai.types import Schema
from google.genai.types import Type
from google.genai.types import FunctionDeclaration, Schema, Type
import pytest
@ -790,13 +788,13 @@ class TestToGeminiSchema:
result = to_gemini_schema({})
assert isinstance(result, Schema)
assert result.type == Type.OBJECT
assert result.properties == {"dummy_DO_NOT_GENERATE": Schema(type="string")}
assert result.properties is None
def test_to_gemini_schema_dict_with_only_object_type(self):
result = to_gemini_schema({"type": "object"})
assert isinstance(result, Schema)
assert result.type == Type.OBJECT
assert result.properties == {"dummy_DO_NOT_GENERATE": Schema(type="string")}
assert result.properties is None
def test_to_gemini_schema_basic_types(self):
openapi_schema = {
@ -814,6 +812,42 @@ class TestToGeminiSchema:
assert gemini_schema.properties["age"].type == Type.INTEGER
assert gemini_schema.properties["is_active"].type == Type.BOOLEAN
def test_to_gemini_schema_array_string_types(self):
openapi_schema = {
"type": "object",
"properties": {
"boolean_field": {"type": "boolean"},
"nonnullable_string": {"type": ["string"]},
"nullable_string": {"type": ["string", "null"]},
"nullable_number": {"type": ["null", "integer"]},
"object_nullable": {"type": "null"},
"multi_types_nullable": {"type": ["string", "null", "integer"]},
"empty_default_object": {},
},
}
gemini_schema = to_gemini_schema(openapi_schema)
assert isinstance(gemini_schema, Schema)
assert gemini_schema.type == Type.OBJECT
assert gemini_schema.properties["boolean_field"].type == Type.BOOLEAN
assert gemini_schema.properties["nonnullable_string"].type == Type.STRING
assert not gemini_schema.properties["nonnullable_string"].nullable
assert gemini_schema.properties["nullable_string"].type == Type.STRING
assert gemini_schema.properties["nullable_string"].nullable
assert gemini_schema.properties["nullable_number"].type == Type.INTEGER
assert gemini_schema.properties["nullable_number"].nullable
assert gemini_schema.properties["object_nullable"].type == Type.OBJECT
assert gemini_schema.properties["object_nullable"].nullable
assert gemini_schema.properties["multi_types_nullable"].type == Type.STRING
assert gemini_schema.properties["multi_types_nullable"].nullable
assert gemini_schema.properties["empty_default_object"].type == Type.OBJECT
assert not gemini_schema.properties["empty_default_object"].nullable
def test_to_gemini_schema_nested_objects(self):
openapi_schema = {
"type": "object",
@ -895,7 +929,15 @@ class TestToGeminiSchema:
def test_to_gemini_schema_nested_dict(self):
openapi_schema = {
"type": "object",
"properties": {"metadata": {"key1": "value1", "key2": 123}},
"properties": {
"metadata": {
"type": "object",
"properties": {
"key1": {"type": "object"},
"key2": {"type": "string"},
},
}
},
}
gemini_schema = to_gemini_schema(openapi_schema)
# Since metadata is not properties nor item, it will call to_gemini_schema recursively.
@ -903,9 +945,15 @@ class TestToGeminiSchema:
assert (
gemini_schema.properties["metadata"].type == Type.OBJECT
) # add object type by default
assert gemini_schema.properties["metadata"].properties == {
"dummy_DO_NOT_GENERATE": Schema(type="string")
}
assert len(gemini_schema.properties["metadata"].properties) == 2
assert (
gemini_schema.properties["metadata"].properties["key1"].type
== Type.OBJECT
)
assert (
gemini_schema.properties["metadata"].properties["key2"].type
== Type.STRING
)
def test_to_gemini_schema_ignore_title_default_format(self):
openapi_schema = {