Skip to content

Commit 06e6fc9

Browse files
GWealecopybara-github
authored andcommitted
feat: wire runtime entrypoints to service factory defaults
This change routes adk run and the FastAPI server through the new session/artifact service factory, keeps the default experience backed by per-agent .adk storage Co-authored-by: George Weale <[email protected]> PiperOrigin-RevId: 836733234
1 parent 5453b5b commit 06e6fc9

File tree

6 files changed

+508
-77
lines changed

6 files changed

+508
-77
lines changed

src/google/adk/cli/cli.py

Lines changed: 57 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,14 @@
1515
from __future__ import annotations
1616

1717
from datetime import datetime
18+
from pathlib import Path
1819
from typing import Optional
1920
from typing import Union
2021

2122
import click
2223
from google.genai import types
2324
from pydantic import BaseModel
2425

25-
from ..agents.base_agent import BaseAgent
2626
from ..agents.llm_agent import LlmAgent
2727
from ..apps.app import App
2828
from ..artifacts.base_artifact_service import BaseArtifactService
@@ -35,8 +35,11 @@
3535
from ..sessions.session import Session
3636
from ..utils.context_utils import Aclosing
3737
from ..utils.env_utils import is_env_enabled
38+
from .service_registry import load_services_module
3839
from .utils import envs
3940
from .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

4245
class 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)

src/google/adk/cli/fast_api.py

Lines changed: 24 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -34,20 +34,19 @@
3434
from starlette.types import Lifespan
3535
from watchdog.observers import Observer
3636

37-
from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
3837
from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
3938
from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
4039
from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
41-
from ..memory.in_memory_memory_service import InMemoryMemoryService
4240
from ..runners import Runner
43-
from ..sessions.in_memory_session_service import InMemorySessionService
4441
from .adk_web_server import AdkWebServer
45-
from .service_registry import get_service_registry
4642
from .service_registry import load_services_module
4743
from .utils import envs
4844
from .utils import evals
4945
from .utils.agent_change_handler import AgentChangeEventHandler
5046
from .utils.agent_loader import AgentLoader
47+
from .utils.service_factory import create_artifact_service_from_options
48+
from .utils.service_factory import create_memory_service_from_options
49+
from .utils.service_factory import create_session_service_from_options
5150

5251
logger = logging.getLogger("google_adk." + __name__)
5352

@@ -74,6 +73,8 @@ def get_fast_api_app(
7473
logo_text: Optional[str] = None,
7574
logo_image_url: Optional[str] = None,
7675
) -> FastAPI:
76+
# Convert to absolute path for consistency
77+
agents_dir = str(Path(agents_dir).resolve())
7778

7879
# Set up eval managers.
7980
if eval_storage_uri:
@@ -91,48 +92,32 @@ def get_fast_api_app(
9192
# Load services.py from agents_dir for custom service registration.
9293
load_services_module(agents_dir)
9394

94-
service_registry = get_service_registry()
95-
9695
# Build the Memory service
97-
if memory_service_uri:
98-
memory_service = service_registry.create_memory_service(
99-
memory_service_uri, agents_dir=agents_dir
96+
try:
97+
memory_service = create_memory_service_from_options(
98+
base_dir=agents_dir,
99+
memory_service_uri=memory_service_uri,
100100
)
101-
if not memory_service:
102-
raise click.ClickException(
103-
"Unsupported memory service URI: %s" % memory_service_uri
104-
)
105-
else:
106-
memory_service = InMemoryMemoryService()
101+
except ValueError as exc:
102+
raise click.ClickException(str(exc)) from exc
107103

108104
# Build the Session service
109-
if session_service_uri:
110-
session_kwargs = session_db_kwargs or {}
111-
session_service = service_registry.create_session_service(
112-
session_service_uri, agents_dir=agents_dir, **session_kwargs
113-
)
114-
if not session_service:
115-
# Fallback to DatabaseSessionService if the service registry doesn't
116-
# support the session service URI scheme.
117-
from ..sessions.database_session_service import DatabaseSessionService
118-
119-
session_service = DatabaseSessionService(
120-
db_url=session_service_uri, **session_kwargs
121-
)
122-
else:
123-
session_service = InMemorySessionService()
105+
session_service = create_session_service_from_options(
106+
base_dir=agents_dir,
107+
session_service_uri=session_service_uri,
108+
session_db_kwargs=session_db_kwargs,
109+
per_agent=True, # Multi-agent mode
110+
)
124111

125112
# Build the Artifact service
126-
if artifact_service_uri:
127-
artifact_service = service_registry.create_artifact_service(
128-
artifact_service_uri, agents_dir=agents_dir
113+
try:
114+
artifact_service = create_artifact_service_from_options(
115+
base_dir=agents_dir,
116+
artifact_service_uri=artifact_service_uri,
117+
per_agent=True, # Multi-agent mode
129118
)
130-
if not artifact_service:
131-
raise click.ClickException(
132-
"Unsupported artifact service URI: %s" % artifact_service_uri
133-
)
134-
else:
135-
artifact_service = InMemoryArtifactService()
119+
except ValueError as exc:
120+
raise click.ClickException(str(exc)) from exc
136121

137122
# Build the Credential service
138123
credential_service = InMemoryCredentialService()

0 commit comments

Comments
 (0)