chore: Move API registry to the integrations folder
Added a deprecation warning in the old tools/api_registry file. Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 878660213
This commit is contained in:
Kathy Wu
committed by
Copybara-Service
Copybara-Service
parent
34c560e66e
commit
45fb53b9e2
@@ -15,7 +15,7 @@
|
||||
import os
|
||||
|
||||
from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.tools.api_registry import ApiRegistry
|
||||
from google.adk.integrations.api_registry import ApiRegistry
|
||||
|
||||
# TODO: Fill in with your GCloud project id and MCP server name
|
||||
PROJECT_ID = "your-google-cloud-project-id"
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .api_registry import ApiRegistry
|
||||
|
||||
__all__ = [
|
||||
'ApiRegistry',
|
||||
]
|
||||
@@ -0,0 +1,140 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from typing import Callable
|
||||
|
||||
from google.adk.agents.readonly_context import ReadonlyContext
|
||||
from google.adk.tools.base_toolset import ToolPredicate
|
||||
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
|
||||
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
|
||||
import google.auth
|
||||
import google.auth.transport.requests
|
||||
import httpx
|
||||
|
||||
API_REGISTRY_URL = "https://cloudapiregistry.googleapis.com"
|
||||
|
||||
|
||||
class ApiRegistry:
|
||||
"""Registry that provides McpToolsets for MCP servers registered in API Registry."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_registry_project_id: str,
|
||||
location: str = "global",
|
||||
header_provider: (
|
||||
Callable[[ReadonlyContext], dict[str, str]] | None
|
||||
) = None,
|
||||
):
|
||||
"""Initialize the API Registry.
|
||||
|
||||
Args:
|
||||
api_registry_project_id: The project ID for the Google Cloud API Registry.
|
||||
location: The location of the API Registry resources.
|
||||
header_provider: Optional function to provide additional headers for MCP
|
||||
server calls.
|
||||
"""
|
||||
self.api_registry_project_id = api_registry_project_id
|
||||
self.location = location
|
||||
self._credentials, _ = google.auth.default()
|
||||
self._mcp_servers: dict[str, dict[str, Any]] = {}
|
||||
self._header_provider = header_provider
|
||||
|
||||
url = f"{API_REGISTRY_URL}/v1beta/projects/{self.api_registry_project_id}/locations/{self.location}/mcpServers"
|
||||
|
||||
try:
|
||||
headers = self._get_auth_headers()
|
||||
headers["Content-Type"] = "application/json"
|
||||
page_token = None
|
||||
with httpx.Client() as client:
|
||||
while True:
|
||||
params = {}
|
||||
if page_token:
|
||||
params["pageToken"] = page_token
|
||||
|
||||
response = client.get(url, headers=headers, params=params)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
mcp_servers_list = data.get("mcpServers", [])
|
||||
for server in mcp_servers_list:
|
||||
server_name = server.get("name", "")
|
||||
if server_name:
|
||||
self._mcp_servers[server_name] = server
|
||||
|
||||
page_token = data.get("nextPageToken")
|
||||
if not page_token:
|
||||
break
|
||||
except (httpx.HTTPError, ValueError) as e:
|
||||
# Handle error in fetching or parsing tool definitions
|
||||
raise RuntimeError(
|
||||
f"Error fetching MCP servers from API Registry: {e}"
|
||||
) from e
|
||||
|
||||
def get_toolset(
|
||||
self,
|
||||
mcp_server_name: str,
|
||||
tool_filter: ToolPredicate | list[str] | None = None,
|
||||
tool_name_prefix: str | None = None,
|
||||
) -> McpToolset:
|
||||
"""Return the MCP Toolset based on the params.
|
||||
|
||||
Args:
|
||||
mcp_server_name: Filter to select the MCP server name to get tools from.
|
||||
tool_filter: Optional filter to select specific tools. Can be a list of
|
||||
tool names or a ToolPredicate function.
|
||||
tool_name_prefix: Optional prefix to prepend to the names of the tools
|
||||
returned by the toolset.
|
||||
|
||||
Returns:
|
||||
McpToolset: A toolset for the MCP server specified.
|
||||
"""
|
||||
server = self._mcp_servers.get(mcp_server_name)
|
||||
if not server:
|
||||
raise ValueError(
|
||||
f"MCP server {mcp_server_name} not found in API Registry."
|
||||
)
|
||||
if not server.get("urls"):
|
||||
raise ValueError(f"MCP server {mcp_server_name} has no URLs.")
|
||||
|
||||
mcp_server_url = server["urls"][0]
|
||||
headers = self._get_auth_headers()
|
||||
|
||||
# Only prepend "https://" if the URL doesn't already have a scheme
|
||||
if not mcp_server_url.startswith(("http://", "https://")):
|
||||
mcp_server_url = "https://" + mcp_server_url
|
||||
|
||||
return McpToolset(
|
||||
connection_params=StreamableHTTPConnectionParams(
|
||||
url=mcp_server_url,
|
||||
headers=headers,
|
||||
),
|
||||
tool_filter=tool_filter,
|
||||
tool_name_prefix=tool_name_prefix,
|
||||
header_provider=self._header_provider,
|
||||
)
|
||||
|
||||
def _get_auth_headers(self) -> dict[str, str]:
|
||||
"""Refreshes credentials and returns authorization headers."""
|
||||
request = google.auth.transport.requests.Request()
|
||||
self._credentials.refresh(request)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._credentials.token}",
|
||||
}
|
||||
# Add quota project header if available in ADC
|
||||
quota_project_id = getattr(self._credentials, "quota_project_id", None)
|
||||
if quota_project_id:
|
||||
headers["x-goog-user-project"] = quota_project_id
|
||||
return headers
|
||||
@@ -14,128 +14,13 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from typing import Callable
|
||||
import warnings
|
||||
|
||||
from google.adk.agents.readonly_context import ReadonlyContext
|
||||
import google.auth
|
||||
import google.auth.transport.requests
|
||||
import httpx
|
||||
from google.adk.integrations.api_registry import ApiRegistry
|
||||
|
||||
from .base_toolset import ToolPredicate
|
||||
from .mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
|
||||
from .mcp_tool.mcp_toolset import McpToolset
|
||||
|
||||
API_REGISTRY_URL = "https://cloudapiregistry.googleapis.com"
|
||||
|
||||
|
||||
class ApiRegistry:
|
||||
"""Registry that provides McpToolsets for MCP servers registered in API Registry."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_registry_project_id: str,
|
||||
location: str = "global",
|
||||
header_provider: (
|
||||
Callable[[ReadonlyContext], dict[str, str]] | None
|
||||
) = None,
|
||||
):
|
||||
"""Initialize the API Registry.
|
||||
|
||||
Args:
|
||||
api_registry_project_id: The project ID for the Google Cloud API Registry.
|
||||
location: The location of the API Registry resources.
|
||||
header_provider: Optional function to provide additional headers for MCP
|
||||
server calls.
|
||||
"""
|
||||
self.api_registry_project_id = api_registry_project_id
|
||||
self.location = location
|
||||
self._credentials, _ = google.auth.default()
|
||||
self._mcp_servers: dict[str, dict[str, Any]] = {}
|
||||
self._header_provider = header_provider
|
||||
|
||||
url = f"{API_REGISTRY_URL}/v1beta/projects/{self.api_registry_project_id}/locations/{self.location}/mcpServers"
|
||||
|
||||
try:
|
||||
headers = self._get_auth_headers()
|
||||
headers["Content-Type"] = "application/json"
|
||||
page_token = None
|
||||
with httpx.Client() as client:
|
||||
while True:
|
||||
params = {}
|
||||
if page_token:
|
||||
params["pageToken"] = page_token
|
||||
|
||||
response = client.get(url, headers=headers, params=params)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
mcp_servers_list = data.get("mcpServers", [])
|
||||
for server in mcp_servers_list:
|
||||
server_name = server.get("name", "")
|
||||
if server_name:
|
||||
self._mcp_servers[server_name] = server
|
||||
|
||||
page_token = data.get("nextPageToken")
|
||||
if not page_token:
|
||||
break
|
||||
except (httpx.HTTPError, ValueError) as e:
|
||||
# Handle error in fetching or parsing tool definitions
|
||||
raise RuntimeError(
|
||||
f"Error fetching MCP servers from API Registry: {e}"
|
||||
) from e
|
||||
|
||||
def get_toolset(
|
||||
self,
|
||||
mcp_server_name: str,
|
||||
tool_filter: ToolPredicate | list[str] | None = None,
|
||||
tool_name_prefix: str | None = None,
|
||||
) -> McpToolset:
|
||||
"""Return the MCP Toolset based on the params.
|
||||
|
||||
Args:
|
||||
mcp_server_name: Filter to select the MCP server name to get tools from.
|
||||
tool_filter: Optional filter to select specific tools. Can be a list of
|
||||
tool names or a ToolPredicate function.
|
||||
tool_name_prefix: Optional prefix to prepend to the names of the tools
|
||||
returned by the toolset.
|
||||
|
||||
Returns:
|
||||
McpToolset: A toolset for the MCP server specified.
|
||||
"""
|
||||
server = self._mcp_servers.get(mcp_server_name)
|
||||
if not server:
|
||||
raise ValueError(
|
||||
f"MCP server {mcp_server_name} not found in API Registry."
|
||||
)
|
||||
if not server.get("urls"):
|
||||
raise ValueError(f"MCP server {mcp_server_name} has no URLs.")
|
||||
|
||||
mcp_server_url = server["urls"][0]
|
||||
headers = self._get_auth_headers()
|
||||
|
||||
# Only prepend "https://" if the URL doesn't already have a scheme
|
||||
if not mcp_server_url.startswith(("http://", "https://")):
|
||||
mcp_server_url = "https://" + mcp_server_url
|
||||
|
||||
return McpToolset(
|
||||
connection_params=StreamableHTTPConnectionParams(
|
||||
url=mcp_server_url,
|
||||
headers=headers,
|
||||
),
|
||||
tool_filter=tool_filter,
|
||||
tool_name_prefix=tool_name_prefix,
|
||||
header_provider=self._header_provider,
|
||||
)
|
||||
|
||||
def _get_auth_headers(self) -> dict[str, str]:
|
||||
"""Refreshes credentials and returns authorization headers."""
|
||||
request = google.auth.transport.requests.Request()
|
||||
self._credentials.refresh(request)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._credentials.token}",
|
||||
}
|
||||
# Add quota project header if available in ADC
|
||||
quota_project_id = getattr(self._credentials, "quota_project_id", None)
|
||||
if quota_project_id:
|
||||
headers["x-goog-user-project"] = quota_project_id
|
||||
return headers
|
||||
warnings.warn(
|
||||
"google.adk.tools.api_registry is moved to"
|
||||
" google.adk.integrations.api_registry",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
+15
-6
@@ -18,8 +18,8 @@ from unittest.mock import create_autospec
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
|
||||
from google.adk.tools import api_registry
|
||||
from google.adk.tools.api_registry import ApiRegistry
|
||||
from google.adk.integrations import api_registry
|
||||
from google.adk.integrations.api_registry import ApiRegistry
|
||||
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
|
||||
import httpx
|
||||
|
||||
@@ -218,7 +218,10 @@ class TestApiRegistry(unittest.IsolatedAsyncioTestCase):
|
||||
)
|
||||
mock_response.raise_for_status.assert_called_once()
|
||||
|
||||
@patch("google.adk.tools.api_registry.McpToolset", autospec=True)
|
||||
@patch(
|
||||
"google.adk.integrations.api_registry.api_registry.McpToolset",
|
||||
autospec=True,
|
||||
)
|
||||
@patch("httpx.Client", autospec=True)
|
||||
async def test_get_toolset_success(self, MockHttpClient, MockMcpToolset):
|
||||
mock_response = MagicMock()
|
||||
@@ -245,7 +248,10 @@ class TestApiRegistry(unittest.IsolatedAsyncioTestCase):
|
||||
)
|
||||
self.assertEqual(toolset, MockMcpToolset.return_value)
|
||||
|
||||
@patch("google.adk.tools.api_registry.McpToolset", autospec=True)
|
||||
@patch(
|
||||
"google.adk.integrations.api_registry.api_registry.McpToolset",
|
||||
autospec=True,
|
||||
)
|
||||
@patch("httpx.Client", autospec=True)
|
||||
async def test_get_toolset_with_quota_project_id_success(
|
||||
self, MockHttpClient, MockMcpToolset
|
||||
@@ -277,7 +283,10 @@ class TestApiRegistry(unittest.IsolatedAsyncioTestCase):
|
||||
)
|
||||
self.assertEqual(toolset, MockMcpToolset.return_value)
|
||||
|
||||
@patch("google.adk.tools.api_registry.McpToolset", autospec=True)
|
||||
@patch(
|
||||
"google.adk.integrations.api_registry.api_registry.McpToolset",
|
||||
autospec=True,
|
||||
)
|
||||
@patch("httpx.Client", autospec=True)
|
||||
async def test_get_toolset_with_filter_and_prefix(
|
||||
self, MockHttpClient, MockMcpToolset
|
||||
@@ -321,7 +330,7 @@ class TestApiRegistry(unittest.IsolatedAsyncioTestCase):
|
||||
with (
|
||||
patch.object(httpx, "Client", autospec=True) as MockHttpClient,
|
||||
patch.object(
|
||||
api_registry, "McpToolset", autospec=True
|
||||
api_registry.api_registry, "McpToolset", autospec=True
|
||||
) as MockMcpToolset,
|
||||
):
|
||||
mock_response = create_autospec(httpx.Response, instance=True)
|
||||
Reference in New Issue
Block a user