fix(streaming): aggregate model info per invocation, not per session

src/agentevals/converter.py

@@ -45,6 +45,10 @@ class ConversionResult:
     trace_id: str
     invocations: list[Invocation] = field(default_factory=list)
     warnings: list[str] = field(default_factory=list)
+    # LLM spans each invocation was built from, parallel to ``invocations``.
+    # Kept so callers can aggregate per-invocation model info (token counts,
+    # models, providers) without walking the whole trace for every invocation.
+    invocation_llm_spans: list[list[Span]] = field(default_factory=list)
 
 
 def convert_trace(trace: Trace, format: str | None = None) -> ConversionResult:
@@ -87,12 +91,39 @@ def _convert_adk_trace(trace: Trace) -> ConversionResult:
 
     for invoke_span in invoke_spans:
         try:
-            invocation = _convert_invoke_span(invoke_span)
+            invocation, llm_spans = _convert_invoke_span(invoke_span)
             result.invocations.append(invocation)
+            result.invocation_llm_spans.append(llm_spans)
         except Exception as exc:
             msg = f"Trace {trace.trace_id}: failed to convert invoke_agent span {invoke_span.span_id}: {exc}"
             logger.warning(msg)
             result.warnings.append(msg)
+            # Orchestrators like SequentialAgent don't call an LLM themselves,
+            # so after pruning the invocation has no LLM descendants and the
+            # converter raises. Dropping the whole invocation would silently
+            # shrink a 3-step trace to 2 rows and bias every per-invocation
+            # count downstream, so we keep it: the empty ``invocation_llm_spans``
+            # slot ensures token totals stay honest, and we fall back
+            # ``user_content`` / ``final_response`` to the previous invocation
+            # so callers still get a renderable row. Without a previous
+            # invocation (the rare first-span-failed case) we emit an empty
+            # Content so the row is still well-formed.
+            prev = result.invocations[-1] if result.invocations else None
+            fallback = prev.user_content if prev is not None else genai_types.Content(
+                role="user", parts=[]
+            )
+            fallback_response = prev.final_response if prev is not None else genai_types.Content(
+                role="model", parts=[]
+            )
+            result.invocations.append(
+                Invocation(
+                    invocation_id=invoke_span.get_tag(ADK_INVOCATION_ID, invoke_span.span_id),
+                    user_content=fallback,
+                    final_response=fallback_response,
+                    creation_timestamp=invoke_span.start_time / 1_000_000.0,
+                )
+            )
+            result.invocation_llm_spans.append([])
 
     return result
 
@@ -127,7 +158,7 @@ def _find_adk_spans(trace: Trace, operation: str) -> list[Span]:
     return matches
 
 
-def _convert_invoke_span(invoke_span: Span) -> Invocation:
+def _convert_invoke_span(invoke_span: Span) -> tuple[Invocation, list[Span]]:
     llm_spans = find_adk_llm_spans_in(invoke_span)
     if not llm_spans:
         raise ValueError(
@@ -148,14 +179,16 @@ def _convert_invoke_span(invoke_span: Span) -> Invocation:
 
     invocation_id = invoke_span.get_tag(ADK_INVOCATION_ID, invoke_span.span_id)
 
-    return Invocation(
+    invocation = Invocation(
         invocation_id=invocation_id,
         user_content=user_content,
         final_response=final_response,
         intermediate_data=intermediate_data,
         creation_timestamp=invoke_span.start_time / 1_000_000.0,
     )
 
+    return invocation, llm_spans
+
 
 def _find_children_by_op(root: Span, op_prefix: str) -> list[Span]:
     results: list[Span] = []

src/agentevals/extraction.py

@@ -482,16 +482,22 @@ def collect(span: Span) -> None:
         elif is_adk_generate_content_llm_span(span):
             generate_content_spans.append(span)
 
-    _walk_descendants(root, collect)
+    _walk_descendants(root, collect, skip_invoke_agents=True)
     call_llm_spans.sort(key=lambda s: s.start_time)
     generate_content_spans.sort(key=lambda s: s.start_time)
     return call_llm_spans or generate_content_spans
 
 
-def _walk_descendants(span: Span, visit) -> None:
+def _walk_descendants(span: Span, visit, skip_invoke_agents: bool = False) -> None:
     for child in span.children:
+        # When collecting the LLM spans that belong to one invocation, a nested
+        # invoke_agent span is a separate invocation: its subtree's LLM spans are
+        # attributed to that child invocation, so walking into it here would
+        # double-count them.
+        if skip_invoke_agents and child.operation_name.startswith("invoke_agent"):
+            continue
         visit(child)
-        _walk_descendants(child, visit)
+        _walk_descendants(child, visit, skip_invoke_agents=skip_invoke_agents)
 
 
 def is_llm_span(span: Span) -> bool:

src/agentevals/genai_converter.py

@@ -93,8 +93,20 @@ def convert_genai_trace(trace: Trace) -> ConversionResult:
                 logger.debug(f"Multi-turn conversation: {len(llm_root_spans)} LLM spans")
                 try:
                     turns = _extract_multiturn_turns(llm_root_spans)
-                    for turn in turns:
+                    if len(turns) == len(llm_root_spans):
+                        # One turn per conversation span: attribute each turn to
+                        # its own span so per-invocation token counts are honest.
+                        per_turn_spans = [[span] for span in llm_root_spans]
+                    else:
+                        # Turns are derived from messages inside the same set of
+                        # conversation spans, so per-turn span attribution is not
+                        # possible at span granularity. Attribute every span to
+                        # the first turn (keeping the session total honest) and
+                        # leave the remaining turns empty to avoid double counting.
+                        per_turn_spans = [list(llm_root_spans)] + [[] for _ in turns[1:]]
+                    for turn, turn_spans in zip(turns, per_turn_spans, strict=True):
                         result.invocations.append(_turn_to_invocation(turn))
+                        result.invocation_llm_spans.append(turn_spans)
                 except Exception as exc:
                     msg = f"Trace {trace.trace_id}: failed to convert multi-turn conversation: {exc}"
                     logger.warning(msg)
@@ -110,14 +122,17 @@ def convert_genai_trace(trace: Trace) -> ConversionResult:
 
     for inv_span in invocation_spans:
         try:
-            turn = _extract_single_turn(inv_span)
+            turn, llm_spans = _extract_single_turn(inv_span)
             result.invocations.append(_turn_to_invocation(turn))
+            result.invocation_llm_spans.append(llm_spans)
         except Exception as exc:
             msg = f"Failed to convert span {inv_span.span_id}: {exc}"
             logger.warning(msg)
             result.warnings.append(msg)
 
-    result.invocations = _deduplicate_invocations(result.invocations)
+    result.invocations, result.invocation_llm_spans = _deduplicate_invocations(
+        result.invocations, result.invocation_llm_spans
+    )
     return result
 
 
@@ -158,7 +173,7 @@ def _find_genai_invocation_spans(trace: Trace) -> list[Span]:
     return candidates
 
 
-def _extract_single_turn(inv_span: Span) -> _ConversationTurn:
+def _extract_single_turn(inv_span: Span) -> tuple[_ConversationTurn, list[Span]]:
     llm_spans = _find_llm_spans(inv_span)
 
     logger.debug(f"Converting invocation span: {inv_span.operation_name}")
@@ -177,7 +192,7 @@ def _extract_single_turn(inv_span: Span) -> _ConversationTurn:
     assistant_text = _extract_assistant_text(llm_spans[-1])
     tool_calls, tool_responses = _extract_tool_calls(tool_spans, llm_spans)
 
-    return _ConversationTurn(
+    turn = _ConversationTurn(
         invocation_id=f"genai-{inv_span.span_id}",
         user_text=user_text,
         assistant_text=assistant_text,
@@ -186,6 +201,8 @@ def _extract_single_turn(inv_span: Span) -> _ConversationTurn:
         start_time=float(inv_span.start_time),
     )
 
+    return turn, llm_spans
+
 
 def _extract_multiturn_turns(llm_spans: list[Span]) -> list[_ConversationTurn]:
     messages_raw = llm_spans[0].get_tag(OTEL_GENAI_INPUT_MESSAGES, "[]")
@@ -254,17 +271,27 @@ def _extract_multiturn_turns(llm_spans: list[Span]) -> list[_ConversationTurn]:
     return turns
 
 
-def _deduplicate_invocations(invocations: list[Invocation]) -> list[Invocation]:
+def _deduplicate_invocations(
+    invocations: list[Invocation],
+    llm_spans: list[list[Span]] | None = None,
+) -> tuple[list[Invocation], list[list[Span]] | None]:
     """Deduplicate invocations with the same user text, keeping the best one.
 
     The OpenAI instrumentor creates separate LLM calls for tool-use loops within
     a single conversation turn. Each call logs the full conversation history, so
     multiple spans produce invocations with the same user text. We keep the last
     one per unique user text — it has the final response (not the intermediate
     tool-call-only response).
+
+    When ``llm_spans`` is provided it is filtered in lockstep with the
+    invocations so the per-invocation span mapping stays aligned, and the
+    dropped invocations' spans are merged into the surviving invocation for the
+    same user text so real token spend is not discarded.
+
+    Always returns the ``(invocations, llm_spans)`` tuple.
     """
     if len(invocations) <= 1:
-        return invocations
+        return invocations, llm_spans
 
     def _user_text(inv: Invocation) -> str:
         if inv.user_content and inv.user_content.parts:
@@ -281,10 +308,29 @@ def _user_text(inv: Invocation) -> str:
             seen[text] = i
 
     if len(seen) + len(always_keep) == len(invocations):
-        return invocations
+        return invocations, llm_spans
 
     keep = always_keep | set(seen.values())
-    return [inv for i, inv in enumerate(invocations) if i in keep]
+    deduped = [inv for i, inv in enumerate(invocations) if i in keep]
+
+    if llm_spans is None:
+        return deduped, None
+
+    kept_positions = [i for i in range(len(invocations)) if i in keep]
+    position_of_kept = {i: pos for pos, i in enumerate(kept_positions)}
+    merged: list[list[Span]] = [
+        list(llm_spans[i]) if i < len(llm_spans) and llm_spans[i] else [] for i in kept_positions
+    ]
+    for i, inv in enumerate(invocations):
+        if i in keep or i >= len(llm_spans) or not llm_spans[i]:
+            continue
+        text = _user_text(inv)
+        if not text.strip():
+            continue
+        survivor = seen.get(text)
+        if survivor is not None:
+            merged[position_of_kept[survivor]].extend(llm_spans[i])
+    return deduped, merged
 
 
 def _turn_to_invocation(turn: _ConversationTurn) -> Invocation:

src/agentevals/streaming/ws_server.py

@@ -23,10 +23,9 @@
 from ..extraction import (
     extract_extended_model_info_from_attrs,
     extract_token_usage_from_attrs,
-    is_llm_span,
     parse_tool_response_content,
 )
-from ..loader.base import Trace
+from ..loader.base import Span
 from ..loader.otlp import OtlpJsonLoader
 from ..trace_attrs import OTEL_GENAI_INPUT_MESSAGES, OTEL_GENAI_REQUEST_MODEL, OTEL_SERVICE_NAME
 from ..utils.log_enrichment import enrich_spans_with_logs
@@ -741,12 +740,10 @@ async def _extract_invocations(self, session: TraceSession) -> list[dict]:
 
             invocations_data = []
 
-            for trace_idx, conv_result in enumerate(conversion_results):
+            for _trace_idx, conv_result in enumerate(conversion_results):
                 if conv_result.warnings:
                     logger.warning("Conversion warnings: %s", conv_result.warnings)
 
-                trace = traces[trace_idx] if trace_idx < len(traces) else None
-
                 for inv_idx, inv in enumerate(conv_result.invocations):
                     user_text = ""
                     if inv.user_content and inv.user_content.parts:
@@ -781,8 +778,17 @@ async def _extract_invocations(self, session: TraceSession) -> list[dict]:
                             )
 
                     model_info = {}
-                    if trace:
-                        model_info = self._extract_model_info_from_trace(trace, inv_idx)
+                    if inv_idx >= len(conv_result.invocation_llm_spans):
+                        logger.warning(
+                            "Index drift: invocation %d has no recorded LLM spans "
+                            "(%d recorded); reporting blank model info",
+                            inv_idx,
+                            len(conv_result.invocation_llm_spans),
+                        )
+                        inv_llm_spans = []
+                    else:
+                        inv_llm_spans = conv_result.invocation_llm_spans[inv_idx]
+                    model_info = self._extract_model_info_from_llm_spans(inv_llm_spans)
 
                     invocations_data.append(
                         {
@@ -805,8 +811,14 @@ async def _extract_invocations(self, session: TraceSession) -> list[dict]:
             logger.exception("Failed to extract invocations")
             return []
 
-    def _extract_model_info_from_trace(self, trace: Trace, invocation_idx: int) -> dict:
-        """Extract model information from LLM spans in the trace."""
+    @staticmethod
+    def _extract_model_info_from_llm_spans(llm_spans: list[Span]) -> dict:
+        """Extract model information from the LLM spans of a single invocation.
+
+        Aggregates only the spans that belong to the invocation, so each
+        invocation shows its own token counts / models / providers instead of
+        the whole-session aggregate.
+        """
         model_info: dict[str, Any] = {}
         models_used: set[str] = set()
         total_input_tokens = 0
@@ -820,10 +832,14 @@ def _extract_model_info_from_trace(self, trace: Trace, invocation_idx: int) -> d
         first_temperature: float | None = None
         first_max_tokens: int | None = None
 
-        llm_spans = [s for s in trace.all_spans if is_llm_span(s) or "call_llm" in s.operation_name]
-        llm_spans.sort(key=lambda s: s.start_time)
+        # The caller already hands over the invocation's own LLM spans, so no
+        # re-filtering is needed here. (Re-filtering would drop a provider
+        # `generate_content`-only trace to nothing even when usage metadata is
+        # present.)
+        spans = list(llm_spans)
+        spans.sort(key=lambda s: s.start_time)
 
-        for span in llm_spans:
+        for span in spans:
             in_toks, out_toks, model = extract_token_usage_from_attrs(span.tags)
             if model and model != "unknown":
                 models_used.add(model)

tests/test_extraction.py

@@ -579,7 +579,7 @@ def test_find_llm_spans_in_ignores_provider_generate_content_without_adk_payload
         ext = AdkExtractor()
         assert ext.find_llm_spans_in(root) == []
 
-    def test_find_llm_spans_in_prefers_call_llm_over_generate_content(self):
+    def test_find_llm_spans_in_prefers_call_llm_spans(self):
         call_llm = _span(op="call_llm gemini", span_id="llm1", start_time=20)
         generate_content = _span(
             op="generate_content gemini",
@@ -589,8 +589,34 @@ def test_find_llm_spans_in_prefers_call_llm_over_generate_content(self):
         )
         root = _span(op="invoke_agent a", children=[generate_content, call_llm])
         ext = AdkExtractor()
+        # When call_llm spans are present they are preferred (matching the
+        # pre-existing behaviour); generate_content spans are a fallback only.
         assert [s.span_id for s in ext.find_llm_spans_in(root)] == ["llm1"]
 
+    def test_find_llm_spans_in_falls_back_to_generate_content(self):
+        generate_content = _span(
+            op="generate_content gemini",
+            tags={ADK_LLM_REQUEST: "{}"},
+            span_id="llm2",
+            start_time=10,
+        )
+        root = _span(op="invoke_agent a", children=[generate_content])
+        ext = AdkExtractor()
+        assert [s.span_id for s in ext.find_llm_spans_in(root)] == ["llm2"]
+
+    def test_find_llm_spans_in_skips_nested_invoke_agent(self):
+        # A coordinator delegates to a sub-agent: the sub-agent's invoke_agent
+        # span nests under the coordinator's, but its LLM spans belong to the
+        # sub-agent invocation and must not be double-counted on the coordinator.
+        sub_llm = _span(op="call_llm gemini", span_id="sub_llm", start_time=10)
+        nested_invoke = _span(
+            op="invoke_agent sub_agent", span_id="sub_invoke", children=[sub_llm]
+        )
+        own_llm = _span(op="call_llm gemini", span_id="own_llm", start_time=20)
+        root = _span(op="invoke_agent coordinator", children=[own_llm, nested_invoke])
+        ext = AdkExtractor()
+        assert [s.span_id for s in ext.find_llm_spans_in(root)] == ["own_llm"]
+
     def test_find_tool_spans_in(self):
         child_llm = _span(op="call_llm gemini", span_id="llm1")
         child_tool = _span(op="execute_tool search", span_id="tool1")