mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
99405d6a8a
commit
712da1bd36
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user