diff --git a/src/google/adk/tools/computer_use/computer_use_toolset.py b/src/google/adk/tools/computer_use/computer_use_toolset.py index 43504f8d..103ba73e 100644 --- a/src/google/adk/tools/computer_use/computer_use_toolset.py +++ b/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 diff --git a/src/google/adk/tools/data_agent/credentials.py b/src/google/adk/tools/data_agent/credentials.py index 3503cfa0..673d5be7 100644 --- a/src/google/adk/tools/data_agent/credentials.py +++ b/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 diff --git a/src/google/adk/tools/data_agent/data_agent_toolset.py b/src/google/adk/tools/data_agent/data_agent_toolset.py index df7d7190..cc721412 100644 --- a/src/google/adk/tools/data_agent/data_agent_toolset.py +++ b/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 diff --git a/src/google/adk/tools/google_api_tool/google_api_toolset.py b/src/google/adk/tools/google_api_tool/google_api_toolset.py index 2e425b40..157fa6a7 100644 --- a/src/google/adk/tools/google_api_tool/google_api_toolset.py +++ b/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() diff --git a/src/google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py b/src/google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py index 131b186c..7fad7cba 100644 --- a/src/google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py +++ b/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()) diff --git a/src/google/adk/tools/retrieval/__init__.py b/src/google/adk/tools/retrieval/__init__.py index 93fabf08..d86a85a9 100644 --- a/src/google/adk/tools/retrieval/__init__.py +++ b/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 diff --git a/src/google/adk/tools/retrieval/files_retrieval.py b/src/google/adk/tools/retrieval/files_retrieval.py index 9db1bc39..306fca8c 100644 --- a/src/google/adk/tools/retrieval/files_retrieval.py +++ b/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, ) diff --git a/src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py b/src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py index a4e43974..25bbe52d 100644 --- a/src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py +++ b/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, diff --git a/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py b/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py index 2509f883..f366a28c 100644 --- a/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py +++ b/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',