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 import inspect
from textwrap import dedent from textwrap import dedent
from typing import Any from typing import Any, Dict, List, Optional, Union
from typing import Dict
from typing import List
from typing import Optional
from typing import Union
from fastapi.encoders import jsonable_encoder from fastapi.encoders import jsonable_encoder
from fastapi.openapi.models import Operation from fastapi.openapi.models import Operation, Parameter, Schema
from fastapi.openapi.models import Parameter
from fastapi.openapi.models import Schema
from ..common.common import ApiParameter from ..common.common import ApiParameter, PydocHelper, to_snake_case
from ..common.common import PydocHelper
from ..common.common import to_snake_case
class OperationParser: class OperationParser:
@ -110,7 +102,8 @@ class OperationParser:
description = request_body.description or '' description = request_body.description or ''
if schema and schema.type == 'object': 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( self.params.append(
ApiParameter( ApiParameter(
original_name=prop_name, original_name=prop_name,

View File

@ -17,6 +17,7 @@ from typing import Dict
from typing import List from typing import List
from typing import Literal from typing import Literal
from typing import Optional from typing import Optional
from typing import Sequence
from typing import Tuple from typing import Tuple
from typing import Union 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: def to_gemini_schema(openapi_schema: Optional[Dict[str, Any]] = None) -> Schema:
"""Converts an OpenAPI schema dictionary to a Gemini Schema object. """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"): if not openapi_schema.get("type"):
openapi_schema["type"] = "object" 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(): for key, value in openapi_schema.items():
snake_case_key = to_snake_case(key) snake_case_key = to_snake_case(key)
# Check if the snake_case_key exists in the Schema model's fields. # 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 # Format: properties[expiration].format: only 'enum' and 'date-time' are
# supported for STRING type # supported for STRING type
continue 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] = { pydantic_schema_data[snake_case_key] = {
k: to_gemini_schema(v) for k, v in value.items() 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' 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): def test_dedupe_param_names(sample_operation):
"""Test _dedupe_param_names method.""" """Test _dedupe_param_names method."""
parser = OperationParser(sample_operation, should_parse=False) parser = OperationParser(sample_operation, should_parse=False)

View File

@ -14,11 +14,9 @@
import json import json
from unittest.mock import MagicMock from unittest.mock import MagicMock, patch
from unittest.mock import patch
from fastapi.openapi.models import MediaType from fastapi.openapi.models import MediaType, Operation
from fastapi.openapi.models import Operation
from fastapi.openapi.models import Parameter as OpenAPIParameter from fastapi.openapi.models import Parameter as OpenAPIParameter
from fastapi.openapi.models import RequestBody from fastapi.openapi.models import RequestBody
from fastapi.openapi.models import Schema as OpenAPISchema 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.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.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.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 (
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import snake_to_lower_camel RestApiTool,
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema snake_to_lower_camel,
to_gemini_schema,
)
from google.adk.tools.tool_context import ToolContext from google.adk.tools.tool_context import ToolContext
from google.genai.types import FunctionDeclaration from google.genai.types import FunctionDeclaration, Schema, Type
from google.genai.types import Schema
from google.genai.types import Type
import pytest import pytest
@ -790,13 +788,13 @@ class TestToGeminiSchema:
result = to_gemini_schema({}) result = to_gemini_schema({})
assert isinstance(result, Schema) assert isinstance(result, Schema)
assert result.type == Type.OBJECT 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): def test_to_gemini_schema_dict_with_only_object_type(self):
result = to_gemini_schema({"type": "object"}) result = to_gemini_schema({"type": "object"})
assert isinstance(result, Schema) assert isinstance(result, Schema)
assert result.type == Type.OBJECT 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): def test_to_gemini_schema_basic_types(self):
openapi_schema = { openapi_schema = {
@ -814,6 +812,42 @@ class TestToGeminiSchema:
assert gemini_schema.properties["age"].type == Type.INTEGER assert gemini_schema.properties["age"].type == Type.INTEGER
assert gemini_schema.properties["is_active"].type == Type.BOOLEAN 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): def test_to_gemini_schema_nested_objects(self):
openapi_schema = { openapi_schema = {
"type": "object", "type": "object",
@ -895,7 +929,15 @@ class TestToGeminiSchema:
def test_to_gemini_schema_nested_dict(self): def test_to_gemini_schema_nested_dict(self):
openapi_schema = { openapi_schema = {
"type": "object", "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) gemini_schema = to_gemini_schema(openapi_schema)
# Since metadata is not properties nor item, it will call to_gemini_schema recursively. # Since metadata is not properties nor item, it will call to_gemini_schema recursively.
@ -903,9 +945,15 @@ class TestToGeminiSchema:
assert ( assert (
gemini_schema.properties["metadata"].type == Type.OBJECT gemini_schema.properties["metadata"].type == Type.OBJECT
) # add object type by default ) # add object type by default
assert gemini_schema.properties["metadata"].properties == { assert len(gemini_schema.properties["metadata"].properties) == 2
"dummy_DO_NOT_GENERATE": Schema(type="string") 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): def test_to_gemini_schema_ignore_title_default_format(self):
openapi_schema = { openapi_schema = {