diff --git a/src/google/adk/cli/cli.py b/src/google/adk/cli/cli.py index a1b63a4c..941f1c28 100644 --- a/src/google/adk/cli/cli.py +++ b/src/google/adk/cli/cli.py @@ -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 //.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) diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index df06b1cf..131213ec 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -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 - ) - 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() + session_service = create_session_service_from_options( + base_dir=agents_dir, + session_service_uri=session_service_uri, + session_db_kwargs=session_db_kwargs, + ) # 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() diff --git a/src/google/adk/cli/utils/local_storage.py b/src/google/adk/cli/utils/local_storage.py index ec7099b8..6fb6a83e 100644 --- a/src/google/adk/cli/utils/local_storage.py +++ b/src/google/adk/cli/utils/local_storage.py @@ -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 diff --git a/src/google/adk/cli/utils/service_factory.py b/src/google/adk/cli/utils/service_factory.py index 60f4ddd3..840c5b1c 100644 --- a/src/google/adk/cli/utils/service_factory.py +++ b/src/google/adk/cli/utils/service_factory.py @@ -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 //.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, diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index 75d56790..263e4704 100755 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -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( diff --git a/tests/unittests/cli/utils/test_local_storage.py b/tests/unittests/cli/utils/test_local_storage.py index 324ce5e0..39bce7a5 100644 --- a/tests/unittests/cli/utils/test_local_storage.py +++ b/tests/unittests/cli/utils/test_local_storage.py @@ -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: diff --git a/tests/unittests/cli/utils/test_service_factory.py b/tests/unittests/cli/utils/test_service_factory.py index 9d9afdd2..a8eb0fdb 100644 --- a/tests/unittests/cli/utils/test_service_factory.py +++ b/tests/unittests/cli/utils/test_service_factory.py @@ -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()