mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Add a service registry to provide a generic way to register custom service implementations to be used in FastAPI server
To register a custom service: - Create a factory function that takes a URI and returns an instance of your custom service. This function will parse any details it needs from the URI. - Register your factory with the global service registry. You need to define a unique URI scheme for your service (e.g., custom). PiperOrigin-RevId: 822310466
This commit is contained in:
Shangjie Chen
committed by
Copybara-Service
Copybara-Service
parent
409df1378f
commit
391628fcdc
@@ -0,0 +1,169 @@
|
||||
# Copyright 2025 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 unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_services():
|
||||
"""Mock all service implementation classes to avoid real instantiation."""
|
||||
with (
|
||||
patch(
|
||||
"google.adk.sessions.vertex_ai_session_service.VertexAiSessionService"
|
||||
) as mock_vertex_session,
|
||||
patch(
|
||||
"google.adk.sessions.database_session_service.DatabaseSessionService"
|
||||
) as mock_db_session,
|
||||
patch(
|
||||
"google.adk.artifacts.gcs_artifact_service.GcsArtifactService"
|
||||
) as mock_gcs_artifact,
|
||||
patch(
|
||||
"google.adk.memory.vertex_ai_rag_memory_service.VertexAiRagMemoryService"
|
||||
) as mock_rag_memory,
|
||||
patch(
|
||||
"google.adk.memory.vertex_ai_memory_bank_service.VertexAiMemoryBankService"
|
||||
) as mock_agentengine_memory,
|
||||
):
|
||||
yield {
|
||||
"vertex_session": mock_vertex_session,
|
||||
"db_session": mock_db_session,
|
||||
"gcs_artifact": mock_gcs_artifact,
|
||||
"rag_memory": mock_rag_memory,
|
||||
"agentengine_memory": mock_agentengine_memory,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registry():
|
||||
from google.adk.cli.service_registry import get_service_registry
|
||||
|
||||
return get_service_registry()
|
||||
|
||||
|
||||
# Session Service Tests
|
||||
def test_create_session_service_sqlite(registry, mock_services):
|
||||
registry.create_session_service("sqlite:///test.db")
|
||||
mock_services["db_session"].assert_called_once_with(
|
||||
db_url="sqlite:///test.db"
|
||||
)
|
||||
|
||||
|
||||
def test_create_session_service_sqlite_with_kwargs(registry, mock_services):
|
||||
registry.create_session_service(
|
||||
"sqlite:///test.db", pool_size=10, agents_dir="foo"
|
||||
)
|
||||
mock_services["db_session"].assert_called_once_with(
|
||||
db_url="sqlite:///test.db", pool_size=10
|
||||
)
|
||||
|
||||
|
||||
def test_create_session_service_postgresql(registry, mock_services):
|
||||
registry.create_session_service("postgresql://user:pass@host/db")
|
||||
mock_services["db_session"].assert_called_once_with(
|
||||
db_url="postgresql://user:pass@host/db"
|
||||
)
|
||||
|
||||
|
||||
@patch("google.adk.cli.utils.envs.load_dotenv_for_agent")
|
||||
def test_create_session_service_agentengine_short(
|
||||
mock_load_dotenv, registry, mock_services, monkeypatch
|
||||
):
|
||||
monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project")
|
||||
monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1")
|
||||
registry.create_session_service(
|
||||
"agentengine://123", agents_dir="/path/to/agents"
|
||||
)
|
||||
mock_services["vertex_session"].assert_called_once_with(
|
||||
project="test-project", location="us-central1", agent_engine_id="123"
|
||||
)
|
||||
mock_load_dotenv.assert_called_once_with("", "/path/to/agents")
|
||||
|
||||
|
||||
def test_create_session_service_agentengine_full(registry, mock_services):
|
||||
uri = "agentengine://projects/p/locations/l/reasoningEngines/123"
|
||||
registry.create_session_service(uri, agents_dir="/path/to/agents")
|
||||
mock_services["vertex_session"].assert_called_once_with(
|
||||
project="p", location="l", agent_engine_id="123"
|
||||
)
|
||||
|
||||
|
||||
# Artifact Service Tests
|
||||
def test_create_artifact_service_gcs(registry, mock_services):
|
||||
registry.create_artifact_service(
|
||||
"gs://my-bucket/path/prefix", agents_dir="foo", other_kwarg="bar"
|
||||
)
|
||||
mock_services["gcs_artifact"].assert_called_once_with(
|
||||
bucket_name="my-bucket", other_kwarg="bar"
|
||||
)
|
||||
|
||||
|
||||
# Memory Service Tests
|
||||
@patch("google.adk.cli.utils.envs.load_dotenv_for_agent")
|
||||
def test_create_memory_service_rag(
|
||||
mock_load_dotenv, registry, mock_services, monkeypatch
|
||||
):
|
||||
monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project")
|
||||
monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1")
|
||||
registry.create_memory_service(
|
||||
"rag://corpus-123", agents_dir="/path/to/agents"
|
||||
)
|
||||
mock_services["rag_memory"].assert_called_once_with(
|
||||
rag_corpus=(
|
||||
"projects/test-project/locations/us-central1/ragCorpora/corpus-123"
|
||||
)
|
||||
)
|
||||
mock_load_dotenv.assert_called_once_with("", "/path/to/agents")
|
||||
|
||||
|
||||
@patch("google.adk.cli.utils.envs.load_dotenv_for_agent")
|
||||
def test_create_memory_service_agentengine_short(
|
||||
mock_load_dotenv, registry, mock_services, monkeypatch
|
||||
):
|
||||
monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project")
|
||||
monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1")
|
||||
registry.create_memory_service(
|
||||
"agentengine://456", agents_dir="/path/to/agents"
|
||||
)
|
||||
mock_services["agentengine_memory"].assert_called_once_with(
|
||||
project="test-project", location="us-central1", agent_engine_id="456"
|
||||
)
|
||||
mock_load_dotenv.assert_called_once_with("", "/path/to/agents")
|
||||
|
||||
|
||||
def test_create_memory_service_agentengine_full(registry, mock_services):
|
||||
uri = "agentengine://projects/p/locations/l/reasoningEngines/456"
|
||||
registry.create_memory_service(uri, agents_dir="/path/to/agents")
|
||||
mock_services["agentengine_memory"].assert_called_once_with(
|
||||
project="p", location="l", agent_engine_id="456"
|
||||
)
|
||||
|
||||
|
||||
# General Tests
|
||||
def test_unsupported_scheme(registry, mock_services):
|
||||
session_service = registry.create_session_service("unsupported://foo")
|
||||
artifact_service = registry.create_artifact_service("unsupported://foo")
|
||||
memory_service = registry.create_memory_service("unsupported://foo")
|
||||
assert session_service is None
|
||||
assert artifact_service is None
|
||||
assert memory_service is None
|
||||
for service in [
|
||||
"vertex_session",
|
||||
"db_session",
|
||||
"gcs_artifact",
|
||||
"rag_memory",
|
||||
"agentengine_memory",
|
||||
]:
|
||||
mock_services[service].assert_not_called()
|
||||
Reference in New Issue
Block a user