mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Populate AppDetails to each Invocation
AppDetails require two pieces of information: 1) Instructions 2) Tools Both these pieces of information are gathered using the llm_request that was passed to the model. This approach, slightly invasive, ensures that we capture the "exact" instructions and tools that were given to the model. PiperOrigin-RevId: 811180648
This commit is contained in:
Ankur Sharma
committed by
Copybara-Service
Copybara-Service
parent
2a2da0fe03
commit
d48679582d
@@ -32,6 +32,7 @@ from ..sessions.base_session_service import BaseSessionService
|
||||
from ..sessions.in_memory_session_service import InMemorySessionService
|
||||
from ..sessions.session import Session
|
||||
from ..utils.context_utils import Aclosing
|
||||
from .app_details import AgentDetails
|
||||
from .app_details import AppDetails
|
||||
from .eval_case import EvalCase
|
||||
from .eval_case import Invocation
|
||||
@@ -39,6 +40,7 @@ from .eval_case import InvocationEvent
|
||||
from .eval_case import InvocationEvents
|
||||
from .eval_case import SessionInput
|
||||
from .eval_set import EvalSet
|
||||
from .request_intercepter_plugin import _RequestIntercepterPlugin
|
||||
|
||||
_USER_AUTHOR = "user"
|
||||
_DEFAULT_AUTHOR = "agent"
|
||||
@@ -180,12 +182,16 @@ class EvaluationGenerator:
|
||||
if callable(reset_func):
|
||||
reset_func()
|
||||
|
||||
request_intercepter_plugin = _RequestIntercepterPlugin(
|
||||
name="request_intercepter_plugin"
|
||||
)
|
||||
async with Runner(
|
||||
app_name=app_name,
|
||||
agent=root_agent,
|
||||
artifact_service=artifact_service,
|
||||
session_service=session_service,
|
||||
memory_service=memory_service,
|
||||
plugins=[request_intercepter_plugin],
|
||||
) as runner:
|
||||
events = []
|
||||
|
||||
@@ -212,30 +218,36 @@ class EvaluationGenerator:
|
||||
|
||||
events.append(event)
|
||||
|
||||
return EvaluationGenerator.convert_events_to_eval_invocations(events)
|
||||
app_details_by_invocation_id = (
|
||||
EvaluationGenerator._get_app_details_by_invocation_id(
|
||||
events, request_intercepter_plugin
|
||||
)
|
||||
)
|
||||
return EvaluationGenerator.convert_events_to_eval_invocations(
|
||||
events, app_details_by_invocation_id
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def convert_events_to_eval_invocations(
|
||||
events: list[Event],
|
||||
app_details_per_invocation: Optional[dict[str, AppDetails]] = None,
|
||||
) -> list[Invocation]:
|
||||
"""Converts a list of events to eval invocations."""
|
||||
# Group Events by invocation id. Events that share the same invocation id
|
||||
# belong to the same invocation.
|
||||
events_by_invocation_id: dict[str, list[Event]] = {}
|
||||
|
||||
for event in events:
|
||||
invocation_id = event.invocation_id
|
||||
|
||||
if invocation_id not in events_by_invocation_id:
|
||||
events_by_invocation_id[invocation_id] = []
|
||||
|
||||
events_by_invocation_id[invocation_id].append(event)
|
||||
events_by_invocation_id = (
|
||||
EvaluationGenerator._collect_events_by_invocation_id(events)
|
||||
)
|
||||
|
||||
invocations = []
|
||||
for invocation_id, events in events_by_invocation_id.items():
|
||||
final_response = None
|
||||
user_content = ""
|
||||
invocation_timestamp = 0
|
||||
app_details = None
|
||||
if (
|
||||
app_details_per_invocation
|
||||
and invocation_id in app_details_per_invocation
|
||||
):
|
||||
app_details = app_details_per_invocation[invocation_id]
|
||||
|
||||
events_to_add = []
|
||||
|
||||
@@ -271,11 +283,61 @@ class EvaluationGenerator:
|
||||
invocation_events=invocation_events
|
||||
),
|
||||
creation_timestamp=invocation_timestamp,
|
||||
app_details=app_details,
|
||||
)
|
||||
)
|
||||
|
||||
return invocations
|
||||
|
||||
@staticmethod
|
||||
def _get_app_details_by_invocation_id(
|
||||
events: list[Event], request_intercepter: _RequestIntercepterPlugin
|
||||
) -> dict[str, AppDetails]:
|
||||
"""Creates an AppDetails object from the list of events."""
|
||||
events_by_invocation_id = (
|
||||
EvaluationGenerator._collect_events_by_invocation_id(events)
|
||||
)
|
||||
app_details_by_invocation_id = {}
|
||||
|
||||
for invocation_id, events in events_by_invocation_id.items():
|
||||
app_details = AppDetails(agent_details={})
|
||||
app_details_by_invocation_id[invocation_id] = app_details
|
||||
|
||||
for event in events:
|
||||
if event.author == _USER_AUTHOR:
|
||||
continue
|
||||
|
||||
llm_request = request_intercepter.get_model_request(event)
|
||||
|
||||
if not llm_request:
|
||||
continue
|
||||
|
||||
if event.author not in app_details.agent_details:
|
||||
agent_name = event.author
|
||||
app_details.agent_details[agent_name] = AgentDetails(
|
||||
name=agent_name,
|
||||
instructions=llm_request.config.system_instruction,
|
||||
tool_declarations=llm_request.config.tools or [],
|
||||
)
|
||||
|
||||
return app_details_by_invocation_id
|
||||
|
||||
@staticmethod
|
||||
def _collect_events_by_invocation_id(events: list[Event]) -> dict[str, Event]:
|
||||
# Group Events by invocation id. Events that share the same invocation id
|
||||
# belong to the same invocation.
|
||||
events_by_invocation_id: dict[str, list[Event]] = {}
|
||||
|
||||
for event in events:
|
||||
invocation_id = event.invocation_id
|
||||
|
||||
if invocation_id not in events_by_invocation_id:
|
||||
events_by_invocation_id[invocation_id] = []
|
||||
|
||||
events_by_invocation_id[invocation_id].append(event)
|
||||
|
||||
return events_by_invocation_id
|
||||
|
||||
@staticmethod
|
||||
def _process_query_with_session(session_data, data):
|
||||
"""Process the queries using the existing session data without invoking the runner."""
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
import uuid
|
||||
|
||||
from typing_extensions import override
|
||||
|
||||
from ..agents.callback_context import CallbackContext
|
||||
from ..models.llm_request import LlmRequest
|
||||
from ..models.llm_response import LlmResponse
|
||||
from ..plugins.base_plugin import BasePlugin
|
||||
|
||||
logger = logging.getLogger("google_adk." + __name__)
|
||||
|
||||
_LLM_REQUEST_ID_KEY = "__llm_request_key__"
|
||||
|
||||
|
||||
class _RequestIntercepterPlugin(BasePlugin):
|
||||
"""A plugin that intercepts requests that are made to the model and couples them with the model response.
|
||||
|
||||
NOTE: This implementation is intended for eval systems internal usage. Do not
|
||||
take direct depdency on it.
|
||||
|
||||
Context behind the creation of this intercepter:
|
||||
Some of the newer AutoRater backed metrics need access the pieces of
|
||||
information that were presented to the model like instructions and the list
|
||||
of available tools.
|
||||
|
||||
We intercept the llm_request using this intercepter and make it available to
|
||||
eval system.
|
||||
|
||||
How is it done?
|
||||
The class maintains a cache of llm_requests that pass through it. Each request
|
||||
is given a unique id. The id is put in custom_metadata field of the response.
|
||||
Eval systems have access to the response and can use the request id to
|
||||
get the llm_request.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__(name=name)
|
||||
self._llm_requests_cache: dict[str, LlmRequest] = {}
|
||||
|
||||
@override
|
||||
async def before_model_callback(
|
||||
self, *, callback_context: CallbackContext, llm_request: LlmRequest
|
||||
) -> Optional[LlmResponse]:
|
||||
# We add the llm_request to the call back context so that we can fetch
|
||||
# it later.
|
||||
request_id = str(uuid.uuid4())
|
||||
self._llm_requests_cache[request_id] = llm_request
|
||||
callback_context.state[_LLM_REQUEST_ID_KEY] = request_id
|
||||
|
||||
@override
|
||||
async def after_model_callback(
|
||||
self, *, callback_context: CallbackContext, llm_response: LlmResponse
|
||||
) -> Optional[LlmResponse]:
|
||||
# Fetch the request_id from the callback_context
|
||||
if callback_context and _LLM_REQUEST_ID_KEY in callback_context.state:
|
||||
if llm_response.custom_metadata is None:
|
||||
llm_response.custom_metadata = {}
|
||||
|
||||
llm_response.custom_metadata[_LLM_REQUEST_ID_KEY] = (
|
||||
callback_context.state[_LLM_REQUEST_ID_KEY]
|
||||
)
|
||||
|
||||
def get_model_request(
|
||||
self, llm_response: LlmResponse
|
||||
) -> Optional[LlmRequest]:
|
||||
"""Fetches the request object, if found."""
|
||||
if (
|
||||
llm_response.custom_metadata
|
||||
and _LLM_REQUEST_ID_KEY in llm_response.custom_metadata
|
||||
):
|
||||
request_id = llm_response.custom_metadata[_LLM_REQUEST_ID_KEY]
|
||||
|
||||
if request_id in self._llm_requests_cache:
|
||||
return self._llm_requests_cache[request_id]
|
||||
else:
|
||||
logger.warning("`%s` not found in llm_request_cache.", request_id)
|
||||
@@ -14,8 +14,14 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest import mock
|
||||
|
||||
from google.adk.evaluation.app_details import AgentDetails
|
||||
from google.adk.evaluation.app_details import AppDetails
|
||||
from google.adk.evaluation.evaluation_generator import EvaluationGenerator
|
||||
from google.adk.evaluation.request_intercepter_plugin import _RequestIntercepterPlugin
|
||||
from google.adk.events.event import Event
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.genai import types
|
||||
|
||||
|
||||
@@ -195,3 +201,128 @@ class TestConvertEventsToEvalInvocation:
|
||||
assert events[1].author == "sub_agent_1"
|
||||
assert events[2].author == "sub_agent_1"
|
||||
assert events[3].author == "sub_agent_2"
|
||||
|
||||
|
||||
class TestGetAppDetailsByInvocationId:
|
||||
"""Test cases for EvaluationGenerator._get_app_details_by_invocation_id method."""
|
||||
|
||||
def test_get_app_details_by_invocation_id_empty(self):
|
||||
"""Tests with an empty list of events."""
|
||||
mock_request_intercepter = mock.MagicMock(spec=_RequestIntercepterPlugin)
|
||||
app_details = EvaluationGenerator._get_app_details_by_invocation_id(
|
||||
[], mock_request_intercepter
|
||||
)
|
||||
assert app_details == {}
|
||||
|
||||
def test_get_app_details_by_invocation_id_no_model_requests(self):
|
||||
"""Tests when request_intercepter returns no model requests."""
|
||||
mock_request_intercepter = mock.MagicMock(spec=_RequestIntercepterPlugin)
|
||||
mock_request_intercepter.get_model_request.return_value = None
|
||||
events = [
|
||||
_build_event("user", [types.Part(text="Hello")], "inv1"),
|
||||
_build_event("agent", [types.Part(text="Hi there!")], "inv1"),
|
||||
]
|
||||
app_details = EvaluationGenerator._get_app_details_by_invocation_id(
|
||||
events, mock_request_intercepter
|
||||
)
|
||||
assert app_details == {"inv1": AppDetails(agent_details={})}
|
||||
mock_request_intercepter.get_model_request.assert_called_once_with(
|
||||
events[1]
|
||||
)
|
||||
|
||||
def test_get_app_details_single_invocation_single_agent(self):
|
||||
"""Tests a single invocation with one agent."""
|
||||
mock_request_intercepter = mock.MagicMock(spec=_RequestIntercepterPlugin)
|
||||
mock_llm_request = LlmRequest(model="test")
|
||||
mock_llm_request.config.system_instruction = "instruction1"
|
||||
mock_llm_request.config.tools = [types.Tool()]
|
||||
mock_request_intercepter.get_model_request.return_value = mock_llm_request
|
||||
|
||||
events = [
|
||||
_build_event("user", [types.Part(text="Hello")], "inv1"),
|
||||
_build_event("agent", [types.Part(text="Hi there!")], "inv1"),
|
||||
]
|
||||
app_details = EvaluationGenerator._get_app_details_by_invocation_id(
|
||||
events, mock_request_intercepter
|
||||
)
|
||||
|
||||
expected_app_details = {
|
||||
"inv1": AppDetails(
|
||||
agent_details={
|
||||
"agent": AgentDetails(
|
||||
name="agent",
|
||||
instructions="instruction1",
|
||||
tool_declarations=[types.Tool()],
|
||||
)
|
||||
}
|
||||
)
|
||||
}
|
||||
assert app_details == expected_app_details
|
||||
mock_request_intercepter.get_model_request.assert_called_once_with(
|
||||
events[1]
|
||||
)
|
||||
|
||||
def test_get_app_details_multiple_invocations_multiple_agents(self):
|
||||
"""Tests multiple invocations with multiple agents."""
|
||||
mock_request_intercepter = mock.MagicMock(spec=_RequestIntercepterPlugin)
|
||||
|
||||
def get_model_request_side_effect(event):
|
||||
mock_llm_request = LlmRequest(model="test")
|
||||
if event.invocation_id == "inv1" and event.author == "agent1":
|
||||
mock_llm_request.config.system_instruction = "instruction1"
|
||||
mock_llm_request.config.tools = [
|
||||
types.Tool(
|
||||
function_declarations=[types.FunctionDeclaration(name="tool1")]
|
||||
)
|
||||
]
|
||||
return mock_llm_request
|
||||
if event.invocation_id == "inv2" and event.author == "agent2":
|
||||
mock_llm_request.config.system_instruction = "instruction2"
|
||||
return mock_llm_request
|
||||
return None
|
||||
|
||||
mock_request_intercepter.get_model_request.side_effect = (
|
||||
get_model_request_side_effect
|
||||
)
|
||||
|
||||
events = [
|
||||
_build_event("user", [types.Part(text="Hello")], "inv1"),
|
||||
_build_event("agent1", [types.Part(text="Hi there!")], "inv1"),
|
||||
_build_event("user", [types.Part(text="Hello again")], "inv2"),
|
||||
_build_event("agent2", [types.Part(text="Hi again!")], "inv2"),
|
||||
_build_event(
|
||||
"agent1", [types.Part(text="Hi again from agent1")], "inv2"
|
||||
), # no request
|
||||
]
|
||||
app_details = EvaluationGenerator._get_app_details_by_invocation_id(
|
||||
events, mock_request_intercepter
|
||||
)
|
||||
|
||||
expected_app_details = {
|
||||
"inv1": AppDetails(
|
||||
agent_details={
|
||||
"agent1": AgentDetails(
|
||||
name="agent1",
|
||||
instructions="instruction1",
|
||||
tool_declarations=[
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name="tool1")
|
||||
]
|
||||
)
|
||||
],
|
||||
)
|
||||
}
|
||||
),
|
||||
"inv2": AppDetails(
|
||||
agent_details={
|
||||
"agent2": AgentDetails(
|
||||
name="agent2",
|
||||
instructions="instruction2",
|
||||
tool_declarations=[],
|
||||
)
|
||||
}
|
||||
),
|
||||
}
|
||||
assert app_details == expected_app_details
|
||||
assert mock_request_intercepter.get_model_request.call_count == 3
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest import mock
|
||||
|
||||
from google.adk.agents.callback_context import CallbackContext
|
||||
from google.adk.evaluation.request_intercepter_plugin import _LLM_REQUEST_ID_KEY
|
||||
from google.adk.evaluation.request_intercepter_plugin import _RequestIntercepterPlugin
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.models.llm_response import LlmResponse
|
||||
from google.genai import types
|
||||
|
||||
|
||||
class TestRequestIntercepterPlugin:
|
||||
|
||||
async def test_intercept_request_and_response(self):
|
||||
plugin = _RequestIntercepterPlugin(name="test_plugin")
|
||||
llm_request = LlmRequest(
|
||||
model="test_model",
|
||||
contents=[
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[types.Part(text="hello")],
|
||||
)
|
||||
],
|
||||
)
|
||||
mock_invocation_context = mock.MagicMock()
|
||||
mock_invocation_context.session.state = {}
|
||||
callback_context = CallbackContext(mock_invocation_context)
|
||||
llm_response = LlmResponse()
|
||||
|
||||
# Test before_model_callback
|
||||
await plugin.before_model_callback(
|
||||
callback_context=callback_context, llm_request=llm_request
|
||||
)
|
||||
assert _LLM_REQUEST_ID_KEY in callback_context.state
|
||||
request_id = callback_context.state[_LLM_REQUEST_ID_KEY]
|
||||
assert isinstance(request_id, str)
|
||||
|
||||
# Test after_model_callback
|
||||
await plugin.after_model_callback(
|
||||
callback_context=callback_context, llm_response=llm_response
|
||||
)
|
||||
assert llm_response.custom_metadata is not None
|
||||
assert _LLM_REQUEST_ID_KEY in llm_response.custom_metadata
|
||||
assert llm_response.custom_metadata[_LLM_REQUEST_ID_KEY] == request_id
|
||||
|
||||
# Test get_model_request
|
||||
retrieved_request = plugin.get_model_request(llm_response)
|
||||
assert retrieved_request == llm_request
|
||||
|
||||
def test_get_model_request_not_found(self):
|
||||
plugin = _RequestIntercepterPlugin(name="test_plugin")
|
||||
llm_response = LlmResponse()
|
||||
assert plugin.get_model_request(llm_response) is None
|
||||
|
||||
llm_response_with_metadata = LlmResponse(
|
||||
custom_metadata={_LLM_REQUEST_ID_KEY: "non_existent_id"}
|
||||
)
|
||||
assert plugin.get_model_request(llm_response_with_metadata) is None
|
||||
Reference in New Issue
Block a user