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
+6 -6
View File
@@ -327,15 +327,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(
@@ -472,15 +472,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(
+121 -10
View File
@@ -28,7 +28,12 @@ from typing import Tuple
import click
from google.adk.agents.base_agent import BaseAgent
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
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
@@ -151,9 +156,9 @@ async def test_run_input_file_outputs(
input_path = tmp_path / "input.json"
input_path.write_text(json.dumps(input_json))
artifact_service = cli.InMemoryArtifactService()
session_service = cli.InMemorySessionService()
credential_service = cli.InMemoryCredentialService()
artifact_service = InMemoryArtifactService()
session_service = InMemorySessionService()
credential_service = InMemoryCredentialService()
dummy_root = BaseAgent(name="root")
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
async def test_run_cli_app_uses_app_name_for_sessions(
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
created_app_names: List[str] = []
original_session_cls = cli.InMemorySessionService
class _SpySessionService(original_session_cls):
class _SpySessionService(InMemorySessionService):
async def create_session(self, *, app_name: str, **kwargs: Any) -> Any:
created_app_names.append(app_name)
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_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
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
async def test_run_interactively_whitespace_and_exit(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""run_interactively should skip blank input, echo once, then exit."""
# 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")
artifact_service = cli.InMemoryArtifactService()
credential_service = cli.InMemoryCredentialService()
artifact_service = InMemoryArtifactService()
credential_service = InMemoryCredentialService()
root_agent = BaseAgent(name="root")
# 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",
)