diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py index db5edb0a..37e36ff9 100644 --- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py +++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py @@ -18,6 +18,7 @@ import json import logging import ssl from typing import Any +from typing import Callable from typing import Dict from typing import Final from typing import List @@ -71,6 +72,9 @@ class OpenAPIToolset(BaseToolset): tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, tool_name_prefix: Optional[str] = None, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None, + header_provider: Optional[ + Callable[[ReadonlyContext], Dict[str, str]] + ] = None, ): """Initializes the OpenAPIToolset. @@ -116,8 +120,14 @@ class OpenAPIToolset(BaseToolset): - ssl.SSLContext: Custom SSL context for advanced configuration This is useful for enterprise environments where requests go through a TLS-intercepting proxy with a custom CA certificate. + header_provider: A callable that returns a dictionary of headers to be + included in API requests. The callable receives the ReadonlyContext as + an argument, allowing dynamic header generation based on the current + context. Useful for adding custom headers like correlation IDs, + authentication tokens, or other request metadata. """ super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix) + self._header_provider = header_provider if not spec_dict: spec_dict = self._load_spec(spec_str, spec_str_type) self._ssl_verify = ssl_verify @@ -189,7 +199,11 @@ class OpenAPIToolset(BaseToolset): tools = [] for o in operations: - tool = RestApiTool.from_parsed_operation(o, ssl_verify=self._ssl_verify) + tool = RestApiTool.from_parsed_operation( + o, + ssl_verify=self._ssl_verify, + header_provider=self._header_provider, + ) logger.info("Parsed tool: %s", tool.name) tools.append(tool) return tools diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py index 60fb1362..5c27b168 100644 --- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py +++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py @@ -16,6 +16,7 @@ from __future__ import annotations import ssl from typing import Any +from typing import Callable from typing import Dict from typing import List from typing import Literal @@ -29,6 +30,7 @@ from google.genai.types import FunctionDeclaration import requests from typing_extensions import override +from ....agents.readonly_context import ReadonlyContext from ....auth.auth_credential import AuthCredential from ....auth.auth_schemes import AuthScheme from ..._gemini_schema_util import _to_gemini_schema @@ -90,6 +92,9 @@ class RestApiTool(BaseTool): auth_credential: Optional[Union[AuthCredential, str]] = None, should_parse_operation=True, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None, + header_provider: Optional[ + Callable[[ReadonlyContext], Dict[str, str]] + ] = None, ): """Initializes the RestApiTool with the given parameters. @@ -122,6 +127,11 @@ class RestApiTool(BaseTool): - False: Disable SSL verification (insecure, not recommended) - str: Path to a CA bundle file or directory for custom CA - ssl.SSLContext: Custom SSL context for advanced configuration + header_provider: A callable that returns a dictionary of headers to be + included in API requests. The callable receives the ReadonlyContext as + an argument, allowing dynamic header generation based on the current + context. Useful for adding custom headers like correlation IDs, + authentication tokens, or other request metadata. """ # Gemini restrict the length of function name to be less than 64 characters self.name = name[:60] @@ -145,6 +155,7 @@ class RestApiTool(BaseTool): self.credential_exchanger = AutoAuthCredentialExchanger() self._default_headers: Dict[str, str] = {} self._ssl_verify = ssl_verify + self._header_provider = header_provider if should_parse_operation: self._operation_parser = OperationParser(self.operation) @@ -153,12 +164,20 @@ class RestApiTool(BaseTool): cls, parsed: ParsedOperation, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None, + header_provider: Optional[ + Callable[[ReadonlyContext], Dict[str, str]] + ] = None, ) -> "RestApiTool": """Initializes the RestApiTool from a ParsedOperation object. Args: parsed: A ParsedOperation object. ssl_verify: SSL certificate verification option. + header_provider: A callable that returns a dictionary of headers to be + included in API requests. The callable receives the ReadonlyContext as + an argument, allowing dynamic header generation based on the current + context. Useful for adding custom headers like correlation IDs, + authentication tokens, or other request metadata. Returns: A RestApiTool object. @@ -178,6 +197,7 @@ class RestApiTool(BaseTool): auth_scheme=parsed.auth_scheme, auth_credential=parsed.auth_credential, ssl_verify=ssl_verify, + header_provider=header_provider, ) generated._operation_parser = operation_parser return generated @@ -450,6 +470,13 @@ class RestApiTool(BaseTool): request_params = self._prepare_request_params(api_params, api_args) if self._ssl_verify is not None: request_params["verify"] = self._ssl_verify + + # Add headers from header_provider if configured + if self._header_provider is not None and tool_context is not None: + provider_headers = self._header_provider(tool_context) + if provider_headers: + request_params.setdefault("headers", {}).update(provider_headers) + response = requests.request(**request_params) # Parse API response diff --git a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_toolset.py b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_toolset.py index f5d9e997..5238a287 100644 --- a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_toolset.py +++ b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_toolset.py @@ -135,9 +135,8 @@ def test_openapi_toolset_configure_auth_on_init(openapi_spec: Dict): auth_scheme=auth_scheme, auth_credential=auth_credential, ) - for tool in toolset._tools: - assert tool.auth_scheme == auth_scheme - assert tool.auth_credential == auth_credential + assert all(tool.auth_scheme == auth_scheme for tool in toolset._tools) + assert all(tool.auth_credential == auth_credential for tool in toolset._tools) @pytest.mark.parametrize( @@ -151,8 +150,7 @@ def test_openapi_toolset_verify_on_init( spec_dict=openapi_spec, ssl_verify=verify_value, ) - for tool in toolset._tools: - assert tool._ssl_verify == verify_value + assert all(tool._ssl_verify == verify_value for tool in toolset._tools) def test_openapi_toolset_configure_verify_all(openapi_spec: Dict[str, Any]): @@ -160,15 +158,13 @@ def test_openapi_toolset_configure_verify_all(openapi_spec: Dict[str, Any]): toolset = OpenAPIToolset(spec_dict=openapi_spec) # Initially verify should be None - for tool in toolset._tools: - assert tool._ssl_verify is None + assert all(tool._ssl_verify is None for tool in toolset._tools) # Configure verify for all tools ca_bundle_path = "/path/to/custom-ca.crt" toolset.configure_ssl_verify_all(ca_bundle_path) - for tool in toolset._tools: - assert tool._ssl_verify == ca_bundle_path + assert all(tool._ssl_verify == ca_bundle_path for tool in toolset._tools) async def test_openapi_toolset_tool_name_prefix(openapi_spec: Dict[str, Any]): @@ -183,10 +179,42 @@ async def test_openapi_toolset_tool_name_prefix(openapi_spec: Dict[str, Any]): assert len(prefixed_tools) == 5 # Verify all tool names are prefixed - for tool in prefixed_tools: - assert tool.name.startswith(f"{prefix}_") + assert all(tool.name.startswith(f"{prefix}_") for tool in prefixed_tools) # Verify specific tool name is prefixed expected_prefixed_name = "my_api_calendar_calendars_insert" prefixed_tool_names = [t.name for t in prefixed_tools] assert expected_prefixed_name in prefixed_tool_names + + +def test_openapi_toolset_header_provider(openapi_spec: Dict[str, Any]): + """Test header_provider parameter is passed to tools.""" + + def my_header_provider(context): + return {"X-Custom-Header": "custom-value", "X-Request-ID": "12345"} + + toolset = OpenAPIToolset( + spec_dict=openapi_spec, + header_provider=my_header_provider, + ) + + # Verify the toolset has the header_provider set + assert toolset._header_provider is my_header_provider + + # Verify all tools have the header_provider + assert all( + tool._header_provider is my_header_provider for tool in toolset._tools + ) + + +def test_openapi_toolset_header_provider_none_by_default( + openapi_spec: Dict[str, Any], +): + """Test that header_provider is None by default.""" + toolset = OpenAPIToolset(spec_dict=openapi_spec) + + # Verify the toolset has no header_provider by default + assert toolset._header_provider is None + + # Verify all tools have no header_provider + assert all(tool._header_provider is None for tool in toolset._tools) diff --git a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py index e91fca9c..560813e6 100644 --- a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py +++ b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py @@ -1036,6 +1036,149 @@ class TestRestApiTool: call_kwargs = mock_request.call_args[1] assert call_kwargs["verify"] == ca_bundle_path + def test_init_with_header_provider( + self, + sample_endpoint, + sample_operation, + ): + """Test that header_provider is stored correctly.""" + + def my_header_provider(context): + return {"X-Custom": "value"} + + tool = RestApiTool( + name="test_tool", + description="Test Tool", + endpoint=sample_endpoint, + operation=sample_operation, + header_provider=my_header_provider, + ) + assert tool._header_provider is my_header_provider + + def test_init_header_provider_none_by_default( + self, + sample_endpoint, + sample_operation, + ): + """Test that header_provider is None by default.""" + tool = RestApiTool( + name="test_tool", + description="Test Tool", + endpoint=sample_endpoint, + operation=sample_operation, + ) + assert tool._header_provider is None + + @pytest.mark.asyncio + async def test_call_with_header_provider( + self, + mock_tool_context, + sample_endpoint, + sample_operation, + sample_auth_scheme, + sample_auth_credential, + ): + """Test that header_provider adds headers to the request.""" + mock_response = mock.create_autospec( + requests.Response, instance=True, spec_set=True + ) + mock_response.json.return_value = {"result": "success"} + + def my_header_provider(context): + return {"X-Custom-Header": "custom-value", "X-Request-ID": "12345"} + + tool = RestApiTool( + name="test_tool", + description="Test Tool", + endpoint=sample_endpoint, + operation=sample_operation, + auth_scheme=sample_auth_scheme, + auth_credential=sample_auth_credential, + header_provider=my_header_provider, + ) + + with patch.object( + requests, "request", return_value=mock_response, autospec=True + ) as mock_request: + await tool.call(args={}, tool_context=mock_tool_context) + + # Verify the headers were added to the request + assert mock_request.called + _, call_kwargs = mock_request.call_args + assert call_kwargs["headers"]["X-Custom-Header"] == "custom-value" + assert call_kwargs["headers"]["X-Request-ID"] == "12345" + + @pytest.mark.asyncio + async def test_call_header_provider_receives_tool_context( + self, + mock_tool_context, + sample_endpoint, + sample_operation, + sample_auth_scheme, + sample_auth_credential, + ): + """Test that header_provider receives the tool_context.""" + mock_response = mock.create_autospec( + requests.Response, instance=True, spec_set=True + ) + mock_response.json.return_value = {"result": "success"} + + received_context = [] + + def my_header_provider(context): + received_context.append(context) + return {"X-Test": "test"} + + tool = RestApiTool( + name="test_tool", + description="Test Tool", + endpoint=sample_endpoint, + operation=sample_operation, + auth_scheme=sample_auth_scheme, + auth_credential=sample_auth_credential, + header_provider=my_header_provider, + ) + + with patch.object( + requests, "request", return_value=mock_response, autospec=True + ): + await tool.call(args={}, tool_context=mock_tool_context) + + # Verify header_provider was called with the tool_context + assert len(received_context) == 1 + assert received_context[0] is mock_tool_context + + @pytest.mark.asyncio + async def test_call_without_header_provider( + self, + mock_tool_context, + sample_endpoint, + sample_operation, + sample_auth_scheme, + sample_auth_credential, + ): + """Test that call works without header_provider.""" + mock_response = mock.create_autospec( + requests.Response, instance=True, spec_set=True + ) + mock_response.json.return_value = {"result": "success"} + + tool = RestApiTool( + name="test_tool", + description="Test Tool", + endpoint=sample_endpoint, + operation=sample_operation, + auth_scheme=sample_auth_scheme, + auth_credential=sample_auth_credential, + ) + + with patch.object( + requests, "request", return_value=mock_response, autospec=True + ): + result = await tool.call(args={}, tool_context=mock_tool_context) + + assert result == {"result": "success"} + def test_snake_to_lower_camel(): assert snake_to_lower_camel("single") == "single"