1515from __future__ import annotations
1616
1717from datetime import datetime
18+ from pathlib import Path
1819from typing import Optional
1920from typing import Union
2021
2122import click
2223from google .genai import types
2324from pydantic import BaseModel
2425
25- from ..agents .base_agent import BaseAgent
2626from ..agents .llm_agent import LlmAgent
2727from ..apps .app import App
2828from ..artifacts .base_artifact_service import BaseArtifactService
3535from ..sessions .session import Session
3636from ..utils .context_utils import Aclosing
3737from ..utils .env_utils import is_env_enabled
38+ from .service_registry import load_services_module
3839from .utils import envs
3940from .utils .agent_loader import AgentLoader
41+ from .utils .service_factory import create_artifact_service_from_options
42+ from .utils .service_factory import create_session_service_from_options
4043
4144
4245class InputFile (BaseModel ):
@@ -66,7 +69,7 @@ async def run_input_file(
6669 )
6770 with open (input_path , 'r' , encoding = 'utf-8' ) as f :
6871 input_file = InputFile .model_validate_json (f .read ())
69- input_file .state ['_time' ] = datetime .now ()
72+ input_file .state ['_time' ] = datetime .now (). isoformat ()
7073
7174 session = await session_service .create_session (
7275 app_name = app_name , user_id = user_id , state = input_file .state
@@ -134,6 +137,8 @@ async def run_cli(
134137 saved_session_file : Optional [str ] = None ,
135138 save_session : bool ,
136139 session_id : Optional [str ] = None ,
140+ session_service_uri : Optional [str ] = None ,
141+ artifact_service_uri : Optional [str ] = None ,
137142) -> None :
138143 """Runs an interactive CLI for a certain agent.
139144
@@ -148,24 +153,47 @@ async def run_cli(
148153 contains a previously saved session, exclusive with input_file.
149154 save_session: bool, whether to save the session on exit.
150155 session_id: Optional[str], the session ID to save the session to on exit.
156+ session_service_uri: Optional[str], custom session service URI.
157+ artifact_service_uri: Optional[str], custom artifact service URI.
151158 """
159+ agent_parent_path = Path (agent_parent_dir ).resolve ()
160+ agent_root = agent_parent_path / agent_folder_name
161+ load_services_module (str (agent_root ))
162+ user_id = 'test_user'
152163
153- artifact_service = InMemoryArtifactService ()
154- session_service = InMemorySessionService ()
155- credential_service = InMemoryCredentialService ()
164+ # Create session and artifact services using factory functions
165+ session_service = create_session_service_from_options (
166+ base_dir = agent_root ,
167+ session_service_uri = session_service_uri ,
168+ )
156169
157- user_id = 'test_user'
158- agent_or_app = AgentLoader (agents_dir = agent_parent_dir ).load_agent (
170+ artifact_service = create_artifact_service_from_options (
171+ base_dir = agent_root ,
172+ artifact_service_uri = artifact_service_uri ,
173+ )
174+
175+ credential_service = InMemoryCredentialService ()
176+ agents_dir = str (agent_parent_path )
177+ agent_or_app = AgentLoader (agents_dir = agents_dir ).load_agent (
159178 agent_folder_name
160179 )
161180 session_app_name = (
162181 agent_or_app .name if isinstance (agent_or_app , App ) else agent_folder_name
163182 )
164- session = await session_service .create_session (
165- app_name = session_app_name , user_id = user_id
166- )
167183 if not is_env_enabled ('ADK_DISABLE_LOAD_DOTENV' ):
168- envs .load_dotenv_for_agent (agent_folder_name , agent_parent_dir )
184+ envs .load_dotenv_for_agent (agent_folder_name , agents_dir )
185+
186+ # Helper function for printing events
187+ def _print_event (event ) -> None :
188+ content = event .content
189+ if not content or not content .parts :
190+ return
191+ text_parts = [part .text for part in content .parts if part .text ]
192+ if not text_parts :
193+ return
194+ author = event .author or 'system'
195+ click .echo (f'[{ author } ]: { "" .join (text_parts )} ' )
196+
169197 if input_file :
170198 session = await run_input_file (
171199 app_name = session_app_name ,
@@ -177,16 +205,22 @@ async def run_cli(
177205 input_path = input_file ,
178206 )
179207 elif saved_session_file :
208+ # Load the saved session from file
180209 with open (saved_session_file , 'r' , encoding = 'utf-8' ) as f :
181210 loaded_session = Session .model_validate_json (f .read ())
182211
212+ # Create a new session in the service, copying state from the file
213+ session = await session_service .create_session (
214+ app_name = session_app_name ,
215+ user_id = user_id ,
216+ state = loaded_session .state if loaded_session else None ,
217+ )
218+
219+ # Append events from the file to the new session and display them
183220 if loaded_session :
184221 for event in loaded_session .events :
185222 await session_service .append_event (session , event )
186- content = event .content
187- if not content or not content .parts or not content .parts [0 ].text :
188- continue
189- click .echo (f'[{ event .author } ]: { content .parts [0 ].text } ' )
223+ _print_event (event )
190224
191225 await run_interactively (
192226 agent_or_app ,
@@ -196,6 +230,9 @@ async def run_cli(
196230 credential_service ,
197231 )
198232 else :
233+ session = await session_service .create_session (
234+ app_name = session_app_name , user_id = user_id
235+ )
199236 click .echo (f'Running agent { agent_or_app .name } , type exit to exit.' )
200237 await run_interactively (
201238 agent_or_app ,
@@ -207,19 +244,17 @@ async def run_cli(
207244
208245 if save_session :
209246 session_id = session_id or input ('Session ID to save: ' )
210- session_path = (
211- f'{ agent_parent_dir } /{ agent_folder_name } /{ session_id } .session.json'
212- )
247+ session_path = agent_root / f'{ session_id } .session.json'
213248
214249 # Fetch the session again to get all the details.
215250 session = await session_service .get_session (
216251 app_name = session .app_name ,
217252 user_id = session .user_id ,
218253 session_id = session .id ,
219254 )
220- with open ( session_path , 'w' , encoding = 'utf-8' ) as f :
221- f . write (
222- session . model_dump_json ( indent = 2 , exclude_none = True , by_alias = True )
223- )
255+ session_path . write_text (
256+ session . model_dump_json ( indent = 2 , exclude_none = True , by_alias = True ),
257+ encoding = 'utf-8' ,
258+ )
224259
225260 print ('Session saved to' , session_path )
0 commit comments