refactor(types): type the computer use, data agent, retrieval and Google API tools for strict mypy

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

Author: GWeale

src/google/adk/tools/computer_use/computer_use_toolset.py

@@ -20,6 +20,7 @@
 import logging
 from typing import Any
 from typing import Callable
+from typing import cast
 from typing import Optional
 from typing import Union
 
@@ -72,7 +73,7 @@ def __init__(
     self._excluded_predefined_functions = excluded_predefined_functions
     self._allow_private_network_access = allow_private_network_access
     self._initialized = False
-    self._tools = None
+    self._tools: Optional[list[ComputerUseTool]] = None
 
   async def _ensure_initialized(self) -> None:
     if not self._initialized:
@@ -99,7 +100,9 @@ def _wrap_method_with_state_binding(
 
     @functools.wraps(method)
     async def wrapper(
-        *args: Any, tool_context: ToolContext = None, **kwargs: Any
+        *args: Any,
+        tool_context: Optional[ToolContext] = None,
+        **kwargs: Any,
     ) -> Any:
       # Prepare computer before each tool call
       # Computers that need session state (e.g., AgentEngineSandboxComputer)
@@ -121,7 +124,7 @@ async def wrapper(
             annotation=ToolContext,
         )
     ]
-    wrapper.__signature__ = orig_sig.replace(parameters=new_params)
+    setattr(wrapper, "__signature__", orig_sig.replace(parameters=new_params))
 
     return wrapper
 
@@ -200,16 +203,14 @@ async def adapt_computer_use_tool(
       logger.warning("Method %s not found in tools_dict", method_name)
       return
 
-    original_tool = llm_request.tools_dict[method_name]
+    original_tool = cast(ComputerUseTool, llm_request.tools_dict[method_name])
 
     # Create the adapted function using the adapter
-    # Handle both sync and async adapter functions
-    if asyncio.iscoroutinefunction(adapter_func):
-      # If adapter_func is async, await it to get the adapted function
-      adapted_func = await adapter_func(original_tool.func)
+    adapted_func_or_awaitable = adapter_func(original_tool.func)
+    if inspect.isawaitable(adapted_func_or_awaitable):
+      adapted_func = await adapted_func_or_awaitable
     else:
-      # If adapter_func is sync, call it directly
-      adapted_func = adapter_func(original_tool.func)
+      adapted_func = adapted_func_or_awaitable
 
     # Get the name from the adapted function
     new_method_name = adapted_func.__name__
@@ -232,7 +233,9 @@ async def adapt_computer_use_tool(
     )
 
   @override
-  async def get_tools(
+  # list is invariant, so the narrower element type is not a compatible
+  # override; widening it to BaseTool would change this public signature.
+  async def get_tools(  # type: ignore[override]
       self,
       readonly_context: Optional[ReadonlyContext] = None,
   ) -> list[ComputerUseTool]:
@@ -306,16 +309,20 @@ async def process_llm_request(
       if not self._tools:
         await self.get_tools()
 
-      for tool in self._tools:
-        llm_request.tools_dict[tool.name] = tool
+      assert self._tools is not None
+      for computer_tool in self._tools:
+        llm_request.tools_dict[computer_tool.name] = computer_tool
 
       # Initialize config if needed
       llm_request.config = llm_request.config or types.GenerateContentConfig()
       llm_request.config.tools = llm_request.config.tools or []
 
       # Check if computer use is already configured
-      for tool in llm_request.config.tools:
-        if isinstance(tool, types.Tool) and tool.computer_use:
+      for configured_tool in llm_request.config.tools:
+        if (
+            isinstance(configured_tool, types.Tool)
+            and configured_tool.computer_use
+        ):
           logger.debug("Computer use already configured in LLM request")
           return
 

src/google/adk/tools/data_agent/credentials.py

@@ -25,7 +25,9 @@ class DataAgentCredentialsConfig(BaseGoogleCredentialsConfig):
 
   def __post_init__(self) -> DataAgentCredentialsConfig:
     """Populate default scope if scopes is None."""
-    super().__post_init__()
+    # pydantic wraps the base @model_validator in a descriptor proxy that mypy
+    # does not treat as callable; it binds to the function normally at runtime.
+    super().__post_init__()  # type: ignore[operator]
 
     if not self.scopes:
       self.scopes = DATA_AGENT_DEFAULT_SCOPE

src/google/adk/tools/data_agent/data_agent_toolset.py

@@ -14,9 +14,9 @@
 
 from __future__ import annotations
 
+from typing import Any
+from typing import Callable
 from typing import List
-from typing import Optional
-from typing import Union
 
 from google.adk.agents.readonly_context import ReadonlyContext
 from typing_extensions import override
@@ -36,9 +36,9 @@ class DataAgentToolset(BaseToolset):
   def __init__(
       self,
       *,
-      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
-      credentials_config: Optional[DataAgentCredentialsConfig] = None,
-      data_agent_tool_config: Optional[DataAgentToolConfig] = None,
+      tool_filter: ToolPredicate | list[str] | None = None,
+      credentials_config: DataAgentCredentialsConfig | None = None,
+      data_agent_tool_config: DataAgentToolConfig | None = None,
   ):
     super().__init__(tool_filter=tool_filter)
     self._credentials_config = credentials_config
@@ -49,8 +49,9 @@ def __init__(
     )
 
   def _is_tool_selected(
-      self, tool: BaseTool, readonly_context: ReadonlyContext
+      self, tool: BaseTool, readonly_context: ReadonlyContext | None
   ) -> bool:
+    # Unlike the base implementation, an empty tool_filter selects no tools.
     if self.tool_filter is None:
       return True
 
@@ -64,9 +65,9 @@ def _is_tool_selected(
 
   @override
   async def get_tools(
-      self, readonly_context: Optional[ReadonlyContext] = None
+      self, readonly_context: ReadonlyContext | None = None
   ) -> List[BaseTool]:
-    funcs = [
+    funcs: list[Callable[..., Any]] = [
         data_agent_tool.list_accessible_data_agents,
         data_agent_tool.get_data_agent_info,
         data_agent_tool.ask_data_agent,
@@ -92,5 +93,5 @@ async def get_tools(
     ]
 
   @override
-  async def close(self):
+  async def close(self) -> None:
     pass

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

@@ -15,10 +15,9 @@
 from __future__ import annotations
 
 import logging
+from typing import Callable
 from typing import Dict
 from typing import List
-from typing import Optional
-from typing import Union
 
 import httpx
 from typing_extensions import override
@@ -62,15 +61,15 @@ def __init__(
       self,
       api_name: str,
       api_version: str,
-      client_id: Optional[str] = None,
-      client_secret: Optional[str] = None,
-      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
-      service_account: Optional[ServiceAccount] = None,
-      tool_name_prefix: Optional[str] = None,
+      client_id: str | None = None,
+      client_secret: str | None = None,
+      tool_filter: ToolPredicate | List[str] | None = None,
+      service_account: ServiceAccount | None = None,
+      tool_name_prefix: str | None = None,
       *,
-      additional_headers: Optional[Dict[str, str]] = None,
-      additional_scopes: Optional[List[str]] = None,
-      discovery_url: Optional[str] = None,
+      additional_headers: Dict[str, str] | None = None,
+      additional_scopes: List[str] | None = None,
+      discovery_url: str | None = None,
   ):
     super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
     self.api_name = api_name
@@ -82,7 +81,8 @@ def __init__(
     self._additional_scopes = additional_scopes
     self._discovery_url = discovery_url
 
-    self._httpx_client_factory = None
+    self._httpx_client_factory: Callable[[], httpx.AsyncClient] | None = None
+    self._mtls_certs: MtlsClientCerts | None = None
     use_client_cert = use_client_cert_effective()
 
     if use_client_cert:
@@ -102,8 +102,10 @@ def client_factory() -> httpx.AsyncClient:
     self._openapi_toolset = self._load_toolset_with_oidc_auth()
 
   @override
-  async def get_tools(
-      self, readonly_context: Optional[ReadonlyContext] = None
+  # list is invariant, so the narrower element type is not a compatible
+  # override; widening it to BaseTool would change this public signature.
+  async def get_tools(  # type: ignore[override]
+      self, readonly_context: ReadonlyContext | None = None
   ) -> List[GoogleApiTool]:
     """Get all tools in the toolset."""
     return [
@@ -118,9 +120,7 @@ async def get_tools(
         if self._is_tool_selected(tool, readonly_context)
     ]
 
-  def set_tool_filter(
-      self, tool_filter: Union[ToolPredicate, List[str]]
-  ) -> None:
+  def set_tool_filter(self, tool_filter: ToolPredicate | List[str]) -> None:
     self.tool_filter = tool_filter
 
   def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
@@ -171,5 +171,5 @@ def configure_sa_auth(self, service_account: ServiceAccount) -> None:
   async def close(self) -> None:
     if self._openapi_toolset:
       await self._openapi_toolset.close()
-    if hasattr(self, '_mtls_certs') and self._mtls_certs:
+    if self._mtls_certs:
       self._mtls_certs.close()

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

@@ -15,12 +15,12 @@
 from __future__ import annotations
 
 import argparse
+from collections.abc import Mapping
 import json
 import logging
 import socket
 from typing import Any
 from typing import Dict
-from typing import List
 
 # Google API client
 from googleapiclient.discovery import build
@@ -50,9 +50,11 @@ def __init__(
     self._api_name = api_name
     self._api_version = api_version
     self._discovery_url = discovery_url
-    self._google_api_resource = None
-    self._google_api_spec = None
-    self._openapi_spec = {
+    self._google_api_resource: object | None = None
+    # Discovery documents are heterogeneous JSON objects, and this attribute
+    # is only populated once the document has been fetched.
+    self._google_api_spec: Any = None
+    self._openapi_spec: dict[str, Any] = {
         "openapi": "3.0.0",
         "info": {},
         "servers": [],
@@ -108,10 +110,12 @@ def fetch_google_api_spec(self) -> None:
         )
 
       # Access the underlying API discovery document
-      self._google_api_spec = self._google_api_resource._rootDesc
-
-      if not self._google_api_spec:
+      root_desc = getattr(self._google_api_resource, "_rootDesc", None)
+      if not isinstance(root_desc, dict) or not root_desc:
         raise ValueError("Failed to retrieve API specification")
+      if not all(isinstance(key, str) for key in root_desc):
+        raise ValueError("API specification keys must be strings")
+      self._google_api_spec = root_desc
 
       logger.info("Successfully fetched %s API specification", self._api_name)
     except HttpError as e:
@@ -200,7 +204,7 @@ def _convert_security_schemes(self) -> None:
     if oauth2:
       # Handle OAuth2
       scopes = oauth2.get("scopes", {})
-      formatted_scopes = {}
+      formatted_scopes: dict[str, str] = {}
 
       for scope, scope_info in scopes.items():
         formatted_scopes[scope] = scope_info.get("description", "")
@@ -244,8 +248,8 @@ def _convert_schemas(self) -> None:
       ] = converted_schema
 
   def _convert_schema_object(
-      self, schema_def: Dict[str, Any]
-  ) -> Dict[str, Any]:
+      self, schema_def: Mapping[str, Any]
+  ) -> dict[str, Any]:
     """Recursively convert a Google API schema object to OpenAPI schema.
 
     Args:
@@ -254,7 +258,7 @@ def _convert_schema_object(
     Returns:
         Converted OpenAPI schema object
     """
-    result = {}
+    result: dict[str, Any] = {}
 
     # Convert the type
     if "type" in schema_def:
@@ -332,7 +336,7 @@ def _convert_schema_object(
     return result
 
   def _convert_resources(
-      self, resources: Dict[str, Any], parent_path: str = ""
+      self, resources: Mapping[str, Any], parent_path: str = ""
   ) -> None:
     """Recursively convert all resources and their methods.
 
@@ -352,7 +356,7 @@ def _convert_resources(
         self._convert_resources(nested_resources, resource_path)
 
   def _convert_methods(
-      self, methods: Dict[str, Any], resource_path: str
+      self, methods: Mapping[str, Any], resource_path: str
   ) -> None:
     """Convert methods for a specific resource path.
 
@@ -382,7 +386,7 @@ def _convert_methods(
           self._convert_operation(method_data, path_params)
       )
 
-  def _extract_path_parameters(self, path: str) -> List[str]:
+  def _extract_path_parameters(self, path: str) -> list[str]:
     """Extract path parameters from a URL path.
 
     Args:
@@ -403,8 +407,8 @@ def _extract_path_parameters(self, path: str) -> List[str]:
     return params
 
   def _convert_operation(
-      self, method_data: Dict[str, Any], path_params: List[str]
-  ) -> Dict[str, Any]:
+      self, method_data: Mapping[str, Any], path_params: list[str]
+  ) -> dict[str, Any]:
     """Convert a Google API method to an OpenAPI operation.
 
     Args:
@@ -414,7 +418,7 @@ def _convert_operation(
     Returns:
         OpenAPI operation object
     """
-    operation = {
+    operation: dict[str, Any] = {
         "operationId": method_data.get("id", ""),
         "summary": method_data.get("description", ""),
         "description": method_data.get("description", ""),
@@ -491,8 +495,8 @@ def _convert_operation(
     return operation
 
   def _convert_parameter_schema(
-      self, param_data: Dict[str, Any]
-  ) -> Dict[str, Any]:
+      self, param_data: Mapping[str, Any]
+  ) -> dict[str, Any]:
     """Convert a parameter definition to an OpenAPI schema.
 
     Args:
@@ -501,7 +505,7 @@ def _convert_parameter_schema(
     Returns:
         OpenAPI schema for the parameter
     """
-    schema = {}
+    schema: dict[str, Any] = {}
 
     # Convert type
     param_type = param_data.get("type", "string")
@@ -536,7 +540,7 @@ def save_openapi_spec(self, output_path: str) -> None:
     logger.info("OpenAPI specification saved to %s", output_path)
 
 
-def main():
+def main() -> int:
   """Command line interface for the converter."""
   parser = argparse.ArgumentParser(
       description=(
@@ -575,4 +579,4 @@ def main():
 
 
 if __name__ == "__main__":
-  main()
+  raise SystemExit(main())