chore: Add abstract type annotation support to AFC

PiperOrigin-RevId: 831462920

Author: sararob

src/google/adk/tools/_automatic_function_calling_util.py

@@ -296,20 +296,55 @@ def from_function_with_options(
 ) -> 'types.FunctionDeclaration':
 
   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,
-    ):
-      # 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
+  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
+          )
+
   declaration = types.FunctionDeclaration(
       name=func.__name__,
       description=func.__doc__,
@@ -324,6 +359,12 @@ 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
 
@@ -372,17 +413,35 @@ 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']
     )
 
-  declaration.response = (
-      _function_parameter_parse_util._parse_schema_from_parameter(
-          variant,
-          return_value,
-          func.__name__,
+  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
+          )
       )
-  )
+      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

src/google/adk/tools/_function_parameter_parse_util.py

@@ -49,6 +49,91 @@
 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:

tests/unittests/tools/test_from_function_with_options.py

@@ -12,6 +12,7 @@
 # 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
 
@@ -192,3 +193,52 @@ def test_function() -> None:
   # 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