refactor(dea): simplify A2A tool metrics extraction and transcript telemetry

- Implement stateless single-pass A2A ConversationToken decoding in gcp_data_engineering_agent.py
- Extract function call parameters, outputs, and failure flags without state locks or diffing
- Format structured per-turn tool telemetry inside conversation_history in dataengineeringagentevaluator.py
- Add typed state attributes to EvalDeaRequest in dataengineeringagentinput.py

Author: James Nguyen

evalbench/dataset/dataengineeringagentinput.py

@@ -1,5 +1,6 @@
 import json
 import copy
+from typing import Any
 
 
 class EvalDeaRequest:
@@ -21,8 +22,11 @@ def __init__(self, raw_dict: dict, job_id: str = "", trace_id: str = ""):
         self.payload_str = json.dumps(raw_dict)
         self.payload = self.payload_str
 
-        self.agent_results = []
-        self.scoring_results = []
+        self.agent_results: list[dict[str, Any]] = []
+        self.scoring_results: list[dict[str, Any]] = []
+
+        self.accumulated_tools: list[str] = []
+        self.this_turn_tool_details: list[dict[str, Any]] = []
 
     @classmethod
     def init_from_proto(cls, proto):

evalbench/evaluator/dataengineeringagentevaluator.py

@@ -200,9 +200,37 @@ def process_scenario(
                 agent_text
             )
 
+            this_turn_tools = getattr(
+                eval_result, "this_turn_tool_details", []
+            )
+
+            tools_by_name = {}
+            for t in this_turn_tools:
+                name = t["name"]
+                params = t["params"]
+                output = t.get("output", {})
+                fail = t.get("fail", 0)
+
+                tool_data = tools_by_name.setdefault(
+                    name, {"parameters": [], "outputs": [], "fail": 0}
+                )
+                tool_data["parameters"].append(params)
+                tool_data["outputs"].append(output)
+                if fail:
+                    tool_data["fail"] = 1
+
+            agent_json = {
+                "response": agent_text,
+                "stats": {
+                    "tools": {
+                        "byName": tools_by_name
+                    }
+                }
+            }
+
             conversation_history.append({
                 "user": current_prompt,
-                "agent": agent_text,
+                "agent": json.dumps(agent_json),
             })
 
             # Simulated User checks conversation plan and generates next prompt
@@ -316,7 +344,9 @@ def _finalize_scenario(
             "prompt": scenario["starting_prompt"],
             "conversation_history": json.dumps(conversation_history, indent=2),
             "scenario": scenario,
-            "accumulated_tools": [],
+            "accumulated_tools": getattr(
+                eval_result, "accumulated_tools", []
+            ),
             "accumulated_skills": [],
             "job_id": job_id,
             "metadata": metadata,

evalbench/generators/models/gcp_data_engineering_agent.py

@@ -1,9 +1,11 @@
 import asyncio
+import base64
 import collections.abc
 import concurrent.futures
 import datetime
 import json
 import logging
+import re
 import threading
 import uuid
 from typing import Any, Coroutine
@@ -197,6 +199,81 @@ def _find_agent_text_recursive(obj: Any) -> str:
     return "\n\n".join(texts)
 
 
+def _find_json_objects(
+    text: str,
+) -> collections.abc.Iterable[dict[str, Any]]:
+    """Finds top-level JSON objects starting with {"id": in text."""
+    decoder = json.JSONDecoder()
+    for match in re.finditer(r'\{"id":', text):
+        start = match.start()
+        try:
+            obj, _ = decoder.raw_decode(text, start)
+            if isinstance(obj, dict):
+                yield obj
+        except json.JSONDecodeError:
+            continue
+
+
+def _extract_tool_details_from_token(
+    token_str: str,
+) -> list[dict[str, Any]]:
+    """Extracts tool execution details from A2A conversation token."""
+    if not token_str:
+        return []
+
+    cleaned_token_str = re.sub(r"[^A-Za-z0-9+/=]", "", token_str)
+
+    try:
+        decoded_bytes = base64.b64decode(cleaned_token_str)
+        decoded_str = decoded_bytes.decode("utf-8", errors="ignore")
+    except (ValueError, UnicodeDecodeError) as e:
+        logger.exception("Failed to decode token: %s", e)
+        return []
+
+    calls = {}
+    responses = {}
+
+    for js in _find_json_objects(decoded_str):
+        content = js.get("content", {})
+        if not isinstance(content, dict):
+            continue
+        parts = content.get("parts")
+        if not isinstance(parts, list):
+            continue
+        for part in parts:
+            if not isinstance(part, dict):
+                continue
+            if "functionCall" in part:
+                fc = part["functionCall"]
+                fc_id = fc.get("id")
+                if fc_id:
+                    calls[fc_id] = {
+                        "id": fc_id,
+                        "name": fc["name"],
+                        "params": fc.get("args", {}),
+                        "output": {},
+                        "fail": 0,
+                    }
+            if "functionResponse" in part:
+                fr = part["functionResponse"]
+                fr_id = fr.get("id")
+                if fr_id:
+                    responses[fr_id] = fr.get("response", {})
+
+    # Match calls and responses to determine failure and store output
+    for fr_id, resp_payload in responses.items():
+        if fr_id in calls:
+            calls[fr_id]["output"] = resp_payload
+            if isinstance(resp_payload, dict):
+                if resp_payload.get("error") or resp_payload.get("errors"):
+                    calls[fr_id]["fail"] = 1
+            elif isinstance(resp_payload, str):
+                if "error" in resp_payload.lower():
+                    calls[fr_id]["fail"] = 1
+
+    return list(calls.values())
+
+
 class DataEngineeringAgentGenerator(QueryGenerator):
     """Data Engineering Agent (DEA) Query Generator using the A2A SDK."""
 
@@ -261,7 +338,12 @@ def generate_internal(self, prompt: EvalDeaRequest) -> EvalDeaRequest:
         )
 
         try:
-            prompt.generated_nl_response = self.run_async(coro)
+            reply_text, new_token = self.run_async(coro)
+            prompt.generated_nl_response = reply_text
+            if new_token:
+                all_tools = _extract_tool_details_from_token(new_token)
+                prompt.this_turn_tool_details = all_tools
+                prompt.accumulated_tools = list({t["name"] for t in all_tools})
         except Exception:
             logger.exception("A2A SDK messaging error")
             raise
@@ -272,7 +354,7 @@ async def _run_client(
         prompt: str,
         conversation_id: str | None,
         target_workspace: str,
-    ) -> str:
+    ) -> tuple[str, str]:
         """Core asynchronous A2A SDK connection loop."""
         # Configure Client in standard Non-Streaming Mode
         config = ClientConfig(
@@ -326,11 +408,13 @@ async def _run_client(
         }
 
         # Handle ConversationToken state memory thread-safely
-        token = ""
+        conversation_token = ""
         with self._token_lock:
-            token = self._conversation_token_cache.get(conversation_id, "")
-        if token:
-            message_req.metadata[CONVERSATION_TOKEN_URI] = token
+            conversation_token = self._conversation_token_cache.get(
+                conversation_id, ""
+            )
+        if conversation_token:
+            message_req.metadata[CONVERSATION_TOKEN_URI] = conversation_token
 
         context = ClientCallContext(
             timeout=300.0,
@@ -340,7 +424,7 @@ async def _run_client(
         )
 
         reply_text = ""
-        new_token = ""
+        new_conversation_token = ""
 
         try:
             async for resp in client.send_message(
@@ -358,19 +442,23 @@ async def _run_client(
                     resp.HasField("task")
                     and CONVERSATION_TOKEN_URI in resp.task.metadata
                 ):
-                    new_token = resp.task.metadata[CONVERSATION_TOKEN_URI]
+                    new_conversation_token = resp.task.metadata[
+                        CONVERSATION_TOKEN_URI
+                    ]
         except Exception as e:
             self._log_api_error_details(e)
             raise
         finally:
             await client.close()
 
         # Cache the new token thread-safely
-        if new_token:
+        if new_conversation_token:
             with self._token_lock:
-                self._conversation_token_cache[conversation_id] = new_token
+                self._conversation_token_cache[
+                    conversation_id
+                ] = new_conversation_token
 
-        return reply_text.strip()
+        return reply_text.strip(), new_conversation_token
 
     @staticmethod
     def run_async(coro: Coroutine[Any, Any, Any]) -> Any: