feat(conformance): Integrates RecordingsPlugin into AdkWebServer to record Llm interactions and tool calls

When start the server with `--extra_plugins=google.adk.cli.plugins.recordings_plugin.RecordingsPlugin`, it will trigger recording with expected state in session.

PiperOrigin-RevId: 808432022
This commit is contained in:
Wei Sun (Jack)
2025-09-18 00:05:06 -07:00
committed by Copybara-Service
parent 99405d6a8a
commit 712da1bd36
4 changed files with 109 additions and 17 deletions
+65 -7
View File
@@ -16,6 +16,7 @@ from __future__ import annotations
import asyncio
from contextlib import asynccontextmanager
import importlib
import logging
import os
import time
@@ -75,6 +76,7 @@ from ..evaluation.eval_set_results_manager import EvalSetResultsManager
from ..evaluation.eval_sets_manager import EvalSetsManager
from ..events.event import Event
from ..memory.base_memory_service import BaseMemoryService
from ..plugins.base_plugin import BasePlugin
from ..runners import Runner
from ..sessions.base_session_service import BaseSessionService
from ..sessions.session import Session
@@ -354,6 +356,7 @@ class AdkWebServer:
eval_sets_manager: EvalSetsManager,
eval_set_results_manager: EvalSetResultsManager,
agents_dir: str,
extra_plugins: Optional[list[str]] = None,
):
self.agent_loader = agent_loader
self.session_service = session_service
@@ -363,39 +366,94 @@ class AdkWebServer:
self.eval_sets_manager = eval_sets_manager
self.eval_set_results_manager = eval_set_results_manager
self.agents_dir = agents_dir
self.extra_plugins = extra_plugins or []
# Internal propeties we want to allow being modified from callbacks.
self.runners_to_clean: set[str] = set()
self.current_app_name_ref: SharedValue[str] = SharedValue(value="")
self.runner_dict = {}
async def get_runner_async(self, app_name: str) -> Runner:
"""Returns the runner for the given app."""
"""Returns the cached runner for the given app."""
# Handle cleanup
if app_name in self.runners_to_clean:
self.runners_to_clean.remove(app_name)
runner = self.runner_dict.pop(app_name, None)
await cleanup.close_runners(list([runner]))
envs.load_dotenv_for_agent(os.path.basename(app_name), self.agents_dir)
# Return cached runner if exists
if app_name in self.runner_dict:
return self.runner_dict[app_name]
# Create new runner
envs.load_dotenv_for_agent(os.path.basename(app_name), self.agents_dir)
agent_or_app = self.agent_loader.load_agent(app_name)
agentic_app = None
# Instantiate extra plugins if configured
extra_plugins_instances = self._instantiate_extra_plugins()
if isinstance(agent_or_app, BaseAgent):
agentic_app = App(
name=app_name,
root_agent=agent_or_app,
plugins=extra_plugins_instances,
)
else:
agentic_app = agent_or_app
runner = Runner(
# Combine existing plugins with extra plugins
all_plugins = (agent_or_app.plugins or []) + extra_plugins_instances
agentic_app = App(
name=agent_or_app.name,
root_agent=agent_or_app.root_agent,
plugins=all_plugins,
)
runner = self._create_runner(agentic_app)
self.runner_dict[app_name] = runner
return runner
def _create_runner(self, agentic_app: App) -> Runner:
"""Create a runner with common services."""
return Runner(
app=agentic_app,
artifact_service=self.artifact_service,
session_service=self.session_service,
memory_service=self.memory_service,
credential_service=self.credential_service,
)
self.runner_dict[app_name] = runner
return runner
def _instantiate_extra_plugins(self) -> list[BasePlugin]:
"""Instantiate extra plugins from the configured list.
Returns:
List of instantiated BasePlugin objects.
"""
extra_plugins_instances = []
for qualified_name in self.extra_plugins:
try:
plugin_obj = self._import_plugin_object(qualified_name)
if isinstance(plugin_obj, BasePlugin):
extra_plugins_instances.append(plugin_obj)
elif issubclass(plugin_obj, BasePlugin):
extra_plugins_instances.append(plugin_obj(name=qualified_name))
except Exception as e:
logger.error("Failed to load plugin %s: %s", qualified_name, e)
return extra_plugins_instances
def _import_plugin_object(self, qualified_name: str) -> Any:
"""Import a plugin object (class or instance) from a fully qualified name.
Args:
qualified_name: Fully qualified name (e.g., 'my_package.my_plugin.MyPlugin')
Returns:
The imported object, which can be either a class or an instance.
Raises:
ImportError: If the module cannot be imported.
AttributeError: If the object doesn't exist in the module.
"""
module_name, obj_name = qualified_name.rsplit(".", 1)
module = importlib.import_module(module_name)
return getattr(module, obj_name)
def get_fast_api_app(
self,
@@ -22,6 +22,7 @@ import logging
from typing import Any
from typing import AsyncGenerator
from typing import Dict
from typing import Literal
from typing import Optional
import httpx
@@ -178,34 +179,59 @@ class AdkWebServerClient:
async def run_agent(
self,
request: RunAgentRequest,
mode: Optional[Literal["record", "replay"]] = None,
test_case_dir: Optional[str] = None,
user_message_index: Optional[int] = None,
) -> AsyncGenerator[Event, None]:
"""Run an agent with streaming Server-Sent Events response.
Args:
request: The RunAgentRequest containing agent execution parameters
mode: Optional conformance mode ("record" or "replay") to trigger recording
test_case_dir: Optional test case directory path for conformance recording
user_message_index: Optional user message index for conformance recording
Yields:
Event objects streamed from the agent execution
Raises:
ValueError: If mode is provided but test_case_dir or user_message_index is None
httpx.HTTPStatusError: If the request fails
json.JSONDecodeError: If event data cannot be parsed
"""
# TODO: Prepare headers for conformance tracking
headers = {}
# Add recording parameters to state_delta for conformance tests
if mode:
if test_case_dir is None or user_message_index is None:
raise ValueError(
"test_case_dir and user_message_index must be provided when mode is"
" specified"
)
# Modify request state_delta in place
if request.state_delta is None:
request.state_delta = {}
if mode == "replay":
request.state_delta["_adk_replay_config"] = {
"dir": str(test_case_dir),
"user_message_index": user_message_index,
}
else: # record mode
request.state_delta["_adk_recordings_config"] = {
"dir": str(test_case_dir),
"user_message_index": user_message_index,
}
async with self._get_client() as client:
async with client.stream(
"POST",
"/run_sse",
json=request.model_dump(by_alias=True, exclude_none=True),
headers=headers,
) as response:
response.raise_for_status()
async for line in response.aiter_lines():
if line.startswith("data:") and (data := line[5:].strip()):
try:
event_data = json.loads(data)
yield Event.model_validate(event_data)
except (json.JSONDecodeError, ValueError) as exc:
logger.warning("Failed to parse event data: %s", exc)
event_data = json.loads(data)
yield Event.model_validate(event_data)
else:
logger.debug("Non data line received: %s", line)
+2
View File
@@ -70,6 +70,7 @@ def get_fast_api_app(
otel_to_cloud: bool = False,
reload_agents: bool = False,
lifespan: Optional[Lifespan[FastAPI]] = None,
extra_plugins: Optional[list[str]] = None,
) -> FastAPI:
# Set up eval managers.
if eval_storage_uri:
@@ -187,6 +188,7 @@ def get_fast_api_app(
eval_sets_manager=eval_sets_manager,
eval_set_results_manager=eval_set_results_manager,
agents_dir=agents_dir,
extra_plugins=extra_plugins,
)
# Callbacks & other optional args for when constructing the FastAPI instance
+8 -2
View File
@@ -28,6 +28,7 @@ def dump_pydantic_to_yaml(
indent: int = 2,
sort_keys: bool = True,
exclude_none: bool = True,
exclude_defaults: bool = True,
) -> None:
"""Dump a Pydantic model to a YAML file with multiline strings using | style.
@@ -38,7 +39,11 @@ def dump_pydantic_to_yaml(
sort_keys: Whether to sort dictionary keys (default: True).
exclude_none: Exclude fields with None values (default: True).
"""
model_dict = model.model_dump(exclude_none=exclude_none, mode='json')
model_dict = model.model_dump(
exclude_none=exclude_none,
exclude_defaults=exclude_defaults,
mode='json',
)
file_path = Path(file_path)
file_path.parent.mkdir(parents=True, exist_ok=True)
@@ -55,7 +60,7 @@ def dump_pydantic_to_yaml(
return super(_MultilineDumper, self).increase_indent(flow, False)
def multiline_str_representer(dumper, data):
if '\n' in data:
if '\n' in data or '"' in data or "'" in data:
return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
return dumper.represent_scalar('tag:yaml.org,2002:str', data)
@@ -70,4 +75,5 @@ def dump_pydantic_to_yaml(
indent=indent,
sort_keys=sort_keys,
default_flow_style=False,
width=1000000, # Essentially disable text wraps
)