feat: add memory_service option to CLI

chore: consolidate ADK service CLI options

PiperOrigin-RevId: 769944881

src/google/adk/cli/cli_deploy.py

@@ -55,7 +55,7 @@ COPY "agents/{app_name}/" "/app/agents/{app_name}/"
 
 EXPOSE {port}
 
-CMD adk {command} --port={port} {host_option} {session_db_option} {trace_to_cloud_option} "/app/agents"
+CMD adk {command} --port={port} {host_option} {service_option} {trace_to_cloud_option} "/app/agents"
 """
 
 _AGENT_ENGINE_APP_TEMPLATE = """
@@ -84,6 +84,32 @@ def _resolve_project(project_in_option: Optional[str]) -> str:
   return project
 
 
+def _get_service_option_by_adk_version(
+    adk_version: str,
+    session_uri: Optional[str],
+    artifact_uri: Optional[str],
+    memory_uri: Optional[str],
+) -> str:
+  """Returns service option string based on adk_version."""
+  if adk_version >= '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 adk_version >= '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}'
+  else:
+    return f'--session_db_url={session_uri}' if session_uri else ''
+
+
 def to_cloud_run(
     *,
     agent_folder: str,
@@ -96,9 +122,10 @@ def to_cloud_run(
     trace_to_cloud: bool,
     with_ui: bool,
     verbosity: str,
-    session_db_url: str,
-    artifact_storage_uri: Optional[str],
     adk_version: str,
+    session_service_uri: Optional[str] = None,
+    artifact_service_uri: Optional[str] = None,
+    memory_service_uri: Optional[str] = None,
 ):
   """Deploys an agent to Google Cloud Run.
 
@@ -126,9 +153,10 @@ def to_cloud_run(
     trace_to_cloud: Whether to enable Cloud Trace.
     with_ui: Whether to deploy with UI.
     verbosity: The verbosity level of the CLI.
-    session_db_url: The database URL to connect the session.
-    artifact_storage_uri: The artifact storage URI to store the artifacts.
     adk_version: The ADK version to use in 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.
   """
   app_name = app_name or os.path.basename(agent_folder)
 
@@ -162,12 +190,12 @@ def to_cloud_run(
         port=port,
         command='web' if with_ui else 'api_server',
         install_agent_deps=install_agent_deps,
-        session_db_option=f'--session_db_url={session_db_url}'
+        service_option=_get_service_option_by_adk_version(
-        if session_db_url
+            adk_version,
-        else '',
+            session_service_uri,
-        artifact_storage_option=f'--artifact_storage_uri={artifact_storage_uri}'
+            artifact_service_uri,
-        if artifact_storage_uri
+            memory_service_uri,
-        else '',
+        ),
         trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '',
         adk_version=adk_version,
         host_option=host_option,

src/google/adk/cli/cli_tools_click.py

@@ -417,28 +417,87 @@ def cli_eval(
       print(eval_result.model_dump_json(indent=2))
 
 
-def fast_api_common_options():
+def adk_services_options():
-  """Decorator to add common fast api options to click commands."""
+  """Decorator to add ADK services options to click commands."""
 
   def decorator(func):
     @click.option(
-        "--session_db_url",
+        "--session_service_uri",
         help=(
-            """Optional. The database URL to store the session.
+            """Optional. The URI of the session service.
           - Use 'agentengine://<agent_engine_resource_id>' to connect to Agent Engine sessions.
           - 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 DB URLs."""
+          - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported database URIs."""
         ),
     )
     @click.option(
-        "--artifact_storage_uri",
+        "--artifact_service_uri",
         type=str,
         help=(
-            "Optional. The artifact storage URI to store the artifacts,"
+            "Optional. The URI of the artifact service,"
             " supported URIs: gs://<bucket name> for GCS artifact service."
         ),
         default=None,
     )
+    @click.option(
+        "--memory_service_uri",
+        type=str,
+        help=(
+            """Optional. The URI of the memory service.
+            - Use 'rag://<rag_corpus_id>' to connect to Vertex AI Rag Memory Service."""
+        ),
+        default=None,
+    )
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+      return func(*args, **kwargs)
+
+    return wrapper
+
+  return decorator
+
+
+def deprecated_adk_services_options():
+  """Depracated ADK services options."""
+
+  def warn(alternative_param, ctx, param, value):
+    if value:
+      click.echo(
+          click.style(
+              f"WARNING: Deprecated option {param.name} is used. Please use"
+              f" {alternative_param} instead.",
+              fg="yellow",
+          ),
+          err=True,
+      )
+    return value
+
+  def decorator(func):
+    @click.option(
+        "--session_db_url",
+        help="Deprecated. Use --session_service_uri instead.",
+        callback=functools.partial(warn, "--session_service_uri"),
+    )
+    @click.option(
+        "--artifact_storage_uri",
+        type=str,
+        help="Deprecated. Use --artifact_service_uri instead.",
+        callback=functools.partial(warn, "--artifact_service_uri"),
+        default=None,
+    )
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+      return func(*args, **kwargs)
+
+    return wrapper
+
+  return decorator
+
+
+def fast_api_common_options():
+  """Decorator to add common fast api options to click commands."""
+
+  def decorator(func):
     @click.option(
         "--host",
         type=str,
@@ -489,6 +548,8 @@ def fast_api_common_options():
 
 @main.command("web")
 @fast_api_common_options()
+@adk_services_options()
+@deprecated_adk_services_options()
 @click.argument(
     "agents_dir",
     type=click.Path(
@@ -498,14 +559,17 @@ def fast_api_common_options():
 )
 def cli_web(
     agents_dir: str,
-    session_db_url: str = "",
-    artifact_storage_uri: Optional[str] = None,
     log_level: str = "INFO",
     allow_origins: Optional[list[str]] = None,
     host: str = "127.0.0.1",
     port: int = 8000,
     trace_to_cloud: bool = False,
     reload: bool = True,
+    session_service_uri: Optional[str] = None,
+    artifact_service_uri: Optional[str] = None,
+    memory_service_uri: Optional[str] = None,
+    session_db_url: Optional[str] = None,  # Deprecated
+    artifact_storage_uri: Optional[str] = None,  # Deprecated
 ):
   """Starts a FastAPI server with Web UI for agents.
 
@@ -514,7 +578,7 @@ def cli_web(
 
   Example:
 
-    adk web --session_db_url=[db_url] --port=[port] path/to/agents_dir
+    adk web --session_service_uri=[uri] --port=[port] path/to/agents_dir
   """
   logs.setup_adk_logger(getattr(logging, log_level.upper()))
 
@@ -540,10 +604,13 @@ def cli_web(
         fg="green",
     )
 
+  session_service_uri = session_service_uri or session_db_url
+  artifact_service_uri = artifact_service_uri or artifact_storage_uri
   app = get_fast_api_app(
       agents_dir=agents_dir,
-      session_db_url=session_db_url,
+      session_service_uri=session_service_uri,
-      artifact_storage_uri=artifact_storage_uri,
+      artifact_service_uri=artifact_service_uri,
+      memory_service_uri=memory_service_uri,
       allow_origins=allow_origins,
       web=True,
       trace_to_cloud=trace_to_cloud,
@@ -571,16 +638,21 @@ def cli_web(
     default=os.getcwd(),
 )
 @fast_api_common_options()
+@adk_services_options()
+@deprecated_adk_services_options()
 def cli_api_server(
     agents_dir: str,
-    session_db_url: str = "",
-    artifact_storage_uri: Optional[str] = None,
     log_level: str = "INFO",
     allow_origins: Optional[list[str]] = None,
     host: str = "127.0.0.1",
     port: int = 8000,
     trace_to_cloud: bool = False,
     reload: bool = True,
+    session_service_uri: Optional[str] = None,
+    artifact_service_uri: Optional[str] = None,
+    memory_service_uri: Optional[str] = None,
+    session_db_url: Optional[str] = None,  # Deprecated
+    artifact_storage_uri: Optional[str] = None,  # Deprecated
 ):
   """Starts a FastAPI server for agents.
 
@@ -589,15 +661,18 @@ def cli_api_server(
 
   Example:
 
-    adk api_server --session_db_url=[db_url] --port=[port] path/to/agents_dir
+    adk api_server --session_service_uri=[uri] --port=[port] path/to/agents_dir
   """
   logs.setup_adk_logger(getattr(logging, log_level.upper()))
 
+  session_service_uri = session_service_uri or session_db_url
+  artifact_service_uri = artifact_service_uri or artifact_storage_uri
   config = uvicorn.Config(
       get_fast_api_app(
           agents_dir=agents_dir,
-          session_db_url=session_db_url,
+          session_service_uri=session_service_uri,
-          artifact_storage_uri=artifact_storage_uri,
+          artifact_service_uri=artifact_service_uri,
+          memory_service_uri=memory_service_uri,
           allow_origins=allow_origins,
           web=False,
           trace_to_cloud=trace_to_cloud,
@@ -689,27 +764,6 @@ def cli_api_server(
     default="WARNING",
     help="Optional. Override the default verbosity level.",
 )
-@click.option(
-    "--session_db_url",
-    help=(
-        """Optional. The database URL to store the session.
-
-  - Use 'agentengine://<agent_engine_resource_id>' to connect to Agent Engine sessions.
-
-  - 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 DB URLs."""
-    ),
-)
-@click.option(
-    "--artifact_storage_uri",
-    type=str,
-    help=(
-        "Optional. The artifact storage URI to store the artifacts, supported"
-        " URIs: gs://<bucket name> for GCS artifact service."
-    ),
-    default=None,
-)
 @click.argument(
     "agent",
     type=click.Path(
@@ -726,6 +780,8 @@ def cli_api_server(
         " version in the dev environment)"
     ),
 )
+@adk_services_options()
+@deprecated_adk_services_options()
 def cli_deploy_cloud_run(
     agent: str,
     project: Optional[str],
@@ -737,9 +793,12 @@ def cli_deploy_cloud_run(
     trace_to_cloud: bool,
     with_ui: bool,
     verbosity: str,
-    session_db_url: str,
-    artifact_storage_uri: Optional[str],
     adk_version: str,
+    session_service_uri: Optional[str] = None,
+    artifact_service_uri: Optional[str] = None,
+    memory_service_uri: Optional[str] = None,
+    session_db_url: Optional[str] = None,  # Deprecated
+    artifact_storage_uri: Optional[str] = None,  # Deprecated
 ):
   """Deploys an agent to Cloud Run.
 
@@ -749,6 +808,8 @@ def cli_deploy_cloud_run(
 
     adk deploy cloud_run --project=[project] --region=[region] path/to/my_agent
   """
+  session_service_uri = session_service_uri or session_db_url
+  artifact_service_uri = artifact_service_uri or artifact_storage_uri
   try:
     cli_deploy.to_cloud_run(
         agent_folder=agent,
@@ -761,9 +822,10 @@ def cli_deploy_cloud_run(
         trace_to_cloud=trace_to_cloud,
         with_ui=with_ui,
         verbosity=verbosity,
-        session_db_url=session_db_url,
-        artifact_storage_uri=artifact_storage_uri,
         adk_version=adk_version,
+        session_service_uri=session_service_uri,
+        artifact_service_uri=artifact_service_uri,
+        memory_service_uri=memory_service_uri,
     )
   except Exception as e:
     click.secho(f"Deploy failed: {e}", fg="red", err=True)

src/google/adk/cli/fast_api.py

@@ -68,6 +68,7 @@ from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManag
 from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
 from ..events.event import Event
 from ..memory.in_memory_memory_service import InMemoryMemoryService
+from ..memory.vertex_ai_rag_memory_service import VertexAiRagMemoryService
 from ..runners import Runner
 from ..sessions.database_session_service import DatabaseSessionService
 from ..sessions.in_memory_session_service import InMemorySessionService
@@ -193,8 +194,9 @@ class GetEventGraphResult(common.BaseModel):
 def get_fast_api_app(
     *,
     agents_dir: str,
-    session_db_url: str = "",
+    session_service_uri: Optional[str] = None,
-    artifact_storage_uri: Optional[str] = None,
+    artifact_service_uri: Optional[str] = None,
+    memory_service_uri: Optional[str] = None,
     allow_origins: Optional[list[str]] = None,
     web: bool,
     trace_to_cloud: bool = False,
@@ -257,14 +259,28 @@ def get_fast_api_app(
   eval_set_results_manager = LocalEvalSetResultsManager(agents_dir=agents_dir)
 
   # Build the Memory service
-  memory_service = InMemoryMemoryService()
+  if memory_service_uri:
+    if memory_service_uri.startswith("rag://"):
+      rag_corpus = memory_service_uri.split("://")[1]
+      if not rag_corpus:
+        raise click.ClickException("Rag corpus can not be empty.")
+      envs.load_dotenv_for_agent("", agents_dir)
+      memory_service = VertexAiRagMemoryService(
+          rag_corpus=f'projects/{os.environ["GOOGLE_CLOUD_PROJECT"]}/locations/{os.environ["GOOGLE_CLOUD_LOCATION"]}/ragCorpora/{rag_corpus}'
+      )
+    else:
+      raise click.ClickException(
+          "Unsupported memory service URI: %s" % memory_service_uri
+      )
+  else:
+    memory_service = InMemoryMemoryService()
 
   # Build the Session service
   agent_engine_id = ""
-  if session_db_url:
+  if session_service_uri:
-    if session_db_url.startswith("agentengine://"):
+    if session_service_uri.startswith("agentengine://"):
       # Create vertex session service
-      agent_engine_id = session_db_url.split("://")[1]
+      agent_engine_id = session_service_uri.split("://")[1]
       if not agent_engine_id:
         raise click.ClickException("Agent engine id can not be empty.")
       envs.load_dotenv_for_agent("", agents_dir)
@@ -273,18 +289,18 @@ def get_fast_api_app(
           os.environ["GOOGLE_CLOUD_LOCATION"],
       )
     else:
-      session_service = DatabaseSessionService(db_url=session_db_url)
+      session_service = DatabaseSessionService(db_url=session_service_uri)
   else:
     session_service = InMemorySessionService()
 
   # Build the Artifact service
-  if artifact_storage_uri:
+  if artifact_service_uri:
-    if artifact_storage_uri.startswith("gs://"):
+    if artifact_service_uri.startswith("gs://"):
-      gcs_bucket = artifact_storage_uri.split("://")[1]
+      gcs_bucket = artifact_service_uri.split("://")[1]
       artifact_service = GcsArtifactService(bucket_name=gcs_bucket)
     else:
       raise click.ClickException(
-          "Unsupported artifact storage URI: %s" % artifact_storage_uri
+          "Unsupported artifact service URI: %s" % artifact_service_uri
       )
   else:
     artifact_service = InMemoryArtifactService()

tests/unittests/cli/test_fast_api.py

@@ -458,7 +458,12 @@ def test_app(
   ):
     # Get the FastAPI app, but don't actually run it
     app = get_fast_api_app(
-        agents_dir=".", web=True, session_db_url="", allow_origins=["*"]
+        agents_dir=".",
+        web=True,
+        session_service_uri="",
+        artifact_service_uri="",
+        memory_service_uri="",
+        allow_origins=["*"],
     )
 
     # Create a TestClient that doesn't start a real server

tests/unittests/cli/utils/test_cli_deploy.py

@@ -87,6 +87,41 @@ def test_resolve_project_from_gcloud(monkeypatch: pytest.MonkeyPatch) -> None:
     mocked_echo.assert_called_once()
 
 
+# _get_service_option_by_adk_version
+def test_get_service_option_by_adk_version() -> None:
+  """It should return the explicit project value untouched."""
+  assert cli_deploy._get_service_option_by_adk_version(
+      adk_version="1.3.0",
+      session_uri="sqlite://",
+      artifact_uri="gs://bucket",
+      memory_uri="rag://",
+  ) == (
+      "--session_service_uri=sqlite:// "
+      "--artifact_service_uri=gs://bucket "
+      "--memory_service_uri=rag://"
+  )
+
+  assert (
+      cli_deploy._get_service_option_by_adk_version(
+          adk_version="1.2.0",
+          session_uri="sqlite://",
+          artifact_uri="gs://bucket",
+          memory_uri="rag://",
+      )
+      == "--session_db_url=sqlite:// --artifact_storage_uri=gs://bucket"
+  )
+
+  assert (
+      cli_deploy._get_service_option_by_adk_version(
+          adk_version="0.5.0",
+          session_uri="sqlite://",
+          artifact_uri="gs://bucket",
+          memory_uri="rag://",
+      )
+      == "--session_db_url=sqlite://"
+  )
+
+
 # to_cloud_run
 @pytest.mark.parametrize("include_requirements", [True, False])
 def test_to_cloud_run_happy_path(
@@ -127,8 +162,9 @@ def test_to_cloud_run_happy_path(
       trace_to_cloud=True,
       with_ui=True,
       verbosity="info",
-      session_db_url="sqlite://",
+      session_service_uri="sqlite://",
-      artifact_storage_uri="gs://bucket",
+      artifact_service_uri="gs://bucket",
+      memory_service_uri="rag://",
       adk_version="0.0.5",
   )
 
@@ -170,9 +206,10 @@ def test_to_cloud_run_cleans_temp_dir(
       trace_to_cloud=False,
       with_ui=False,
       verbosity="info",
-      session_db_url=None,
+      adk_version="1.0.0",
-      artifact_storage_uri=None,
+      session_service_uri=None,
-      adk_version="0.0.5",
+      artifact_service_uri=None,
+      memory_service_uri=None,
   )
 
   assert deleted["path"] == tmp_dir