Convert Session To Invocation.

PiperOrigin-RevId: 758356942
This commit is contained in:
Ankur Sharma 2025-05-13 13:39:17 -07:00 committed by Copybara-Service
parent 4188673b0f
commit 82b8e304ec
2 changed files with 93 additions and 4 deletions

View File

@ -12,11 +12,17 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from typing import Any from typing import Any, Tuple
from deprecated import deprecated
from google.genai import types as genai_types
from ...evaluation.eval_case import IntermediateData
from ...evaluation.eval_case import Invocation
from ...sessions.session import Session from ...sessions.session import Session
@deprecated(reason='Use convert_session_to_eval_invocations instead.')
def convert_session_to_eval_format(session: Session) -> list[dict[str, Any]]: def convert_session_to_eval_format(session: Session) -> list[dict[str, Any]]:
"""Converts a session data into eval format. """Converts a session data into eval format.
@ -91,3 +97,79 @@ def convert_session_to_eval_format(session: Session) -> list[dict[str, Any]]:
}) })
return eval_case return eval_case
def convert_session_to_eval_invocations(session: Session) -> list[Invocation]:
"""Converts a session data into a list of Invocation.
Args:
session: The session that should be converted.
Returns:
list: A list of invocation.
"""
invocations: list[Invocation] = []
events = session.events if session and session.events else []
for event in events:
if event.author == 'user':
if not event.content or not event.content.parts:
continue
# The content present in this event is the user content.
user_content = event.content
invocation_id = event.invocation_id
invocaton_timestamp = event.timestamp
# Find the corresponding tool usage or response for the query
tool_uses: list[genai_types.FunctionCall] = []
intermediate_responses: list[Tuple[str, list[genai_types.Part]]] = []
# Check subsequent events to extract tool uses or responses for this turn.
for subsequent_event in events[events.index(event) + 1 :]:
event_author = subsequent_event.author or 'agent'
if event_author == 'user':
# We found an event where the author was the user. This means that a
# new turn has started. So close this turn here.
break
if not subsequent_event.content or not subsequent_event.content.parts:
continue
intermediate_response_parts = []
for subsequent_part in subsequent_event.content.parts:
# Some events have both function call and reference
if subsequent_part.function_call:
tool_uses.append(subsequent_part.function_call)
elif subsequent_part.text:
# Also keep track of all the natural language responses that
# agent (or sub agents) generated.
intermediate_response_parts.append(subsequent_part)
if intermediate_response_parts:
# Only add an entry if there any intermediate entries.
intermediate_responses.append(
(event_author, intermediate_response_parts)
)
# If we are here then either we are done reading all the events or we
# encountered an event that had content authored by the end-user.
# This, basically means an end of turn.
# We assume that the last natural language intermediate response is the
# final response from the agent/model. We treat that as a reference.
invocations.append(
Invocation(
user_content=user_content,
invocation_id=invocation_id,
creation_timestamp=invocaton_timestamp,
intermediate_data=IntermediateData(
tool_uses=tool_uses,
intermediate_responses=intermediate_responses[:-1],
),
final_response=genai_types.Content(
parts=intermediate_responses[-1][1]
),
)
)
return invocations

View File

@ -13,7 +13,7 @@
# limitations under the License. # limitations under the License.
from typing import Any, Optional from typing import Any, Optional, Tuple
from google.genai import types as genai_types from google.genai import types as genai_types
from pydantic import BaseModel from pydantic import BaseModel
@ -26,9 +26,16 @@ class IntermediateData(BaseModel):
tool_uses: list[genai_types.FunctionCall] tool_uses: list[genai_types.FunctionCall]
"""Tool use trajectory in chronological order.""" """Tool use trajectory in chronological order."""
intermediate_responses: list[genai_types.Part] intermediate_responses: list[Tuple[str, list[genai_types.Part]]]
"""Intermediate responses generated by sub-agents to convey progress or status """Intermediate responses generated by sub-agents to convey progress or status
in a multi-agent system, distinct from the final response.""" in a multi-agent system, distinct from the final response.
This is expressed as a Tuple of:
- Author: Usually the sub-agent name that generated the intermediate
response.
- A list of Parts that comprise of the response.
"""
class Invocation(BaseModel): class Invocation(BaseModel):