chore: fix mypy strict type errors in adk tools

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 944705634

Author: GWeale

src/google/adk/tools/_automatic_function_calling_util.py

@@ -59,7 +59,7 @@
 }
 
 
-def _get_fields_dict(func: Callable) -> Dict:
+def _get_fields_dict(func: Callable[..., Any]) -> Dict[str, Any]:
   param_signature = dict(inspect.signature(func).parameters)
   fields_dict = {
       name: (
@@ -94,7 +94,7 @@ def _get_fields_dict(func: Callable) -> Dict:
   return fields_dict
 
 
-def _annotate_nullable_fields(schema: Dict):
+def _annotate_nullable_fields(schema: Dict[str, Any]) -> None:
   for _, property_schema in schema.get('properties', {}).items():
     # for Optional[T], the pydantic schema is:
     # {
@@ -117,7 +117,7 @@ def _annotate_nullable_fields(schema: Dict):
         break
 
 
-def _annotate_required_fields(schema: Dict):
+def _annotate_required_fields(schema: Dict[str, Any]) -> None:
   required = [
       field_name
       for field_name, field_schema in schema.get('properties', {}).items()
@@ -126,7 +126,7 @@ def _annotate_required_fields(schema: Dict):
   schema['required'] = required
 
 
-def _remove_any_of(schema: Dict):
+def _remove_any_of(schema: Dict[str, Any]) -> None:
   for _, property_schema in schema.get('properties', {}).items():
     union_types = property_schema.pop('anyOf', None)
     # Take the first non-null type.
@@ -136,17 +136,17 @@ def _remove_any_of(schema: Dict):
           property_schema.update(type_)
 
 
-def _remove_default(schema: Dict):
+def _remove_default(schema: Dict[str, Any]) -> None:
   for _, property_schema in schema.get('properties', {}).items():
     property_schema.pop('default', None)
 
 
-def _remove_nullable(schema: Dict):
+def _remove_nullable(schema: Dict[str, Any]) -> None:
   for _, property_schema in schema.get('properties', {}).items():
     property_schema.pop('nullable', None)
 
 
-def _remove_title(schema: Dict):
+def _remove_title(schema: Dict[str, Any]) -> None:
   for _, property_schema in schema.get('properties', {}).items():
     property_schema.pop('title', None)
 
@@ -162,7 +162,9 @@ def _get_pydantic_schema(func: Callable) -> Dict:
   return pydantic.create_model(func.__name__, **fields_dict).model_json_schema()
 
 
-def _process_pydantic_schema(vertexai: bool, schema: Dict) -> Dict:
+def _process_pydantic_schema(
+    vertexai: bool, schema: Dict[str, Any]
+) -> Dict[str, Any]:
   _annotate_nullable_fields(schema)
   _annotate_required_fields(schema)
   if not vertexai:
@@ -173,7 +175,9 @@ def _process_pydantic_schema(vertexai: bool, schema: Dict) -> Dict:
   return schema
 
 
-def _map_pydantic_type_to_property_schema(property_schema: Dict):
+def _map_pydantic_type_to_property_schema(
+    property_schema: Dict[str, Any],
+) -> None:
   if 'type' in property_schema:
     property_schema['type'] = _py_type_2_schema_type.get(
         property_schema['type'], 'TYPE_UNSPECIFIED'
@@ -190,20 +194,20 @@ def _map_pydantic_type_to_property_schema(property_schema: Dict):
       property_schema['type'] = type_['type']
 
 
-def _map_pydantic_type_to_schema_type(schema: Dict):
+def _map_pydantic_type_to_schema_type(schema: Dict[str, Any]) -> None:
   for _, property_schema in schema.get('properties', {}).items():
     _map_pydantic_type_to_property_schema(property_schema)
 
 
-def _get_return_type(func: Callable) -> Any:
+def _get_return_type(func: Callable[..., Any]) -> Any:
   return _py_type_2_schema_type.get(
       inspect.signature(func).return_annotation.__name__,
       inspect.signature(func).return_annotation.__name__,
   )
 
 
 def build_function_declaration(
-    func: Union[Callable, BaseModel],
+    func: Union[Callable[..., Any], BaseModel],
     ignore_params: Optional[list[str]] = None,
     variant: GoogleLLMVariant = GoogleLLMVariant.GEMINI_API,
 ) -> types.FunctionDeclaration:
@@ -222,45 +226,41 @@ def build_function_declaration(
 
   # ========== ADK defined function tool declaration (old behavior) ==========
   signature = inspect.signature(func)
-  should_update_signature = False
-  new_func = None
   if not ignore_params:
     ignore_params = []
-  for name, _ in signature.parameters.items():
-    if name in ignore_params:
-      should_update_signature = True
-      break
-  if should_update_signature:
-    new_params = [
-        param
+  should_update_signature = any(
+      name in ignore_params for name in signature.parameters
+  )
+  if not should_update_signature:
+    return from_function_with_options(func, variant)
+
+  if isinstance(func, type):
+    fields = {
+        name: (param.annotation, param.default)
         for name, param in signature.parameters.items()
         if name not in ignore_params
-    ]
-    if isinstance(func, type):
-      fields = {
-          name: (param.annotation, param.default)
-          for name, param in signature.parameters.items()
-          if name not in ignore_params
-      }
-      new_func = create_model(func.__name__, **fields)
-    else:
-      new_sig = signature.replace(parameters=new_params)
-      new_func = FunctionType(
-          func.__code__,
-          func.__globals__,
-          func.__name__,
-          func.__defaults__,
-          func.__closure__,
-      )
-      new_func.__signature__ = new_sig
-      new_func.__doc__ = func.__doc__
-      new_func.__annotations__ = func.__annotations__
-
-  return (
-      from_function_with_options(func, variant)
-      if not should_update_signature
-      else from_function_with_options(new_func, variant)
+    }
+    return from_function_with_options(
+        create_model(func.__name__, **fields), variant
+    )
+
+  new_params = [
+      param
+      for name, param in signature.parameters.items()
+      if name not in ignore_params
+  ]
+  new_sig = signature.replace(parameters=new_params)
+  new_func = FunctionType(
+      func.__code__,
+      func.__globals__,
+      func.__name__,
+      func.__defaults__,
+      func.__closure__,
   )
+  setattr(new_func, '__signature__', new_sig)
+  new_func.__doc__ = func.__doc__
+  new_func.__annotations__ = func.__annotations__
+  return from_function_with_options(new_func, variant)
 
 
 def build_function_declaration_for_langchain(
@@ -316,7 +316,7 @@ def build_function_declaration_util(
 
 
 def from_function_with_options(
-    func: Callable,
+    func: Callable[..., Any],
     variant: GoogleLLMVariant = GoogleLLMVariant.GEMINI_API,
 ) -> 'types.FunctionDeclaration':
 

src/google/adk/tools/_function_parameter_parse_util.py

@@ -151,7 +151,7 @@ def _raise_for_unsupported_param(
   ) from exception
 
 
-def _raise_for_invalid_enum_value(param: inspect.Parameter):
+def _raise_for_invalid_enum_value(param: inspect.Parameter) -> None:
   """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 [
@@ -193,14 +193,14 @@ def _is_builtin_primitive_or_compound(
   return annotation in _py_builtin_type_to_schema_type.keys()
 
 
-def _raise_for_any_of_if_mldev(schema: types.Schema):
+def _raise_for_any_of_if_mldev(schema: types.Schema) -> None:
   if schema.any_of:
     raise ValueError(
         'AnyOf is not supported in function declaration schema for Google AI.'
     )
 
 
-def _update_for_default_if_mldev(schema: types.Schema):
+def _update_for_default_if_mldev(schema: types.Schema) -> None:
   if schema.default is not None:
     # TODO: Remove this workaround once mldev supports default value.
     schema.default = None
@@ -212,7 +212,7 @@ def _update_for_default_if_mldev(schema: types.Schema):
 
 def _raise_if_schema_unsupported(
     variant: GoogleLLMVariant, schema: types.Schema
-):
+) -> None:
   if variant == GoogleLLMVariant.GEMINI_API:
     _raise_for_any_of_if_mldev(schema)
     # _update_for_default_if_mldev(schema) # No need of this since GEMINI now supports default value

src/google/adk/tools/bigtable/metadata_tool.py

@@ -16,6 +16,7 @@
 
 import enum
 import logging
+from typing import Any
 
 from google.auth.credentials import Credentials
 from google.cloud.bigtable import enums
@@ -25,7 +26,7 @@
 logger = logging.getLogger(f"google_adk.{__name__}")
 
 
-def list_instances(project_id: str, credentials: Credentials) -> dict:
+def list_instances(project_id: str, credentials: Credentials) -> dict[str, Any]:
   """List Bigtable instance ids in a Google Cloud project.
 
   Args:
@@ -86,7 +87,7 @@ def list_instances(project_id: str, credentials: Credentials) -> dict:
 
 def get_instance_info(
     project_id: str, instance_id: str, credentials: Credentials
-) -> dict:
+) -> dict[str, Any]:
   """Get metadata information about a Bigtable instance.
 
   Args:
@@ -128,7 +129,7 @@ def get_instance_info(
 
 def list_tables(
     project_id: str, instance_id: str, credentials: Credentials
-) -> dict:
+) -> dict[str, Any]:
   """List tables and their metadata in a Bigtable instance.
 
   Args:
@@ -182,7 +183,7 @@ def get_table_info(
     instance_id: str,
     table_id: str,
     credentials: Credentials,
-) -> dict:
+) -> dict[str, Any]:
   """Get metadata information about a Bigtable table.
 
   Args:
@@ -241,7 +242,7 @@ def _enum_name_from_value(
 
 def list_clusters(
     project_id: str, instance_id: str, credentials: Credentials
-) -> dict:
+) -> dict[str, Any]:
   """List clusters and their metadata in a Bigtable instance.
 
   Args:
@@ -315,7 +316,7 @@ def get_cluster_info(
     instance_id: str,
     cluster_id: str,
     credentials: Credentials,
-) -> dict:
+) -> dict[str, Any]:
   """Get detailed metadata information about a Bigtable cluster.
 
   Args:

src/google/adk/tools/bigtable/query_tool.py

@@ -43,7 +43,7 @@ async def execute_sql(
     parameters: Dict[str, Any] | None = None,
     parameter_types: Dict[str, Any] | None = None,
     _view_parameters: Dict[str, Any] | None = None,
-) -> dict:
+) -> Dict[str, Any]:
   """Execute a GoogleSQL query from a Bigtable table.
 
   Args:
@@ -83,7 +83,7 @@ async def execute_sql(
   """
   del tool_context  # Unused for now
 
-  def _execute_sql():
+  def _execute_sql() -> Dict[str, Any]:
     try:
       bt_client = client.get_bigtable_data_client(
           project=project_id, credentials=credentials
@@ -122,7 +122,7 @@ def _execute_sql():
       finally:
         eqi.close()
 
-      result = {"status": "SUCCESS", "rows": rows}
+      result: Dict[str, Any] = {"status": "SUCCESS", "rows": rows}
       if truncated:
         result["result_is_likely_truncated"] = True
       return result

src/google/adk/tools/enterprise_search_tool.py

@@ -41,7 +41,7 @@ class EnterpriseWebSearchTool(BaseTool):
 
   """
 
-  def __init__(self):
+  def __init__(self) -> None:
     """Initializes the Enterprise Web Search tool."""
     # Name and description are not used because this is a model built-in tool.
     super().__init__(

src/google/adk/tools/exit_loop_tool.py

@@ -17,7 +17,7 @@
 from .tool_context import ToolContext
 
 
-def exit_loop(tool_context: ToolContext):
+def exit_loop(tool_context: ToolContext) -> None:
   """Exits the loop.
 
   Call this function only when you are instructed to do so.

src/google/adk/tools/google_api_tool/google_api_toolset.py

@@ -118,7 +118,9 @@ async def get_tools(
         if self._is_tool_selected(tool, readonly_context)
     ]
 
-  def set_tool_filter(self, tool_filter: Union[ToolPredicate, List[str]]):
+  def set_tool_filter(
+      self, tool_filter: Union[ToolPredicate, List[str]]
+  ) -> None:
     self.tool_filter = tool_filter
 
   def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
@@ -158,15 +160,15 @@ def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
         httpx_client_factory=self._httpx_client_factory,
     )
 
-  def configure_auth(self, client_id: str, client_secret: str):
+  def configure_auth(self, client_id: str, client_secret: str) -> None:
     self._client_id = client_id
     self._client_secret = client_secret
 
-  def configure_sa_auth(self, service_account: ServiceAccount):
+  def configure_sa_auth(self, service_account: ServiceAccount) -> None:
     self._service_account = service_account
 
   @override
-  async def close(self):
+  async def close(self) -> None:
     if self._openapi_toolset:
       await self._openapi_toolset.close()
     if hasattr(self, '_mtls_certs') and self._mtls_certs:

src/google/adk/tools/google_maps_grounding_tool.py

@@ -39,7 +39,7 @@ class GoogleMapsGroundingTool(BaseTool):
   GOOGLE_GENAI_USE_ENTERPRISE=TRUE)
   """
 
-  def __init__(self):
+  def __init__(self) -> None:
     # Name and description are not used because this is a model built-in tool.
     super().__init__(name='google_maps', description='google_maps')
 

src/google/adk/tools/load_mcp_resource_tool.py

@@ -106,7 +106,7 @@ async def process_llm_request(
 
   async def _append_resources_to_llm_request(
       self, *, tool_context: ToolContext, llm_request: LlmRequest
-  ):
+  ) -> None:
     try:
       resource_names = await self._mcp_toolset.list_resources()
       if resource_names:

src/google/adk/tools/load_memory_tool.py

@@ -56,7 +56,7 @@ class LoadMemoryTool(FunctionTool):
   NOTE: Currently this tool only uses text part from the memory.
   """
 
-  def __init__(self):
+  def __init__(self) -> None:
     super().__init__(load_memory)
 
   @override