mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: avoid local .adk storage in Cloud Run/GKE
Default session and artifact services to in-memory when running in Cloud Run/Kubernetes (or when agents_dir isn’t writable) to prevent startup failures from attempting to create .adk under read-only/unwritable container paths (e.g. /app/agents/.adk). Local development defaults are unchanged. - ADK_FORCE_LOCAL_STORAGE=1 to always use .adk defaults - ADK_DISABLE_LOCAL_STORAGE=1 to always avoid local storage If local artifact initialization raises PermissionError, fall back to in-memory and log a warning Close #3907 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 853315459
This commit is contained in:
committed by
Copybara-Service
parent
d5bd8d939e
commit
b30c2f4e13
@@ -137,6 +137,7 @@ async def run_cli(
|
||||
session_id: Optional[str] = None,
|
||||
session_service_uri: Optional[str] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = True,
|
||||
) -> None:
|
||||
"""Runs an interactive CLI for a certain agent.
|
||||
|
||||
@@ -153,6 +154,7 @@ async def run_cli(
|
||||
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.
|
||||
use_local_storage: bool, whether to use local .adk storage by default.
|
||||
"""
|
||||
agent_parent_path = Path(agent_parent_dir).resolve()
|
||||
agent_root = agent_parent_path / agent_folder_name
|
||||
@@ -169,17 +171,19 @@ async def run_cli(
|
||||
if isinstance(agent_or_app, App) and agent_or_app.name != agent_folder_name:
|
||||
app_name_to_dir = {agent_or_app.name: agent_folder_name}
|
||||
|
||||
# Create session and artifact services using factory functions
|
||||
# Sessions persist under <agents_dir>/<agent>/.adk/session.db by default.
|
||||
# Create session and artifact services using factory functions.
|
||||
# Sessions persist under <agents_dir>/<agent>/.adk/session.db when enabled.
|
||||
session_service = create_session_service_from_options(
|
||||
base_dir=agent_parent_path,
|
||||
session_service_uri=session_service_uri,
|
||||
app_name_to_dir=app_name_to_dir,
|
||||
use_local_storage=use_local_storage,
|
||||
)
|
||||
|
||||
artifact_service = create_artifact_service_from_options(
|
||||
base_dir=agent_root,
|
||||
artifact_service_uri=artifact_service_uri,
|
||||
use_local_storage=use_local_storage,
|
||||
)
|
||||
|
||||
credential_service = InMemoryCredentialService()
|
||||
|
||||
@@ -26,6 +26,7 @@ from packaging.version import parse
|
||||
|
||||
_IS_WINDOWS = os.name == 'nt'
|
||||
_GCLOUD_CMD = 'gcloud.cmd' if _IS_WINDOWS else 'gcloud'
|
||||
_LOCAL_STORAGE_FLAG_MIN_VERSION: Final[str] = '1.21.0'
|
||||
|
||||
_DOCKERFILE_TEMPLATE: Final[str] = """
|
||||
FROM python:3.11-slim
|
||||
@@ -442,26 +443,38 @@ def _get_service_option_by_adk_version(
|
||||
session_uri: Optional[str],
|
||||
artifact_uri: Optional[str],
|
||||
memory_uri: Optional[str],
|
||||
use_local_storage: Optional[bool] = None,
|
||||
) -> str:
|
||||
"""Returns service option string based on adk_version."""
|
||||
parsed_version = parse(adk_version)
|
||||
options: list[str] = []
|
||||
|
||||
if parsed_version >= parse('1.3.0'):
|
||||
session_option = (
|
||||
f'--session_service_uri={session_uri}' if session_uri else ''
|
||||
)
|
||||
artifact_option = (
|
||||
f'--artifact_service_uri={artifact_uri}' if artifact_uri else ''
|
||||
)
|
||||
memory_option = f'--memory_service_uri={memory_uri}' if memory_uri else ''
|
||||
return f'{session_option} {artifact_option} {memory_option}'
|
||||
elif parsed_version >= parse('1.2.0'):
|
||||
session_option = f'--session_db_url={session_uri}' if session_uri else ''
|
||||
artifact_option = (
|
||||
f'--artifact_storage_uri={artifact_uri}' if artifact_uri else ''
|
||||
)
|
||||
return f'{session_option} {artifact_option}'
|
||||
if session_uri:
|
||||
options.append(f'--session_service_uri={session_uri}')
|
||||
if artifact_uri:
|
||||
options.append(f'--artifact_service_uri={artifact_uri}')
|
||||
if memory_uri:
|
||||
options.append(f'--memory_service_uri={memory_uri}')
|
||||
else:
|
||||
return f'--session_db_url={session_uri}' if session_uri else ''
|
||||
if session_uri:
|
||||
options.append(f'--session_db_url={session_uri}')
|
||||
if parsed_version >= parse('1.2.0') and artifact_uri:
|
||||
options.append(f'--artifact_storage_uri={artifact_uri}')
|
||||
|
||||
if use_local_storage is not None and parsed_version >= parse(
|
||||
_LOCAL_STORAGE_FLAG_MIN_VERSION
|
||||
):
|
||||
# Only valid when session/artifact URIs are unset; otherwise the CLI
|
||||
# rejects the combination to avoid confusing precedence.
|
||||
if session_uri is None and artifact_uri is None:
|
||||
options.append((
|
||||
'--use_local_storage'
|
||||
if use_local_storage
|
||||
else '--no_use_local_storage'
|
||||
))
|
||||
|
||||
return ' '.join(options)
|
||||
|
||||
|
||||
def to_cloud_run(
|
||||
@@ -482,6 +495,7 @@ def to_cloud_run(
|
||||
session_service_uri: Optional[str] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
memory_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = False,
|
||||
a2a: bool = False,
|
||||
extra_gcloud_args: Optional[tuple[str, ...]] = None,
|
||||
):
|
||||
@@ -517,8 +531,12 @@ def to_cloud_run(
|
||||
session_service_uri: The URI of the session service.
|
||||
artifact_service_uri: The URI of the artifact service.
|
||||
memory_service_uri: The URI of the memory service.
|
||||
use_local_storage: Whether to use local .adk storage in the container.
|
||||
"""
|
||||
app_name = app_name or os.path.basename(agent_folder)
|
||||
if parse(adk_version) >= parse('1.3.0') and not use_local_storage:
|
||||
session_service_uri = session_service_uri or 'memory://'
|
||||
artifact_service_uri = artifact_service_uri or 'memory://'
|
||||
|
||||
click.echo(f'Start generating Cloud Run source files in {temp_folder}')
|
||||
|
||||
@@ -559,6 +577,7 @@ def to_cloud_run(
|
||||
session_service_uri,
|
||||
artifact_service_uri,
|
||||
memory_service_uri,
|
||||
use_local_storage,
|
||||
),
|
||||
trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '',
|
||||
allow_origins_option=allow_origins_option,
|
||||
@@ -944,6 +963,7 @@ def to_gke(
|
||||
session_service_uri: Optional[str] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
memory_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = False,
|
||||
a2a: bool = False,
|
||||
):
|
||||
"""Deploys an agent to Google Kubernetes Engine(GKE).
|
||||
@@ -969,6 +989,7 @@ def to_gke(
|
||||
session_service_uri: The URI of the session service.
|
||||
artifact_service_uri: The URI of the artifact service.
|
||||
memory_service_uri: The URI of the memory service.
|
||||
use_local_storage: Whether to use local .adk storage in the container.
|
||||
"""
|
||||
click.secho(
|
||||
'\n🚀 Starting ADK Agent Deployment to GKE...', fg='cyan', bold=True
|
||||
@@ -982,6 +1003,9 @@ def to_gke(
|
||||
click.echo('--------------------------------------------------\n')
|
||||
|
||||
app_name = app_name or os.path.basename(agent_folder)
|
||||
if parse(adk_version) >= parse('1.3.0') and not use_local_storage:
|
||||
session_service_uri = session_service_uri or 'memory://'
|
||||
artifact_service_uri = artifact_service_uri or 'memory://'
|
||||
|
||||
click.secho('STEP 1: Preparing build environment...', bold=True)
|
||||
click.echo(f' - Using temporary directory: {temp_folder}')
|
||||
@@ -1024,6 +1048,7 @@ def to_gke(
|
||||
session_service_uri,
|
||||
artifact_service_uri,
|
||||
memory_service_uri,
|
||||
use_local_storage,
|
||||
),
|
||||
trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '',
|
||||
allow_origins_option=allow_origins_option,
|
||||
|
||||
@@ -36,7 +36,6 @@ from . import cli_create
|
||||
from . import cli_deploy
|
||||
from .. import version
|
||||
from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
|
||||
from ..sessions.migration import migration_runner
|
||||
from .cli import run_cli
|
||||
from .fast_api import get_fast_api_app
|
||||
from .utils import envs
|
||||
@@ -368,24 +367,26 @@ def validate_exclusive(ctx, param, value):
|
||||
return value
|
||||
|
||||
|
||||
def adk_services_options():
|
||||
def adk_services_options(*, default_use_local_storage: bool = True):
|
||||
"""Decorator to add ADK services options to click commands."""
|
||||
|
||||
def decorator(func):
|
||||
@click.option(
|
||||
"--session_service_uri",
|
||||
help=textwrap.dedent(
|
||||
"""\
|
||||
help=textwrap.dedent("""\
|
||||
Optional. The URI of the session service.
|
||||
- Leave unset to use the in-memory session service (default).
|
||||
If set, ADK uses this service.
|
||||
|
||||
If unset, ADK chooses a default session service (see
|
||||
--use_local_storage).
|
||||
- Use 'agentengine://<agent_engine>' to connect to Agent Engine
|
||||
sessions. <agent_engine> can either be the full qualified resource
|
||||
name 'projects/abc/locations/us-central1/reasoningEngines/123' or
|
||||
the resource id '123'.
|
||||
- Use 'memory://' to run with the in-memory session service.
|
||||
- Use 'sqlite://<path_to_sqlite_file>' to connect to a SQLite DB.
|
||||
- See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported database URIs."""
|
||||
),
|
||||
- See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls
|
||||
for supported database URIs."""),
|
||||
)
|
||||
@click.option(
|
||||
"--artifact_service_uri",
|
||||
@@ -393,13 +394,29 @@ def adk_services_options():
|
||||
help=textwrap.dedent(
|
||||
"""\
|
||||
Optional. The URI of the artifact service.
|
||||
- Leave unset to store artifacts under '.adk/artifacts' locally.
|
||||
If set, ADK uses this service.
|
||||
|
||||
If unset, ADK chooses a default artifact service (see
|
||||
--use_local_storage).
|
||||
- Use 'gs://<bucket_name>' to connect to the GCS artifact service.
|
||||
- Use 'memory://' to force the in-memory artifact service.
|
||||
- Use 'file://<path>' to store artifacts in a custom local directory."""
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
@click.option(
|
||||
"--use_local_storage/--no_use_local_storage",
|
||||
default=default_use_local_storage,
|
||||
show_default=True,
|
||||
help=(
|
||||
"Optional. Whether to use local .adk storage when "
|
||||
"--session_service_uri and --artifact_service_uri are unset. "
|
||||
"Cannot be combined with explicit service URIs. When the agents "
|
||||
"directory isn't writable (common in Cloud Run/Kubernetes), ADK "
|
||||
"falls back to in-memory unless overridden by "
|
||||
"ADK_FORCE_LOCAL_STORAGE=1 or ADK_DISABLE_LOCAL_STORAGE=1."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--memory_service_uri",
|
||||
type=str,
|
||||
@@ -415,6 +432,17 @@ def adk_services_options():
|
||||
)
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
ctx = click.get_current_context(silent=True)
|
||||
if ctx is not None:
|
||||
use_local_storage_source = ctx.get_parameter_source("use_local_storage")
|
||||
if use_local_storage_source != ParameterSource.DEFAULT and (
|
||||
kwargs.get("session_service_uri") is not None
|
||||
or kwargs.get("artifact_service_uri") is not None
|
||||
):
|
||||
raise click.UsageError(
|
||||
"--use_local_storage/--no_use_local_storage cannot be used with "
|
||||
"--session_service_uri or --artifact_service_uri."
|
||||
)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
@@ -423,7 +451,7 @@ def adk_services_options():
|
||||
|
||||
|
||||
@main.command("run", cls=HelpfulCommand)
|
||||
@adk_services_options()
|
||||
@adk_services_options(default_use_local_storage=True)
|
||||
@click.option(
|
||||
"--save_session",
|
||||
type=bool,
|
||||
@@ -481,6 +509,7 @@ def cli_run(
|
||||
session_service_uri: Optional[str] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
memory_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = True,
|
||||
):
|
||||
"""Runs an interactive CLI for a certain agent.
|
||||
|
||||
@@ -513,6 +542,7 @@ def cli_run(
|
||||
session_id=session_id,
|
||||
session_service_uri=session_service_uri,
|
||||
artifact_service_uri=artifact_service_uri,
|
||||
use_local_storage=use_local_storage,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1113,7 +1143,7 @@ def fast_api_common_options():
|
||||
@main.command("web")
|
||||
@fast_api_common_options()
|
||||
@web_options()
|
||||
@adk_services_options()
|
||||
@adk_services_options(default_use_local_storage=True)
|
||||
@deprecated_adk_services_options()
|
||||
@click.argument(
|
||||
"agents_dir",
|
||||
@@ -1136,6 +1166,7 @@ def cli_web(
|
||||
session_service_uri: Optional[str] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
memory_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = True,
|
||||
session_db_url: Optional[str] = None, # Deprecated
|
||||
artifact_storage_uri: Optional[str] = None, # Deprecated
|
||||
a2a: bool = False,
|
||||
@@ -1184,6 +1215,7 @@ def cli_web(
|
||||
session_service_uri=session_service_uri,
|
||||
artifact_service_uri=artifact_service_uri,
|
||||
memory_service_uri=memory_service_uri,
|
||||
use_local_storage=use_local_storage,
|
||||
eval_storage_uri=eval_storage_uri,
|
||||
allow_origins=allow_origins,
|
||||
web=True,
|
||||
@@ -1221,7 +1253,7 @@ def cli_web(
|
||||
default=os.getcwd(),
|
||||
)
|
||||
@fast_api_common_options()
|
||||
@adk_services_options()
|
||||
@adk_services_options(default_use_local_storage=True)
|
||||
@deprecated_adk_services_options()
|
||||
def cli_api_server(
|
||||
agents_dir: str,
|
||||
@@ -1237,6 +1269,7 @@ def cli_api_server(
|
||||
session_service_uri: Optional[str] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
memory_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = True,
|
||||
session_db_url: Optional[str] = None, # Deprecated
|
||||
artifact_storage_uri: Optional[str] = None, # Deprecated
|
||||
a2a: bool = False,
|
||||
@@ -1262,6 +1295,7 @@ def cli_api_server(
|
||||
session_service_uri=session_service_uri,
|
||||
artifact_service_uri=artifact_service_uri,
|
||||
memory_service_uri=memory_service_uri,
|
||||
use_local_storage=use_local_storage,
|
||||
eval_storage_uri=eval_storage_uri,
|
||||
allow_origins=allow_origins,
|
||||
web=False,
|
||||
@@ -1403,7 +1437,7 @@ def cli_api_server(
|
||||
multiple=True,
|
||||
)
|
||||
# TODO: Add eval_storage_uri option back when evals are supported in Cloud Run.
|
||||
@adk_services_options()
|
||||
@adk_services_options(default_use_local_storage=False)
|
||||
@deprecated_adk_services_options()
|
||||
@click.pass_context
|
||||
def cli_deploy_cloud_run(
|
||||
@@ -1424,6 +1458,7 @@ def cli_deploy_cloud_run(
|
||||
session_service_uri: Optional[str] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
memory_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = False,
|
||||
session_db_url: Optional[str] = None, # Deprecated
|
||||
artifact_storage_uri: Optional[str] = None, # Deprecated
|
||||
a2a: bool = False,
|
||||
@@ -1501,6 +1536,7 @@ def cli_deploy_cloud_run(
|
||||
session_service_uri=session_service_uri,
|
||||
artifact_service_uri=artifact_service_uri,
|
||||
memory_service_uri=memory_service_uri,
|
||||
use_local_storage=use_local_storage,
|
||||
a2a=a2a,
|
||||
extra_gcloud_args=tuple(gcloud_args),
|
||||
)
|
||||
@@ -1508,47 +1544,6 @@ def cli_deploy_cloud_run(
|
||||
click.secho(f"Deploy failed: {e}", fg="red", err=True)
|
||||
|
||||
|
||||
@main.group()
|
||||
def migrate():
|
||||
"""ADK migration commands."""
|
||||
pass
|
||||
|
||||
|
||||
@migrate.command("session", cls=HelpfulCommand)
|
||||
@click.option(
|
||||
"--source_db_url",
|
||||
required=True,
|
||||
help=(
|
||||
"SQLAlchemy URL of source database in database session service, e.g."
|
||||
" sqlite:///source.db."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--dest_db_url",
|
||||
required=True,
|
||||
help=(
|
||||
"SQLAlchemy URL of destination database in database session service,"
|
||||
" e.g. sqlite:///dest.db."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--log_level",
|
||||
type=LOG_LEVELS,
|
||||
default="INFO",
|
||||
help="Optional. Set the logging level",
|
||||
)
|
||||
def cli_migrate_session(
|
||||
*, source_db_url: str, dest_db_url: str, log_level: str
|
||||
):
|
||||
"""Migrates a session database to the latest schema version."""
|
||||
logs.setup_adk_logger(getattr(logging, log_level.upper()))
|
||||
try:
|
||||
migration_runner.upgrade(source_db_url, dest_db_url)
|
||||
click.secho("Migration check and upgrade process finished.", fg="green")
|
||||
except Exception as e:
|
||||
click.secho(f"Migration failed: {e}", fg="red", err=True)
|
||||
|
||||
|
||||
@deploy.command("agent_engine")
|
||||
@click.option(
|
||||
"--api_key",
|
||||
@@ -1843,7 +1838,7 @@ def cli_deploy_agent_engine(
|
||||
" version in the dev environment)"
|
||||
),
|
||||
)
|
||||
@adk_services_options()
|
||||
@adk_services_options(default_use_local_storage=False)
|
||||
@click.argument(
|
||||
"agent",
|
||||
type=click.Path(
|
||||
@@ -1866,6 +1861,7 @@ def cli_deploy_gke(
|
||||
session_service_uri: Optional[str] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
memory_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = False,
|
||||
):
|
||||
"""Deploys an agent to GKE.
|
||||
|
||||
@@ -1894,6 +1890,7 @@ def cli_deploy_gke(
|
||||
session_service_uri=session_service_uri,
|
||||
artifact_service_uri=artifact_service_uri,
|
||||
memory_service_uri=memory_service_uri,
|
||||
use_local_storage=use_local_storage,
|
||||
)
|
||||
except Exception as e:
|
||||
click.secho(f"Deploy failed: {e}", fg="red", err=True)
|
||||
|
||||
@@ -76,6 +76,7 @@ def get_fast_api_app(
|
||||
session_db_kwargs: Optional[Mapping[str, Any]] = None,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
memory_service_uri: Optional[str] = None,
|
||||
use_local_storage: bool = True,
|
||||
eval_storage_uri: Optional[str] = None,
|
||||
allow_origins: Optional[list[str]] = None,
|
||||
web: bool,
|
||||
@@ -122,6 +123,7 @@ def get_fast_api_app(
|
||||
base_dir=agents_dir,
|
||||
session_service_uri=session_service_uri,
|
||||
session_db_kwargs=session_db_kwargs,
|
||||
use_local_storage=use_local_storage,
|
||||
)
|
||||
|
||||
# Build the Artifact service
|
||||
@@ -130,6 +132,7 @@ def get_fast_api_app(
|
||||
base_dir=agents_dir,
|
||||
artifact_service_uri=artifact_service_uri,
|
||||
strict_uri=True,
|
||||
use_local_storage=use_local_storage,
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise click.ClickException(str(exc)) from exc
|
||||
|
||||
@@ -13,7 +13,9 @@
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import errno
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Optional
|
||||
@@ -21,12 +23,111 @@ 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 ...utils.env_utils import is_env_enabled
|
||||
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__)
|
||||
|
||||
_DISABLE_LOCAL_STORAGE_ENV = "ADK_DISABLE_LOCAL_STORAGE"
|
||||
_FORCE_LOCAL_STORAGE_ENV = "ADK_FORCE_LOCAL_STORAGE"
|
||||
_LOCAL_STORAGE_ERRNOS = frozenset({
|
||||
errno.EACCES,
|
||||
errno.EPERM,
|
||||
errno.EROFS,
|
||||
})
|
||||
|
||||
_CLOUD_RUN_SERVICE_ENV = "K_SERVICE"
|
||||
_KUBERNETES_HOST_ENV = "KUBERNETES_SERVICE_HOST"
|
||||
|
||||
|
||||
def _is_cloud_run() -> bool:
|
||||
"""Returns True when running in Cloud Run."""
|
||||
return bool(os.environ.get(_CLOUD_RUN_SERVICE_ENV))
|
||||
|
||||
|
||||
def _is_kubernetes() -> bool:
|
||||
"""Returns True when running in Kubernetes (including GKE)."""
|
||||
return bool(os.environ.get(_KUBERNETES_HOST_ENV))
|
||||
|
||||
|
||||
def _is_dir_writable(path: Path) -> bool:
|
||||
"""Returns True if the directory exists and is writable/executable."""
|
||||
try:
|
||||
if not path.exists() or not path.is_dir():
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
return os.access(path, os.W_OK | os.X_OK)
|
||||
|
||||
|
||||
def _resolve_use_local_storage(
|
||||
*,
|
||||
base_path: Path,
|
||||
requested: bool,
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Resolves effective local storage setting with safe defaults."""
|
||||
if is_env_enabled(_DISABLE_LOCAL_STORAGE_ENV):
|
||||
warning_message = (
|
||||
"Local storage is disabled by %s; using in-memory services. "
|
||||
"Set --session_service_uri/--artifact_service_uri for production "
|
||||
"deployments."
|
||||
) % _DISABLE_LOCAL_STORAGE_ENV
|
||||
return False, warning_message
|
||||
|
||||
if is_env_enabled(_FORCE_LOCAL_STORAGE_ENV):
|
||||
if not _is_dir_writable(base_path):
|
||||
warning_message = (
|
||||
"Local storage is forced by %s, but %s is not writable; "
|
||||
"using in-memory services."
|
||||
) % (_FORCE_LOCAL_STORAGE_ENV, base_path)
|
||||
return False, warning_message
|
||||
return True, None
|
||||
|
||||
if not requested:
|
||||
return False, None
|
||||
|
||||
if _is_cloud_run() or _is_kubernetes():
|
||||
warning_message = (
|
||||
"Detected Cloud Run/Kubernetes runtime; using in-memory services "
|
||||
"instead of local .adk storage. Set %s=1 to force local storage."
|
||||
) % _FORCE_LOCAL_STORAGE_ENV
|
||||
return False, warning_message
|
||||
|
||||
if not _is_dir_writable(base_path):
|
||||
warning_message = (
|
||||
"Agents directory %s is not writable; using in-memory services "
|
||||
"instead of local .adk storage. Set %s=1 to force local storage."
|
||||
) % (base_path, _FORCE_LOCAL_STORAGE_ENV)
|
||||
return False, warning_message
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
def _create_in_memory_session_service(
|
||||
warning_message: str | None = None,
|
||||
*warning_args: object,
|
||||
) -> BaseSessionService:
|
||||
"""Creates an in-memory session service, optionally logging a warning."""
|
||||
if warning_message is not None:
|
||||
logger.warning(warning_message, *warning_args)
|
||||
from ...sessions.in_memory_session_service import InMemorySessionService
|
||||
|
||||
return InMemorySessionService()
|
||||
|
||||
|
||||
def _create_in_memory_artifact_service(
|
||||
warning_message: str | None = None,
|
||||
*warning_args: object,
|
||||
) -> BaseArtifactService:
|
||||
"""Creates an in-memory artifact service, optionally logging a warning."""
|
||||
if warning_message is not None:
|
||||
logger.warning(warning_message, *warning_args)
|
||||
from ...artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||||
|
||||
return InMemoryArtifactService()
|
||||
|
||||
|
||||
def create_session_service_from_options(
|
||||
*,
|
||||
@@ -34,6 +135,7 @@ def create_session_service_from_options(
|
||||
session_service_uri: Optional[str] = None,
|
||||
session_db_kwargs: Optional[dict[str, Any]] = None,
|
||||
app_name_to_dir: Optional[dict[str, str]] = None,
|
||||
use_local_storage: bool = True,
|
||||
) -> BaseSessionService:
|
||||
"""Creates a session service based on CLI/web options."""
|
||||
base_path = Path(base_dir)
|
||||
@@ -64,12 +166,36 @@ def create_session_service_from_options(
|
||||
)
|
||||
return DatabaseSessionService(db_url=session_service_uri, **fallback_kwargs)
|
||||
|
||||
# Default to per-agent local SQLite storage in <agents_root>/<agent>/.adk/.
|
||||
return create_local_session_service(
|
||||
base_dir=base_path,
|
||||
per_agent=True,
|
||||
app_name_to_dir=app_name_to_dir,
|
||||
effective_use_local_storage, auto_warning = _resolve_use_local_storage(
|
||||
base_path=base_path,
|
||||
requested=use_local_storage,
|
||||
)
|
||||
if not effective_use_local_storage:
|
||||
if auto_warning is not None:
|
||||
return _create_in_memory_session_service(auto_warning)
|
||||
return _create_in_memory_session_service(
|
||||
"Local session storage is disabled; using in-memory session service. "
|
||||
"Set --session_service_uri for production deployments."
|
||||
)
|
||||
|
||||
# Default to per-agent local SQLite storage in <agents_root>/<agent>/.adk/.
|
||||
try:
|
||||
return create_local_session_service(
|
||||
base_dir=base_path,
|
||||
per_agent=True,
|
||||
app_name_to_dir=app_name_to_dir,
|
||||
)
|
||||
except OSError as exc:
|
||||
if exc.errno not in _LOCAL_STORAGE_ERRNOS and not isinstance(
|
||||
exc, PermissionError
|
||||
):
|
||||
raise
|
||||
return _create_in_memory_session_service(
|
||||
"Failed to initialize local session storage under %s (%r); "
|
||||
"falling back to in-memory session service.",
|
||||
base_path,
|
||||
exc,
|
||||
)
|
||||
|
||||
|
||||
def create_memory_service_from_options(
|
||||
@@ -102,6 +228,7 @@ def create_artifact_service_from_options(
|
||||
base_dir: Path | str,
|
||||
artifact_service_uri: Optional[str] = None,
|
||||
strict_uri: bool = False,
|
||||
use_local_storage: bool = True,
|
||||
) -> BaseArtifactService:
|
||||
"""Creates an artifact service based on CLI/web options."""
|
||||
base_path = Path(base_dir)
|
||||
@@ -118,13 +245,34 @@ def create_artifact_service_from_options(
|
||||
raise ValueError(
|
||||
f"Unsupported artifact service URI: {artifact_service_uri}"
|
||||
)
|
||||
logger.warning(
|
||||
return _create_in_memory_artifact_service(
|
||||
"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
|
||||
|
||||
return create_local_artifact_service(base_dir=base_path)
|
||||
effective_use_local_storage, auto_warning = _resolve_use_local_storage(
|
||||
base_path=base_path,
|
||||
requested=use_local_storage,
|
||||
)
|
||||
if not effective_use_local_storage:
|
||||
if auto_warning is not None:
|
||||
return _create_in_memory_artifact_service(auto_warning)
|
||||
return _create_in_memory_artifact_service(
|
||||
"Local artifact storage is disabled; using in-memory artifact service. "
|
||||
"Set --artifact_service_uri for production deployments."
|
||||
)
|
||||
|
||||
try:
|
||||
return create_local_artifact_service(base_dir=base_path)
|
||||
except OSError as exc:
|
||||
if exc.errno not in _LOCAL_STORAGE_ERRNOS and not isinstance(
|
||||
exc, PermissionError
|
||||
):
|
||||
raise
|
||||
return _create_in_memory_artifact_service(
|
||||
"Failed to initialize local artifact storage under %s (%r); "
|
||||
"falling back to in-memory artifact service.",
|
||||
base_path,
|
||||
exc,
|
||||
)
|
||||
|
||||
@@ -293,7 +293,10 @@ async def test_run_cli_save_session(
|
||||
|
||||
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)
|
||||
service = create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(service, FileArtifactService)
|
||||
expected_root = Path(tmp_path) / ".adk" / "artifacts"
|
||||
assert service.root_dir == expected_root
|
||||
@@ -304,7 +307,10 @@ def test_create_artifact_service_uses_shared_root(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Artifact service should use a single file artifact service."""
|
||||
service = create_artifact_service_from_options(base_dir=tmp_path)
|
||||
service = create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(service, FileArtifactService)
|
||||
expected_root = Path(tmp_path) / ".adk" / "artifacts"
|
||||
assert service.root_dir == expected_root
|
||||
|
||||
@@ -128,13 +128,15 @@ def test_resolve_project_from_gcloud_fails(
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"adk_version, session_uri, artifact_uri, memory_uri, expected",
|
||||
"adk_version, session_uri, artifact_uri, memory_uri, use_local_storage, "
|
||||
"expected",
|
||||
[
|
||||
(
|
||||
"1.3.0",
|
||||
"sqlite://s",
|
||||
"gs://a",
|
||||
"rag://m",
|
||||
None,
|
||||
(
|
||||
"--session_service_uri=sqlite://s --artifact_service_uri=gs://a"
|
||||
" --memory_service_uri=rag://m"
|
||||
@@ -145,6 +147,7 @@ def test_resolve_project_from_gcloud_fails(
|
||||
"sqlite://s",
|
||||
"gs://a",
|
||||
"rag://m",
|
||||
None,
|
||||
"--session_db_url=sqlite://s --artifact_storage_uri=gs://a",
|
||||
),
|
||||
(
|
||||
@@ -152,6 +155,7 @@ def test_resolve_project_from_gcloud_fails(
|
||||
"sqlite://s",
|
||||
"gs://a",
|
||||
"rag://m",
|
||||
None,
|
||||
"--session_db_url=sqlite://s",
|
||||
),
|
||||
(
|
||||
@@ -159,16 +163,49 @@ def test_resolve_project_from_gcloud_fails(
|
||||
"sqlite://s",
|
||||
None,
|
||||
None,
|
||||
"--session_service_uri=sqlite://s ",
|
||||
None,
|
||||
"--session_service_uri=sqlite://s",
|
||||
),
|
||||
(
|
||||
"1.3.0",
|
||||
None,
|
||||
"gs://a",
|
||||
"rag://m",
|
||||
" --artifact_service_uri=gs://a --memory_service_uri=rag://m",
|
||||
None,
|
||||
"--artifact_service_uri=gs://a --memory_service_uri=rag://m",
|
||||
),
|
||||
(
|
||||
"1.2.0",
|
||||
None,
|
||||
"gs://a",
|
||||
None,
|
||||
None,
|
||||
"--artifact_storage_uri=gs://a",
|
||||
),
|
||||
(
|
||||
"1.21.0",
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
False,
|
||||
"--no_use_local_storage",
|
||||
),
|
||||
(
|
||||
"1.21.0",
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
True,
|
||||
"--use_local_storage",
|
||||
),
|
||||
(
|
||||
"1.21.0",
|
||||
"sqlite://s",
|
||||
"gs://a",
|
||||
None,
|
||||
False,
|
||||
"--session_service_uri=sqlite://s --artifact_service_uri=gs://a",
|
||||
),
|
||||
("1.2.0", None, "gs://a", None, " --artifact_storage_uri=gs://a"),
|
||||
],
|
||||
)
|
||||
def test_get_service_option_by_adk_version(
|
||||
@@ -176,6 +213,7 @@ def test_get_service_option_by_adk_version(
|
||||
session_uri: str | None,
|
||||
artifact_uri: str | None,
|
||||
memory_uri: str | None,
|
||||
use_local_storage: bool | None,
|
||||
expected: str,
|
||||
) -> None:
|
||||
"""It should return the correct service URI flags for a given ADK version."""
|
||||
@@ -184,6 +222,7 @@ def test_get_service_option_by_adk_version(
|
||||
session_uri=session_uri,
|
||||
artifact_uri=artifact_uri,
|
||||
memory_uri=memory_uri,
|
||||
use_local_storage=use_local_storage,
|
||||
)
|
||||
assert actual.rstrip() == expected.rstrip()
|
||||
|
||||
|
||||
@@ -16,13 +16,17 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock
|
||||
|
||||
from google.adk.artifacts.file_artifact_service import FileArtifactService
|
||||
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||||
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
|
||||
|
||||
|
||||
@@ -52,6 +56,7 @@ async def test_create_session_service_defaults_to_per_agent_sqlite(
|
||||
agent_dir.mkdir()
|
||||
service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
|
||||
assert isinstance(service, PerAgentDatabaseSessionService)
|
||||
@@ -71,6 +76,7 @@ async def test_create_session_service_respects_app_name_mapping(
|
||||
service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
app_name_to_dir={logical_name: "agent_folder"},
|
||||
use_local_storage=True,
|
||||
)
|
||||
|
||||
assert isinstance(service, PerAgentDatabaseSessionService)
|
||||
@@ -173,3 +179,175 @@ def test_create_memory_service_raises_on_unknown_scheme(
|
||||
base_dir=tmp_path,
|
||||
memory_service_uri="unknown://foo",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_session_service_defaults_to_in_memory_when_disabled(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=False,
|
||||
)
|
||||
|
||||
assert isinstance(service, InMemorySessionService)
|
||||
session = await service.create_session(app_name="agent_a", user_id="user")
|
||||
assert session.app_name == "agent_a"
|
||||
assert not (tmp_path / "agent_a" / ".adk").exists()
|
||||
|
||||
|
||||
def test_create_artifact_service_defaults_to_in_memory_when_disabled(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
service = service_factory.create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=False,
|
||||
)
|
||||
|
||||
assert isinstance(service, InMemoryArtifactService)
|
||||
assert not (tmp_path / ".adk").exists()
|
||||
|
||||
|
||||
def test_create_session_service_fallbacks_to_in_memory_on_permission_error(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
def _raise_permission_error(*_args, **_kwargs):
|
||||
raise PermissionError("nope")
|
||||
|
||||
monkeypatch.setattr(
|
||||
service_factory, "create_local_session_service", _raise_permission_error
|
||||
)
|
||||
|
||||
service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
|
||||
assert isinstance(service, InMemorySessionService)
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="chmod behavior differs on Windows")
|
||||
def test_create_services_default_to_in_memory_when_agents_dir_unwritable(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
agents_dir = tmp_path / "agents"
|
||||
agents_dir.mkdir()
|
||||
try:
|
||||
agents_dir.chmod(0o555)
|
||||
if os.access(agents_dir, os.W_OK | os.X_OK):
|
||||
pytest.skip("Test cannot make directory unwritable in this environment.")
|
||||
|
||||
session_service = service_factory.create_session_service_from_options(
|
||||
base_dir=agents_dir,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(session_service, InMemorySessionService)
|
||||
|
||||
artifact_service = service_factory.create_artifact_service_from_options(
|
||||
base_dir=agents_dir,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(artifact_service, InMemoryArtifactService)
|
||||
finally:
|
||||
agents_dir.chmod(0o755)
|
||||
|
||||
|
||||
def test_adk_disable_local_storage_env_forces_in_memory(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("ADK_DISABLE_LOCAL_STORAGE", "1")
|
||||
|
||||
session_service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(session_service, InMemorySessionService)
|
||||
|
||||
artifact_service = service_factory.create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(artifact_service, InMemoryArtifactService)
|
||||
|
||||
|
||||
def test_cloud_run_env_defaults_to_in_memory(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("K_SERVICE", "adk-service")
|
||||
|
||||
session_service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(session_service, InMemorySessionService)
|
||||
|
||||
artifact_service = service_factory.create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(artifact_service, InMemoryArtifactService)
|
||||
|
||||
|
||||
def test_kubernetes_env_defaults_to_in_memory(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "10.0.0.1")
|
||||
|
||||
session_service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(session_service, InMemorySessionService)
|
||||
|
||||
artifact_service = service_factory.create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
assert isinstance(artifact_service, InMemoryArtifactService)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_adk_force_local_storage_env_overrides_flag(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("ADK_FORCE_LOCAL_STORAGE", "1")
|
||||
agent_dir = tmp_path / "agent_a"
|
||||
agent_dir.mkdir()
|
||||
|
||||
session_service = service_factory.create_session_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=False,
|
||||
)
|
||||
assert isinstance(session_service, PerAgentDatabaseSessionService)
|
||||
await session_service.create_session(app_name="agent_a", user_id="user")
|
||||
assert (agent_dir / ".adk" / "session.db").exists()
|
||||
|
||||
artifact_service = service_factory.create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=False,
|
||||
)
|
||||
assert isinstance(artifact_service, FileArtifactService)
|
||||
|
||||
|
||||
def test_create_artifact_service_fallbacks_to_in_memory_on_permission_error(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
def _raise_permission_error(*_args, **_kwargs):
|
||||
raise PermissionError("nope")
|
||||
|
||||
monkeypatch.setattr(
|
||||
service_factory, "create_local_artifact_service", _raise_permission_error
|
||||
)
|
||||
|
||||
service = service_factory.create_artifact_service_from_options(
|
||||
base_dir=tmp_path,
|
||||
use_local_storage=True,
|
||||
)
|
||||
|
||||
assert isinstance(service, InMemoryArtifactService)
|
||||
|
||||
Reference in New Issue
Block a user