feat: wire runtime entrypoints to service factory defaults

This change routes adk run and the FastAPI server through the new session/artifact service factory, keeps the default experience backed by per-agent .adk storage

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 836733234
This commit is contained in:
George Weale
2025-11-25 10:47:44 -08:00
committed by Copybara-Service
parent 5453b5bfde
commit 06e6fc9132
6 changed files with 509 additions and 78 deletions
+57 -22
View File
@@ -15,6 +15,7 @@
from __future__ import annotations from __future__ import annotations
from datetime import datetime from datetime import datetime
from pathlib import Path
from typing import Optional from typing import Optional
from typing import Union from typing import Union
@@ -22,7 +23,6 @@ import click
from google.genai import types from google.genai import types
from pydantic import BaseModel from pydantic import BaseModel
from ..agents.base_agent import BaseAgent
from ..agents.llm_agent import LlmAgent from ..agents.llm_agent import LlmAgent
from ..apps.app import App from ..apps.app import App
from ..artifacts.base_artifact_service import BaseArtifactService from ..artifacts.base_artifact_service import BaseArtifactService
@@ -35,8 +35,11 @@ from ..sessions.in_memory_session_service import InMemorySessionService
from ..sessions.session import Session from ..sessions.session import Session
from ..utils.context_utils import Aclosing from ..utils.context_utils import Aclosing
from ..utils.env_utils import is_env_enabled from ..utils.env_utils import is_env_enabled
from .service_registry import load_services_module
from .utils import envs from .utils import envs
from .utils.agent_loader import AgentLoader from .utils.agent_loader import AgentLoader
from .utils.service_factory import create_artifact_service_from_options
from .utils.service_factory import create_session_service_from_options
class InputFile(BaseModel): class InputFile(BaseModel):
@@ -66,7 +69,7 @@ async def run_input_file(
) )
with open(input_path, 'r', encoding='utf-8') as f: with open(input_path, 'r', encoding='utf-8') as f:
input_file = InputFile.model_validate_json(f.read()) input_file = InputFile.model_validate_json(f.read())
input_file.state['_time'] = datetime.now() input_file.state['_time'] = datetime.now().isoformat()
session = await session_service.create_session( session = await session_service.create_session(
app_name=app_name, user_id=user_id, state=input_file.state app_name=app_name, user_id=user_id, state=input_file.state
@@ -134,6 +137,8 @@ async def run_cli(
saved_session_file: Optional[str] = None, saved_session_file: Optional[str] = None,
save_session: bool, save_session: bool,
session_id: Optional[str] = None, session_id: Optional[str] = None,
session_service_uri: Optional[str] = None,
artifact_service_uri: Optional[str] = None,
) -> None: ) -> None:
"""Runs an interactive CLI for a certain agent. """Runs an interactive CLI for a certain agent.
@@ -148,24 +153,47 @@ async def run_cli(
contains a previously saved session, exclusive with input_file. contains a previously saved session, exclusive with input_file.
save_session: bool, whether to save the session on exit. save_session: bool, whether to save the session on exit.
session_id: Optional[str], the session ID to save the session to on exit. session_id: Optional[str], the session ID to save the session to on exit.
session_service_uri: Optional[str], custom session service URI.
artifact_service_uri: Optional[str], custom artifact service URI.
""" """
agent_parent_path = Path(agent_parent_dir).resolve()
artifact_service = InMemoryArtifactService() agent_root = agent_parent_path / agent_folder_name
session_service = InMemorySessionService() load_services_module(str(agent_root))
credential_service = InMemoryCredentialService()
user_id = 'test_user' user_id = 'test_user'
agent_or_app = AgentLoader(agents_dir=agent_parent_dir).load_agent(
# Create session and artifact services using factory functions
session_service = create_session_service_from_options(
base_dir=agent_root,
session_service_uri=session_service_uri,
)
artifact_service = create_artifact_service_from_options(
base_dir=agent_root,
artifact_service_uri=artifact_service_uri,
)
credential_service = InMemoryCredentialService()
agents_dir = str(agent_parent_path)
agent_or_app = AgentLoader(agents_dir=agents_dir).load_agent(
agent_folder_name agent_folder_name
) )
session_app_name = ( session_app_name = (
agent_or_app.name if isinstance(agent_or_app, App) else agent_folder_name agent_or_app.name if isinstance(agent_or_app, App) else agent_folder_name
) )
session = await session_service.create_session(
app_name=session_app_name, user_id=user_id
)
if not is_env_enabled('ADK_DISABLE_LOAD_DOTENV'): if not is_env_enabled('ADK_DISABLE_LOAD_DOTENV'):
envs.load_dotenv_for_agent(agent_folder_name, agent_parent_dir) envs.load_dotenv_for_agent(agent_folder_name, agents_dir)
# Helper function for printing events
def _print_event(event) -> None:
content = event.content
if not content or not content.parts:
return
text_parts = [part.text for part in content.parts if part.text]
if not text_parts:
return
author = event.author or 'system'
click.echo(f'[{author}]: {"".join(text_parts)}')
if input_file: if input_file:
session = await run_input_file( session = await run_input_file(
app_name=session_app_name, app_name=session_app_name,
@@ -177,16 +205,22 @@ async def run_cli(
input_path=input_file, input_path=input_file,
) )
elif saved_session_file: elif saved_session_file:
# Load the saved session from file
with open(saved_session_file, 'r', encoding='utf-8') as f: with open(saved_session_file, 'r', encoding='utf-8') as f:
loaded_session = Session.model_validate_json(f.read()) loaded_session = Session.model_validate_json(f.read())
# Create a new session in the service, copying state from the file
session = await session_service.create_session(
app_name=session_app_name,
user_id=user_id,
state=loaded_session.state if loaded_session else None,
)
# Append events from the file to the new session and display them
if loaded_session: if loaded_session:
for event in loaded_session.events: for event in loaded_session.events:
await session_service.append_event(session, event) await session_service.append_event(session, event)
content = event.content _print_event(event)
if not content or not content.parts or not content.parts[0].text:
continue
click.echo(f'[{event.author}]: {content.parts[0].text}')
await run_interactively( await run_interactively(
agent_or_app, agent_or_app,
@@ -196,6 +230,9 @@ async def run_cli(
credential_service, credential_service,
) )
else: else:
session = await session_service.create_session(
app_name=session_app_name, user_id=user_id
)
click.echo(f'Running agent {agent_or_app.name}, type exit to exit.') click.echo(f'Running agent {agent_or_app.name}, type exit to exit.')
await run_interactively( await run_interactively(
agent_or_app, agent_or_app,
@@ -207,9 +244,7 @@ async def run_cli(
if save_session: if save_session:
session_id = session_id or input('Session ID to save: ') session_id = session_id or input('Session ID to save: ')
session_path = ( session_path = agent_root / f'{session_id}.session.json'
f'{agent_parent_dir}/{agent_folder_name}/{session_id}.session.json'
)
# Fetch the session again to get all the details. # Fetch the session again to get all the details.
session = await session_service.get_session( session = await session_service.get_session(
@@ -217,9 +252,9 @@ async def run_cli(
user_id=session.user_id, user_id=session.user_id,
session_id=session.id, session_id=session.id,
) )
with open(session_path, 'w', encoding='utf-8') as f: session_path.write_text(
f.write( session.model_dump_json(indent=2, exclude_none=True, by_alias=True),
session.model_dump_json(indent=2, exclude_none=True, by_alias=True) encoding='utf-8',
) )
print('Session saved to', session_path) print('Session saved to', session_path)
+23 -38
View File
@@ -34,20 +34,19 @@ from opentelemetry.sdk.trace import TracerProvider
from starlette.types import Lifespan from starlette.types import Lifespan
from watchdog.observers import Observer from watchdog.observers import Observer
from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
from ..memory.in_memory_memory_service import InMemoryMemoryService
from ..runners import Runner from ..runners import Runner
from ..sessions.in_memory_session_service import InMemorySessionService
from .adk_web_server import AdkWebServer from .adk_web_server import AdkWebServer
from .service_registry import get_service_registry
from .service_registry import load_services_module from .service_registry import load_services_module
from .utils import envs from .utils import envs
from .utils import evals from .utils import evals
from .utils.agent_change_handler import AgentChangeEventHandler from .utils.agent_change_handler import AgentChangeEventHandler
from .utils.agent_loader import AgentLoader 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__) logger = logging.getLogger("google_adk." + __name__)
@@ -74,6 +73,8 @@ def get_fast_api_app(
logo_text: Optional[str] = None, logo_text: Optional[str] = None,
logo_image_url: Optional[str] = None, logo_image_url: Optional[str] = None,
) -> FastAPI: ) -> FastAPI:
# Convert to absolute path for consistency
agents_dir = str(Path(agents_dir).resolve())
# Set up eval managers. # Set up eval managers.
if eval_storage_uri: if eval_storage_uri:
@@ -91,48 +92,32 @@ def get_fast_api_app(
# Load services.py from agents_dir for custom service registration. # Load services.py from agents_dir for custom service registration.
load_services_module(agents_dir) load_services_module(agents_dir)
service_registry = get_service_registry()
# Build the Memory service # Build the Memory service
if memory_service_uri: try:
memory_service = service_registry.create_memory_service( memory_service = create_memory_service_from_options(
memory_service_uri, agents_dir=agents_dir base_dir=agents_dir,
memory_service_uri=memory_service_uri,
) )
if not memory_service: except ValueError as exc:
raise click.ClickException( raise click.ClickException(str(exc)) from exc
"Unsupported memory service URI: %s" % memory_service_uri
)
else:
memory_service = InMemoryMemoryService()
# Build the Session service # Build the Session service
if session_service_uri: session_service = create_session_service_from_options(
session_kwargs = session_db_kwargs or {} base_dir=agents_dir,
session_service = service_registry.create_session_service( session_service_uri=session_service_uri,
session_service_uri, agents_dir=agents_dir, **session_kwargs session_db_kwargs=session_db_kwargs,
per_agent=True, # Multi-agent mode
) )
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 # Build the Artifact service
if artifact_service_uri: try:
artifact_service = service_registry.create_artifact_service( artifact_service = create_artifact_service_from_options(
artifact_service_uri, agents_dir=agents_dir base_dir=agents_dir,
artifact_service_uri=artifact_service_uri,
per_agent=True, # Multi-agent mode
) )
if not artifact_service: except ValueError as exc:
raise click.ClickException( raise click.ClickException(str(exc)) from exc
"Unsupported artifact service URI: %s" % artifact_service_uri
)
else:
artifact_service = InMemoryArtifactService()
# Build the Credential service # Build the Credential service
credential_service = InMemoryCredentialService() credential_service = InMemoryCredentialService()
+138
View File
@@ -0,0 +1,138 @@
# 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 __future__ import annotations
import logging
from pathlib import Path
from typing import Any
from typing import Optional
from ...artifacts.base_artifact_service import BaseArtifactService
from ...memory.base_memory_service import BaseMemoryService
from ...sessions.base_session_service import BaseSessionService
from ..service_registry import get_service_registry
from .local_storage import create_local_artifact_service
logger = logging.getLogger("google_adk." + __name__)
def create_session_service_from_options(
*,
base_dir: Path | str,
session_service_uri: Optional[str] = None,
session_db_kwargs: Optional[dict[str, Any]] = None,
per_agent: bool = False,
) -> BaseSessionService:
"""Creates a session service based on CLI/web options."""
base_path = Path(base_dir)
registry = get_service_registry()
kwargs: dict[str, Any] = {
"agents_dir": str(base_path),
"per_agent": per_agent,
}
if session_db_kwargs:
kwargs.update(session_db_kwargs)
if session_service_uri:
if per_agent:
logger.warning(
"per_agent is not supported with remote session service URIs,"
" ignoring"
)
logger.info("Using session service URI: %s", session_service_uri)
service = registry.create_session_service(session_service_uri, **kwargs)
if service is not None:
return service
# Fallback to DatabaseSessionService if the registry doesn't support the
# session service URI scheme. This keeps support for SQLAlchemy-compatible
# databases like AlloyDB or Cloud Spanner without explicit registration.
from ...sessions.database_session_service import DatabaseSessionService
fallback_kwargs = dict(kwargs)
fallback_kwargs.pop("agents_dir", None)
fallback_kwargs.pop("per_agent", None)
logger.info(
"Falling back to DatabaseSessionService for URI: %s",
session_service_uri,
)
return DatabaseSessionService(db_url=session_service_uri, **fallback_kwargs)
logger.info("Using in-memory session service")
from ...sessions.in_memory_session_service import InMemorySessionService
return InMemorySessionService()
def create_memory_service_from_options(
*,
base_dir: Path | str,
memory_service_uri: Optional[str] = None,
) -> BaseMemoryService:
"""Creates a memory service based on CLI/web options."""
base_path = Path(base_dir)
registry = get_service_registry()
if memory_service_uri:
logger.info("Using memory service URI: %s", memory_service_uri)
service = registry.create_memory_service(
memory_service_uri,
agents_dir=str(base_path),
)
if service is None:
raise ValueError(f"Unsupported memory service URI: {memory_service_uri}")
return service
logger.info("Using in-memory memory service")
from ...memory.in_memory_memory_service import InMemoryMemoryService
return InMemoryMemoryService()
def create_artifact_service_from_options(
*,
base_dir: Path | str,
artifact_service_uri: Optional[str] = None,
per_agent: bool = False,
) -> BaseArtifactService:
"""Creates an artifact service based on CLI/web options."""
base_path = Path(base_dir)
registry = get_service_registry()
if artifact_service_uri:
if per_agent:
logger.warning(
"per_agent is not supported with remote artifact service URIs,"
" ignoring"
)
logger.info("Using artifact service URI: %s", artifact_service_uri)
service = registry.create_artifact_service(
artifact_service_uri,
agents_dir=str(base_path),
per_agent=per_agent,
)
if service is None:
logger.warning(
"Unsupported artifact service URI: %s, falling back to in-memory",
artifact_service_uri,
)
from ...artifacts.in_memory_artifact_service import InMemoryArtifactService
return InMemoryArtifactService()
return service
if per_agent:
logger.info("Using shared file artifact service rooted at %s", base_dir)
return create_local_artifact_service(base_dir=base_path, per_agent=per_agent)
+6 -6
View File
@@ -327,15 +327,15 @@ def test_app(
with ( with (
patch("signal.signal", return_value=None), patch("signal.signal", return_value=None),
patch( patch(
"google.adk.cli.fast_api.InMemorySessionService", "google.adk.cli.fast_api.create_session_service_from_options",
return_value=mock_session_service, return_value=mock_session_service,
), ),
patch( patch(
"google.adk.cli.fast_api.InMemoryArtifactService", "google.adk.cli.fast_api.create_artifact_service_from_options",
return_value=mock_artifact_service, return_value=mock_artifact_service,
), ),
patch( patch(
"google.adk.cli.fast_api.InMemoryMemoryService", "google.adk.cli.fast_api.create_memory_service_from_options",
return_value=mock_memory_service, return_value=mock_memory_service,
), ),
patch( patch(
@@ -472,15 +472,15 @@ def test_app_with_a2a(
with ( with (
patch("signal.signal", return_value=None), patch("signal.signal", return_value=None),
patch( patch(
"google.adk.cli.fast_api.InMemorySessionService", "google.adk.cli.fast_api.create_session_service_from_options",
return_value=mock_session_service, return_value=mock_session_service,
), ),
patch( patch(
"google.adk.cli.fast_api.InMemoryArtifactService", "google.adk.cli.fast_api.create_artifact_service_from_options",
return_value=mock_artifact_service, return_value=mock_artifact_service,
), ),
patch( patch(
"google.adk.cli.fast_api.InMemoryMemoryService", "google.adk.cli.fast_api.create_memory_service_from_options",
return_value=mock_memory_service, return_value=mock_memory_service,
), ),
patch( patch(
+121 -10
View File
@@ -28,7 +28,12 @@ from typing import Tuple
import click import click
from google.adk.agents.base_agent import BaseAgent from google.adk.agents.base_agent import BaseAgent
from google.adk.apps.app import App from google.adk.apps.app import App
from google.adk.artifacts.file_artifact_service import FileArtifactService
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
import google.adk.cli.cli as cli import google.adk.cli.cli as cli
from google.adk.cli.utils.service_factory import create_artifact_service_from_options
from google.adk.sessions.in_memory_session_service import InMemorySessionService
import pytest import pytest
@@ -151,9 +156,9 @@ async def test_run_input_file_outputs(
input_path = tmp_path / "input.json" input_path = tmp_path / "input.json"
input_path.write_text(json.dumps(input_json)) input_path.write_text(json.dumps(input_json))
artifact_service = cli.InMemoryArtifactService() artifact_service = InMemoryArtifactService()
session_service = cli.InMemorySessionService() session_service = InMemorySessionService()
credential_service = cli.InMemoryCredentialService() credential_service = InMemoryCredentialService()
dummy_root = BaseAgent(name="root") dummy_root = BaseAgent(name="root")
session = await cli.run_input_file( session = await cli.run_input_file(
@@ -189,6 +194,34 @@ async def test_run_cli_with_input_file(fake_agent, tmp_path: Path) -> None:
) )
@pytest.mark.asyncio
async def test_run_cli_loads_services_module(
fake_agent, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""run_cli should load custom services from the agents directory."""
parent_dir, folder_name = fake_agent
input_json = {"state": {}, "queries": ["ping"]}
input_path = tmp_path / "input.json"
input_path.write_text(json.dumps(input_json))
loaded_dirs: list[str] = []
monkeypatch.setattr(
cli, "load_services_module", lambda path: loaded_dirs.append(path)
)
agent_root = parent_dir / folder_name
await cli.run_cli(
agent_parent_dir=str(parent_dir),
agent_folder_name=folder_name,
input_file=str(input_path),
saved_session_file=None,
save_session=False,
)
assert loaded_dirs == [str(agent_root.resolve())]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_run_cli_app_uses_app_name_for_sessions( async def test_run_cli_app_uses_app_name_for_sessions(
fake_app_agent, tmp_path: Path, monkeypatch: pytest.MonkeyPatch fake_app_agent, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
@@ -197,15 +230,20 @@ async def test_run_cli_app_uses_app_name_for_sessions(
parent_dir, folder_name, app_name = fake_app_agent parent_dir, folder_name, app_name = fake_app_agent
created_app_names: List[str] = [] created_app_names: List[str] = []
original_session_cls = cli.InMemorySessionService class _SpySessionService(InMemorySessionService):
class _SpySessionService(original_session_cls):
async def create_session(self, *, app_name: str, **kwargs: Any) -> Any: async def create_session(self, *, app_name: str, **kwargs: Any) -> Any:
created_app_names.append(app_name) created_app_names.append(app_name)
return await super().create_session(app_name=app_name, **kwargs) return await super().create_session(app_name=app_name, **kwargs)
monkeypatch.setattr(cli, "InMemorySessionService", _SpySessionService) spy_session_service = _SpySessionService()
def _session_factory(**_: Any) -> InMemorySessionService:
return spy_session_service
monkeypatch.setattr(
cli, "create_session_service_from_options", _session_factory
)
input_json = {"state": {}, "queries": ["ping"]} input_json = {"state": {}, "queries": ["ping"]}
input_path = tmp_path / "input_app.json" input_path = tmp_path / "input_app.json"
@@ -253,16 +291,89 @@ async def test_run_cli_save_session(
assert "id" in data and "events" in data assert "id" in data and "events" in data
def test_create_artifact_service_defaults_to_file(tmp_path: Path) -> None:
"""Service factory should default to FileArtifactService when URI is unset."""
service = create_artifact_service_from_options(base_dir=tmp_path)
assert isinstance(service, FileArtifactService)
expected_root = Path(tmp_path) / ".adk" / "artifacts"
assert service.root_dir == expected_root
assert expected_root.exists()
def test_create_artifact_service_per_agent_uses_shared_root(
tmp_path: Path,
) -> None:
"""Multi-agent mode should still use a single file artifact service."""
service = create_artifact_service_from_options(
base_dir=tmp_path, per_agent=True
)
assert isinstance(service, FileArtifactService)
expected_root = Path(tmp_path) / ".adk" / "artifacts"
assert service.root_dir == expected_root
assert expected_root.exists()
def test_create_artifact_service_respects_memory_uri(tmp_path: Path) -> None:
"""Service factory should honor memory:// URIs."""
service = create_artifact_service_from_options(
base_dir=tmp_path, artifact_service_uri="memory://"
)
assert isinstance(service, InMemoryArtifactService)
def test_create_artifact_service_accepts_file_uri(tmp_path: Path) -> None:
"""Service factory should allow custom local roots via file:// URIs."""
custom_root = tmp_path / "custom_artifacts"
service = create_artifact_service_from_options(
base_dir=tmp_path, artifact_service_uri=custom_root.as_uri()
)
assert isinstance(service, FileArtifactService)
assert service.root_dir == custom_root
assert custom_root.exists()
def test_create_artifact_service_file_uri_rejects_per_agent(tmp_path: Path):
"""file:// URIs are incompatible with per-agent mode."""
custom_root = tmp_path / "custom"
with pytest.raises(ValueError, match="multi-agent"):
create_artifact_service_from_options(
base_dir=tmp_path,
artifact_service_uri=custom_root.as_uri(),
per_agent=True,
)
@pytest.mark.asyncio
async def test_run_cli_accepts_memory_scheme(
fake_agent, tmp_path: Path
) -> None:
"""run_cli should allow configuring in-memory services via memory:// URIs."""
parent_dir, folder_name = fake_agent
input_json = {"state": {}, "queries": []}
input_path = tmp_path / "noop.json"
input_path.write_text(json.dumps(input_json))
await cli.run_cli(
agent_parent_dir=str(parent_dir),
agent_folder_name=folder_name,
input_file=str(input_path),
saved_session_file=None,
save_session=False,
session_service_uri="memory://",
artifact_service_uri="memory://",
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_run_interactively_whitespace_and_exit( async def test_run_interactively_whitespace_and_exit(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None: ) -> None:
"""run_interactively should skip blank input, echo once, then exit.""" """run_interactively should skip blank input, echo once, then exit."""
# make a session that belongs to dummy agent # make a session that belongs to dummy agent
session_service = cli.InMemorySessionService() session_service = InMemorySessionService()
sess = await session_service.create_session(app_name="dummy", user_id="u") sess = await session_service.create_session(app_name="dummy", user_id="u")
artifact_service = cli.InMemoryArtifactService() artifact_service = InMemoryArtifactService()
credential_service = cli.InMemoryCredentialService() credential_service = InMemoryCredentialService()
root_agent = BaseAgent(name="root") root_agent = BaseAgent(name="root")
# fake user input: blank -> 'hello' -> 'exit' # fake user input: blank -> 'hello' -> 'exit'
@@ -0,0 +1,162 @@
# 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.
"""Tests for service factory helpers."""
from __future__ import annotations
from pathlib import Path
from unittest.mock import Mock
import google.adk.cli.utils.service_factory as service_factory
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
from google.adk.sessions.database_session_service import DatabaseSessionService
from google.adk.sessions.in_memory_session_service import InMemorySessionService
import pytest
def test_create_session_service_uses_registry(tmp_path: Path, monkeypatch):
registry = Mock()
expected = object()
registry.create_session_service.return_value = expected
monkeypatch.setattr(service_factory, "get_service_registry", lambda: registry)
result = service_factory.create_session_service_from_options(
base_dir=tmp_path,
session_service_uri="sqlite:///test.db",
)
assert result is expected
registry.create_session_service.assert_called_once_with(
"sqlite:///test.db",
agents_dir=str(tmp_path),
per_agent=False,
)
def test_create_session_service_per_agent_uri(tmp_path: Path, monkeypatch):
registry = Mock()
expected = object()
registry.create_session_service.return_value = expected
monkeypatch.setattr(service_factory, "get_service_registry", lambda: registry)
result = service_factory.create_session_service_from_options(
base_dir=tmp_path,
session_service_uri="memory://",
per_agent=True,
)
assert result is expected
registry.create_session_service.assert_called_once_with(
"memory://", agents_dir=str(tmp_path), per_agent=True
)
@pytest.mark.parametrize("per_agent", [True, False])
def test_create_session_service_defaults_to_memory(
tmp_path: Path, per_agent: bool
):
service = service_factory.create_session_service_from_options(
base_dir=tmp_path,
per_agent=per_agent,
)
assert isinstance(service, InMemorySessionService)
def test_create_session_service_fallbacks_to_database(
tmp_path: Path, monkeypatch
):
registry = Mock()
registry.create_session_service.return_value = None
monkeypatch.setattr(service_factory, "get_service_registry", lambda: registry)
service = service_factory.create_session_service_from_options(
base_dir=tmp_path,
session_service_uri="sqlite+aiosqlite:///:memory:",
session_db_kwargs={"echo": True},
)
assert isinstance(service, DatabaseSessionService)
assert service.db_engine.url.drivername == "sqlite+aiosqlite"
assert service.db_engine.echo is True
registry.create_session_service.assert_called_once_with(
"sqlite+aiosqlite:///:memory:",
agents_dir=str(tmp_path),
per_agent=False,
echo=True,
)
@pytest.mark.parametrize("per_agent", [True, False])
def test_create_artifact_service_uses_registry(
tmp_path: Path, monkeypatch, per_agent: bool
):
registry = Mock()
expected = object()
registry.create_artifact_service.return_value = expected
monkeypatch.setattr(service_factory, "get_service_registry", lambda: registry)
result = service_factory.create_artifact_service_from_options(
base_dir=tmp_path,
artifact_service_uri="gs://bucket/path",
per_agent=per_agent,
)
assert result is expected
registry.create_artifact_service.assert_called_once_with(
"gs://bucket/path",
agents_dir=str(tmp_path),
per_agent=per_agent,
)
def test_create_memory_service_uses_registry(tmp_path: Path, monkeypatch):
registry = Mock()
expected = object()
registry.create_memory_service.return_value = expected
monkeypatch.setattr(service_factory, "get_service_registry", lambda: registry)
result = service_factory.create_memory_service_from_options(
base_dir=tmp_path,
memory_service_uri="rag://my-corpus",
)
assert result is expected
registry.create_memory_service.assert_called_once_with(
"rag://my-corpus",
agents_dir=str(tmp_path),
)
def test_create_memory_service_defaults_to_in_memory(tmp_path: Path):
service = service_factory.create_memory_service_from_options(
base_dir=tmp_path
)
assert isinstance(service, InMemoryMemoryService)
def test_create_memory_service_raises_on_unknown_scheme(
tmp_path: Path, monkeypatch
):
registry = Mock()
registry.create_memory_service.return_value = None
monkeypatch.setattr(service_factory, "get_service_registry", lambda: registry)
with pytest.raises(ValueError):
service_factory.create_memory_service_from_options(
base_dir=tmp_path,
memory_service_uri="unknown://foo",
)