1717import sys
1818from typing import Any
1919from 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
2522import google .auth
2623import google .auth .transport .requests
2724import 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