google-adk 0.5.0__py3-none-any.whl → 1.1.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 (139) hide show
  1. google/adk/agents/base_agent.py +76 -30
  2. google/adk/agents/callback_context.py +2 -6
  3. google/adk/agents/llm_agent.py +122 -30
  4. google/adk/agents/loop_agent.py +1 -1
  5. google/adk/agents/parallel_agent.py +7 -0
  6. google/adk/agents/readonly_context.py +8 -0
  7. google/adk/agents/run_config.py +1 -1
  8. google/adk/agents/sequential_agent.py +31 -0
  9. google/adk/agents/transcription_entry.py +4 -2
  10. google/adk/artifacts/gcs_artifact_service.py +1 -1
  11. google/adk/artifacts/in_memory_artifact_service.py +1 -1
  12. google/adk/auth/auth_credential.py +10 -2
  13. google/adk/auth/auth_preprocessor.py +7 -1
  14. google/adk/auth/auth_tool.py +3 -4
  15. google/adk/cli/agent_graph.py +5 -5
  16. google/adk/cli/browser/index.html +4 -4
  17. google/adk/cli/browser/{main-ULN5R5I5.js → main-PKDNKWJE.js} +59 -60
  18. google/adk/cli/browser/polyfills-B6TNHZQ6.js +17 -0
  19. google/adk/cli/cli.py +10 -9
  20. google/adk/cli/cli_deploy.py +7 -2
  21. google/adk/cli/cli_eval.py +109 -115
  22. google/adk/cli/cli_tools_click.py +179 -67
  23. google/adk/cli/fast_api.py +248 -197
  24. google/adk/cli/utils/agent_loader.py +137 -0
  25. google/adk/cli/utils/cleanup.py +40 -0
  26. google/adk/cli/utils/common.py +23 -0
  27. google/adk/cli/utils/evals.py +83 -0
  28. google/adk/cli/utils/logs.py +8 -5
  29. google/adk/code_executors/__init__.py +3 -1
  30. google/adk/code_executors/built_in_code_executor.py +52 -0
  31. google/adk/code_executors/code_execution_utils.py +2 -1
  32. google/adk/code_executors/container_code_executor.py +0 -1
  33. google/adk/code_executors/vertex_ai_code_executor.py +6 -8
  34. google/adk/evaluation/__init__.py +1 -1
  35. google/adk/evaluation/agent_evaluator.py +168 -128
  36. google/adk/evaluation/eval_case.py +104 -0
  37. google/adk/evaluation/eval_metrics.py +74 -0
  38. google/adk/evaluation/eval_result.py +86 -0
  39. google/adk/evaluation/eval_set.py +39 -0
  40. google/adk/evaluation/eval_set_results_manager.py +47 -0
  41. google/adk/evaluation/eval_sets_manager.py +43 -0
  42. google/adk/evaluation/evaluation_generator.py +88 -113
  43. google/adk/evaluation/evaluator.py +58 -0
  44. google/adk/evaluation/local_eval_set_results_manager.py +113 -0
  45. google/adk/evaluation/local_eval_sets_manager.py +264 -0
  46. google/adk/evaluation/response_evaluator.py +106 -1
  47. google/adk/evaluation/trajectory_evaluator.py +84 -2
  48. google/adk/events/event.py +6 -1
  49. google/adk/events/event_actions.py +6 -1
  50. google/adk/examples/base_example_provider.py +1 -0
  51. google/adk/examples/example_util.py +3 -2
  52. google/adk/flows/llm_flows/_code_execution.py +9 -1
  53. google/adk/flows/llm_flows/audio_transcriber.py +4 -3
  54. google/adk/flows/llm_flows/base_llm_flow.py +58 -21
  55. google/adk/flows/llm_flows/contents.py +3 -1
  56. google/adk/flows/llm_flows/functions.py +9 -8
  57. google/adk/flows/llm_flows/instructions.py +18 -80
  58. google/adk/flows/llm_flows/single_flow.py +2 -2
  59. google/adk/memory/__init__.py +1 -1
  60. google/adk/memory/_utils.py +23 -0
  61. google/adk/memory/base_memory_service.py +23 -21
  62. google/adk/memory/in_memory_memory_service.py +57 -25
  63. google/adk/memory/memory_entry.py +37 -0
  64. google/adk/memory/vertex_ai_rag_memory_service.py +38 -15
  65. google/adk/models/anthropic_llm.py +16 -9
  66. google/adk/models/base_llm.py +2 -1
  67. google/adk/models/base_llm_connection.py +2 -0
  68. google/adk/models/gemini_llm_connection.py +11 -11
  69. google/adk/models/google_llm.py +12 -2
  70. google/adk/models/lite_llm.py +80 -23
  71. google/adk/models/llm_response.py +16 -3
  72. google/adk/models/registry.py +1 -1
  73. google/adk/runners.py +98 -42
  74. google/adk/sessions/__init__.py +1 -1
  75. google/adk/sessions/_session_util.py +2 -1
  76. google/adk/sessions/base_session_service.py +6 -33
  77. google/adk/sessions/database_session_service.py +57 -67
  78. google/adk/sessions/in_memory_session_service.py +106 -24
  79. google/adk/sessions/session.py +3 -0
  80. google/adk/sessions/vertex_ai_session_service.py +44 -51
  81. google/adk/telemetry.py +7 -2
  82. google/adk/tools/__init__.py +4 -7
  83. google/adk/tools/_memory_entry_utils.py +30 -0
  84. google/adk/tools/agent_tool.py +10 -10
  85. google/adk/tools/apihub_tool/apihub_toolset.py +55 -74
  86. google/adk/tools/apihub_tool/clients/apihub_client.py +10 -3
  87. google/adk/tools/apihub_tool/clients/secret_client.py +1 -0
  88. google/adk/tools/application_integration_tool/application_integration_toolset.py +111 -85
  89. google/adk/tools/application_integration_tool/clients/connections_client.py +28 -1
  90. google/adk/tools/application_integration_tool/clients/integration_client.py +7 -5
  91. google/adk/tools/application_integration_tool/integration_connector_tool.py +69 -26
  92. google/adk/tools/base_toolset.py +96 -0
  93. google/adk/tools/bigquery/__init__.py +28 -0
  94. google/adk/tools/bigquery/bigquery_credentials.py +216 -0
  95. google/adk/tools/bigquery/bigquery_tool.py +116 -0
  96. google/adk/tools/{built_in_code_execution_tool.py → enterprise_search_tool.py} +17 -11
  97. google/adk/tools/function_parameter_parse_util.py +9 -2
  98. google/adk/tools/function_tool.py +33 -3
  99. google/adk/tools/get_user_choice_tool.py +1 -0
  100. google/adk/tools/google_api_tool/__init__.py +24 -70
  101. google/adk/tools/google_api_tool/google_api_tool.py +12 -6
  102. google/adk/tools/google_api_tool/{google_api_tool_set.py → google_api_toolset.py} +57 -55
  103. google/adk/tools/google_api_tool/google_api_toolsets.py +108 -0
  104. google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py +40 -42
  105. google/adk/tools/google_search_tool.py +2 -2
  106. google/adk/tools/langchain_tool.py +96 -49
  107. google/adk/tools/load_memory_tool.py +14 -5
  108. google/adk/tools/mcp_tool/__init__.py +3 -2
  109. google/adk/tools/mcp_tool/conversion_utils.py +6 -2
  110. google/adk/tools/mcp_tool/mcp_session_manager.py +80 -69
  111. google/adk/tools/mcp_tool/mcp_tool.py +35 -32
  112. google/adk/tools/mcp_tool/mcp_toolset.py +99 -194
  113. google/adk/tools/openapi_tool/auth/credential_exchangers/base_credential_exchanger.py +1 -3
  114. google/adk/tools/openapi_tool/auth/credential_exchangers/service_account_exchanger.py +6 -7
  115. google/adk/tools/openapi_tool/common/common.py +5 -1
  116. google/adk/tools/openapi_tool/openapi_spec_parser/__init__.py +7 -2
  117. google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py +27 -7
  118. google/adk/tools/openapi_tool/openapi_spec_parser/operation_parser.py +36 -32
  119. google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py +11 -1
  120. google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py +1 -1
  121. google/adk/tools/preload_memory_tool.py +27 -18
  122. google/adk/tools/retrieval/__init__.py +1 -1
  123. google/adk/tools/retrieval/vertex_ai_rag_retrieval.py +1 -1
  124. google/adk/tools/toolbox_toolset.py +107 -0
  125. google/adk/tools/transfer_to_agent_tool.py +0 -1
  126. google/adk/utils/__init__.py +13 -0
  127. google/adk/utils/instructions_utils.py +131 -0
  128. google/adk/version.py +1 -1
  129. {google_adk-0.5.0.dist-info → google_adk-1.1.0.dist-info}/METADATA +18 -19
  130. google_adk-1.1.0.dist-info/RECORD +200 -0
  131. google/adk/agents/remote_agent.py +0 -50
  132. google/adk/cli/browser/polyfills-FFHMD2TL.js +0 -18
  133. google/adk/cli/fast_api.py.orig +0 -728
  134. google/adk/tools/google_api_tool/google_api_tool_sets.py +0 -112
  135. google/adk/tools/toolbox_tool.py +0 -46
  136. google_adk-0.5.0.dist-info/RECORD +0 -180
  137. {google_adk-0.5.0.dist-info → google_adk-1.1.0.dist-info}/WHEEL +0 -0
  138. {google_adk-0.5.0.dist-info → google_adk-1.1.0.dist-info}/entry_points.txt +0 -0
  139. {google_adk-0.5.0.dist-info → google_adk-1.1.0.dist-info}/licenses/LICENSE +0 -0
@@ -1,728 +0,0 @@
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 asyncio
16
- import importlib
17
- import json
18
- import logging
19
- import os
20
- from pathlib import Path
21
- import re
22
- import sys
23
- import traceback
24
- import typing
25
- from typing import Any
26
- from typing import List
27
- from typing import Literal
28
- from typing import Optional
29
-
30
- import click
31
- from fastapi import FastAPI
32
- from fastapi import HTTPException
33
- from fastapi import Query
34
- from fastapi import Response
35
- from fastapi.middleware.cors import CORSMiddleware
36
- from fastapi.responses import FileResponse
37
- from fastapi.responses import RedirectResponse
38
- from fastapi.responses import StreamingResponse
39
- from fastapi.staticfiles import StaticFiles
40
- from fastapi.websockets import WebSocket
41
- from fastapi.websockets import WebSocketDisconnect
42
- from google.genai import types
43
- from opentelemetry import trace
44
- from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
45
- from opentelemetry.sdk.trace import export
46
- from opentelemetry.sdk.trace import ReadableSpan
47
- from opentelemetry.sdk.trace import TracerProvider
48
- from pydantic import BaseModel
49
- from pydantic import ValidationError
50
-
51
- from ..agents import RunConfig
52
- from ..agents.live_request_queue import LiveRequest
53
- from ..agents.live_request_queue import LiveRequestQueue
54
- from ..agents.llm_agent import Agent
55
- from ..agents.run_config import StreamingMode
56
- from ..artifacts import InMemoryArtifactService
57
- from ..events.event import Event
58
- from ..runners import Runner
59
- from ..sessions.database_session_service import DatabaseSessionService
60
- from ..sessions.in_memory_session_service import InMemorySessionService
61
- from ..sessions.session import Session
62
- from ..sessions.vertex_ai_session_service import VertexAiSessionService
63
- from .cli_eval import EVAL_SESSION_ID_PREFIX
64
- from .cli_eval import EvalMetric
65
- from .cli_eval import EvalMetricResult
66
- from .cli_eval import EvalStatus
67
- from .utils import create_empty_state
68
- from .utils import envs
69
- from .utils import evals
70
-
71
- logger = logging.getLogger(__name__)
72
-
73
- _EVAL_SET_FILE_EXTENSION = ".evalset.json"
74
-
75
-
76
- class ApiServerSpanExporter(export.SpanExporter):
77
-
78
- def __init__(self, trace_dict):
79
- self.trace_dict = trace_dict
80
-
81
- def export(
82
- self, spans: typing.Sequence[ReadableSpan]
83
- ) -> export.SpanExportResult:
84
- for span in spans:
85
- if span.name == "call_llm" or span.name == "send_data":
86
- attributes = dict(span.attributes)
87
- attributes["trace_id"] = span.get_span_context().trace_id
88
- attributes["span_id"] = span.get_span_context().span_id
89
- self.trace_dict[attributes["gcp.vertex.agent.event_id"]] = attributes
90
- return export.SpanExportResult.SUCCESS
91
-
92
- def force_flush(self, timeout_millis: int = 30000) -> bool:
93
- return True
94
-
95
-
96
- class AgentRunRequest(BaseModel):
97
- app_name: str
98
- user_id: str
99
- session_id: str
100
- new_message: types.Content
101
- streaming: bool = False
102
-
103
-
104
- class AddSessionToEvalSetRequest(BaseModel):
105
- eval_id: str
106
- session_id: str
107
- user_id: str
108
-
109
-
110
- class RunEvalRequest(BaseModel):
111
- eval_ids: list[str] # if empty, then all evals in the eval set are run.
112
- eval_metrics: list[EvalMetric]
113
-
114
-
115
- class RunEvalResult(BaseModel):
116
- eval_set_id: str
117
- eval_id: str
118
- final_eval_status: EvalStatus
119
- eval_metric_results: list[tuple[EvalMetric, EvalMetricResult]]
120
- session_id: str
121
-
122
-
123
- def get_fast_api_app(
124
- *,
125
- agent_dir: str,
126
- session_db_url: str = "",
127
- allow_origins: Optional[list[str]] = None,
128
- web: bool,
129
- ) -> FastAPI:
130
- # InMemory tracing dict.
131
- trace_dict: dict[str, Any] = {}
132
-
133
- # Set up tracing in the FastAPI server.
134
- provider = TracerProvider()
135
- provider.add_span_processor(
136
- export.SimpleSpanProcessor(ApiServerSpanExporter(trace_dict))
137
- )
138
- if os.environ.get("ADK_TRACE_TO_CLOUD", "0") == "1":
139
- processor = export.BatchSpanProcessor(
140
- CloudTraceSpanExporter(
141
- project_id=os.environ.get("GOOGLE_CLOUD_PROJECT", "")
142
- )
143
- )
144
- provider.add_span_processor(processor)
145
-
146
- trace.set_tracer_provider(provider)
147
-
148
- # Run the FastAPI server.
149
- app = FastAPI()
150
- origins = ["http://localhost:4200"]
151
- app.add_middleware(
152
- CORSMiddleware,
153
- allow_origins=origins,
154
- allow_credentials=True,
155
- allow_methods=["*"],
156
- allow_headers=["*"],
157
- )
158
-
159
- if allow_origins:
160
- app.add_middleware(
161
- CORSMiddleware,
162
- allow_origins=allow_origins,
163
- allow_credentials=True,
164
- allow_methods=["*"],
165
- allow_headers=["*"],
166
- )
167
-
168
- if agent_dir not in sys.path:
169
- sys.path.append(agent_dir)
170
-
171
- runner_dict = {}
172
- root_agent_dict = {}
173
-
174
- # Build the Artifact service
175
- artifact_service = InMemoryArtifactService()
176
-
177
- # Build the Session service
178
- agent_engine_id = ""
179
- if session_db_url:
180
- if session_db_url.startswith("agentengine://"):
181
- # Create vertex session service
182
- agent_engine_id = session_db_url.split("://")[1]
183
- if not agent_engine_id:
184
- raise click.ClickException("Agent engine id can not be empty.")
185
- envs.load_dotenv_for_agent("", agent_dir)
186
- session_service = VertexAiSessionService(
187
- os.environ["GOOGLE_CLOUD_PROJECT"],
188
- os.environ["GOOGLE_CLOUD_LOCATION"],
189
- )
190
- else:
191
- session_service = DatabaseSessionService(db_url=session_db_url)
192
- else:
193
- session_service = InMemorySessionService()
194
-
195
- @app.get("/list-apps")
196
- def list_apps() -> list[str]:
197
- base_path = Path.cwd() / agent_dir
198
- if not base_path.exists():
199
- raise HTTPException(status_code=404, detail="Path not found")
200
- if not base_path.is_dir():
201
- raise HTTPException(status_code=400, detail="Not a directory")
202
- agent_names = [
203
- x
204
- for x in os.listdir(base_path)
205
- if os.path.isdir(os.path.join(base_path, x))
206
- and not x.startswith(".")
207
- and x != "__pycache__"
208
- ]
209
- agent_names.sort()
210
- return agent_names
211
-
212
- @app.get("/debug/trace/{event_id}")
213
- def get_trace_dict(event_id: str) -> Any:
214
- event_dict = trace_dict.get(event_id, None)
215
- if event_dict is None:
216
- raise HTTPException(status_code=404, detail="Trace not found")
217
- return event_dict
218
-
219
- @app.get(
220
- "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
221
- response_model_exclude_none=True,
222
- )
223
- def get_session(app_name: str, user_id: str, session_id: str) -> Session:
224
- # Connect to managed session if agent_engine_id is set.
225
- app_name = agent_engine_id if agent_engine_id else app_name
226
- session = session_service.get_session(app_name, user_id, session_id)
227
- if not session:
228
- raise HTTPException(status_code=404, detail="Session not found")
229
- return session
230
-
231
- @app.get(
232
- "/apps/{app_name}/users/{user_id}/sessions",
233
- response_model_exclude_none=True,
234
- )
235
- def list_sessions(app_name: str, user_id: str) -> list[Session]:
236
- # Connect to managed session if agent_engine_id is set.
237
- app_name = agent_engine_id if agent_engine_id else app_name
238
- return [
239
- session
240
- for session in session_service.list_sessions(app_name, user_id).sessions
241
- # Remove sessions that were generated as a part of Eval.
242
- if not session.id.startswith(EVAL_SESSION_ID_PREFIX)
243
- ]
244
-
245
- @app.post(
246
- "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
247
- response_model_exclude_none=True,
248
- )
249
- def create_session(
250
- app_name: str,
251
- user_id: str,
252
- session_id: str,
253
- state: Optional[dict[str, Any]] = None,
254
- ) -> Session:
255
- # Connect to managed session if agent_engine_id is set.
256
- app_name = agent_engine_id if agent_engine_id else app_name
257
- if session_service.get_session(app_name, user_id, session_id) is not None:
258
- logger.warning("Session already exists: %s", session_id)
259
- raise HTTPException(
260
- status_code=400, detail=f"Session already exists: {session_id}"
261
- )
262
-
263
- logger.info("New session created: %s", session_id)
264
- return session_service.create_session(
265
- app_name, user_id, state, session_id=session_id
266
- )
267
-
268
- @app.post(
269
- "/apps/{app_name}/users/{user_id}/sessions",
270
- response_model_exclude_none=True,
271
- )
272
- def create_session(
273
- app_name: str,
274
- user_id: str,
275
- state: Optional[dict[str, Any]] = None,
276
- ) -> Session:
277
- # Connect to managed session if agent_engine_id is set.
278
- app_name = agent_engine_id if agent_engine_id else app_name
279
-
280
- logger.info("New session created")
281
- return session_service.create_session(app_name, user_id, state)
282
-
283
- def _get_eval_set_file_path(app_name, agent_dir, eval_set_id) -> str:
284
- return os.path.join(
285
- agent_dir,
286
- app_name,
287
- eval_set_id + _EVAL_SET_FILE_EXTENSION,
288
- )
289
-
290
- @app.post(
291
- "/apps/{app_name}/eval_sets/{eval_set_id}",
292
- response_model_exclude_none=True,
293
- )
294
- def create_eval_set(
295
- app_name: str,
296
- eval_set_id: str,
297
- ):
298
- """Creates an eval set, given the id."""
299
- pattern = r"^[a-zA-Z0-9_]+$"
300
- if not bool(re.fullmatch(pattern, eval_set_id)):
301
- raise HTTPException(
302
- status_code=400,
303
- detail=(
304
- f"Invalid eval set id. Eval set id should have the `{pattern}`"
305
- " format"
306
- ),
307
- )
308
- # Define the file path
309
- new_eval_set_path = _get_eval_set_file_path(
310
- app_name, agent_dir, eval_set_id
311
- )
312
-
313
- logger.info("Creating eval set file `%s`", new_eval_set_path)
314
-
315
- if not os.path.exists(new_eval_set_path):
316
- # Write the JSON string to the file
317
- logger.info("Eval set file doesn't exist, we will create a new one.")
318
- with open(new_eval_set_path, "w") as f:
319
- empty_content = json.dumps([], indent=2)
320
- f.write(empty_content)
321
-
322
- @app.get(
323
- "/apps/{app_name}/eval_sets",
324
- response_model_exclude_none=True,
325
- )
326
- def list_eval_sets(app_name: str) -> list[str]:
327
- """Lists all eval sets for the given app."""
328
- eval_set_file_path = os.path.join(agent_dir, app_name)
329
- eval_sets = []
330
- for file in os.listdir(eval_set_file_path):
331
- if file.endswith(_EVAL_SET_FILE_EXTENSION):
332
- eval_sets.append(
333
- os.path.basename(file).removesuffix(_EVAL_SET_FILE_EXTENSION)
334
- )
335
-
336
- return sorted(eval_sets)
337
-
338
- @app.post(
339
- "/apps/{app_name}/eval_sets/{eval_set_id}/add_session",
340
- response_model_exclude_none=True,
341
- )
342
- def add_session_to_eval_set(
343
- app_name: str, eval_set_id: str, req: AddSessionToEvalSetRequest
344
- ):
345
- pattern = r"^[a-zA-Z0-9_]+$"
346
- if not bool(re.fullmatch(pattern, req.eval_id)):
347
- raise HTTPException(
348
- status_code=400,
349
- detail=f"Invalid eval id. Eval id should have the `{pattern}` format",
350
- )
351
-
352
- # Get the session
353
- session = session_service.get_session(app_name, req.user_id, req.session_id)
354
- assert session, "Session not found."
355
- # Load the eval set file data
356
- eval_set_file_path = _get_eval_set_file_path(
357
- app_name, agent_dir, eval_set_id
358
- )
359
- with open(eval_set_file_path, "r") as file:
360
- eval_set_data = json.load(file) # Load JSON into a list
361
-
362
- if [x for x in eval_set_data if x["name"] == req.eval_id]:
363
- raise HTTPException(
364
- status_code=400,
365
- detail=(
366
- f"Eval id `{req.eval_id}` already exists in `{eval_set_id}`"
367
- " eval set."
368
- ),
369
- )
370
-
371
- # Convert the session data to evaluation format
372
- test_data = evals.convert_session_to_eval_format(session)
373
-
374
- # Populate the session with initial session state.
375
- initial_session_state = create_empty_state(_get_root_agent(app_name))
376
-
377
- eval_set_data.append({
378
- "name": req.eval_id,
379
- "data": test_data,
380
- "initial_session": {
381
- "state": initial_session_state,
382
- "app_name": app_name,
383
- "user_id": req.user_id,
384
- },
385
- })
386
- # Serialize the test data to JSON and write to the eval set file.
387
- with open(eval_set_file_path, "w") as f:
388
- f.write(json.dumps(eval_set_data, indent=2))
389
-
390
- @app.get(
391
- "/apps/{app_name}/eval_sets/{eval_set_id}/evals",
392
- response_model_exclude_none=True,
393
- )
394
- def list_evals_in_eval_set(
395
- app_name: str,
396
- eval_set_id: str,
397
- ) -> list[str]:
398
- """Lists all evals in an eval set."""
399
- # Load the eval set file data
400
- eval_set_file_path = _get_eval_set_file_path(
401
- app_name, agent_dir, eval_set_id
402
- )
403
- with open(eval_set_file_path, "r") as file:
404
- eval_set_data = json.load(file) # Load JSON into a list
405
-
406
- return sorted([x["name"] for x in eval_set_data])
407
-
408
- @app.post(
409
- "/apps/{app_name}/eval_sets/{eval_set_id}/run_eval",
410
- response_model_exclude_none=True,
411
- )
412
- def run_eval(
413
- app_name: str, eval_set_id: str, req: RunEvalRequest
414
- ) -> list[RunEvalResult]:
415
- from .cli_eval import run_evals
416
-
417
- """Runs an eval given the details in the eval request."""
418
- # Create a mapping from eval set file to all the evals that needed to be
419
- # run.
420
- eval_set_file_path = _get_eval_set_file_path(
421
- app_name, agent_dir, eval_set_id
422
- )
423
- eval_set_to_evals = {eval_set_file_path: req.eval_ids}
424
-
425
- if not req.eval_ids:
426
- logger.info(
427
- "Eval ids to run list is empty. We will all evals in the eval set."
428
- )
429
- root_agent = _get_root_agent(app_name)
430
- eval_results = list(
431
- run_evals(
432
- eval_set_to_evals,
433
- root_agent,
434
- getattr(root_agent, "reset_data", None),
435
- req.eval_metrics,
436
- session_service=session_service,
437
- artifact_service=artifact_service,
438
- )
439
- )
440
-
441
- run_eval_results = []
442
- for eval_result in eval_results:
443
- run_eval_results.append(
444
- RunEvalResult(
445
- app_name=app_name,
446
- eval_set_id=eval_set_id,
447
- eval_id=eval_result.eval_id,
448
- final_eval_status=eval_result.final_eval_status,
449
- eval_metric_results=eval_result.eval_metric_results,
450
- session_id=eval_result.session_id,
451
- )
452
- )
453
- return run_eval_results
454
-
455
- @app.delete("/apps/{app_name}/users/{user_id}/sessions/{session_id}")
456
- def delete_session(app_name: str, user_id: str, session_id: str):
457
- # Connect to managed session if agent_engine_id is set.
458
- app_name = agent_engine_id if agent_engine_id else app_name
459
- session_service.delete_session(app_name, user_id, session_id)
460
-
461
- @app.get(
462
- "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}",
463
- response_model_exclude_none=True,
464
- )
465
- def load_artifact(
466
- app_name: str,
467
- user_id: str,
468
- session_id: str,
469
- artifact_name: str,
470
- version: Optional[int] = Query(None),
471
- ) -> Optional[types.Part]:
472
- artifact = artifact_service.load_artifact(
473
- app_name, user_id, session_id, artifact_name, version
474
- )
475
- if not artifact:
476
- raise HTTPException(status_code=404, detail="Artifact not found")
477
- return artifact
478
-
479
- @app.get(
480
- "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}/versions/{version_id}",
481
- response_model_exclude_none=True,
482
- )
483
- def load_artifact_version(
484
- app_name: str,
485
- user_id: str,
486
- session_id: str,
487
- artifact_name: str,
488
- version_id: int,
489
- ) -> Optional[types.Part]:
490
- artifact = artifact_service.load_artifact(
491
- app_name, user_id, session_id, artifact_name, version_id
492
- )
493
- if not artifact:
494
- raise HTTPException(status_code=404, detail="Artifact not found")
495
- return artifact
496
-
497
- @app.get(
498
- "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts",
499
- response_model_exclude_none=True,
500
- )
501
- def list_artifact_names(
502
- app_name: str, user_id: str, session_id: str
503
- ) -> list[str]:
504
- return artifact_service.list_artifact_keys(app_name, user_id, session_id)
505
-
506
- @app.get(
507
- "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}/versions",
508
- response_model_exclude_none=True,
509
- )
510
- def list_artifact_versions(
511
- app_name: str, user_id: str, session_id: str, artifact_name: str
512
- ) -> list[int]:
513
- return artifact_service.list_versions(
514
- app_name, user_id, session_id, artifact_name
515
- )
516
-
517
- @app.delete(
518
- "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}",
519
- )
520
- def delete_artifact(
521
- app_name: str, user_id: str, session_id: str, artifact_name: str
522
- ):
523
- artifact_service.delete_artifact(
524
- app_name, user_id, session_id, artifact_name
525
- )
526
-
527
- @app.post("/run", response_model_exclude_none=True)
528
- async def agent_run(req: AgentRunRequest) -> list[Event]:
529
- # Connect to managed session if agent_engine_id is set.
530
- app_id = agent_engine_id if agent_engine_id else req.app_name
531
- session = session_service.get_session(app_id, req.user_id, req.session_id)
532
- if not session:
533
- raise HTTPException(status_code=404, detail="Session not found")
534
- runner = _get_runner(req.app_name)
535
- events = [
536
- event
537
- async for event in runner.run_async(
538
- user_id=req.user_id,
539
- session_id=req.session_id,
540
- new_message=req.new_message,
541
- )
542
- ]
543
- logger.info("Generated %s events in agent run: %s", len(events), events)
544
- return events
545
-
546
- @app.post("/run_sse")
547
- async def agent_run_sse(req: AgentRunRequest) -> StreamingResponse:
548
- # Connect to managed session if agent_engine_id is set.
549
- app_id = agent_engine_id if agent_engine_id else req.app_name
550
- # SSE endpoint
551
- session = session_service.get_session(app_id, req.user_id, req.session_id)
552
- if not session:
553
- raise HTTPException(status_code=404, detail="Session not found")
554
-
555
- # Convert the events to properly formatted SSE
556
- async def event_generator():
557
- try:
558
- stream_mode = StreamingMode.SSE if req.streaming else StreamingMode.NONE
559
- runner = _get_runner(req.app_name)
560
- async for event in runner.run_async(
561
- user_id=req.user_id,
562
- session_id=req.session_id,
563
- new_message=req.new_message,
564
- run_config=RunConfig(streaming_mode=stream_mode),
565
- ):
566
- # Format as SSE data
567
- sse_event = event.model_dump_json(exclude_none=True, by_alias=True)
568
- logger.info("Generated event in agent run streaming: %s", sse_event)
569
- yield f"data: {sse_event}\n\n"
570
- except Exception as e:
571
- logger.exception("Error in event_generator: %s", e)
572
- # You might want to yield an error event here
573
- yield f'data: {{"error": "{str(e)}"}}\n\n'
574
-
575
- # Returns a streaming response with the proper media type for SSE
576
- return StreamingResponse(
577
- event_generator(),
578
- media_type="text/event-stream",
579
- )
580
-
581
- @app.get(
582
- "/apps/{app_name}/users/{user_id}/sessions/{session_id}/events/{event_id}/graph",
583
- response_model_exclude_none=True,
584
- )
585
- def get_event_graph(
586
- app_name: str, user_id: str, session_id: str, event_id: str
587
- ):
588
- # Connect to managed session if agent_engine_id is set.
589
- app_id = agent_engine_id if agent_engine_id else app_name
590
- session = session_service.get_session(app_id, user_id, session_id)
591
- session_events = session.events if session else []
592
- event = next((x for x in session_events if x.id == event_id), None)
593
- if event:
594
- from . import agent_graph
595
-
596
- function_calls = event.get_function_calls()
597
- function_responses = event.get_function_responses()
598
- root_agent = _get_root_agent(app_name)
599
- image_bytes = None
600
- if function_calls:
601
- function_call_highlights = []
602
- for function_call in function_calls:
603
- from_name = event.author
604
- to_name = function_call.name
605
- function_call_highlights.append((from_name, to_name))
606
- image_bytes = agent_graph.get_agent_graph(
607
- root_agent, function_call_highlights, True
608
- )
609
- elif function_responses:
610
- function_responses_highlights = []
611
- for function_response in function_responses:
612
- from_name = function_response.name
613
- to_name = event.author
614
- function_responses_highlights.append((from_name, to_name))
615
- image_bytes = agent_graph.get_agent_graph(
616
- root_agent, function_responses_highlights, True
617
- )
618
- else:
619
- from_name = event.author
620
- to_name = ""
621
- image_bytes = agent_graph.get_agent_graph(
622
- root_agent, [(from_name, to_name)], True
623
- )
624
- return Response(content=image_bytes, media_type="image/png")
625
- return None
626
-
627
- @app.websocket("/run_live")
628
- async def agent_live_run(
629
- websocket: WebSocket,
630
- app_name: str,
631
- user_id: str,
632
- session_id: str,
633
- modalities: List[Literal["TEXT", "AUDIO"]] = Query(
634
- default=["TEXT", "AUDIO"]
635
- ), # Only allows "TEXT" or "AUDIO"
636
- ) -> None:
637
- await websocket.accept()
638
-
639
- # Connect to managed session if agent_engine_id is set.
640
- app_id = agent_engine_id if agent_engine_id else app_name
641
- session = session_service.get_session(app_id, user_id, session_id)
642
- if not session:
643
- # Accept first so that the client is aware of connection establishment,
644
- # then close with a specific code.
645
- await websocket.close(code=1002, reason="Session not found")
646
- return
647
-
648
- live_request_queue = LiveRequestQueue()
649
-
650
- async def forward_events():
651
- runner = _get_runner(app_name)
652
- async for event in runner.run_live(
653
- session=session, live_request_queue=live_request_queue
654
- ):
655
- await websocket.send_text(
656
- event.model_dump_json(exclude_none=True, by_alias=True)
657
- )
658
-
659
- async def process_messages():
660
- try:
661
- while True:
662
- data = await websocket.receive_text()
663
- # Validate and send the received message to the live queue.
664
- live_request_queue.send(LiveRequest.model_validate_json(data))
665
- except ValidationError as ve:
666
- logger.error("Validation error in process_messages: %s", ve)
667
-
668
- # Run both tasks concurrently and cancel all if one fails.
669
- tasks = [
670
- asyncio.create_task(forward_events()),
671
- asyncio.create_task(process_messages()),
672
- ]
673
- done, pending = await asyncio.wait(
674
- tasks, return_when=asyncio.FIRST_EXCEPTION
675
- )
676
- try:
677
- # This will re-raise any exception from the completed tasks.
678
- for task in done:
679
- task.result()
680
- except WebSocketDisconnect:
681
- logger.info("Client disconnected during process_messages.")
682
- except Exception as e:
683
- logger.exception("Error during live websocket communication: %s", e)
684
- traceback.print_exc()
685
- finally:
686
- for task in pending:
687
- task.cancel()
688
-
689
- def _get_root_agent(app_name: str) -> Agent:
690
- """Returns the root agent for the given app."""
691
- if app_name in root_agent_dict:
692
- return root_agent_dict[app_name]
693
- envs.load_dotenv_for_agent(os.path.basename(app_name), agent_dir)
694
- agent_module = importlib.import_module(app_name)
695
- root_agent: Agent = agent_module.agent.root_agent
696
- root_agent_dict[app_name] = root_agent
697
- return root_agent
698
-
699
- def _get_runner(app_name: str) -> Runner:
700
- """Returns the runner for the given app."""
701
- if app_name in runner_dict:
702
- return runner_dict[app_name]
703
- root_agent = _get_root_agent(app_name)
704
- runner = Runner(
705
- app_name=agent_engine_id if agent_engine_id else app_name,
706
- agent=root_agent,
707
- artifact_service=artifact_service,
708
- session_service=session_service,
709
- )
710
- runner_dict[app_name] = runner
711
- return runner
712
-
713
- if web:
714
- BASE_DIR = Path(__file__).parent.resolve()
715
- ANGULAR_DIST_PATH = BASE_DIR / "browser"
716
-
717
- @app.get("/")
718
- async def redirect_to_dev_ui():
719
- return RedirectResponse("/dev-ui")
720
-
721
- @app.get("/dev-ui")
722
- async def dev_ui():
723
- return FileResponse(BASE_DIR / "browser/index.html")
724
-
725
- app.mount(
726
- "/", StaticFiles(directory=ANGULAR_DIST_PATH, html=True), name="static"
727
- )
728
- return app