mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
chore: Add abstract type annotation support to AFC
PiperOrigin-RevId: 831545133
This commit is contained in:
Google Team Member
committed by
Copybara-Service
Copybara-Service
parent
69627b699f
commit
6bb0b7417e
@@ -296,55 +296,20 @@ def from_function_with_options(
|
||||
) -> 'types.FunctionDeclaration':
|
||||
|
||||
parameters_properties = {}
|
||||
parameters_json_schema = {}
|
||||
annotation_under_future = typing.get_type_hints(func)
|
||||
try:
|
||||
for name, param in inspect.signature(func).parameters.items():
|
||||
if param.kind in (
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
inspect.Parameter.KEYWORD_ONLY,
|
||||
inspect.Parameter.POSITIONAL_ONLY,
|
||||
):
|
||||
param = _function_parameter_parse_util._handle_params_as_deferred_annotations(
|
||||
param, annotation_under_future, name
|
||||
)
|
||||
|
||||
schema = _function_parameter_parse_util._parse_schema_from_parameter(
|
||||
variant, param, func.__name__
|
||||
)
|
||||
parameters_properties[name] = schema
|
||||
except ValueError:
|
||||
# If the function has complex parameter types that fail in _parse_schema_from_parameter,
|
||||
# we try to generate a json schema for the parameter using pydantic.TypeAdapter.
|
||||
parameters_properties = {}
|
||||
for name, param in inspect.signature(func).parameters.items():
|
||||
if param.kind in (
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
inspect.Parameter.KEYWORD_ONLY,
|
||||
inspect.Parameter.POSITIONAL_ONLY,
|
||||
):
|
||||
try:
|
||||
if param.annotation == inspect.Parameter.empty:
|
||||
param = param.replace(annotation=Any)
|
||||
|
||||
param = _function_parameter_parse_util._handle_params_as_deferred_annotations(
|
||||
param, annotation_under_future, name
|
||||
)
|
||||
|
||||
_function_parameter_parse_util._raise_for_invalid_enum_value(param)
|
||||
|
||||
json_schema_dict = _function_parameter_parse_util._generate_json_schema_for_parameter(
|
||||
param
|
||||
)
|
||||
|
||||
parameters_json_schema[name] = types.Schema.model_validate(
|
||||
json_schema_dict
|
||||
)
|
||||
except Exception as e:
|
||||
_function_parameter_parse_util._raise_for_unsupported_param(
|
||||
param, func.__name__, e
|
||||
)
|
||||
for name, param in inspect.signature(func).parameters.items():
|
||||
if param.kind in (
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
inspect.Parameter.KEYWORD_ONLY,
|
||||
inspect.Parameter.POSITIONAL_ONLY,
|
||||
):
|
||||
# This snippet catches the case when type hints are stored as strings
|
||||
if isinstance(param.annotation, str):
|
||||
param = param.replace(annotation=typing.get_type_hints(func)[name])
|
||||
|
||||
schema = _function_parameter_parse_util._parse_schema_from_parameter(
|
||||
variant, param, func.__name__
|
||||
)
|
||||
parameters_properties[name] = schema
|
||||
declaration = types.FunctionDeclaration(
|
||||
name=func.__name__,
|
||||
description=func.__doc__,
|
||||
@@ -359,12 +324,6 @@ def from_function_with_options(
|
||||
declaration.parameters
|
||||
)
|
||||
)
|
||||
elif parameters_json_schema:
|
||||
declaration.parameters = types.Schema(
|
||||
type='OBJECT',
|
||||
properties=parameters_json_schema,
|
||||
)
|
||||
|
||||
if variant == GoogleLLMVariant.GEMINI_API:
|
||||
return declaration
|
||||
|
||||
@@ -413,35 +372,17 @@ def from_function_with_options(
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
annotation=return_annotation,
|
||||
)
|
||||
# This snippet catches the case when type hints are stored as strings
|
||||
if isinstance(return_value.annotation, str):
|
||||
return_value = return_value.replace(
|
||||
annotation=typing.get_type_hints(func)['return']
|
||||
)
|
||||
|
||||
response_schema: Optional[types.Schema] = None
|
||||
response_json_schema: Optional[Union[Dict[str, Any], types.Schema]] = None
|
||||
try:
|
||||
response_schema = (
|
||||
_function_parameter_parse_util._parse_schema_from_parameter(
|
||||
variant,
|
||||
return_value,
|
||||
func.__name__,
|
||||
)
|
||||
)
|
||||
except ValueError:
|
||||
try:
|
||||
response_json_schema = (
|
||||
_function_parameter_parse_util._generate_json_schema_for_parameter(
|
||||
return_value
|
||||
)
|
||||
declaration.response = (
|
||||
_function_parameter_parse_util._parse_schema_from_parameter(
|
||||
variant,
|
||||
return_value,
|
||||
func.__name__,
|
||||
)
|
||||
response_json_schema = types.Schema.model_validate(response_json_schema)
|
||||
except Exception as e:
|
||||
_function_parameter_parse_util._raise_for_unsupported_param(
|
||||
return_value, func.__name__, e
|
||||
)
|
||||
if response_schema:
|
||||
declaration.response = response_schema
|
||||
elif response_json_schema:
|
||||
declaration.response = response_json_schema
|
||||
)
|
||||
return declaration
|
||||
|
||||
@@ -49,91 +49,6 @@ _py_builtin_type_to_schema_type = {
|
||||
logger = logging.getLogger('google_adk.' + __name__)
|
||||
|
||||
|
||||
def _handle_params_as_deferred_annotations(
|
||||
param: inspect.Parameter, annotation_under_future: dict[str, Any], name: str
|
||||
) -> inspect.Parameter:
|
||||
"""Catches the case when type hints are stored as strings."""
|
||||
if isinstance(param.annotation, str):
|
||||
param = param.replace(annotation=annotation_under_future[name])
|
||||
return param
|
||||
|
||||
|
||||
def _add_unevaluated_items_to_fixed_len_tuple_schema(
|
||||
json_schema: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Adds 'unevaluatedItems': False to schemas for fixed-length tuples.
|
||||
|
||||
For example, the schema for a parameter of type `tuple[float, float]` would
|
||||
be:
|
||||
{
|
||||
"type": "array",
|
||||
"prefixItems": [
|
||||
{
|
||||
"type": "number"
|
||||
},
|
||||
{
|
||||
"type": "number"
|
||||
},
|
||||
],
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
"unevaluatedItems": False
|
||||
}
|
||||
|
||||
"""
|
||||
if (
|
||||
json_schema.get('maxItems')
|
||||
and (
|
||||
json_schema.get('prefixItems')
|
||||
and len(json_schema['prefixItems']) == json_schema['maxItems']
|
||||
)
|
||||
and json_schema.get('type') == 'array'
|
||||
):
|
||||
json_schema['unevaluatedItems'] = False
|
||||
return json_schema
|
||||
|
||||
|
||||
def _raise_for_unsupported_param(
|
||||
param: inspect.Parameter,
|
||||
func_name: str,
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
raise ValueError(
|
||||
f'Failed to parse the parameter {param} of function {func_name} for'
|
||||
' automatic function calling.Automatic function calling works best with'
|
||||
' simpler function signature schema, consider manually parsing your'
|
||||
f' function declaration for function {func_name}.'
|
||||
) from exception
|
||||
|
||||
|
||||
def _raise_for_invalid_enum_value(param: inspect.Parameter):
|
||||
"""Raises an error if the default value is not a valid enum value."""
|
||||
if inspect.isclass(param.annotation) and issubclass(param.annotation, Enum):
|
||||
if param.default is not inspect.Parameter.empty and param.default not in [
|
||||
e.value for e in param.annotation
|
||||
]:
|
||||
raise ValueError(
|
||||
f'Default value {param.default} is not a valid enum value for'
|
||||
f' {param.annotation}.'
|
||||
)
|
||||
|
||||
|
||||
def _generate_json_schema_for_parameter(
|
||||
param: inspect.Parameter,
|
||||
) -> dict[str, Any]:
|
||||
"""Generates a JSON schema for a parameter using pydantic.TypeAdapter."""
|
||||
|
||||
param_schema_adapter = pydantic.TypeAdapter(
|
||||
param.annotation,
|
||||
config=pydantic.ConfigDict(arbitrary_types_allowed=True),
|
||||
)
|
||||
json_schema_dict = param_schema_adapter.json_schema()
|
||||
json_schema_dict = _add_unevaluated_items_to_fixed_len_tuple_schema(
|
||||
json_schema_dict
|
||||
)
|
||||
return json_schema_dict
|
||||
|
||||
|
||||
def _is_builtin_primitive_or_compound(
|
||||
annotation: inspect.Parameter.annotation,
|
||||
) -> bool:
|
||||
|
||||
@@ -12,7 +12,6 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
from typing import Dict
|
||||
|
||||
@@ -193,52 +192,3 @@ def test_from_function_with_options_no_params():
|
||||
# VERTEX_AI should have response schema for None return
|
||||
assert declaration.response is not None
|
||||
assert declaration.response.type == types.Type.NULL
|
||||
|
||||
|
||||
def test_from_function_with_collections_type_parameter():
|
||||
"""Test from_function_with_options with collections type parameter."""
|
||||
|
||||
def test_function(
|
||||
artifact_key: str,
|
||||
input_edit_ids: Sequence[str],
|
||||
) -> str:
|
||||
"""Saves a sequence of edit IDs."""
|
||||
return f'Saved {len(input_edit_ids)} edit IDs for artifact {artifact_key}'
|
||||
|
||||
declaration = _automatic_function_calling_util.from_function_with_options(
|
||||
test_function, GoogleLLMVariant.VERTEX_AI
|
||||
)
|
||||
|
||||
assert declaration.name == 'test_function'
|
||||
assert declaration.parameters.type == types.Type.OBJECT
|
||||
assert (
|
||||
declaration.parameters.properties['artifact_key'].type
|
||||
== types.Type.STRING
|
||||
)
|
||||
assert (
|
||||
declaration.parameters.properties['input_edit_ids'].type
|
||||
== types.Type.ARRAY
|
||||
)
|
||||
assert (
|
||||
declaration.parameters.properties['input_edit_ids'].items.type
|
||||
== types.Type.STRING
|
||||
)
|
||||
assert declaration.response.type == types.Type.STRING
|
||||
|
||||
|
||||
def test_from_function_with_collections_return_type():
|
||||
"""Test from_function_with_options with collections return type."""
|
||||
|
||||
def test_function(
|
||||
names: list[str],
|
||||
) -> Sequence[str]:
|
||||
"""Returns a sequence of names."""
|
||||
return names
|
||||
|
||||
declaration = _automatic_function_calling_util.from_function_with_options(
|
||||
test_function, GoogleLLMVariant.VERTEX_AI
|
||||
)
|
||||
|
||||
assert declaration.name == 'test_function'
|
||||
assert declaration.response.type == types.Type.ARRAY
|
||||
assert declaration.response.items.type == types.Type.STRING
|
||||
|
||||
Reference in New Issue
Block a user