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

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

@@ -20,6 +20,7 @@ import inspect
 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 @@ class ComputerUseToolset(BaseToolset):
     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 @@ class ComputerUseToolset(BaseToolset):
 
     @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 @@ class ComputerUseToolset(BaseToolset):
             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 @@ class ComputerUseToolset(BaseToolset):
       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 @@ class ComputerUseToolset(BaseToolset):
     )
 
   @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 @@ class ComputerUseToolset(BaseToolset):
       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 @@ class DataAgentToolset(BaseToolset):
     )
 
   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 @@ class DataAgentToolset(BaseToolset):
 
   @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 @@ class DataAgentToolset(BaseToolset):
     ]
 
   @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 @@ class GoogleApiToolset(BaseToolset):
       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 @@ class GoogleApiToolset(BaseToolset):
     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 @@ class GoogleApiToolset(BaseToolset):
     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 @@ class GoogleApiToolset(BaseToolset):
         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 @@ class GoogleApiToolset(BaseToolset):
   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 @@ class GoogleApiToOpenApiConverter:
     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 @@ class GoogleApiToOpenApiConverter:
         )
 
       # 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 @@ class GoogleApiToOpenApiConverter:
     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 @@ class GoogleApiToOpenApiConverter:
       ] = 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 @@ class GoogleApiToOpenApiConverter:
     Returns:
         Converted OpenAPI schema object
     """
-    result = {}
+    result: dict[str, Any] = {}
 
     # Convert the type
     if "type" in schema_def:
@@ -332,7 +336,7 @@ class GoogleApiToOpenApiConverter:
     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 @@ class GoogleApiToOpenApiConverter:
         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 @@ class GoogleApiToOpenApiConverter:
           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 @@ class GoogleApiToOpenApiConverter:
     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 @@ class GoogleApiToOpenApiConverter:
     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 @@ class GoogleApiToOpenApiConverter:
     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 @@ class GoogleApiToOpenApiConverter:
     Returns:
         OpenAPI schema for the parameter
     """
-    schema = {}
+    schema: dict[str, Any] = {}
 
     # Convert type
     param_type = param_data.get("type", "string")
@@ -536,7 +540,7 @@ class GoogleApiToOpenApiConverter:
     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())

src/google/adk/tools/retrieval/__init__.py

@@ -12,8 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TYPE_CHECKING
+
 from .base_retrieval_tool import BaseRetrievalTool
 
+if TYPE_CHECKING:
+  from .files_retrieval import FilesRetrieval as FilesRetrieval
+  from .llama_index_retrieval import LlamaIndexRetrieval as LlamaIndexRetrieval
+  from .vertex_ai_rag_retrieval import VertexAiRagRetrieval as VertexAiRagRetrieval
+
 __all__ = [
     "BaseRetrievalTool",
     "FilesRetrieval",
@@ -22,7 +29,7 @@ __all__ = [
 ]
 
 
-def __getattr__(name: str):
+def __getattr__(name: str) -> object:
   if name == "FilesRetrieval":
     try:
       from .files_retrieval import FilesRetrieval

src/google/adk/tools/retrieval/files_retrieval.py

@@ -17,7 +17,9 @@
 from __future__ import annotations
 
 import logging
+from typing import cast
 from typing import Optional
+from typing import Protocol
 
 from llama_index.core import SimpleDirectoryReader
 from llama_index.core import VectorStoreIndex
@@ -28,6 +30,14 @@ from .llama_index_retrieval import LlamaIndexRetrieval
 logger = logging.getLogger("google_adk." + __name__)
 
 
+class _EmbeddingFactory(Protocol):
+
+  def __call__(
+      self, *, model_name: str, embed_batch_size: int
+  ) -> BaseEmbedding:
+    ...
+
+
 def _get_default_embedding_model() -> BaseEmbedding:
   """Get the default Google Gemini embedding model.
 
@@ -40,7 +50,8 @@ def _get_default_embedding_model() -> BaseEmbedding:
   try:
     from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
 
-    return GoogleGenAIEmbedding(
+    factory = cast(_EmbeddingFactory, GoogleGenAIEmbedding)
+    return factory(
         model_name="gemini-embedding-2-preview",
         embed_batch_size=1,
     )

src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py

@@ -19,11 +19,13 @@ from __future__ import annotations
 import asyncio
 import logging
 from typing import Any
+from typing import cast
 from typing import TYPE_CHECKING
 
 from google.genai import types
 from typing_extensions import override
 
+from ...models.llm_request import LlmRequest
 from ...utils.model_name_utils import is_gemini_model
 from ...utils.model_name_utils import is_gemini_model_id_check_disabled
 from ..tool_context import ToolContext
@@ -31,7 +33,6 @@ from .base_retrieval_tool import BaseRetrievalTool
 
 if TYPE_CHECKING:
   from ...dependencies.vertexai import rag
-  from ...models import LlmRequest
 
 logger = logging.getLogger('google_adk.' + __name__)
 
@@ -44,15 +45,21 @@ class VertexAiRagRetrieval(BaseRetrievalTool):
       *,
       name: str,
       description: str,
-      rag_corpora: list[str] = None,
-      rag_resources: list[rag.RagResource] = None,
-      similarity_top_k: int = None,
-      vector_distance_threshold: float = None,
+      rag_corpora: list[str] | None = None,
+      rag_resources: list[rag.RagResource] | None = None,
+      similarity_top_k: int | None = None,
+      vector_distance_threshold: float | None = None,
   ):
     super().__init__(name=name, description=description)
+    # VertexRagStore validates from attributes, so it rebuilds each resource as
+    # its own type and the originals are unrecoverable from it. retrieval_query
+    # needs the vertexai ones, so keep them.
+    self._rag_resources = rag_resources
     self.vertex_rag_store = types.VertexRagStore(
         rag_corpora=rag_corpora,
-        rag_resources=rag_resources,
+        rag_resources=cast(
+            'list[types.VertexRagStoreRagResource] | None', rag_resources
+        ),
         similarity_top_k=similarity_top_k,
         vector_distance_threshold=vector_distance_threshold,
     )
@@ -95,10 +102,14 @@ class VertexAiRagRetrieval(BaseRetrievalTool):
   ) -> Any:
     from ...dependencies.vertexai import rag
 
+    query = args.get('query')
+    if not isinstance(query, str):
+      raise ValueError("Vertex AI RAG retrieval requires a string 'query'.")
+
     response = await asyncio.to_thread(
         rag.retrieval_query,
-        text=args['query'],
-        rag_resources=self.vertex_rag_store.rag_resources,
+        text=query,
+        rag_resources=self._rag_resources,
         rag_corpora=self.vertex_rag_store.rag_corpora,
         similarity_top_k=self.vertex_rag_store.similarity_top_k,
         vector_distance_threshold=self.vertex_rag_store.vector_distance_threshold,

tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py

@@ -16,6 +16,8 @@ from google.adk.agents.llm_agent import Agent
 from google.adk.tools.function_tool import FunctionTool
 from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
 from google.genai import types
+import pytest
+from vertexai.preview import rag
 
 from ... import testing_utils
 
@@ -24,6 +26,47 @@ def noop_tool(x: str) -> str:
   return x
 
 
+def test_vertex_rag_resources_are_converted_for_gemini():
+  resource = rag.RagResource(
+      rag_corpus='projects/p/locations/l/ragCorpora/c',
+      rag_file_ids=['file-1'],
+  )
+
+  retrieval = VertexAiRagRetrieval(
+      name='rag_retrieval',
+      description='rag_retrieval',
+      rag_resources=[resource],
+  )
+
+  assert retrieval.vertex_rag_store.rag_resources == [
+      types.VertexRagStoreRagResource(
+          rag_corpus='projects/p/locations/l/ragCorpora/c',
+          rag_file_ids=['file-1'],
+      )
+  ]
+
+
+@pytest.mark.asyncio
+async def test_retrieval_query_gets_the_original_rag_resources(mocker):
+  resource = rag.RagResource(
+      rag_corpus='projects/p/locations/l/ragCorpora/c',
+      rag_file_ids=['file-1'],
+  )
+  retrieval = VertexAiRagRetrieval(
+      name='rag_retrieval',
+      description='rag_retrieval',
+      rag_resources=[resource],
+  )
+  retrieval_query = mocker.patch(
+      'google.adk.dependencies.vertexai.rag.retrieval_query'
+  )
+  retrieval_query.return_value.contexts.contexts = []
+
+  await retrieval.run_async(args={'query': 'q'}, tool_context=mocker.Mock())
+
+  assert retrieval_query.call_args.kwargs['rag_resources'] == [resource]
+
+
 def test_vertex_rag_retrieval_for_non_gemini():
   responses = [
       'response1',