Skip to content

Commit 253ac87

Browse files
wukathcopybara-github
authored andcommitted
fix: nit fixes for API Registry - updating imports / type hints and formatting
Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 845980315
1 parent 0088b0f commit 253ac87

File tree

1 file changed

+12
-15
lines changed

1 file changed

+12
-15
lines changed

src/google/adk/tools/api_registry.py

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,8 @@
1717
import sys
1818
from typing import Any
1919
from typing import Callable
20-
from typing import Dict
21-
from typing import List
22-
from typing import Optional
23-
from typing import Union
2420

21+
from google.adk.agents.readonly_context import ReadonlyContext
2522
import google.auth
2623
import google.auth.transport.requests
2724
import httpx
@@ -40,9 +37,9 @@ def __init__(
4037
self,
4138
api_registry_project_id: str,
4239
location: str = "global",
43-
header_provider: Optional[
44-
Callable[[ReadonlyContext], Dict[str, str]]
45-
] = None,
40+
header_provider: (
41+
Callable[[ReadonlyContext], dict[str, str]] | None
42+
) = None,
4643
):
4744
"""Initialize the API Registry.
4845
@@ -55,7 +52,7 @@ def __init__(
5552
self.api_registry_project_id = api_registry_project_id
5653
self.location = location
5754
self._credentials, _ = google.auth.default()
58-
self._mcp_servers: Dict[str, Dict[str, Any]] = {}
55+
self._mcp_servers: dict[str, dict[str, Any]] = {}
5956
self._header_provider = header_provider
6057

6158
url = f"{API_REGISTRY_URL}/v1beta/projects/{self.api_registry_project_id}/locations/{self.location}/mcpServers"
@@ -79,14 +76,13 @@ def __init__(
7976
def get_toolset(
8077
self,
8178
mcp_server_name: str,
82-
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
83-
tool_name_prefix: Optional[str] = None,
79+
tool_filter: ToolPredicate | list[str] | None = None,
80+
tool_name_prefix: str | None = None,
8481
) -> McpToolset:
8582
"""Return the MCP Toolset based on the params.
8683
8784
Args:
88-
mcp_server_name: Filter to select the MCP server name to get tools
89-
from.
85+
mcp_server_name: Filter to select the MCP server name to get tools from.
9086
tool_filter: Optional filter to select specific tools. Can be a list of
9187
tool names or a ToolPredicate function.
9288
tool_name_prefix: Optional prefix to prepend to the names of the tools
@@ -116,14 +112,15 @@ def get_toolset(
116112
header_provider=self._header_provider,
117113
)
118114

119-
def _get_auth_headers(self) -> Dict[str, str]:
115+
def _get_auth_headers(self) -> dict[str, str]:
120116
"""Refreshes credentials and returns authorization headers."""
121117
request = google.auth.transport.requests.Request()
122118
self._credentials.refresh(request)
123119
headers = {
124120
"Authorization": f"Bearer {self._credentials.token}",
125121
}
126122
# Add quota project header if available in ADC
127-
if self._credentials.quota_project_id:
128-
headers["x-goog-user-project"] = self._credentials.quota_project_id
123+
quota_project_id = getattr(self._credentials, "quota_project_id", None)
124+
if quota_project_id:
125+
headers["x-goog-user-project"] = quota_project_id
129126
return headers

0 commit comments

Comments
 (0)