You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
feat!: Introduction of ADK folder for local session and artifact storage
Default CLI session storage to SQLite instead of in-memory Previously, adk run and adk web used in-memory session storage by default, causing sessions to be lost on restart. Now sessions persist to .adk/session.db automatically. To use in-memory storage, pass --session-service-uri memory:// Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 838975328
This commit is contained in:
committed by
Copybara-Service
parent
dd827af2ee
commit
ed9da3fa45
@@ -26,12 +26,10 @@ from pydantic import BaseModel
|
||||
from ..agents.llm_agent import LlmAgent
|
||||
from ..apps.app import App
|
||||
from ..artifacts.base_artifact_service import BaseArtifactService
|
||||
from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||||
from ..auth.credential_service.base_credential_service import BaseCredentialService
|
||||
from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
|
||||
from ..runners import Runner
|
||||
from ..sessions.base_session_service import BaseSessionService
|
||||
from ..sessions.in_memory_session_service import InMemorySessionService
|
||||
from ..sessions.session import Session
|
||||
from ..utils.context_utils import Aclosing
|
||||
from ..utils.env_utils import is_env_enabled
|
||||
@@ -162,8 +160,9 @@ async def run_cli(
|
||||
user_id = 'test_user'
|
||||
|
||||
# Create session and artifact services using factory functions
|
||||
# Sessions persist under <agents_dir>/<agent>/.adk/session.db by default.
|
||||
session_service = create_session_service_from_options(
|
||||
base_dir=agent_root,
|
||||
base_dir=agent_parent_path,
|
||||
session_service_uri=session_service_uri,
|
||||
)
|
||||
|
||||
|
||||
@@ -57,6 +57,32 @@ def create_local_database_session_service(
|
||||
return SqliteSessionService(db_path=str(session_db_path))
|
||||
|
||||
|
||||
def create_local_session_service(
|
||||
*,
|
||||
base_dir: Path | str,
|
||||
per_agent: bool = False,
|
||||
) -> BaseSessionService:
|
||||
"""Creates a local SQLite-backed session service.
|
||||
|
||||
Args:
|
||||
base_dir: The base directory for the agent(s).
|
||||
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.
|
||||
|
||||
Returns:
|
||||
A BaseSessionService instance backed by SQLite.
|
||||
"""
|
||||
if per_agent:
|
||||
logger.info(
|
||||
"Using per-agent session storage rooted at %s",
|
||||
base_dir,
|
||||
)
|
||||
return PerAgentDatabaseSessionService(agents_root=base_dir)
|
||||
|
||||
return create_local_database_session_service(base_dir=base_dir)
|
||||
|
||||
|
||||
def create_local_artifact_service(
|
||||
*, base_dir: Path | str
|
||||
) -> BaseArtifactService:
|
||||
|
||||
@@ -23,6 +23,7 @@ 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
|
||||
from .local_storage import create_local_session_service
|
||||
|
||||
logger = logging.getLogger("google_adk." + __name__)
|
||||
|
||||
@@ -62,10 +63,8 @@ def create_session_service_from_options(
|
||||
)
|
||||
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()
|
||||
# Default to per-agent local SQLite storage in <agents_root>/<agent>/.adk/.
|
||||
return create_local_session_service(base_dir=base_path, per_agent=True)
|
||||
|
||||
|
||||
def create_memory_service_from_options(
|
||||
|
||||
@@ -19,10 +19,10 @@ from __future__ import annotations
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock
|
||||
|
||||
from google.adk.cli.utils.local_storage import PerAgentDatabaseSessionService
|
||||
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
|
||||
|
||||
|
||||
@@ -44,12 +44,20 @@ def test_create_session_service_uses_registry(tmp_path: Path, monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def test_create_session_service_defaults_to_memory(tmp_path: Path):
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_session_service_defaults_to_per_agent_sqlite(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
agent_dir = tmp_path / "agent_a"
|
||||
agent_dir.mkdir()
|
||||
service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
)
|
||||
|
||||
assert isinstance(service, InMemorySessionService)
|
||||
assert isinstance(service, PerAgentDatabaseSessionService)
|
||||
session = await service.create_session(app_name="agent_a", user_id="user")
|
||||
assert session.app_name == "agent_a"
|
||||
assert (agent_dir / ".adk" / "session.db").exists()
|
||||
|
||||
|
||||
def test_create_session_service_fallbacks_to_database(
|
||||
|
||||
Reference in New Issue
Block a user