feat: Add support for model endpoints in Agent Registry

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 893585328

Author: wukath

contributing/samples/agent_registry_agent/agent.py

@@ -18,6 +18,7 @@
 
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.integrations.agent_registry import AgentRegistry
+from google.adk.models.google_llm import Gemini
 
 # Project and location can be set via environment variables:
 # GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION
@@ -27,6 +28,8 @@
 # Initialize Agent Registry client
 registry = AgentRegistry(project_id=project_id, location=location)
 
+# List agents, MCP servers, and endpoints resource names from the registry.
+# They can be used to initialize the agent, toolset, and model below.
 print(f"Listing agents in {project_id}/{location}...")
 agents = registry.list_agents()
 for agent in agents.get("agents", []):
@@ -37,6 +40,11 @@
 for server in mcp_servers.get("mcpServers", []):
   print(f"- MCP Server: {server.get('displayName')} ({server.get('name')})")
 
+print(f"\nListing endpoints in {project_id}/{location}...")
+endpoints = registry.list_endpoints()
+for endpoint in endpoints.get("endpoints", []):
+  print(f"- Endpoint: {endpoint.get('displayName')} ({endpoint.get('name')})")
+
 # Example of using a specific agent or MCP server from the registry:
 # (Note: These names should be full resource names as returned by list methods)
 
@@ -52,8 +60,19 @@
     f"projects/{project_id}/locations/{location}/mcpServers/MCP_SERVER_NAME"
 )
 
+# 3. Getting a specific model endpoint configuration
+# This returns a string like:
+# "projects/adk12345/locations/us-central1/publishers/google/models/gemini-2.5-flash"
+# TODO: Replace ENDPOINT_NAME with your endpoint name
+model_name = registry.get_model_name(
+    f"projects/{project_id}/locations/{location}/endpoints/ENDPOINT_NAME"
+)
+
+# Initialize the model using the resolved model name from registry.
+gemini_model = Gemini(model=model_name)
+
 root_agent = LlmAgent(
-    model="gemini-2.5-flash",
+    model=gemini_model,
     name="discovery_agent",
     instruction=(
         "You have access to tools and sub-agents discovered via Registry."

src/google/adk/integrations/agent_registry/agent_registry.py

@@ -24,8 +24,10 @@
 from typing import Callable
 from typing import Dict
 from typing import List
+from typing import Mapping
 from typing import Optional
 from typing import Sequence
+from typing import TypedDict
 from typing import Union
 from urllib.parse import parse_qs
 from urllib.parse import urlparse
@@ -109,6 +111,26 @@ class _ProtocolType(str, Enum):
   CUSTOM = "CUSTOM"
 
 
+class Interface(TypedDict, total=False):
+  """Details for a single connection interface."""
+
+  url: str
+  protocolBinding: str
+
+
+class Endpoint(TypedDict, total=False):
+  """Full metadata for a registered Endpoint."""
+
+  name: str
+  endpointId: str
+  displayName: str
+  description: str
+  interfaces: List[Interface]
+  createTime: str
+  updateTime: str
+  attributes: Dict[str, Any]
+
+
 class AgentRegistry:
   """Client for interacting with the Google Cloud Agent Registry service.
 
@@ -194,7 +216,7 @@ def _make_request(
 
   def _get_connection_uri(
       self,
-      resource_details: Dict[str, Any],
+      resource_details: Mapping[str, Any],
       protocol_type: Optional[_ProtocolType] = None,
       protocol_binding: Optional[A2ATransport] = None,
   ) -> Optional[str]:
@@ -273,6 +295,56 @@ def get_mcp_toolset(self, mcp_server_name: str) -> McpToolset:
         header_provider=self._header_provider,
     )
 
+  # --- Endpoint Methods ---
+
+  def list_endpoints(
+      self,
+      filter_str: Optional[str] = None,
+      page_size: Optional[int] = None,
+      page_token: Optional[str] = None,
+  ) -> Dict[str, Any]:
+    """Fetches a list of Endpoints."""
+    params = {}
+    if filter_str:
+      params["filter"] = filter_str
+    if page_size:
+      params["pageSize"] = str(page_size)
+    if page_token:
+      params["pageToken"] = page_token
+    return self._make_request("endpoints", params=params)
+
+  def get_endpoint(self, name: str) -> Endpoint:
+    """Retrieves details of a specific Endpoint."""
+    return self._make_request(name)  # type: ignore
+
+  def get_model_name(self, endpoint_name: str) -> str:
+    """Retrieves and parses an endpoint into a model resource name.
+
+    Args:
+      endpoint_name: The full resource name of the endpoint.
+
+    Returns:
+      The resolved model resource name string (e.g.
+      projects/.../locations/.../publishers/google/models/...).
+    """
+    endpoint_details = self.get_endpoint(endpoint_name)
+    uri = self._get_connection_uri(endpoint_details)
+    if not uri:
+      raise ValueError(
+          f"Connection URI not found for endpoint: {endpoint_name}"
+      )
+
+    uri = re.sub(r":\w+$", "", uri)
+
+    if uri.startswith("projects/"):
+      return uri
+
+    match = re.search(r"(projects/.+)", uri)
+    if match:
+      return match.group(1)
+
+    return uri
+
   # --- Agent Methods ---
 
   def list_agents(

src/google/adk/models/anthropic_llm.py

@@ -23,6 +23,7 @@
 import json
 import logging
 import os
+import re
 from typing import Any
 from typing import AsyncGenerator
 from typing import Iterable
@@ -364,10 +365,23 @@ class AnthropicLlm(BaseLlm):
   def supported_models(cls) -> list[str]:
     return [r"claude-3-.*", r"claude-.*-4.*"]
 
+  def _resolve_model_name(self, model: Optional[str]) -> str:
+    if not model:
+      return self.model
+    if model.startswith("projects/"):
+      match = re.search(
+          r"projects/[^/]+/locations/[^/]+/(?:publishers/anthropic/models|endpoints)/([^/:]+)",
+          model,
+      )
+      if match:
+        return match.group(1)
+    return model
+
   @override
   async def generate_content_async(
       self, llm_request: LlmRequest, stream: bool = False
   ) -> AsyncGenerator[LlmResponse, None]:
+    model_to_use = self._resolve_model_name(llm_request.model)
     messages = [
         content_to_message_param(content)
         for content in llm_request.contents or []
@@ -390,7 +404,7 @@ async def generate_content_async(
 
     if not stream:
       message = await self._anthropic_client.messages.create(
-          model=llm_request.model,
+          model=model_to_use,
           system=llm_request.config.system_instruction,
           messages=messages,
           tools=tools,
@@ -416,8 +430,9 @@ async def _generate_content_streaming(
     Yields partial LlmResponse objects as content arrives, followed by
     a final aggregated LlmResponse with all content.
     """
+    model_to_use = self._resolve_model_name(llm_request.model)
     raw_stream = await self._anthropic_client.messages.create(
-        model=llm_request.model,
+        model=model_to_use,
         system=llm_request.config.system_instruction,
         messages=messages,
         tools=tools,
@@ -511,17 +526,26 @@ class Claude(AnthropicLlm):
   @cached_property
   @override
   def _anthropic_client(self) -> AsyncAnthropicVertex:
-    if (
-        "GOOGLE_CLOUD_PROJECT" not in os.environ
-        or "GOOGLE_CLOUD_LOCATION" not in os.environ
-    ):
+    project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
+    location = os.environ.get("GOOGLE_CLOUD_LOCATION")
+
+    if self.model.startswith("projects/"):
+      match = re.search(
+          r"projects/([^/]+)/locations/([^/]+)/",
+          self.model,
+      )
+      if match:
+        project_id = match.group(1)
+        location = match.group(2)
+
+    if not project_id or not location:
       raise ValueError(
           "GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION must be set for using"
           " Anthropic on Vertex."
       )
 
     return AsyncAnthropicVertex(
-        project_id=os.environ["GOOGLE_CLOUD_PROJECT"],
-        region=os.environ["GOOGLE_CLOUD_LOCATION"],
+        project_id=project_id,
+        region=location,
         default_headers=get_tracking_headers(),
     )

src/google/adk/models/google_llm.py

@@ -304,13 +304,19 @@ def api_client(self) -> Client:
     """
     from google.genai import Client
 
-    return Client(
-        http_options=types.HttpOptions(
+    base_url = self.base_url
+
+    kwargs: dict[str, Any] = {
+        'http_options': types.HttpOptions(
             headers=self._tracking_headers(),
             retry_options=self.retry_options,
-            base_url=self.base_url,
+            base_url=base_url,
         )
-    )
+    }
+    if self.model.startswith('projects/'):
+      kwargs['vertexai'] = True
+
+    return Client(**kwargs)
 
   @cached_property
   def _api_backend(self) -> GoogleLLMVariant:
@@ -336,11 +342,19 @@ def _live_api_version(self) -> str:
   def _live_api_client(self) -> Client:
     from google.genai import Client
 
-    return Client(
-        http_options=types.HttpOptions(
-            headers=self._tracking_headers(), api_version=self._live_api_version
+    base_url = self.base_url
+
+    kwargs: dict[str, Any] = {
+        'http_options': types.HttpOptions(
+            headers=self._tracking_headers(),
+            api_version=self._live_api_version,
+            base_url=base_url,
         )
-    )
+    }
+    if self.model.startswith('projects/'):
+      kwargs['vertexai'] = True
+
+    return Client(**kwargs)
 
   @contextlib.asynccontextmanager
   async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:

tests/unittests/integrations/agent_registry/test_agent_registry.py

@@ -272,6 +272,37 @@ def test_get_mcp_server(self, mock_httpx, registry):
     server = registry.get_mcp_server("test-mcp")
     assert server == {"name": "test-mcp"}
 
+  @patch("httpx.Client")
+  def test_list_endpoints(self, mock_httpx, registry):
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"endpoints": []}
+    mock_response.raise_for_status = MagicMock()
+    mock_httpx.return_value.__enter__.return_value.get.return_value = (
+        mock_response
+    )
+
+    # Mock auth refresh
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    endpoints = registry.list_endpoints()
+    assert endpoints == {"endpoints": []}
+
+  @patch("httpx.Client")
+  def test_get_endpoint(self, mock_httpx, registry):
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"name": "test-endpoint"}
+    mock_response.raise_for_status = MagicMock()
+    mock_httpx.return_value.__enter__.return_value.get.return_value = (
+        mock_response
+    )
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    server = registry.get_endpoint("test-endpoint")
+    assert server == {"name": "test-endpoint"}
+
   @patch("httpx.Client")
   def test_get_mcp_toolset(self, mock_httpx, registry):
     mock_response = MagicMock()
@@ -420,3 +451,41 @@ def test_make_request_raises_generic_exception(self, mock_httpx, registry):
 
     with pytest.raises(RuntimeError, match="API request failed: Generic error"):
       registry._make_request("test-path")
+
+  @patch.object(AgentRegistry, "get_endpoint")
+  def test_get_model_name_starts_with_projects(
+      self, mock_get_endpoint, registry
+  ):
+    mock_get_endpoint.return_value = {
+        "interfaces": [{"url": "projects/p1/locations/l1/models/m1"}]
+    }
+    model_name = registry.get_model_name("test-endpoint")
+    assert model_name == "projects/p1/locations/l1/models/m1"
+
+  @patch.object(AgentRegistry, "get_endpoint")
+  def test_get_model_name_contains_projects(self, mock_get_endpoint, registry):
+    mock_get_endpoint.return_value = {
+        "interfaces": [{
+            "url": (
+                "https://vertexai.googleapis.com/v1/projects/p1/locations/l1/models/m1"
+            )
+        }]
+    }
+    model_name = registry.get_model_name("test-endpoint")
+    assert model_name == "projects/p1/locations/l1/models/m1"
+
+  @patch.object(AgentRegistry, "get_endpoint")
+  def test_get_model_name_strips_suffix(self, mock_get_endpoint, registry):
+    mock_get_endpoint.return_value = {
+        "interfaces": [{"url": "projects/p1/locations/l1/models/m1:predict"}]
+    }
+    model_name = registry.get_model_name("test-endpoint")
+    assert model_name == "projects/p1/locations/l1/models/m1"
+
+  @patch.object(AgentRegistry, "get_endpoint")
+  def test_get_model_name_raises_value_error_if_no_uri(
+      self, mock_get_endpoint, registry
+  ):
+    mock_get_endpoint.return_value = {}
+    with pytest.raises(ValueError, match="Connection URI not found"):
+      registry.get_model_name("test-endpoint")