|
| 1 | +# Copyright 2025 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import os |
| 18 | +from typing import Any |
| 19 | +from typing import Dict |
| 20 | +from typing import Protocol |
| 21 | +from urllib.parse import urlparse |
| 22 | + |
| 23 | +from ..artifacts.base_artifact_service import BaseArtifactService |
| 24 | +from ..memory.base_memory_service import BaseMemoryService |
| 25 | +from ..sessions.base_session_service import BaseSessionService |
| 26 | + |
| 27 | + |
| 28 | +def _load_gcp_config( |
| 29 | + agents_dir: str | None, service_name: str |
| 30 | +) -> tuple[str, str]: |
| 31 | + """Loads GCP project and location from environment.""" |
| 32 | + if not agents_dir: |
| 33 | + raise ValueError(f"agents_dir must be provided for {service_name}") |
| 34 | + |
| 35 | + from .utils import envs |
| 36 | + |
| 37 | + envs.load_dotenv_for_agent("", agents_dir) |
| 38 | + |
| 39 | + project = os.environ.get("GOOGLE_CLOUD_PROJECT") |
| 40 | + location = os.environ.get("GOOGLE_CLOUD_LOCATION") |
| 41 | + |
| 42 | + if not project or not location: |
| 43 | + raise ValueError("GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_LOCATION not set.") |
| 44 | + |
| 45 | + return project, location |
| 46 | + |
| 47 | + |
| 48 | +def _parse_agent_engine_kwargs( |
| 49 | + uri_part: str, agents_dir: str | None |
| 50 | +) -> dict[str, Any]: |
| 51 | + """Helper to parse agent engine resource name.""" |
| 52 | + if not uri_part: |
| 53 | + raise ValueError( |
| 54 | + "Agent engine resource name or resource id can not be empty." |
| 55 | + ) |
| 56 | + if "/" in uri_part: |
| 57 | + parts = uri_part.split("/") |
| 58 | + if not ( |
| 59 | + len(parts) == 6 |
| 60 | + and parts[0] == "projects" |
| 61 | + and parts[2] == "locations" |
| 62 | + and parts[4] == "reasoningEngines" |
| 63 | + ): |
| 64 | + raise ValueError( |
| 65 | + "Agent engine resource name is mal-formatted. It should be of" |
| 66 | + " format :" |
| 67 | + " projects/{project_id}/locations/{location}/reasoningEngines/{resource_id}" |
| 68 | + ) |
| 69 | + project = parts[1] |
| 70 | + location = parts[3] |
| 71 | + agent_engine_id = parts[5] |
| 72 | + else: |
| 73 | + project, location = _load_gcp_config( |
| 74 | + agents_dir, "short-form agent engine IDs" |
| 75 | + ) |
| 76 | + agent_engine_id = uri_part |
| 77 | + return { |
| 78 | + "project": project, |
| 79 | + "location": location, |
| 80 | + "agent_engine_id": agent_engine_id, |
| 81 | + } |
| 82 | + |
| 83 | + |
| 84 | +class ServiceFactory(Protocol): |
| 85 | + """Protocol for service factory functions.""" |
| 86 | + |
| 87 | + def __call__( |
| 88 | + self, uri: str, **kwargs |
| 89 | + ) -> BaseSessionService | BaseArtifactService | BaseMemoryService: |
| 90 | + ... |
| 91 | + |
| 92 | + |
| 93 | +class ServiceRegistry: |
| 94 | + """Registry for custom service URI schemes.""" |
| 95 | + |
| 96 | + def __init__(self): |
| 97 | + self._session_factories: Dict[str, ServiceFactory] = {} |
| 98 | + self._artifact_factories: Dict[str, ServiceFactory] = {} |
| 99 | + self._memory_factories: Dict[str, ServiceFactory] = {} |
| 100 | + |
| 101 | + def register_session_service( |
| 102 | + self, scheme: str, factory: ServiceFactory |
| 103 | + ) -> None: |
| 104 | + """Register a factory for a custom session service URI scheme. |
| 105 | +
|
| 106 | + Args: |
| 107 | + scheme: URI scheme (e.g., 'custom') |
| 108 | + factory: Callable that takes (uri, **kwargs) and returns |
| 109 | + BaseSessionService |
| 110 | + """ |
| 111 | + self._session_factories[scheme] = factory |
| 112 | + |
| 113 | + def register_artifact_service( |
| 114 | + self, scheme: str, factory: ServiceFactory |
| 115 | + ) -> None: |
| 116 | + """Register a factory for a custom artifact service URI scheme.""" |
| 117 | + self._artifact_factories[scheme] = factory |
| 118 | + |
| 119 | + def register_memory_service( |
| 120 | + self, scheme: str, factory: ServiceFactory |
| 121 | + ) -> None: |
| 122 | + """Register a factory for a custom memory service URI scheme.""" |
| 123 | + self._memory_factories[scheme] = factory |
| 124 | + |
| 125 | + def create_session_service( |
| 126 | + self, uri: str, **kwargs |
| 127 | + ) -> BaseSessionService | None: |
| 128 | + """Create session service from URI using registered factories.""" |
| 129 | + scheme = urlparse(uri).scheme |
| 130 | + if scheme and scheme in self._session_factories: |
| 131 | + return self._session_factories[scheme](uri, **kwargs) |
| 132 | + return None |
| 133 | + |
| 134 | + def create_artifact_service( |
| 135 | + self, uri: str, **kwargs |
| 136 | + ) -> BaseArtifactService | None: |
| 137 | + """Create artifact service from URI using registered factories.""" |
| 138 | + scheme = urlparse(uri).scheme |
| 139 | + if scheme and scheme in self._artifact_factories: |
| 140 | + return self._artifact_factories[scheme](uri, **kwargs) |
| 141 | + return None |
| 142 | + |
| 143 | + def create_memory_service( |
| 144 | + self, uri: str, **kwargs |
| 145 | + ) -> BaseMemoryService | None: |
| 146 | + """Create memory service from URI using registered factories.""" |
| 147 | + scheme = urlparse(uri).scheme |
| 148 | + if scheme and scheme in self._memory_factories: |
| 149 | + return self._memory_factories[scheme](uri, **kwargs) |
| 150 | + return None |
| 151 | + |
| 152 | + |
| 153 | +def _register_builtin_services(registry: ServiceRegistry) -> None: |
| 154 | + """Register built-in service implementations.""" |
| 155 | + |
| 156 | + # -- Session Services -- |
| 157 | + def agentengine_session_factory(uri: str, **kwargs): |
| 158 | + from ..sessions.vertex_ai_session_service import VertexAiSessionService |
| 159 | + |
| 160 | + parsed = urlparse(uri) |
| 161 | + params = _parse_agent_engine_kwargs( |
| 162 | + parsed.netloc + parsed.path, kwargs.get("agents_dir") |
| 163 | + ) |
| 164 | + return VertexAiSessionService(**params) |
| 165 | + |
| 166 | + def database_session_factory(uri: str, **kwargs): |
| 167 | + from ..sessions.database_session_service import DatabaseSessionService |
| 168 | + |
| 169 | + kwargs_copy = kwargs.copy() |
| 170 | + kwargs_copy.pop("agents_dir", None) |
| 171 | + return DatabaseSessionService(db_url=uri, **kwargs_copy) |
| 172 | + |
| 173 | + registry.register_session_service("agentengine", agentengine_session_factory) |
| 174 | + for scheme in ["sqlite", "postgresql", "mysql"]: |
| 175 | + registry.register_session_service(scheme, database_session_factory) |
| 176 | + |
| 177 | + # -- Artifact Services -- |
| 178 | + def gcs_artifact_factory(uri: str, **kwargs): |
| 179 | + from ..artifacts.gcs_artifact_service import GcsArtifactService |
| 180 | + |
| 181 | + kwargs_copy = kwargs.copy() |
| 182 | + kwargs_copy.pop("agents_dir", None) |
| 183 | + parsed_uri = urlparse(uri) |
| 184 | + bucket_name = parsed_uri.netloc |
| 185 | + return GcsArtifactService(bucket_name=bucket_name, **kwargs_copy) |
| 186 | + |
| 187 | + registry.register_artifact_service("gs", gcs_artifact_factory) |
| 188 | + |
| 189 | + # -- Memory Services -- |
| 190 | + def rag_memory_factory(uri: str, **kwargs): |
| 191 | + from ..memory.vertex_ai_rag_memory_service import VertexAiRagMemoryService |
| 192 | + |
| 193 | + rag_corpus = urlparse(uri).netloc |
| 194 | + if not rag_corpus: |
| 195 | + raise ValueError("Rag corpus can not be empty.") |
| 196 | + agents_dir = kwargs.get("agents_dir") |
| 197 | + project, location = _load_gcp_config(agents_dir, "RAG memory service") |
| 198 | + return VertexAiRagMemoryService( |
| 199 | + rag_corpus=( |
| 200 | + f"projects/{project}/locations/{location}/ragCorpora/{rag_corpus}" |
| 201 | + ) |
| 202 | + ) |
| 203 | + |
| 204 | + def agentengine_memory_factory(uri: str, **kwargs): |
| 205 | + from ..memory.vertex_ai_memory_bank_service import VertexAiMemoryBankService |
| 206 | + |
| 207 | + parsed = urlparse(uri) |
| 208 | + params = _parse_agent_engine_kwargs( |
| 209 | + parsed.netloc + parsed.path, kwargs.get("agents_dir") |
| 210 | + ) |
| 211 | + return VertexAiMemoryBankService(**params) |
| 212 | + |
| 213 | + registry.register_memory_service("rag", rag_memory_factory) |
| 214 | + registry.register_memory_service("agentengine", agentengine_memory_factory) |
| 215 | + |
| 216 | + |
| 217 | +# Global registry instance |
| 218 | +_global_registry = ServiceRegistry() |
| 219 | +_register_builtin_services(_global_registry) |
| 220 | + |
| 221 | + |
| 222 | +def get_service_registry() -> ServiceRegistry: |
| 223 | + """Get the global service registry instance.""" |
| 224 | + return _global_registry |
0 commit comments