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,8 +91,9 @@ 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)
@@ -127,7 +132,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 +153,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/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_returns_both_call_llm_and_generate_content(self):
         call_llm = _span(op="call_llm gemini", span_id="llm1", start_time=20)
         generate_content = _span(
             op="generate_content gemini",
@@ -589,7 +589,10 @@ 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()
-        assert [s.span_id for s in ext.find_llm_spans_in(root)] == ["llm1"]
+        # Both span kinds are returned (sorted by start time) so a provider
+        # generate_content span's request temperature / response model are not
+        # dropped from modelInfo when call_llm spans are also present.
+        assert [s.span_id for s in ext.find_llm_spans_in(root)] == ["llm2", "llm1"]
 
     def test_find_tool_spans_in(self):
         child_llm = _span(op="call_llm gemini", span_id="llm1")

tests/test_genai_converter.py

@@ -616,15 +616,15 @@ def test_no_dedup_when_all_unique(self):
             self._make_invocation("Q2", "A2"),
             self._make_invocation("Q3", "A3"),
         ]
-        result = _deduplicate_invocations(invocations)
+        result, _ = _deduplicate_invocations(invocations)
         assert len(result) == 3
 
     def test_dedup_keeps_last_duplicate(self):
         invocations = [
             self._make_invocation("Roll a die", "tool_call"),
             self._make_invocation("Roll a die", "I rolled a 3!"),
         ]
-        result = _deduplicate_invocations(invocations)
+        result, _ = _deduplicate_invocations(invocations)
         assert len(result) == 1
         assert result[0].final_response.parts[0].text == "I rolled a 3!"
 
@@ -635,23 +635,49 @@ def test_dedup_preserves_order(self):
             self._make_invocation("Q2", "A2-intermediate"),
             self._make_invocation("Q2", "A2-final"),
         ]
-        result = _deduplicate_invocations(invocations)
+        result, _ = _deduplicate_invocations(invocations)
         assert len(result) == 2
         assert result[0].final_response.parts[0].text == "A1-final"
         assert result[1].final_response.parts[0].text == "A2-final"
 
+    def test_dedup_merges_dropped_spans_into_survivor(self):
+        """Dropped duplicate invocations keep their LLM spans (real spend)."""
+        invocations = [
+            self._make_invocation("Roll a die", "tool_call"),
+            self._make_invocation("Roll a die", "I rolled a 3!"),
+        ]
+        spans = [
+            [_make_genai_llm_span("s1", model="m", input_tokens=50, output_tokens=5)],
+            [_make_genai_llm_span("s2", model="m", input_tokens=100, output_tokens=10)],
+        ]
+        deduped, merged = _deduplicate_invocations(invocations, spans)
+        assert len(deduped) == 1
+        assert deduped[0].final_response.parts[0].text == "I rolled a 3!"
+        # The dropped invocation's span (s1) is merged into the survivor (s2),
+        # so the tool-call spend is not discarded.
+        assert [s.span_id for s in merged[0]] == ["s2", "s1"]
+
+    def test_dedup_no_spans_returns_none_second(self):
+        invocations = [
+            self._make_invocation("Roll a die", "tool_call"),
+            self._make_invocation("Roll a die", "I rolled a 3!"),
+        ]
+        deduped, merged = _deduplicate_invocations(invocations)
+        assert len(deduped) == 1
+        assert merged is None
+
     def test_single_invocation_no_change(self):
         invocations = [self._make_invocation("Q1", "A1")]
-        result = _deduplicate_invocations(invocations)
+        result, _ = _deduplicate_invocations(invocations)
         assert len(result) == 1
 
     def test_empty_list(self):
-        result = _deduplicate_invocations([])
+        result, _ = _deduplicate_invocations([])
         assert result == []
 
 
 class TestTrimCumulativeOutput:
-    """Tests for _trim_cumulative_output — stripping historical tool calls."""
+    """Tests for _trim_cumulative_output 鈥?stripping historical tool calls."""
 
     def test_single_user_message_no_trimming(self):
         span = _make_genai_llm_span(