diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index 3e437a7772..35a27e55cd 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -25,6 +25,7 @@ from typing import AsyncIterator from typing import Awaitable from typing import Callable +from typing import Iterator from typing import Literal from typing import Mapping @@ -44,6 +45,7 @@ from starlette.types import Lifespan from watchdog.observers import Observer +from ..apps.app import App from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService from ..runners import Runner from ..telemetry._agent_engine import get_propagated_context @@ -66,6 +68,28 @@ ) +class AgentAppMapping(Mapping[str, str]): + + def __init__(self, agent_loader: Any) -> None: + self.agent_loader = agent_loader + + def __getitem__(self, app_name: str) -> str: + for folder_name in self.agent_loader.list_agents(): + try: + loaded = self.agent_loader.load_agent(folder_name) + if isinstance(loaded, App) and loaded.name == app_name: + return str(folder_name) + except Exception: + pass + return app_name + + def __iter__(self) -> Iterator[str]: + return iter([]) + + def __len__(self) -> int: + return 1 + + class _QueryRequest(BaseModel): input: dict[str, Any] | None = None class_method: str | None = None @@ -240,12 +264,14 @@ def get_fast_api_app( except ValueError as exc: raise click.ClickException(str(exc)) from exc + app_name_to_dir = AgentAppMapping(agent_loader) # Build the Session service session_service = create_session_service_from_options( base_dir=agents_dir, session_service_uri=session_service_uri, session_db_kwargs=session_db_kwargs, use_local_storage=use_local_storage, + app_name_to_dir=app_name_to_dir, ) # Build the Artifact service diff --git a/src/google/adk/cli/utils/local_storage.py b/src/google/adk/cli/utils/local_storage.py index 7de62a77dc..4b3d1367b0 100644 --- a/src/google/adk/cli/utils/local_storage.py +++ b/src/google/adk/cli/utils/local_storage.py @@ -141,7 +141,7 @@ def __init__( 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._app_name_to_dir = app_name_to_dir or {} self._services: dict[str, BaseSessionService] = {} self._service_lock = asyncio.Lock() diff --git a/src/google/adk/cli/utils/service_factory.py b/src/google/adk/cli/utils/service_factory.py index 8b5e7610ad..ca8294fae9 100644 --- a/src/google/adk/cli/utils/service_factory.py +++ b/src/google/adk/cli/utils/service_factory.py @@ -18,6 +18,8 @@ import os from pathlib import Path from typing import Any +from typing import Mapping +from typing import Optional from urllib.parse import parse_qsl from urllib.parse import urlsplit from urllib.parse import urlunsplit @@ -172,7 +174,7 @@ def create_session_service_from_options( base_dir: Path | str, session_service_uri: str | None = None, session_db_kwargs: dict[str, Any] | None = None, - app_name_to_dir: dict[str, str] | None = None, + app_name_to_dir: Optional[Mapping[str, str]] = None, use_local_storage: bool = True, ) -> BaseSessionService: """Creates a session service based on CLI/web options.""" diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index ec0fdedceb..04c671e447 100644 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -31,8 +31,10 @@ from google.adk.agents.base_agent import BaseAgent from google.adk.agents.llm_agent import LlmAgent from google.adk.agents.run_config import RunConfig +from google.adk.apps.app import App from google.adk.artifacts.base_artifact_service import ArtifactVersion from google.adk.cli import fast_api as fast_api_module +from google.adk.cli.fast_api import AgentAppMapping from google.adk.cli.fast_api import get_fast_api_app from google.adk.errors.input_validation_error import InputValidationError from google.adk.errors.session_not_found_error import SessionNotFoundError @@ -4452,3 +4454,68 @@ def test_create_eval_set_legacy_route_creates_eval_set( if __name__ == "__main__": pytest.main(["-xvs", __file__]) + + +def test_dynamic_app_map_finds_matching_app(): + mock_agent_loader = MagicMock() + mock_agent_loader.list_agents.return_value = ["folder_a", "folder_b"] + + def mock_load_agent(folder): + mock_app = MagicMock(spec=App) + if folder == "folder_b": + mock_app.name = "target_app" + else: + mock_app.name = "other_app" + return mock_app + + mock_agent_loader.load_agent.side_effect = mock_load_agent + + app_map = AgentAppMapping(mock_agent_loader) + + # Should correctly map the internal app name to its physical folder + assert app_map["target_app"] == "folder_b" + + +def test_dynamic_app_map_fallback_when_not_found(): + mock_agent_loader = MagicMock() + mock_agent_loader.list_agents.return_value = ["folder_a"] + + def mock_load_agent(folder): + mock_app = MagicMock(spec=App) + mock_app.name = "other_app" + return mock_app + + mock_agent_loader.load_agent.side_effect = mock_load_agent + + app_map = AgentAppMapping(mock_agent_loader) + + # If the app isn't found, it should return the exact string it was given + assert app_map["missing_app"] == "missing_app" + + +def test_dynamic_app_map_ignores_load_errors(): + mock_agent_loader = MagicMock() + mock_agent_loader.list_agents.return_value = ["broken_folder", "good_folder"] + + def mock_load_agent(folder): + if folder == "broken_folder": + raise ValueError("Agent is completely broken") + mock_app = MagicMock(spec=App) + mock_app.name = "good_app" + return mock_app + + mock_agent_loader.load_agent.side_effect = mock_load_agent + + app_map = AgentAppMapping(mock_agent_loader) + + # Should ignore the ValueError from broken_folder and still find good_app + assert app_map["good_app"] == "good_folder" + + +def test_dynamic_app_map_is_truthy(): + mock_agent_loader = MagicMock() + app_map = AgentAppMapping(mock_agent_loader) + + # Must evaluate to True so that `app_map or {}` evaluates to app_map + assert bool(app_map) is True + assert (app_map or {}) is app_map