Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions src/google/adk/cli/fast_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/google/adk/cli/utils/local_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
4 changes: 3 additions & 1 deletion src/google/adk/cli/utils/service_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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."""
Expand Down
67 changes: 67 additions & 0 deletions tests/unittests/cli/test_fast_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Loading