google-adk 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. google/adk/agents/base_agent.py +4 -4
  2. google/adk/agents/invocation_context.py +1 -1
  3. google/adk/agents/remote_agent.py +1 -1
  4. google/adk/agents/run_config.py +1 -1
  5. google/adk/auth/auth_preprocessor.py +2 -2
  6. google/adk/auth/auth_tool.py +1 -1
  7. google/adk/cli/browser/index.html +2 -2
  8. google/adk/cli/browser/{main-SLIAU2JL.js → main-ZBO76GRM.js} +65 -81
  9. google/adk/cli/cli_create.py +279 -0
  10. google/adk/cli/cli_deploy.py +4 -0
  11. google/adk/cli/cli_eval.py +2 -2
  12. google/adk/cli/cli_tools_click.py +67 -7
  13. google/adk/cli/fast_api.py +51 -16
  14. google/adk/cli/utils/envs.py +0 -3
  15. google/adk/cli/utils/evals.py +2 -2
  16. google/adk/evaluation/evaluation_generator.py +4 -4
  17. google/adk/evaluation/response_evaluator.py +15 -3
  18. google/adk/events/event.py +3 -3
  19. google/adk/flows/llm_flows/_nl_planning.py +10 -4
  20. google/adk/flows/llm_flows/contents.py +1 -1
  21. google/adk/models/lite_llm.py +51 -34
  22. google/adk/planners/plan_re_act_planner.py +2 -2
  23. google/adk/runners.py +1 -1
  24. google/adk/sessions/database_session_service.py +84 -23
  25. google/adk/sessions/state.py +1 -1
  26. google/adk/telemetry.py +2 -2
  27. google/adk/tools/application_integration_tool/clients/integration_client.py +3 -2
  28. google/adk/tools/base_tool.py +1 -1
  29. google/adk/tools/function_parameter_parse_util.py +2 -2
  30. google/adk/tools/google_api_tool/__init__.py +74 -1
  31. google/adk/tools/google_api_tool/google_api_tool_sets.py +91 -34
  32. google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py +3 -1
  33. google/adk/tools/load_memory_tool.py +25 -2
  34. google/adk/tools/mcp_tool/mcp_session_manager.py +176 -0
  35. google/adk/tools/mcp_tool/mcp_tool.py +15 -2
  36. google/adk/tools/mcp_tool/mcp_toolset.py +31 -37
  37. google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py +1 -1
  38. google/adk/tools/toolbox_tool.py +1 -1
  39. google/adk/version.py +1 -1
  40. google_adk-0.2.0.dist-info/METADATA +212 -0
  41. {google_adk-0.1.0.dist-info → google_adk-0.2.0.dist-info}/RECORD +44 -42
  42. google_adk-0.1.0.dist-info/METADATA +0 -160
  43. {google_adk-0.1.0.dist-info → google_adk-0.2.0.dist-info}/WHEEL +0 -0
  44. {google_adk-0.1.0.dist-info → google_adk-0.2.0.dist-info}/entry_points.txt +0 -0
  45. {google_adk-0.1.0.dist-info → google_adk-0.2.0.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,279 @@
1
+ # Copyright 2025 Google LLC
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import os
16
+ import subprocess
17
+ from typing import Optional
18
+ from typing import Tuple
19
+
20
+ import click
21
+
22
+ _INIT_PY_TEMPLATE = """\
23
+ from . import agent
24
+ """
25
+
26
+ _AGENT_PY_TEMPLATE = """\
27
+ from google.adk.agents import Agent
28
+
29
+ root_agent = Agent(
30
+ model='{model_name}',
31
+ name='root_agent',
32
+ description='A helpful assistant for user questions.',
33
+ instruction='Answer user questions to the best of your knowledge',
34
+ )
35
+ """
36
+
37
+
38
+ _GOOGLE_API_MSG = """
39
+ Don't have API Key? Create one in AI Studio: https://aistudio.google.com/apikey
40
+ """
41
+
42
+ _GOOGLE_CLOUD_SETUP_MSG = """
43
+ You need an existing Google Cloud account and project, check out this link for details:
44
+ https://google.github.io/adk-docs/get-started/quickstart/#gemini---google-cloud-vertex-ai
45
+ """
46
+
47
+ _OTHER_MODEL_MSG = """
48
+ Please see below guide to configure other models:
49
+ https://google.github.io/adk-docs/agents/models
50
+ """
51
+
52
+ _SUCCESS_MSG = """
53
+ Agent created in {agent_folder}:
54
+ - .env
55
+ - __init__.py
56
+ - agent.py
57
+ """
58
+
59
+
60
+ def _get_gcp_project_from_gcloud() -> str:
61
+ """Uses gcloud to get default project."""
62
+ try:
63
+ result = subprocess.run(
64
+ ["gcloud", "config", "get-value", "project"],
65
+ capture_output=True,
66
+ text=True,
67
+ check=True,
68
+ )
69
+ return result.stdout.strip()
70
+ except (subprocess.CalledProcessError, FileNotFoundError):
71
+ return ""
72
+
73
+
74
+ def _get_gcp_region_from_gcloud() -> str:
75
+ """Uses gcloud to get default region."""
76
+ try:
77
+ result = subprocess.run(
78
+ ["gcloud", "config", "get-value", "compute/region"],
79
+ capture_output=True,
80
+ text=True,
81
+ check=True,
82
+ )
83
+ return result.stdout.strip()
84
+ except (subprocess.CalledProcessError, FileNotFoundError):
85
+ return ""
86
+
87
+
88
+ def _prompt_str(
89
+ prompt_prefix: str,
90
+ *,
91
+ prior_msg: Optional[str] = None,
92
+ default_value: Optional[str] = None,
93
+ ) -> str:
94
+ if prior_msg:
95
+ click.secho(prior_msg, fg="green")
96
+ while True:
97
+ value: str = click.prompt(
98
+ prompt_prefix, default=default_value or None, type=str
99
+ )
100
+ if value and value.strip():
101
+ return value.strip()
102
+
103
+
104
+ def _prompt_for_google_cloud(
105
+ google_cloud_project: Optional[str],
106
+ ) -> str:
107
+ """Prompts user for Google Cloud project ID."""
108
+ google_cloud_project = (
109
+ google_cloud_project
110
+ or os.environ.get("GOOGLE_CLOUD_PROJECT", None)
111
+ or _get_gcp_project_from_gcloud()
112
+ )
113
+
114
+ google_cloud_project = _prompt_str(
115
+ "Enter Google Cloud project ID", default_value=google_cloud_project
116
+ )
117
+
118
+ return google_cloud_project
119
+
120
+
121
+ def _prompt_for_google_cloud_region(
122
+ google_cloud_region: Optional[str],
123
+ ) -> str:
124
+ """Prompts user for Google Cloud region."""
125
+ google_cloud_region = (
126
+ google_cloud_region
127
+ or os.environ.get("GOOGLE_CLOUD_LOCATION", None)
128
+ or _get_gcp_region_from_gcloud()
129
+ )
130
+
131
+ google_cloud_region = _prompt_str(
132
+ "Enter Google Cloud region",
133
+ default_value=google_cloud_region or "us-central1",
134
+ )
135
+ return google_cloud_region
136
+
137
+
138
+ def _prompt_for_google_api_key(
139
+ google_api_key: Optional[str],
140
+ ) -> str:
141
+ """Prompts user for Google API key."""
142
+ google_api_key = google_api_key or os.environ.get("GOOGLE_API_KEY", None)
143
+
144
+ google_api_key = _prompt_str(
145
+ "Enter Google API key",
146
+ prior_msg=_GOOGLE_API_MSG,
147
+ default_value=google_api_key,
148
+ )
149
+ return google_api_key
150
+
151
+
152
+ def _generate_files(
153
+ agent_folder: str,
154
+ *,
155
+ google_api_key: Optional[str] = None,
156
+ google_cloud_project: Optional[str] = None,
157
+ google_cloud_region: Optional[str] = None,
158
+ model: Optional[str] = None,
159
+ ):
160
+ """Generates a folder name for the agent."""
161
+ os.makedirs(agent_folder, exist_ok=True)
162
+
163
+ dotenv_file_path = os.path.join(agent_folder, ".env")
164
+ init_file_path = os.path.join(agent_folder, "__init__.py")
165
+ agent_file_path = os.path.join(agent_folder, "agent.py")
166
+
167
+ with open(dotenv_file_path, "w", encoding="utf-8") as f:
168
+ lines = []
169
+ if google_api_key:
170
+ lines.append("GOOGLE_GENAI_USE_VERTEXAI=0")
171
+ elif google_cloud_project and google_cloud_region:
172
+ lines.append("GOOGLE_GENAI_USE_VERTEXAI=1")
173
+ if google_api_key:
174
+ lines.append(f"GOOGLE_API_KEY={google_api_key}")
175
+ if google_cloud_project:
176
+ lines.append(f"GOOGLE_CLOUD_PROJECT={google_cloud_project}")
177
+ if google_cloud_region:
178
+ lines.append(f"GOOGLE_CLOUD_LOCATION={google_cloud_region}")
179
+ f.write("\n".join(lines))
180
+
181
+ with open(init_file_path, "w", encoding="utf-8") as f:
182
+ f.write(_INIT_PY_TEMPLATE)
183
+
184
+ with open(agent_file_path, "w", encoding="utf-8") as f:
185
+ f.write(_AGENT_PY_TEMPLATE.format(model_name=model))
186
+
187
+ click.secho(
188
+ _SUCCESS_MSG.format(agent_folder=agent_folder),
189
+ fg="green",
190
+ )
191
+
192
+
193
+ def _prompt_for_model() -> str:
194
+ model_choice = click.prompt(
195
+ """\
196
+ Choose a model for the root agent:
197
+ 1. gemini-2.0-flash-001
198
+ 2. Other models (fill later)
199
+ Choose model""",
200
+ type=click.Choice(["1", "2"]),
201
+ )
202
+ if model_choice == "1":
203
+ return "gemini-2.0-flash-001"
204
+ else:
205
+ click.secho(_OTHER_MODEL_MSG, fg="green")
206
+ return "<FILL_IN_MODEL>"
207
+
208
+
209
+ def _prompt_to_choose_backend(
210
+ google_api_key: Optional[str],
211
+ google_cloud_project: Optional[str],
212
+ google_cloud_region: Optional[str],
213
+ ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
214
+ """Prompts user to choose backend.
215
+
216
+ Returns:
217
+ A tuple of (google_api_key, google_cloud_project, google_cloud_region).
218
+ """
219
+ backend_choice = click.prompt(
220
+ "1. Google AI\n2. Vertex AI\nChoose a backend",
221
+ type=click.Choice(["1", "2"]),
222
+ )
223
+ if backend_choice == "1":
224
+ google_api_key = _prompt_for_google_api_key(google_api_key)
225
+ elif backend_choice == "2":
226
+ click.secho(_GOOGLE_CLOUD_SETUP_MSG, fg="green")
227
+ google_cloud_project = _prompt_for_google_cloud(google_cloud_project)
228
+ google_cloud_region = _prompt_for_google_cloud_region(google_cloud_region)
229
+ return google_api_key, google_cloud_project, google_cloud_region
230
+
231
+
232
+ def run_cmd(
233
+ agent_name: str,
234
+ *,
235
+ model: Optional[str],
236
+ google_api_key: Optional[str],
237
+ google_cloud_project: Optional[str],
238
+ google_cloud_region: Optional[str],
239
+ ):
240
+ """Runs `adk create` command to create agent template.
241
+
242
+ Args:
243
+ agent_name: str, The name of the agent.
244
+ google_api_key: Optional[str], The Google API key for using Google AI as
245
+ backend.
246
+ google_cloud_project: Optional[str], The Google Cloud project for using
247
+ VertexAI as backend.
248
+ google_cloud_region: Optional[str], The Google Cloud region for using
249
+ VertexAI as backend.
250
+ """
251
+ agent_folder = os.path.join(os.getcwd(), agent_name)
252
+ # check folder doesn't exist or it's empty. Otherwise, throw
253
+ if os.path.exists(agent_folder) and os.listdir(agent_folder):
254
+ # Prompt user whether to override existing files using click
255
+ if not click.confirm(
256
+ f"Non-empty folder already exist: '{agent_folder}'\n"
257
+ "Override existing content?",
258
+ default=False,
259
+ ):
260
+ raise click.Abort()
261
+
262
+ if not model:
263
+ model = _prompt_for_model()
264
+
265
+ if not google_api_key and not (google_cloud_project and google_cloud_region):
266
+ if model.startswith("gemini"):
267
+ google_api_key, google_cloud_project, google_cloud_region = (
268
+ _prompt_to_choose_backend(
269
+ google_api_key, google_cloud_project, google_cloud_region
270
+ )
271
+ )
272
+
273
+ _generate_files(
274
+ agent_folder,
275
+ google_api_key=google_api_key,
276
+ google_cloud_project=google_cloud_project,
277
+ google_cloud_region=google_cloud_region,
278
+ model=model,
279
+ )
@@ -84,6 +84,7 @@ def to_cloud_run(
84
84
  port: int,
85
85
  trace_to_cloud: bool,
86
86
  with_ui: bool,
87
+ verbosity: str,
87
88
  ):
88
89
  """Deploys an agent to Google Cloud Run.
89
90
 
@@ -110,6 +111,7 @@ def to_cloud_run(
110
111
  port: The port of the ADK api server.
111
112
  trace_to_cloud: Whether to enable Cloud Trace.
112
113
  with_ui: Whether to deploy with UI.
114
+ verbosity: The verbosity level of the CLI.
113
115
  """
114
116
  app_name = app_name or os.path.basename(agent_folder)
115
117
 
@@ -169,6 +171,8 @@ def to_cloud_run(
169
171
  *region_options,
170
172
  '--port',
171
173
  str(port),
174
+ '--verbosity',
175
+ verbosity,
172
176
  '--labels',
173
177
  'created-by=adk',
174
178
  ],
@@ -112,14 +112,14 @@ def get_evaluation_criteria_or_default(
112
112
 
113
113
 
114
114
  def get_root_agent(agent_module_file_path: str) -> Agent:
115
- """Returns root agent given the agetn module."""
115
+ """Returns root agent given the agent module."""
116
116
  agent_module = _get_agent_module(agent_module_file_path)
117
117
  root_agent = agent_module.agent.root_agent
118
118
  return root_agent
119
119
 
120
120
 
121
121
  def try_get_reset_func(agent_module_file_path: str) -> Any:
122
- """Returns reset function for the agent, if present, given the agetn module."""
122
+ """Returns reset function for the agent, if present, given the agent module."""
123
123
  agent_module = _get_agent_module(agent_module_file_path)
124
124
  reset_func = getattr(agent_module.agent, "reset_data", None)
125
125
  return reset_func
@@ -24,6 +24,7 @@ import click
24
24
  from fastapi import FastAPI
25
25
  import uvicorn
26
26
 
27
+ from . import cli_create
27
28
  from . import cli_deploy
28
29
  from .cli import run_cli
29
30
  from .cli_eval import MISSING_EVAL_DEPENDENCIES_MESSAGE
@@ -42,10 +43,59 @@ def main():
42
43
 
43
44
  @main.group()
44
45
  def deploy():
45
- """Deploy Agent."""
46
+ """Deploys agent to hosted environments."""
46
47
  pass
47
48
 
48
49
 
50
+ @main.command("create")
51
+ @click.option(
52
+ "--model",
53
+ type=str,
54
+ help="Optional. The model used for the root agent.",
55
+ )
56
+ @click.option(
57
+ "--api_key",
58
+ type=str,
59
+ help=(
60
+ "Optional. The API Key needed to access the model, e.g. Google AI API"
61
+ " Key."
62
+ ),
63
+ )
64
+ @click.option(
65
+ "--project",
66
+ type=str,
67
+ help="Optional. The Google Cloud Project for using VertexAI as backend.",
68
+ )
69
+ @click.option(
70
+ "--region",
71
+ type=str,
72
+ help="Optional. The Google Cloud Region for using VertexAI as backend.",
73
+ )
74
+ @click.argument("app_name", type=str, required=True)
75
+ def cli_create_cmd(
76
+ app_name: str,
77
+ model: Optional[str],
78
+ api_key: Optional[str],
79
+ project: Optional[str],
80
+ region: Optional[str],
81
+ ):
82
+ """Creates a new app in the current folder with prepopulated agent template.
83
+
84
+ APP_NAME: required, the folder of the agent source code.
85
+
86
+ Example:
87
+
88
+ adk create path/to/my_app
89
+ """
90
+ cli_create.run_cmd(
91
+ app_name,
92
+ model=model,
93
+ google_api_key=api_key,
94
+ google_cloud_project=project,
95
+ google_cloud_region=region,
96
+ )
97
+
98
+
49
99
  @main.command("run")
50
100
  @click.option(
51
101
  "--save_session",
@@ -62,7 +112,7 @@ def deploy():
62
112
  ),
63
113
  )
64
114
  def cli_run(agent: str, save_session: bool):
65
- """Run an interactive CLI for a certain agent.
115
+ """Runs an interactive CLI for a certain agent.
66
116
 
67
117
  AGENT: The path to the agent source code folder.
68
118
 
@@ -244,7 +294,7 @@ def cli_eval(
244
294
  type=click.Path(
245
295
  exists=True, dir_okay=True, file_okay=False, resolve_path=True
246
296
  ),
247
- default=os.getcwd(),
297
+ default=os.getcwd,
248
298
  )
249
299
  def cli_web(
250
300
  agents_dir: str,
@@ -255,7 +305,7 @@ def cli_web(
255
305
  port: int = 8000,
256
306
  trace_to_cloud: bool = False,
257
307
  ):
258
- """Start a FastAPI server with Web UI for agents.
308
+ """Starts a FastAPI server with Web UI for agents.
259
309
 
260
310
  AGENTS_DIR: The directory of agents, where each sub-directory is a single
261
311
  agent, containing at least `__init__.py` and `agent.py` files.
@@ -274,7 +324,7 @@ def cli_web(
274
324
  @asynccontextmanager
275
325
  async def _lifespan(app: FastAPI):
276
326
  click.secho(
277
- f"""\
327
+ f"""
278
328
  +-----------------------------------------------------------------------------+
279
329
  | ADK Web Server started |
280
330
  | |
@@ -285,7 +335,7 @@ def cli_web(
285
335
  )
286
336
  yield # Startup is done, now app is running
287
337
  click.secho(
288
- """\
338
+ """
289
339
  +-----------------------------------------------------------------------------+
290
340
  | ADK Web Server shutting down... |
291
341
  +-----------------------------------------------------------------------------+
@@ -378,7 +428,7 @@ def cli_api_server(
378
428
  port: int = 8000,
379
429
  trace_to_cloud: bool = False,
380
430
  ):
381
- """Start a FastAPI server for agents.
431
+ """Starts a FastAPI server for agents.
382
432
 
383
433
  AGENTS_DIR: The directory of agents, where each sub-directory is a single
384
434
  agent, containing at least `__init__.py` and `agent.py` files.
@@ -483,6 +533,14 @@ def cli_api_server(
483
533
  " (default: a timestamped folder in the system temp directory)."
484
534
  ),
485
535
  )
536
+ @click.option(
537
+ "--verbosity",
538
+ type=click.Choice(
539
+ ["debug", "info", "warning", "error", "critical"], case_sensitive=False
540
+ ),
541
+ default="WARNING",
542
+ help="Optional. Override the default verbosity level.",
543
+ )
486
544
  @click.argument(
487
545
  "agent",
488
546
  type=click.Path(
@@ -499,6 +557,7 @@ def cli_deploy_cloud_run(
499
557
  port: int,
500
558
  trace_to_cloud: bool,
501
559
  with_ui: bool,
560
+ verbosity: str,
502
561
  ):
503
562
  """Deploys an agent to Cloud Run.
504
563
 
@@ -519,6 +578,7 @@ def cli_deploy_cloud_run(
519
578
  port=port,
520
579
  trace_to_cloud=trace_to_cloud,
521
580
  with_ui=with_ui,
581
+ verbosity=verbosity,
522
582
  )
523
583
  except Exception as e:
524
584
  click.secho(f"Deploy failed: {e}", fg="red", err=True)
@@ -13,7 +13,9 @@
13
13
  # limitations under the License.
14
14
 
15
15
  import asyncio
16
+ from contextlib import asynccontextmanager
16
17
  import importlib
18
+ import inspect
17
19
  import json
18
20
  import logging
19
21
  import os
@@ -28,6 +30,7 @@ from typing import Literal
28
30
  from typing import Optional
29
31
 
30
32
  import click
33
+ from click import Tuple
31
34
  from fastapi import FastAPI
32
35
  from fastapi import HTTPException
33
36
  from fastapi import Query
@@ -56,6 +59,7 @@ from ..agents.llm_agent import Agent
56
59
  from ..agents.run_config import StreamingMode
57
60
  from ..artifacts import InMemoryArtifactService
58
61
  from ..events.event import Event
62
+ from ..memory.in_memory_memory_service import InMemoryMemoryService
59
63
  from ..runners import Runner
60
64
  from ..sessions.database_session_service import DatabaseSessionService
61
65
  from ..sessions.in_memory_session_service import InMemorySessionService
@@ -144,7 +148,7 @@ def get_fast_api_app(
144
148
  export.SimpleSpanProcessor(ApiServerSpanExporter(trace_dict))
145
149
  )
146
150
  if trace_to_cloud:
147
- envs.load_dotenv()
151
+ envs.load_dotenv_for_agent("", agent_dir)
148
152
  if project_id := os.environ.get("GOOGLE_CLOUD_PROJECT", None):
149
153
  processor = export.BatchSpanProcessor(
150
154
  CloudTraceSpanExporter(project_id=project_id)
@@ -158,8 +162,22 @@ def get_fast_api_app(
158
162
 
159
163
  trace.set_tracer_provider(provider)
160
164
 
165
+ exit_stacks = []
166
+
167
+ @asynccontextmanager
168
+ async def internal_lifespan(app: FastAPI):
169
+ if lifespan:
170
+ async with lifespan(app) as lifespan_context:
171
+ yield
172
+
173
+ if exit_stacks:
174
+ for stack in exit_stacks:
175
+ await stack.aclose()
176
+ else:
177
+ yield
178
+
161
179
  # Run the FastAPI server.
162
- app = FastAPI(lifespan=lifespan)
180
+ app = FastAPI(lifespan=internal_lifespan)
163
181
 
164
182
  if allow_origins:
165
183
  app.add_middleware(
@@ -178,6 +196,7 @@ def get_fast_api_app(
178
196
 
179
197
  # Build the Artifact service
180
198
  artifact_service = InMemoryArtifactService()
199
+ memory_service = InMemoryMemoryService()
181
200
 
182
201
  # Build the Session service
183
202
  agent_engine_id = ""
@@ -355,7 +374,7 @@ def get_fast_api_app(
355
374
  "/apps/{app_name}/eval_sets/{eval_set_id}/add_session",
356
375
  response_model_exclude_none=True,
357
376
  )
358
- def add_session_to_eval_set(
377
+ async def add_session_to_eval_set(
359
378
  app_name: str, eval_set_id: str, req: AddSessionToEvalSetRequest
360
379
  ):
361
380
  pattern = r"^[a-zA-Z0-9_]+$"
@@ -390,7 +409,9 @@ def get_fast_api_app(
390
409
  test_data = evals.convert_session_to_eval_format(session)
391
410
 
392
411
  # Populate the session with initial session state.
393
- initial_session_state = create_empty_state(_get_root_agent(app_name))
412
+ initial_session_state = create_empty_state(
413
+ await _get_root_agent_async(app_name)
414
+ )
394
415
 
395
416
  eval_set_data.append({
396
417
  "name": req.eval_id,
@@ -427,7 +448,7 @@ def get_fast_api_app(
427
448
  "/apps/{app_name}/eval_sets/{eval_set_id}/run_eval",
428
449
  response_model_exclude_none=True,
429
450
  )
430
- def run_eval(
451
+ async def run_eval(
431
452
  app_name: str, eval_set_id: str, req: RunEvalRequest
432
453
  ) -> list[RunEvalResult]:
433
454
  from .cli_eval import run_evals
@@ -444,7 +465,7 @@ def get_fast_api_app(
444
465
  logger.info(
445
466
  "Eval ids to run list is empty. We will all evals in the eval set."
446
467
  )
447
- root_agent = _get_root_agent(app_name)
468
+ root_agent = await _get_root_agent_async(app_name)
448
469
  eval_results = list(
449
470
  run_evals(
450
471
  eval_set_to_evals,
@@ -574,7 +595,7 @@ def get_fast_api_app(
574
595
  )
575
596
  if not session:
576
597
  raise HTTPException(status_code=404, detail="Session not found")
577
- runner = _get_runner(req.app_name)
598
+ runner = await _get_runner_async(req.app_name)
578
599
  events = [
579
600
  event
580
601
  async for event in runner.run_async(
@@ -601,7 +622,7 @@ def get_fast_api_app(
601
622
  async def event_generator():
602
623
  try:
603
624
  stream_mode = StreamingMode.SSE if req.streaming else StreamingMode.NONE
604
- runner = _get_runner(req.app_name)
625
+ runner = await _get_runner_async(req.app_name)
605
626
  async for event in runner.run_async(
606
627
  user_id=req.user_id,
607
628
  session_id=req.session_id,
@@ -627,7 +648,7 @@ def get_fast_api_app(
627
648
  "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}/graph",
628
649
  response_model_exclude_none=True,
629
650
  )
630
- def get_event_graph(
651
+ async def get_event_graph(
631
652
  app_name: str, user_id: str, session_id: str, event_id: str
632
653
  ):
633
654
  # Connect to managed session if agent_engine_id is set.
@@ -644,7 +665,7 @@ def get_fast_api_app(
644
665
 
645
666
  function_calls = event.get_function_calls()
646
667
  function_responses = event.get_function_responses()
647
- root_agent = _get_root_agent(app_name)
668
+ root_agent = await _get_root_agent_async(app_name)
648
669
  dot_graph = None
649
670
  if function_calls:
650
671
  function_call_highlights = []
@@ -701,7 +722,7 @@ def get_fast_api_app(
701
722
  live_request_queue = LiveRequestQueue()
702
723
 
703
724
  async def forward_events():
704
- runner = _get_runner(app_name)
725
+ runner = await _get_runner_async(app_name)
705
726
  async for event in runner.run_live(
706
727
  session=session, live_request_queue=live_request_queue
707
728
  ):
@@ -739,26 +760,40 @@ def get_fast_api_app(
739
760
  for task in pending:
740
761
  task.cancel()
741
762
 
742
- def _get_root_agent(app_name: str) -> Agent:
763
+ async def _get_root_agent_async(app_name: str) -> Agent:
743
764
  """Returns the root agent for the given app."""
744
765
  if app_name in root_agent_dict:
745
766
  return root_agent_dict[app_name]
746
- envs.load_dotenv_for_agent(os.path.basename(app_name), agent_dir)
747
767
  agent_module = importlib.import_module(app_name)
748
- root_agent: Agent = agent_module.agent.root_agent
768
+ if getattr(agent_module.agent, "root_agent"):
769
+ root_agent = agent_module.agent.root_agent
770
+ else:
771
+ raise ValueError(f'Unable to find "root_agent" from {app_name}.')
772
+
773
+ # Handle an awaitable root agent and await for the actual agent.
774
+ if inspect.isawaitable(root_agent):
775
+ try:
776
+ agent, exit_stack = await root_agent
777
+ exit_stacks.append(exit_stack)
778
+ root_agent = agent
779
+ except Exception as e:
780
+ raise RuntimeError(f"error getting root agent, {e}") from e
781
+
749
782
  root_agent_dict[app_name] = root_agent
750
783
  return root_agent
751
784
 
752
- def _get_runner(app_name: str) -> Runner:
785
+ async def _get_runner_async(app_name: str) -> Runner:
753
786
  """Returns the runner for the given app."""
787
+ envs.load_dotenv_for_agent(os.path.basename(app_name), agent_dir)
754
788
  if app_name in runner_dict:
755
789
  return runner_dict[app_name]
756
- root_agent = _get_root_agent(app_name)
790
+ root_agent = await _get_root_agent_async(app_name)
757
791
  runner = Runner(
758
792
  app_name=agent_engine_id if agent_engine_id else app_name,
759
793
  agent=root_agent,
760
794
  artifact_service=artifact_service,
761
795
  session_service=session_service,
796
+ memory_service=memory_service,
762
797
  )
763
798
  runner_dict[app_name] = runner
764
799
  return runner
@@ -50,8 +50,5 @@ def load_dotenv_for_agent(
50
50
  agent_name,
51
51
  dotenv_file_path,
52
52
  )
53
- logger.info(
54
- 'Reloaded %s file for %s at %s', filename, agent_name, dotenv_file_path
55
- )
56
53
  else:
57
54
  logger.info('No %s file found for %s', filename, agent_name)
@@ -66,7 +66,7 @@ def convert_session_to_eval_format(session: Session) -> list[dict[str, Any]]:
66
66
  'tool_input': tool_input,
67
67
  })
68
68
  elif subsequent_part.text:
69
- # Also keep track of all the natural langauge responses that
69
+ # Also keep track of all the natural language responses that
70
70
  # agent (or sub agents) generated.
71
71
  intermediate_agent_responses.append(
72
72
  {'author': event_author, 'text': subsequent_part.text}
@@ -75,7 +75,7 @@ def convert_session_to_eval_format(session: Session) -> list[dict[str, Any]]:
75
75
  # If we are here then either we are done reading all the events or we
76
76
  # encountered an event that had content authored by the end-user.
77
77
  # This, basically means an end of turn.
78
- # We assume that the last natural langauge intermediate response is the
78
+ # We assume that the last natural language intermediate response is the
79
79
  # final response from the agent/model. We treat that as a reference.
80
80
  eval_case.append({
81
81
  'query': query,