diff --git a/src/agentevals/converter.py b/src/agentevals/converter.py index 7543369..68075dc 100644 --- a/src/agentevals/converter.py +++ b/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,7 +153,7 @@ 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, @@ -156,6 +161,8 @@ def _convert_invoke_span(invoke_span: Span) -> Invocation: 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] = [] diff --git a/src/agentevals/genai_converter.py b/src/agentevals/genai_converter.py index b635519..78a33c6 100644 --- a/src/agentevals/genai_converter.py +++ b/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,7 +271,10 @@ 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 @@ -262,9 +282,16 @@ def _deduplicate_invocations(invocations: list[Invocation]) -> list[Invocation]: 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: diff --git a/src/agentevals/streaming/ws_server.py b/src/agentevals/streaming/ws_server.py index 741b145..7f57585 100644 --- a/src/agentevals/streaming/ws_server.py +++ b/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) diff --git a/tests/test_extraction.py b/tests/test_extraction.py index 96544f7..6f125bb 100644 --- a/tests/test_extraction.py +++ b/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") diff --git a/tests/test_genai_converter.py b/tests/test_genai_converter.py index 5da61c5..bc1ad6c 100644 --- a/tests/test_genai_converter.py +++ b/tests/test_genai_converter.py @@ -616,7 +616,7 @@ 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): @@ -624,7 +624,7 @@ def test_dedup_keeps_last_duplicate(self): 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( diff --git a/tests/test_model_info.py b/tests/test_model_info.py new file mode 100644 index 0000000..79ae3b6 --- /dev/null +++ b/tests/test_model_info.py @@ -0,0 +1,262 @@ +"""Regression tests for per-invocation model info extraction. + +Covers https://github.com/agentevals-dev/agentevals/issues/204: +`_extract_model_info_from_trace` aggregated LLM spans across the whole trace +for every invocation, so each invocation reported the session-wide totals. +Now the converter records which LLM spans each invocation was built from and +ws_server aggregates only those spans. +""" + +import asyncio +import json + +from agentevals.converter import convert_trace +from agentevals.loader.base import Span, Trace +from agentevals.streaming.session import TraceSession +from agentevals.streaming.ws_server import StreamingTraceManager + + +def _adk_llm_span( + span_id: str, model: str, input_tokens: int, output_tokens: int, start_time: int, parent: str = "invoke" +) -> Span: + """Build an ADK call_llm span with distinct usage metadata and user text.""" + return Span( + trace_id="t1", + span_id=span_id, + parent_span_id=parent, + operation_name="call_llm", + start_time=start_time, + duration=1000, + tags={ + "otel.scope.name": "gcp.vertex.agent", + "gcp.vertex.agent.llm_request": json.dumps( + { + "model": model, + "contents": [ + {"role": "user", "parts": [{"text": f"hello from {span_id}"}]}, + ], + } + ), + "gcp.vertex.agent.llm_response": json.dumps( + { + "content": {"parts": [{"text": f"answer from {span_id}"}], "role": "model"}, + "usage_metadata": { + "prompt_token_count": input_tokens, + "candidates_token_count": output_tokens, + }, + } + ), + "gen_ai.provider.name": "vertex_ai", + "gen_ai.response.finish_reasons": "stop", + }, + ) + + +def _two_invocation_adk_trace() -> Trace: + """ADK trace with two invoke_agent spans, each owning its own LLM spans.""" + invoke1 = Span( + trace_id="t1", + span_id="invoke1", + parent_span_id=None, + operation_name="invoke_agent agent_a", + start_time=1000, + duration=20000, + tags={"otel.scope.name": "gcp.vertex.agent", "gen_ai.operation.name": "invoke_agent"}, + ) + invoke2 = Span( + trace_id="t1", + span_id="invoke2", + parent_span_id=None, + operation_name="invoke_agent agent_b", + start_time=30000, + duration=20000, + tags={"otel.scope.name": "gcp.vertex.agent", "gen_ai.operation.name": "invoke_agent"}, + ) + + llm1 = _adk_llm_span("llm1", "model-a", 100, 20, 2000) + llm2a = _adk_llm_span("llm2a", "model-b", 300, 50, 31000) + llm2b = _adk_llm_span("llm2b", "model-b", 400, 60, 32000) + + llm1.parent_span_id = "invoke1" + llm2a.parent_span_id = "invoke2" + llm2b.parent_span_id = "invoke2" + invoke1.children.append(llm1) + invoke2.children.extend([llm2a, llm2b]) + + return Trace( + trace_id="t1", + root_spans=[invoke1, invoke2], + all_spans=[invoke1, llm1, invoke2, llm2a, llm2b], + ) + + +# --------------------------------------------------------------------------- +# OTLP span dict helpers (for driving a real TraceSession through +# StreamingTraceManager._extract_invocations, the wiring fixed in #204) +# --------------------------------------------------------------------------- + + +def _otlp_attr(key: str, value: str) -> dict: + return {"key": key, "value": {"stringValue": value}} + + +def _otlp_span( + span_id: str, + name: str, + start_ns: int, + end_ns: int, + attrs: dict, + parent: str | None = None, + trace_id: str = "t1", +) -> dict: + span = { + "traceId": trace_id, + "spanId": span_id, + "name": name, + "startTimeUnixNano": start_ns, + "endTimeUnixNano": end_ns, + "attributes": [_otlp_attr(k, str(v)) for k, v in attrs.items()], + } + if parent: + span["parentSpanId"] = parent + return span + + +def _adk_llm_otlp_span( + span_id: str, model: str, input_tokens: int, output_tokens: int, start_ns: int, parent: str +) -> dict: + return _otlp_span( + span_id, + "call_llm", + start_ns, + start_ns + 1_000_000, + { + "otel.scope.name": "gcp.vertex.agent", + "gcp.vertex.agent.llm_request": json.dumps( + { + "model": model, + "contents": [{"role": "user", "parts": [{"text": f"hello from {span_id}"}]}], + } + ), + "gcp.vertex.agent.llm_response": json.dumps( + { + "content": {"parts": [{"text": f"answer from {span_id}"}], "role": "model"}, + "usage_metadata": { + "prompt_token_count": input_tokens, + "candidates_token_count": output_tokens, + }, + } + ), + "gen_ai.provider.name": "vertex_ai", + "gen_ai.response.finish_reasons": "stop", + }, + parent=parent, + ) + + +class TestPerInvocationSpans: + def test_conversion_tracks_each_invocations_own_llm_spans(self): + result = convert_trace(_two_invocation_adk_trace()) + assert len(result.invocations) == 2 + assert len(result.invocation_llm_spans) == 2 + assert [s.span_id for s in result.invocation_llm_spans[0]] == ["llm1"] + assert [s.span_id for s in result.invocation_llm_spans[1]] == ["llm2a", "llm2b"] + + def test_model_info_is_per_invocation_not_session_wide(self): + """Drive a real TraceSession through ``_extract_invocations``. + + This exercises the converter wiring that #204 actually broke, rather + than calling the aggregator directly on hand-picked spans. + """ + manager = StreamingTraceManager() + + invoke1 = _otlp_span( + "invoke1", + "invoke_agent agent_a", + 1_000_000_000, + 21_000_000_000, + {"otel.scope.name": "gcp.vertex.agent", "gen_ai.operation.name": "invoke_agent"}, + ) + invoke2 = _otlp_span( + "invoke2", + "invoke_agent agent_b", + 30_000_000_000, + 50_000_000_000, + {"otel.scope.name": "gcp.vertex.agent", "gen_ai.operation.name": "invoke_agent"}, + ) + llm1 = _adk_llm_otlp_span("llm1", "model-a", 100, 20, 2_000_000_000, parent="invoke1") + llm2a = _adk_llm_otlp_span("llm2a", "model-b", 300, 50, 31_000_000_000, parent="invoke2") + llm2b = _adk_llm_otlp_span("llm2b", "model-b", 400, 60, 32_000_000_000, parent="invoke2") + + session = TraceSession( + session_id="s1", + trace_id="t1", + eval_set_id=None, + spans=[invoke1, llm1, invoke2, llm2a, llm2b], + logs=[], + ) + + data = asyncio.run(manager._extract_invocations(session)) + + assert len(data) == 2 + info_a = data[0]["modelInfo"] + info_b = data[1]["modelInfo"] + assert info_a["inputTokens"] == 100 + assert info_a["outputTokens"] == 20 + assert info_b["inputTokens"] == 700 # 300 + 400, only invocation B's spans + assert info_b["outputTokens"] == 110 # 50 + 60 + + # The bug made every invocation report identical session-wide totals. + assert info_a["inputTokens"] != info_b["inputTokens"] + assert info_a["outputTokens"] != info_b["outputTokens"] + + def test_empty_spans_yield_empty_model_info(self): + manager = StreamingTraceManager() + assert manager._extract_model_info_from_llm_spans([]) == {} + + def test_nested_sub_agent_spans_are_pinned(self): + """Pin the current behavior for nested ``invoke_agent`` spans. + + A coordinator ``invoke_agent`` span that nests a specialist + ``invoke_agent`` span is not yet pruned: ``find_adk_llm_spans_in`` + walks the whole subtree, so the specialist's LLM spans are attributed + to both invocations. This is a known limitation (see PR discussion); + the test documents the current behavior so a future fix is observable. + """ + coordinator = Span( + trace_id="t1", + span_id="coord", + parent_span_id=None, + operation_name="invoke_agent coordinator", + start_time=1000, + duration=30000, + tags={"otel.scope.name": "gcp.vertex.agent", "gen_ai.operation.name": "invoke_agent"}, + ) + specialist = Span( + trace_id="t1", + span_id="spec", + parent_span_id="coord", + operation_name="invoke_agent specialist", + start_time=2000, + duration=20000, + tags={"otel.scope.name": "gcp.vertex.agent", "gen_ai.operation.name": "invoke_agent"}, + ) + llm_root = _adk_llm_span("llm_root", "model-a", 600, 50, 3000, parent="coord") + llm_sub = _adk_llm_span("llm_sub", "model-b", 500, 40, 4000, parent="spec") + coordinator.children.extend([llm_root, specialist]) + specialist.children.append(llm_sub) + + trace = Trace( + trace_id="t1", + root_spans=[coordinator], + all_spans=[coordinator, llm_root, specialist, llm_sub], + ) + + result = convert_trace(trace) + + # Both the coordinator and the nested specialist are treated as + # invocations, and the specialist's LLM spans currently appear under + # both (known limitation, pinned here). + assert len(result.invocations) == 2 + assert [s.span_id for s in result.invocation_llm_spans[0]] == ["llm_root", "llm_sub"] + assert [s.span_id for s in result.invocation_llm_spans[1]] == ["llm_sub"]