mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Change service creation and add app name mapping for sessions
This change refactors how session, memory, and artifact services are created in the fast_api server, using the shared service_factory. Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 839997110
This commit is contained in:
committed by
Copybara-Service
parent
e9182e5eb4
commit
0a07a667e9
@@ -159,11 +159,22 @@ async def run_cli(
|
||||
load_services_module(str(agent_root))
|
||||
user_id = 'test_user'
|
||||
|
||||
agents_dir = str(agent_parent_path)
|
||||
agent_loader = AgentLoader(agents_dir=agents_dir)
|
||||
agent_or_app = agent_loader.load_agent(agent_folder_name)
|
||||
session_app_name = (
|
||||
agent_or_app.name if isinstance(agent_or_app, App) else agent_folder_name
|
||||
)
|
||||
app_name_to_dir = None
|
||||
if isinstance(agent_or_app, App) and agent_or_app.name != agent_folder_name:
|
||||
app_name_to_dir = {agent_or_app.name: agent_folder_name}
|
||||
|
||||
# Create session and artifact services using factory functions
|
||||
# Sessions persist under <agents_dir>/<agent>/.adk/session.db by default.
|
||||
session_service = create_session_service_from_options(
|
||||
base_dir=agent_parent_path,
|
||||
session_service_uri=session_service_uri,
|
||||
app_name_to_dir=app_name_to_dir,
|
||||
)
|
||||
|
||||
artifact_service = create_artifact_service_from_options(
|
||||
@@ -172,13 +183,6 @@ async def run_cli(
|
||||
)
|
||||
|
||||
credential_service = InMemoryCredentialService()
|
||||
agents_dir = str(agent_parent_path)
|
||||
agent_or_app = AgentLoader(agents_dir=agents_dir).load_agent(
|
||||
agent_folder_name
|
||||
)
|
||||
session_app_name = (
|
||||
agent_or_app.name if isinstance(agent_or_app, App) else agent_folder_name
|
||||
)
|
||||
if not is_env_enabled('ADK_DISABLE_LOAD_DOTENV'):
|
||||
envs.load_dotenv_for_agent(agent_folder_name, agents_dir)
|
||||
|
||||
|
||||
@@ -35,28 +35,24 @@ from opentelemetry.sdk.trace import TracerProvider
|
||||
from starlette.types import Lifespan
|
||||
from watchdog.observers import Observer
|
||||
|
||||
from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||||
from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
|
||||
from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
|
||||
from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
|
||||
from ..memory.in_memory_memory_service import InMemoryMemoryService
|
||||
from ..runners import Runner
|
||||
from ..sessions.in_memory_session_service import InMemorySessionService
|
||||
from .adk_web_server import AdkWebServer
|
||||
from .service_registry import get_service_registry
|
||||
from .service_registry import load_services_module
|
||||
from .utils import envs
|
||||
from .utils import evals
|
||||
from .utils.agent_change_handler import AgentChangeEventHandler
|
||||
from .utils.agent_loader import AgentLoader
|
||||
from .utils.service_factory import create_artifact_service_from_options
|
||||
from .utils.service_factory import create_memory_service_from_options
|
||||
from .utils.service_factory import create_session_service_from_options
|
||||
|
||||
logger = logging.getLogger("google_adk." + __name__)
|
||||
|
||||
_LAZY_SERVICE_IMPORTS: dict[str, str] = {
|
||||
"AgentLoader": ".utils.agent_loader",
|
||||
"InMemoryArtifactService": "..artifacts.in_memory_artifact_service",
|
||||
"InMemoryMemoryService": "..memory.in_memory_memory_service",
|
||||
"InMemorySessionService": "..sessions.in_memory_session_service",
|
||||
"LocalEvalSetResultsManager": "..evaluation.local_eval_set_results_manager",
|
||||
"LocalEvalSetsManager": "..evaluation.local_eval_sets_manager",
|
||||
}
|
||||
@@ -112,48 +108,31 @@ def get_fast_api_app(
|
||||
# Load services.py from agents_dir for custom service registration.
|
||||
load_services_module(agents_dir)
|
||||
|
||||
service_registry = get_service_registry()
|
||||
|
||||
# Build the Memory service
|
||||
if memory_service_uri:
|
||||
memory_service = service_registry.create_memory_service(
|
||||
memory_service_uri, agents_dir=agents_dir
|
||||
try:
|
||||
memory_service = create_memory_service_from_options(
|
||||
base_dir=agents_dir,
|
||||
memory_service_uri=memory_service_uri,
|
||||
)
|
||||
if not memory_service:
|
||||
raise click.ClickException(
|
||||
"Unsupported memory service URI: %s" % memory_service_uri
|
||||
)
|
||||
else:
|
||||
memory_service = InMemoryMemoryService()
|
||||
except ValueError as exc:
|
||||
raise click.ClickException(str(exc)) from exc
|
||||
|
||||
# Build the Session service
|
||||
if session_service_uri:
|
||||
session_kwargs = session_db_kwargs or {}
|
||||
session_service = service_registry.create_session_service(
|
||||
session_service_uri, agents_dir=agents_dir, **session_kwargs
|
||||
session_service = create_session_service_from_options(
|
||||
base_dir=agents_dir,
|
||||
session_service_uri=session_service_uri,
|
||||
session_db_kwargs=session_db_kwargs,
|
||||
)
|
||||
if not session_service:
|
||||
# Fallback to DatabaseSessionService if the service registry doesn't
|
||||
# support the session service URI scheme.
|
||||
from ..sessions.database_session_service import DatabaseSessionService
|
||||
|
||||
session_service = DatabaseSessionService(
|
||||
db_url=session_service_uri, **session_kwargs
|
||||
)
|
||||
else:
|
||||
session_service = InMemorySessionService()
|
||||
|
||||
# Build the Artifact service
|
||||
if artifact_service_uri:
|
||||
artifact_service = service_registry.create_artifact_service(
|
||||
artifact_service_uri, agents_dir=agents_dir
|
||||
try:
|
||||
artifact_service = create_artifact_service_from_options(
|
||||
base_dir=agents_dir,
|
||||
artifact_service_uri=artifact_service_uri,
|
||||
strict_uri=True,
|
||||
)
|
||||
if not artifact_service:
|
||||
raise click.ClickException(
|
||||
"Unsupported artifact service URI: %s" % artifact_service_uri
|
||||
)
|
||||
else:
|
||||
artifact_service = InMemoryArtifactService()
|
||||
except ValueError as exc:
|
||||
raise click.ClickException(str(exc)) from exc
|
||||
|
||||
# Build the Credential service
|
||||
credential_service = InMemoryCredentialService()
|
||||
|
||||
@@ -17,6 +17,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Mapping
|
||||
from typing import Optional
|
||||
|
||||
from typing_extensions import override
|
||||
@@ -61,6 +62,7 @@ def create_local_session_service(
|
||||
*,
|
||||
base_dir: Path | str,
|
||||
per_agent: bool = False,
|
||||
app_name_to_dir: Optional[Mapping[str, str]] = None,
|
||||
) -> BaseSessionService:
|
||||
"""Creates a local SQLite-backed session service.
|
||||
|
||||
@@ -69,6 +71,8 @@ def create_local_session_service(
|
||||
per_agent: If True, creates a PerAgentDatabaseSessionService that stores
|
||||
sessions in each agent's .adk folder. If False, creates a single
|
||||
SqliteSessionService at base_dir/.adk/session.db.
|
||||
app_name_to_dir: Optional mapping from logical app name to on-disk agent
|
||||
folder name. Only used when per_agent is True; defaults to identity.
|
||||
|
||||
Returns:
|
||||
A BaseSessionService instance backed by SQLite.
|
||||
@@ -78,7 +82,10 @@ def create_local_session_service(
|
||||
"Using per-agent session storage rooted at %s",
|
||||
base_dir,
|
||||
)
|
||||
return PerAgentDatabaseSessionService(agents_root=base_dir)
|
||||
return PerAgentDatabaseSessionService(
|
||||
agents_root=base_dir,
|
||||
app_name_to_dir=app_name_to_dir,
|
||||
)
|
||||
|
||||
return create_local_database_session_service(base_dir=base_dir)
|
||||
|
||||
@@ -108,23 +115,26 @@ class PerAgentDatabaseSessionService(BaseSessionService):
|
||||
self,
|
||||
*,
|
||||
agents_root: Path | str,
|
||||
app_name_to_dir: Optional[Mapping[str, str]] = None,
|
||||
):
|
||||
self._agents_root = Path(agents_root).resolve()
|
||||
self._app_name_to_dir = dict(app_name_to_dir or {})
|
||||
self._services: dict[str, BaseSessionService] = {}
|
||||
self._service_lock = asyncio.Lock()
|
||||
|
||||
async def _get_service(self, app_name: str) -> BaseSessionService:
|
||||
async with self._service_lock:
|
||||
service = self._services.get(app_name)
|
||||
storage_name = self._app_name_to_dir.get(app_name, app_name)
|
||||
service = self._services.get(storage_name)
|
||||
if service is not None:
|
||||
return service
|
||||
folder = dot_adk_folder_for_agent(
|
||||
agents_root=self._agents_root, app_name=app_name
|
||||
agents_root=self._agents_root, app_name=storage_name
|
||||
)
|
||||
service = create_local_database_session_service(
|
||||
base_dir=folder.agent_dir,
|
||||
)
|
||||
self._services[app_name] = service
|
||||
self._services[storage_name] = service
|
||||
return service
|
||||
|
||||
@override
|
||||
|
||||
@@ -33,6 +33,7 @@ def create_session_service_from_options(
|
||||
base_dir: Path | str,
|
||||
session_service_uri: Optional[str] = None,
|
||||
session_db_kwargs: Optional[dict[str, Any]] = None,
|
||||
app_name_to_dir: Optional[dict[str, str]] = None,
|
||||
) -> BaseSessionService:
|
||||
"""Creates a session service based on CLI/web options."""
|
||||
base_path = Path(base_dir)
|
||||
@@ -64,7 +65,11 @@ def create_session_service_from_options(
|
||||
return DatabaseSessionService(db_url=session_service_uri, **fallback_kwargs)
|
||||
|
||||
# Default to per-agent local SQLite storage in <agents_root>/<agent>/.adk/.
|
||||
return create_local_session_service(base_dir=base_path, per_agent=True)
|
||||
return create_local_session_service(
|
||||
base_dir=base_path,
|
||||
per_agent=True,
|
||||
app_name_to_dir=app_name_to_dir,
|
||||
)
|
||||
|
||||
|
||||
def create_memory_service_from_options(
|
||||
@@ -96,6 +101,7 @@ def create_artifact_service_from_options(
|
||||
*,
|
||||
base_dir: Path | str,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
strict_uri: bool = False,
|
||||
) -> BaseArtifactService:
|
||||
"""Creates an artifact service based on CLI/web options."""
|
||||
base_path = Path(base_dir)
|
||||
@@ -108,6 +114,10 @@ def create_artifact_service_from_options(
|
||||
agents_dir=str(base_path),
|
||||
)
|
||||
if service is None:
|
||||
if strict_uri:
|
||||
raise ValueError(
|
||||
f"Unsupported artifact service URI: {artifact_service_uri}"
|
||||
)
|
||||
logger.warning(
|
||||
"Unsupported artifact service URI: %s, falling back to in-memory",
|
||||
artifact_service_uri,
|
||||
|
||||
@@ -416,15 +416,15 @@ def test_app(
|
||||
with (
|
||||
patch("signal.signal", return_value=None),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.InMemorySessionService",
|
||||
"google.adk.cli.fast_api.create_session_service_from_options",
|
||||
return_value=mock_session_service,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.InMemoryArtifactService",
|
||||
"google.adk.cli.fast_api.create_artifact_service_from_options",
|
||||
return_value=mock_artifact_service,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.InMemoryMemoryService",
|
||||
"google.adk.cli.fast_api.create_memory_service_from_options",
|
||||
return_value=mock_memory_service,
|
||||
),
|
||||
patch(
|
||||
@@ -556,15 +556,15 @@ def test_app_with_a2a(
|
||||
with (
|
||||
patch("signal.signal", return_value=None),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.InMemorySessionService",
|
||||
"google.adk.cli.fast_api.create_session_service_from_options",
|
||||
return_value=mock_session_service,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.InMemoryArtifactService",
|
||||
"google.adk.cli.fast_api.create_artifact_service_from_options",
|
||||
return_value=mock_artifact_service,
|
||||
),
|
||||
patch(
|
||||
"google.adk.cli.fast_api.InMemoryMemoryService",
|
||||
"google.adk.cli.fast_api.create_memory_service_from_options",
|
||||
return_value=mock_memory_service,
|
||||
),
|
||||
patch(
|
||||
|
||||
@@ -17,6 +17,7 @@ from __future__ import annotations
|
||||
from pathlib import Path
|
||||
|
||||
from google.adk.cli.utils.local_storage import create_local_database_session_service
|
||||
from google.adk.cli.utils.local_storage import create_local_session_service
|
||||
from google.adk.cli.utils.local_storage import PerAgentDatabaseSessionService
|
||||
from google.adk.sessions.sqlite_session_service import SqliteSessionService
|
||||
import pytest
|
||||
@@ -48,6 +49,29 @@ async def test_per_agent_session_service_creates_scoped_dot_adk(
|
||||
assert agent_b_sessions.sessions[0].app_name == "agent_b"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_agent_session_service_respects_app_name_alias(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
folder_name = "agent_folder"
|
||||
logical_name = "custom_app"
|
||||
(tmp_path / folder_name).mkdir()
|
||||
|
||||
service = create_local_session_service(
|
||||
base_dir=tmp_path,
|
||||
per_agent=True,
|
||||
app_name_to_dir={logical_name: folder_name},
|
||||
)
|
||||
|
||||
session = await service.create_session(
|
||||
app_name=logical_name,
|
||||
user_id="user",
|
||||
)
|
||||
|
||||
assert session.app_name == logical_name
|
||||
assert (tmp_path / folder_name / ".adk" / "session.db").exists()
|
||||
|
||||
|
||||
def test_create_local_database_session_service_returns_sqlite(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
|
||||
@@ -60,6 +60,25 @@ async def test_create_session_service_defaults_to_per_agent_sqlite(
|
||||
assert (agent_dir / ".adk" / "session.db").exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_session_service_respects_app_name_mapping(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
agent_dir = tmp_path / "agent_folder"
|
||||
logical_name = "custom_app"
|
||||
agent_dir.mkdir()
|
||||
|
||||
service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
app_name_to_dir={logical_name: "agent_folder"},
|
||||
)
|
||||
|
||||
assert isinstance(service, PerAgentDatabaseSessionService)
|
||||
session = await service.create_session(app_name=logical_name, user_id="user")
|
||||
assert session.app_name == logical_name
|
||||
assert (agent_dir / ".adk" / "session.db").exists()
|
||||
|
||||
|
||||
def test_create_session_service_fallbacks_to_database(
|
||||
tmp_path: Path, monkeypatch
|
||||
):
|
||||
@@ -101,6 +120,21 @@ def test_create_artifact_service_uses_registry(tmp_path: Path, monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def test_create_artifact_service_raises_on_unknown_scheme_when_strict(
|
||||
tmp_path: Path, monkeypatch
|
||||
):
|
||||
registry = Mock()
|
||||
registry.create_artifact_service.return_value = None
|
||||
monkeypatch.setattr(service_factory, "get_service_registry", lambda: registry)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
service_factory.create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
artifact_service_uri="unknown://foo",
|
||||
strict_uri=True,
|
||||
)
|
||||
|
||||
|
||||
def test_create_memory_service_uses_registry(tmp_path: Path, monkeypatch):
|
||||
registry = Mock()
|
||||
expected = object()
|
||||
|
||||
Reference in New Issue
Block a user