feat: Add --session_id option to adk run:

--session_id : The session ID to save the session to on exit when --save_session is set to true. User will be prompted to enter a session ID if not set.

PiperOrigin-RevId: 756335619
This commit is contained in:
Google Team Member 2025-05-08 09:06:38 -07:00 committed by Copybara-Service
parent 41564de689
commit a61d20e3df
2 changed files with 14 additions and 1 deletions

View File

@ -105,6 +105,7 @@ async def run_cli(
input_file: Optional[str] = None, input_file: Optional[str] = None,
saved_session_file: Optional[str] = None, saved_session_file: Optional[str] = None,
save_session: bool, save_session: bool,
session_id: Optional[str] = None,
) -> None: ) -> None:
"""Runs an interactive CLI for a certain agent. """Runs an interactive CLI for a certain agent.
@ -118,6 +119,7 @@ async def run_cli(
saved_session_file: Optional[str], the absolute path to the json file that saved_session_file: Optional[str], the absolute path to the json file that
contains a previously saved session, exclusive with input_file. contains a previously saved session, exclusive with input_file.
save_session: bool, whether to save the session on exit. save_session: bool, whether to save the session on exit.
session_id: Optional[str], the session ID to save the session to on exit.
""" """
if agent_parent_dir not in sys.path: if agent_parent_dir not in sys.path:
sys.path.append(agent_parent_dir) sys.path.append(agent_parent_dir)
@ -175,7 +177,7 @@ async def run_cli(
) )
if save_session: if save_session:
session_id = input('Session ID to save: ') session_id = session_id or input('Session ID to save: ')
session_path = f'{agent_module_path}/{session_id}.session.json' session_path = f'{agent_module_path}/{session_id}.session.json'
# Fetch the session again to get all the details. # Fetch the session again to get all the details.

View File

@ -122,6 +122,15 @@ def validate_exclusive(ctx, param, value):
default=False, default=False,
help="Optional. Whether to save the session to a json file on exit.", help="Optional. Whether to save the session to a json file on exit.",
) )
@click.option(
"--session_id",
type=str,
help=(
"Optional. The session ID to save the session to on exit when"
" --save_session is set to true. User will be prompted to enter a"
" session ID if not set."
),
)
@click.option( @click.option(
"--replay", "--replay",
type=click.Path( type=click.Path(
@ -156,6 +165,7 @@ def validate_exclusive(ctx, param, value):
def cli_run( def cli_run(
agent: str, agent: str,
save_session: bool, save_session: bool,
session_id: Optional[str],
replay: Optional[str], replay: Optional[str],
resume: Optional[str], resume: Optional[str],
): ):
@ -179,6 +189,7 @@ def cli_run(
input_file=replay, input_file=replay,
saved_session_file=resume, saved_session_file=resume,
save_session=save_session, save_session=save_session,
session_id=session_id,
) )
) )