google-adk 1.2.1__py3-none-any.whl → 1.4.0__py3-none-any.whl

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():
-  """Decorator to add common fast api options to click commands."""
+def adk_services_options():
+  """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,
-      artifact_storage_uri=artifact_storage_uri,
+      session_service_uri=session_service_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,
-          artifact_storage_uri=artifact_storage_uri,
+          session_service_uri=session_service_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)
@@ -773,12 +835,18 @@ def cli_deploy_cloud_run(
 @click.option(
     "--project",
     type=str,
-    help="Required. Google Cloud project to deploy the agent.",
+    help=(
+        "Required. Google Cloud project to deploy the agent. It will override"
+        " GOOGLE_CLOUD_PROJECT in the .env file (if it exists)."
+    ),
 )
 @click.option(
     "--region",
     type=str,
-    help="Required. Google Cloud region to deploy the agent.",
+    help=(
+        "Required. Google Cloud region to deploy the agent. It will override"
+        " GOOGLE_CLOUD_LOCATION in the .env file (if it exists)."
+    ),
 )
 @click.option(
     "--staging_bucket",
@@ -793,6 +861,20 @@ def cli_deploy_cloud_run(
     default=False,
     help="Optional. Whether to enable Cloud Trace for Agent Engine.",
 )
+@click.option(
+    "--display_name",
+    type=str,
+    show_default=True,
+    default="",
+    help="Optional. Display name of the agent in Agent Engine.",
+)
+@click.option(
+    "--description",
+    type=str,
+    show_default=True,
+    default="",
+    help="Optional. Description of the agent in Agent Engine.",
+)
 @click.option(
     "--adk_app",
     type=str,
@@ -847,6 +929,8 @@ def cli_deploy_agent_engine(
     region: str,
     staging_bucket: str,
     trace_to_cloud: bool,
+    display_name: str,
+    description: str,
     adk_app: str,
     temp_folder: str,
     env_file: str,
@@ -854,27 +938,12 @@ def cli_deploy_agent_engine(
 ):
   """Deploys an agent to Agent Engine.
 
-  Args:
-    agent (str): Required. The path to the agent to be deloyed.
-    project (str): Required. Google Cloud project to deploy the agent.
-    region (str): Required. Google Cloud region to deploy the agent.
-    staging_bucket (str): Required. GCS bucket for staging the deployment
-      artifacts.
-    trace_to_cloud (bool): Required. Whether to enable Cloud Trace.
-    adk_app (str): Required. Python file for defining the ADK application.
-    temp_folder (str): Required. The folder for the generated Agent Engine
-      files. If the folder already exists, its contents will be replaced.
-    env_file (str): Required. The filepath to the `.env` file for environment
-      variables. If it is an empty string, the `.env` file in the `agent`
-      directory will be used if it exists.
-    requirements_file (str): Required. The filepath to the `requirements.txt`
-      file to use. If it is an empty string, the `requirements.txt` file in the
-      `agent` directory will be used if exists.
+  AGENT: The path to the agent source code folder.
 
   Example:
 
     adk deploy agent_engine --project=[project] --region=[region]
-      --staging_bucket=[staging_bucket] path/to/my_agent
+      --staging_bucket=[staging_bucket] --display_name=[app_name] path/to/my_agent
   """
   try:
     cli_deploy.to_agent_engine(
@@ -883,6 +952,8 @@ def cli_deploy_agent_engine(
         region=region,
         staging_bucket=staging_bucket,
         trace_to_cloud=trace_to_cloud,
+        display_name=display_name,
+        description=description,
         adk_app=adk_app,
         temp_folder=temp_folder,
         env_file=env_file,

google/adk/cli/fast_api.py

@@ -57,6 +57,7 @@ from ..agents.llm_agent import Agent
 from ..agents.run_config import StreamingMode
 from ..artifacts.gcs_artifact_service import GcsArtifactService
 from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
+from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
 from ..errors.not_found_error import NotFoundError
 from ..evaluation.eval_case import EvalCase
 from ..evaluation.eval_case import SessionInput
@@ -68,6 +69,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 +195,9 @@ class GetEventGraphResult(common.BaseModel):
 def get_fast_api_app(
     *,
     agents_dir: str,
-    session_db_url: str = "",
-    artifact_storage_uri: Optional[str] = None,
+    session_service_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,38 +260,55 @@ 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_db_url.startswith("agentengine://"):
+  if session_service_uri:
+    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)
       session_service = VertexAiSessionService(
-          os.environ["GOOGLE_CLOUD_PROJECT"],
-          os.environ["GOOGLE_CLOUD_LOCATION"],
+          project=os.environ["GOOGLE_CLOUD_PROJECT"],
+          location=os.environ["GOOGLE_CLOUD_LOCATION"],
+          agent_engine_id=agent_engine_id,
       )
     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_storage_uri.startswith("gs://"):
-      gcs_bucket = artifact_storage_uri.split("://")[1]
+  if artifact_service_uri:
+    if artifact_service_uri.startswith("gs://"):
+      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()
 
+  # Build  the Credential service
+  credential_service = InMemoryCredentialService()
+
   # initialize Agent Loader
   agent_loader = AgentLoader(agents_dir)
 
@@ -341,8 +361,6 @@ def get_fast_api_app(
   async def get_session(
       app_name: str, user_id: str, session_id: str
   ) -> Session:
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else app_name
     session = await session_service.get_session(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
@@ -355,8 +373,6 @@ def get_fast_api_app(
       response_model_exclude_none=True,
   )
   async def list_sessions(app_name: str, user_id: str) -> list[Session]:
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else app_name
     list_sessions_response = await session_service.list_sessions(
         app_name=app_name, user_id=user_id
     )
@@ -377,8 +393,6 @@ def get_fast_api_app(
       session_id: str,
       state: Optional[dict[str, Any]] = None,
   ) -> Session:
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else app_name
     if (
         await session_service.get_session(
             app_name=app_name, user_id=user_id, session_id=session_id
@@ -402,14 +416,19 @@ def get_fast_api_app(
       app_name: str,
       user_id: str,
       state: Optional[dict[str, Any]] = None,
+      events: Optional[list[Event]] = None,
   ) -> Session:
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else app_name
     logger.info("New session created")
-    return await session_service.create_session(
+    session = await session_service.create_session(
         app_name=app_name, user_id=user_id, state=state
     )
 
+    if events:
+      for event in events:
+        await session_service.append_event(session=session, event=event)
+
+    return session
+
   def _get_eval_set_file_path(app_name, agents_dir, eval_set_id) -> str:
     return os.path.join(
         agents_dir,
@@ -644,8 +663,6 @@ def get_fast_api_app(
 
   @app.delete("/apps/{app_name}/users/{user_id}/sessions/{session_id}")
   async def delete_session(app_name: str, user_id: str, session_id: str):
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else app_name
     await session_service.delete_session(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
@@ -661,7 +678,6 @@ def get_fast_api_app(
       artifact_name: str,
       version: Optional[int] = Query(None),
   ) -> Optional[types.Part]:
-    app_name = agent_engine_id if agent_engine_id else app_name
     artifact = await artifact_service.load_artifact(
         app_name=app_name,
         user_id=user_id,
@@ -684,7 +700,6 @@ def get_fast_api_app(
       artifact_name: str,
       version_id: int,
   ) -> Optional[types.Part]:
-    app_name = agent_engine_id if agent_engine_id else app_name
     artifact = await artifact_service.load_artifact(
         app_name=app_name,
         user_id=user_id,
@@ -703,7 +718,6 @@ def get_fast_api_app(
   async def list_artifact_names(
       app_name: str, user_id: str, session_id: str
   ) -> list[str]:
-    app_name = agent_engine_id if agent_engine_id else app_name
     return await artifact_service.list_artifact_keys(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
@@ -715,7 +729,6 @@ def get_fast_api_app(
   async def list_artifact_versions(
       app_name: str, user_id: str, session_id: str, artifact_name: str
   ) -> list[int]:
-    app_name = agent_engine_id if agent_engine_id else app_name
     return await artifact_service.list_versions(
         app_name=app_name,
         user_id=user_id,
@@ -729,7 +742,6 @@ def get_fast_api_app(
   async def delete_artifact(
       app_name: str, user_id: str, session_id: str, artifact_name: str
   ):
-    app_name = agent_engine_id if agent_engine_id else app_name
     await artifact_service.delete_artifact(
         app_name=app_name,
         user_id=user_id,
@@ -739,10 +751,8 @@ def get_fast_api_app(
 
   @app.post("/run", response_model_exclude_none=True)
   async def agent_run(req: AgentRunRequest) -> list[Event]:
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else req.app_name
     session = await session_service.get_session(
-        app_name=app_name, user_id=req.user_id, session_id=req.session_id
+        app_name=req.app_name, user_id=req.user_id, session_id=req.session_id
     )
     if not session:
       raise HTTPException(status_code=404, detail="Session not found")
@@ -760,11 +770,9 @@ def get_fast_api_app(
 
   @app.post("/run_sse")
   async def agent_run_sse(req: AgentRunRequest) -> StreamingResponse:
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else req.app_name
     # SSE endpoint
     session = await session_service.get_session(
-        app_name=app_name, user_id=req.user_id, session_id=req.session_id
+        app_name=req.app_name, user_id=req.user_id, session_id=req.session_id
     )
     if not session:
       raise HTTPException(status_code=404, detail="Session not found")
@@ -802,8 +810,6 @@ def get_fast_api_app(
   async def get_event_graph(
       app_name: str, user_id: str, session_id: str, event_id: str
   ):
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else app_name
     session = await session_service.get_session(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
@@ -859,8 +865,6 @@ def get_fast_api_app(
   ) -> None:
     await websocket.accept()
 
-    # Connect to managed session if agent_engine_id is set.
-    app_name = agent_engine_id if agent_engine_id else app_name
     session = await session_service.get_session(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
@@ -924,11 +928,12 @@ def get_fast_api_app(
       return runner_dict[app_name]
     root_agent = agent_loader.load_agent(app_name)
     runner = Runner(
-        app_name=agent_engine_id if agent_engine_id else app_name,
+        app_name=app_name,
         agent=root_agent,
         artifact_service=artifact_service,
         session_service=session_service,
         memory_service=memory_service,
+        credential_service=credential_service,
     )
     runner_dict[app_name] = runner
     return runner

google/adk/cli/utils/agent_loader.py

@@ -76,7 +76,12 @@ class AgentLoader:
         e.msg = f"Fail to load '{agent_name}' module. " + e.msg
         raise e
     except Exception as e:
-      e.msg = f"Fail to load '{agent_name}' module. " + e.msg
+      if hasattr(e, "msg"):
+        e.msg = f"Fail to load '{agent_name}' module. " + e.msg
+        raise e
+      e.args = (
+          f"Fail to load '{agent_name}' module. {e.args[0] if e.args else ''}",
+      ) + e.args[1:]
       raise e
 
     return None
@@ -110,7 +115,15 @@ class AgentLoader:
         e.msg = f"Fail to load '{agent_name}.agent' module. " + e.msg
         raise e
     except Exception as e:
-      e.msg = f"Fail to load '{agent_name}.agent' module. " + e.msg
+      if hasattr(e, "msg"):
+        e.msg = f"Fail to load '{agent_name}.agent' module. " + e.msg
+        raise e
+      e.args = (
+          (
+              f"Fail to load '{agent_name}.agent' module."
+              f" {e.args[0] if e.args else ''}"
+          ),
+      ) + e.args[1:]
       raise e
 
     return None

google/adk/code_executors/container_code_executor.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import atexit
+import logging
 import os
 from typing import Optional
 
@@ -27,6 +30,7 @@ from .base_code_executor import BaseCodeExecutor
 from .code_execution_utils import CodeExecutionInput
 from .code_execution_utils import CodeExecutionResult
 
+logger = logging.getLogger('google_adk.' + __name__)
 DEFAULT_IMAGE_TAG = 'adk-code-executor:latest'
 
 
@@ -151,13 +155,13 @@ class ContainerCodeExecutor(BaseCodeExecutor):
     if not os.path.exists(self.docker_path):
       raise FileNotFoundError(f'Invalid Docker path: {self.docker_path}')
 
-    print('Building Docker image...')
+    logger.info('Building Docker image...')
     self._client.images.build(
         path=self.docker_path,
         tag=self.image,
         rm=True,
     )
-    print(f'Docker image: {self.image} built.')
+    logger.info('Docker image: %s built.', self.image)
 
   def _verify_python_installation(self):
     """Verifies the container has python3 installed."""
@@ -173,13 +177,13 @@ class ContainerCodeExecutor(BaseCodeExecutor):
     if self.docker_path:
       self._build_docker_image()
 
-    print('Starting container for ContainerCodeExecutor...')
+    logger.info('Starting container for ContainerCodeExecutor...')
     self._container = self._client.containers.run(
         image=self.image,
         detach=True,
         tty=True,
     )
-    print(f'Container {self._container.id} started.')
+    logger.info('Container %s started.', self._container.id)
 
     # Verify the container is able to run python3.
     self._verify_python_installation()
@@ -189,7 +193,7 @@ class ContainerCodeExecutor(BaseCodeExecutor):
     if not self._container:
       return
 
-    print('[Cleanup] Stopping the container...')
+    logger.info('[Cleanup] Stopping the container...')
     self._container.stop()
     self._container.remove()
-    print(f'Container {self._container.id} stopped and removed.')
+    logger.info('Container %s stopped and removed.', self._container.id)