From 90892ae542b490a5e2d3610bfd0fbb7191c85505 Mon Sep 17 00:00:00 2001 From: jariy17 Date: Tue, 22 Sep 2026 20:15:38 +0000 Subject: [PATCH] fix(evaluation): expose shared span helpers and evaluator level lookup --- docs/examples/evaluation_span_helpers.md | 44 +++++ src/bedrock_agentcore/evaluation/__init__.py | 4 + src/bedrock_agentcore/evaluation/client.py | 57 ++----- .../runner/on_demand/on_demand_runner.py | 38 +---- src/bedrock_agentcore/evaluation/spans.py | 50 ++++++ .../runner/on_demand/test_runner.py | 18 -- .../evaluation/test_client.py | 93 +++-------- .../evaluation/test_spans.py | 158 ++++++++++++++++++ 8 files changed, 297 insertions(+), 165 deletions(-) create mode 100644 docs/examples/evaluation_span_helpers.md create mode 100644 src/bedrock_agentcore/evaluation/spans.py create mode 100644 tests/bedrock_agentcore/evaluation/test_spans.py diff --git a/docs/examples/evaluation_span_helpers.md b/docs/examples/evaluation_span_helpers.md new file mode 100644 index 00000000..ee8ebf5f --- /dev/null +++ b/docs/examples/evaluation_span_helpers.md @@ -0,0 +1,44 @@ +# Evaluation span helpers + +Use public helpers to identify tool spans and select evaluation targets from +collected span dictionaries: + +```python +from bedrock_agentcore.evaluation.spans import is_tool_span, tool_span_ids, trace_ids + +spans = [ + { + "traceId": "trace-1", + "spanId": "span-1", + "attributes": {"gen_ai.operation.name": "execute_tool"}, + }, +] + +assert is_tool_span(spans[0]) +assert tool_span_ids(spans, trace_id="trace-1") == ["span-1"] +assert trace_ids(spans) == ["trace-1"] +``` + +These helpers are also exported from `bedrock_agentcore.evaluation`. +They recognize OTel GenAI, OpenInference, and Traceloop tool attributes and +require no AWS client or credentials. Missing or non-dictionary attributes +are treated as non-tool spans. + +`tool_span_ids` preserves input order and duplicate IDs, skipping missing or +empty span IDs. Omitting `trace_id`, or passing an empty string, includes all +traces. `trace_ids` skips missing or empty trace IDs and returns each ID once, +in order of first appearance. + +To look up an evaluator's level, use `EvaluationClient`: + +```python +from bedrock_agentcore.evaluation import EvaluationClient + +client = EvaluationClient(region_name="us-west-2") +level = client.get_evaluator_level("Builtin.Helpfulness") +``` + +This method calls the control plane API and caches the result per evaluator +for the lifetime of the client. If the lookup fails or the response omits the +level, it logs a warning and caches `SESSION` as the fallback, matching the +lookup behavior used by `client.run()`. diff --git a/src/bedrock_agentcore/evaluation/__init__.py b/src/bedrock_agentcore/evaluation/__init__.py index fece1a1a..3394ec0b 100644 --- a/src/bedrock_agentcore/evaluation/__init__.py +++ b/src/bedrock_agentcore/evaluation/__init__.py @@ -60,6 +60,7 @@ from bedrock_agentcore.evaluation.span_to_adot_serializer import ( convert_strands_to_adot, ) +from bedrock_agentcore.evaluation.spans import is_tool_span, tool_span_ids, trace_ids from bedrock_agentcore.evaluation.utils.cloudwatch_span_helper import ( fetch_spans_from_cloudwatch, ) @@ -113,6 +114,9 @@ "convert_strands_to_adot", "create_strands_evaluator", "fetch_spans_from_cloudwatch", + "is_tool_span", + "tool_span_ids", + "trace_ids", ] _STRANDS_EVALS_EXTRAS = { diff --git a/src/bedrock_agentcore/evaluation/client.py b/src/bedrock_agentcore/evaluation/client.py index 1102f519..0b9c3c10 100644 --- a/src/bedrock_agentcore/evaluation/client.py +++ b/src/bedrock_agentcore/evaluation/client.py @@ -8,6 +8,7 @@ from botocore.config import Config from pydantic import BaseModel +import bedrock_agentcore.evaluation.spans as span_helpers from bedrock_agentcore._utils.config import WaitConfig from bedrock_agentcore._utils.polling import wait_until, wait_until_deleted from bedrock_agentcore._utils.snake_case import accept_snake_case_kwargs, convert_kwargs @@ -215,7 +216,7 @@ def run( # Add reference inputs (ground truth) if provided if reference_inputs: - all_trace_ids = self._extract_trace_ids(spans) + all_trace_ids = span_helpers.trace_ids(spans) ref_inputs = self._build_reference_inputs( session_id, reference_inputs, all_trace_ids, target_trace_id=trace_id ) @@ -225,7 +226,7 @@ def run( # Steps 2-4: For each evaluator, look up level, build targets, call API all_results: List[Dict[str, Any]] = [] for evaluator_id in evaluator_ids: - level = self._get_evaluator_level(evaluator_id) + level = self.get_evaluator_level(evaluator_id) logger.info("Evaluating with %s (level=%s)", evaluator_id, level) requests = self._build_requests_for_level(evaluator_id, level, base_input, spans, trace_id) if len(requests) > 1: @@ -241,8 +242,17 @@ def run( ) return all_results - def _get_evaluator_level(self, evaluator_id: str) -> str: - """Look up evaluator level with caching. Falls back to SESSION.""" + def get_evaluator_level(self, evaluator_id: str) -> str: + """Look up an evaluator's level using the control plane client. + + Args: + evaluator_id: The built-in or custom evaluator identifier. + + Returns: + The evaluator level: SESSION, TRACE, or TOOL_CALL. Returns SESSION + if the lookup fails or the response omits the level. Results, + including fallback values, are cached for this client's lifetime. + """ if evaluator_id not in self._evaluator_level_cache: try: response = self._cp_client.get_evaluator(evaluatorId=evaluator_id) @@ -276,7 +286,7 @@ def _build_requests_for_level( if level == "TRACE": if trace_id: return [{**base_input, "evaluationTarget": {"traceIds": [trace_id]}}] - trace_ids = self._extract_trace_ids(spans) + trace_ids = span_helpers.trace_ids(spans) logger.debug("Extracted %d unique trace ID(s) for evaluator %s", len(trace_ids), evaluator_id) if not trace_ids: logger.warning("No trace IDs found for trace-level evaluator %s, skipping", evaluator_id) @@ -287,7 +297,7 @@ def _build_requests_for_level( ] if level == "TOOL_CALL": - tool_span_ids = self._extract_tool_span_ids(spans, trace_id=trace_id) + tool_span_ids = span_helpers.tool_span_ids(spans, trace_id=trace_id) logger.debug("Extracted %d tool span ID(s) for evaluator %s", len(tool_span_ids), evaluator_id) if not tool_span_ids: logger.warning("No tool span IDs found for tool-level evaluator %s, skipping", evaluator_id) @@ -299,41 +309,6 @@ def _build_requests_for_level( raise ValueError(f"Unknown evaluator level: {level}") - @staticmethod - def _extract_trace_ids(spans: list) -> List[str]: - """Extract unique trace IDs from spans, ordered by appearance.""" - return list(dict.fromkeys(span.get("traceId") for span in spans if span.get("traceId"))) - - @staticmethod - def _is_tool_span(span: dict) -> bool: - """Check if a span represents a tool execution (supports Strands, LangGraph, and Traceloop).""" - attrs = span.get("attributes", {}) - if not isinstance(attrs, dict): - return False - return ( - attrs.get("gen_ai.operation.name") == "execute_tool" - or attrs.get("openinference.span.kind") == "TOOL" - or attrs.get("traceloop.span.kind") == "tool" - ) - - @staticmethod - def _extract_tool_span_ids(spans: list, trace_id: Optional[str] = None) -> List[str]: - """Extract span IDs for tool execution spans. - - Args: - spans: List of span dicts. - trace_id: If provided, only include tool spans with this trace ID. - """ - tool_span_ids: List[str] = [] - for span in spans: - if EvaluationClient._is_tool_span(span): - if trace_id and span.get("traceId") != trace_id: - continue - span_id = span.get("spanId") - if span_id: - tool_span_ids.append(span_id) - return tool_span_ids - @staticmethod def _build_reference_inputs( session_id: str, diff --git a/src/bedrock_agentcore/evaluation/runner/on_demand/on_demand_runner.py b/src/bedrock_agentcore/evaluation/runner/on_demand/on_demand_runner.py index 5b36b047..18f5bf71 100644 --- a/src/bedrock_agentcore/evaluation/runner/on_demand/on_demand_runner.py +++ b/src/bedrock_agentcore/evaluation/runner/on_demand/on_demand_runner.py @@ -9,6 +9,7 @@ import boto3 from botocore.config import Config +import bedrock_agentcore.evaluation.spans as span_helpers from bedrock_agentcore.evaluation.agent_span_collector import AgentSpanCollector from ..dataset_types import Dataset, PredefinedScenario, Scenario, SimulatedScenario @@ -228,8 +229,8 @@ def _build_evaluate_requests( List of (evaluator_id, request_dict) tuples. """ base: Dict[str, Any] = {"evaluationInput": {"sessionSpans": spans}} - trace_ids = self._extract_trace_ids(spans) - tool_span_ids = self._extract_tool_span_ids(spans) + trace_ids = span_helpers.trace_ids(spans) + tool_span_ids = span_helpers.tool_span_ids(spans) reference_inputs = self._build_reference_inputs(scenario, trace_ids, session_id) if reference_inputs: @@ -268,27 +269,6 @@ def _build_evaluate_requests( return requests - def _extract_trace_ids(self, spans: list) -> List[str]: - """Extract unique trace IDs from spans, ordered by appearance.""" - trace_ids = [] - seen = set() - for span in spans: - trace_id = span.get("traceId") - if trace_id and trace_id not in seen: - trace_ids.append(trace_id) - seen.add(trace_id) - return trace_ids - - def _extract_tool_span_ids(self, spans: list) -> List[str]: - """Extract span IDs for tool execution spans (supports Strands and LangGraph).""" - tool_span_ids = [] - for span in spans: - if self._is_tool_span(span): - span_id = span.get("spanId") - if span_id: - tool_span_ids.append(span_id) - return tool_span_ids - def _build_reference_inputs( self, scenario: Scenario, @@ -421,15 +401,3 @@ def _batch(items: List[str], size: int) -> Iterator[List[str]]: """Yield successive batches of the given size from items.""" for i in range(0, len(items), size): yield items[i : i + size] - - @staticmethod - def _is_tool_span(span: Dict) -> bool: - """Check if a span represents a tool execution (supports Strands and LangGraph).""" - attrs = span.get("attributes", {}) - if not isinstance(attrs, dict): - return False - return ( - attrs.get("gen_ai.operation.name") == "execute_tool" - or attrs.get("openinference.span.kind") == "TOOL" - or attrs.get("traceloop.span.kind") == "tool" - ) diff --git a/src/bedrock_agentcore/evaluation/spans.py b/src/bedrock_agentcore/evaluation/spans.py new file mode 100644 index 00000000..22d27679 --- /dev/null +++ b/src/bedrock_agentcore/evaluation/spans.py @@ -0,0 +1,50 @@ +"""Helpers for identifying tool spans and extracting evaluation target IDs.""" + +from typing import Any, Dict, List, Optional + + +def is_tool_span(span: Dict[str, Any]) -> bool: + """Return whether a span represents a tool execution. + + Recognizes OTel GenAI (``gen_ai.operation.name=execute_tool``), + OpenInference (``openinference.span.kind=TOOL``), and Traceloop + (``traceloop.span.kind=tool``) attributes. Missing or non-dict + attributes are treated as non-tool spans. + + Args: + span: A span dict containing an optional attributes dict. + """ + attrs = span.get("attributes", {}) + if not isinstance(attrs, dict): + return False + return ( + attrs.get("gen_ai.operation.name") == "execute_tool" + or attrs.get("openinference.span.kind") == "TOOL" + or attrs.get("traceloop.span.kind") == "tool" + ) + + +def tool_span_ids(spans: List[Dict[str, Any]], trace_id: Optional[str] = None) -> List[str]: + """Return tool span IDs in input order, preserving duplicates. + + Args: + spans: Span dicts with attributes, spanId, and optionally traceId. + Spans with missing or empty spanId values are skipped. + trace_id: Restrict results to this trace. None or an empty string + includes tool spans from all traces. + """ + return [ + span["spanId"] + for span in spans + if is_tool_span(span) and span.get("spanId") and (not trace_id or span.get("traceId") == trace_id) + ] + + +def trace_ids(spans: List[Dict[str, Any]]) -> List[str]: + """Return unique trace IDs ordered by first appearance. + + Args: + spans: Span dicts containing traceId values. Missing or empty + traceId values are skipped. + """ + return list(dict.fromkeys(span.get("traceId") for span in spans if span.get("traceId"))) diff --git a/tests/bedrock_agentcore/evaluation/runner/on_demand/test_runner.py b/tests/bedrock_agentcore/evaluation/runner/on_demand/test_runner.py index 79306e10..54c06ce2 100644 --- a/tests/bedrock_agentcore/evaluation/runner/on_demand/test_runner.py +++ b/tests/bedrock_agentcore/evaluation/runner/on_demand/test_runner.py @@ -172,24 +172,6 @@ def test_batch_single(self): def test_batch_empty(self): assert list(OnDemandEvaluationDatasetRunner._batch([], 10)) == [] - def test_is_tool_span_strands(self): - span = {"attributes": {"gen_ai.operation.name": "execute_tool"}} - assert OnDemandEvaluationDatasetRunner._is_tool_span(span) is True - - def test_is_tool_span_langgraph_openinference(self): - span = {"attributes": {"openinference.span.kind": "TOOL"}} - assert OnDemandEvaluationDatasetRunner._is_tool_span(span) is True - - def test_is_tool_span_langgraph_otel(self): - assert OnDemandEvaluationDatasetRunner._is_tool_span({"attributes": {"traceloop.span.kind": "tool"}}) is True - - def test_is_tool_span_not_tool(self): - assert OnDemandEvaluationDatasetRunner._is_tool_span({"attributes": {"gen_ai.operation.name": "chat"}}) is False - - def test_is_tool_span_no_attributes(self): - assert OnDemandEvaluationDatasetRunner._is_tool_span({}) is False - assert OnDemandEvaluationDatasetRunner._is_tool_span({"attributes": None}) is False - # --- Helpers --- def _make_dataset(self): diff --git a/tests/bedrock_agentcore/evaluation/test_client.py b/tests/bedrock_agentcore/evaluation/test_client.py index 00910a31..461fa254 100644 --- a/tests/bedrock_agentcore/evaluation/test_client.py +++ b/tests/bedrock_agentcore/evaluation/test_client.py @@ -225,28 +225,41 @@ def test_custom_look_back_time(self, client): assert duration == timedelta(hours=2) -# --- _get_evaluator_level tests --- +# --- get_evaluator_level tests --- class TestGetEvaluatorLevel: - def test_returns_level_from_api(self, client): - client._cp_client.get_evaluator.return_value = {"level": "TRACE"} - assert client._get_evaluator_level("eval-1") == "TRACE" + @pytest.mark.parametrize("level", ["SESSION", "TRACE", "TOOL_CALL"]) + def test_returns_level_from_api(self, client, level): + client._cp_client.get_evaluator.return_value = {"level": level} + assert client.get_evaluator_level("eval-1") == level + client._cp_client.get_evaluator.assert_called_once_with(evaluatorId="eval-1") + + def test_missing_level_falls_back_to_session(self, client): + client._cp_client.get_evaluator.return_value = {} + assert client.get_evaluator_level("eval-1") == "SESSION" + + def test_cache_is_per_evaluator(self, client): + client._cp_client.get_evaluator.side_effect = [{"level": "TRACE"}, {"level": "TOOL_CALL"}] + assert client.get_evaluator_level("eval-1") == "TRACE" + assert client.get_evaluator_level("eval-2") == "TOOL_CALL" + assert client.get_evaluator_level("eval-1") == "TRACE" + assert client._cp_client.get_evaluator.call_count == 2 def test_caches_level(self, client): client._cp_client.get_evaluator.return_value = {"level": "TRACE"} - client._get_evaluator_level("eval-1") - client._get_evaluator_level("eval-1") + client.get_evaluator_level("eval-1") + client.get_evaluator_level("eval-1") client._cp_client.get_evaluator.assert_called_once() def test_falls_back_to_session_on_error(self, client): client._cp_client.get_evaluator.side_effect = RuntimeError("not found") - assert client._get_evaluator_level("eval-1") == "SESSION" + assert client.get_evaluator_level("eval-1") == "SESSION" def test_caches_fallback(self, client): client._cp_client.get_evaluator.side_effect = RuntimeError("not found") - client._get_evaluator_level("eval-1") - client._get_evaluator_level("eval-1") + client.get_evaluator_level("eval-1") + client.get_evaluator_level("eval-1") client._cp_client.get_evaluator.assert_called_once() @@ -297,68 +310,6 @@ def test_trace_level_batching(self, client): assert len(requests[1]["evaluationTarget"]["traceIds"]) == 2 -# --- Static helper tests --- - - -class TestExtractTraceIds: - def test_extracts_unique_ordered(self): - ids = EvaluationClient._extract_trace_ids(SAMPLE_SPANS) - assert ids == ["trace-1", "trace-2"] - - def test_empty_spans(self): - assert EvaluationClient._extract_trace_ids([]) == [] - - def test_skips_missing_trace_id(self): - spans = [{"spanId": "s1"}, {"traceId": "t1", "spanId": "s2"}] - assert EvaluationClient._extract_trace_ids(spans) == ["t1"] - - -class TestExtractToolSpanIds: - def test_extracts_tool_spans(self): - ids = EvaluationClient._extract_tool_span_ids(SAMPLE_SPANS) - assert ids == ["span-2", "span-3", "span-5"] - - def test_ignores_non_tool_spans(self): - spans = [ - {"name": "Agent.invoke", "kind": "SPAN_KIND_SERVER", "spanId": "s1"}, - {"name": "LLM.call", "kind": "SPAN_KIND_INTERNAL", "spanId": "s2"}, - ] - assert EvaluationClient._extract_tool_span_ids(spans) == [] - - def test_empty_spans(self): - assert EvaluationClient._extract_tool_span_ids([]) == [] - - def test_filters_by_trace_id(self): - ids = EvaluationClient._extract_tool_span_ids(SAMPLE_SPANS, trace_id="trace-1") - assert ids == ["span-2", "span-3"] - - def test_filters_by_trace_id_no_match(self): - ids = EvaluationClient._extract_tool_span_ids(SAMPLE_SPANS, trace_id="trace-999") - assert ids == [] - - def test_extracts_langgraph_tool_spans(self): - spans = [ - {"spanId": "s1", "traceId": "t1", "attributes": {"openinference.span.kind": "TOOL"}}, - {"spanId": "s2", "traceId": "t1", "attributes": {"openinference.span.kind": "LLM"}}, - ] - assert EvaluationClient._extract_tool_span_ids(spans) == ["s1"] - - def test_extracts_traceloop_tool_spans(self): - spans = [ - {"spanId": "s1", "traceId": "t1", "attributes": {"traceloop.span.kind": "tool"}}, - {"spanId": "s2", "traceId": "t1", "attributes": {"traceloop.span.kind": "workflow"}}, - ] - assert EvaluationClient._extract_tool_span_ids(spans) == ["s1"] - - def test_ignores_span_without_tool_attributes(self): - spans = [ - {"spanId": "s1", "traceId": "t1", "attributes": {"gen_ai.operation.name": "invoke_agent"}}, - {"spanId": "s2", "traceId": "t1", "attributes": {"some.other.attr": "value"}}, - {"spanId": "s3", "traceId": "t1"}, - ] - assert EvaluationClient._extract_tool_span_ids(spans) == [] - - # --- trace_id tests --- diff --git a/tests/bedrock_agentcore/evaluation/test_spans.py b/tests/bedrock_agentcore/evaluation/test_spans.py new file mode 100644 index 00000000..4b24253f --- /dev/null +++ b/tests/bedrock_agentcore/evaluation/test_spans.py @@ -0,0 +1,158 @@ +"""Tests for public evaluation span helpers.""" + +import pytest + +from bedrock_agentcore.evaluation.spans import is_tool_span, tool_span_ids, trace_ids + + +@pytest.mark.parametrize( + "span,expected", + [ + ({"attributes": {"gen_ai.operation.name": "execute_tool"}}, True), + ({"attributes": {"openinference.span.kind": "TOOL"}}, True), + ({"attributes": {"traceloop.span.kind": "tool"}}, True), + ({"attributes": {"gen_ai.operation.name": "chat", "openinference.span.kind": "TOOL"}}, True), + ({"attributes": {"gen_ai.operation.name": "chat"}}, False), + ({"attributes": {"openinference.span.kind": "tool"}}, False), + ({}, False), + ({"attributes": None}, False), + ({"attributes": []}, False), + ({"attributes": "execute_tool"}, False), + ], +) +def test_is_tool_span(span, expected): + assert is_tool_span(span) is expected + + +@pytest.mark.parametrize("name", ["is_tool_span", "tool_span_ids", "trace_ids"]) +def test_public_package_exports(name): + from bedrock_agentcore import evaluation + from bedrock_agentcore.evaluation import spans + + assert name in evaluation.__all__ + assert getattr(evaluation, name) is getattr(spans, name) + + +@pytest.mark.parametrize("trace_id", [None, "", "t1", "missing"]) +def test_tool_span_ids_missing_ids_duplicates_and_trace_filter(trace_id): + spans = [ + {"spanId": "s1", "traceId": "t1", "attributes": {"gen_ai.operation.name": "execute_tool"}}, + {"spanId": "s1", "traceId": "t1", "attributes": {"openinference.span.kind": "TOOL"}}, + {"spanId": "s2", "attributes": {"traceloop.span.kind": "tool"}}, + {"spanId": "", "attributes": {"traceloop.span.kind": "tool"}}, + {"spanId": None, "attributes": {"traceloop.span.kind": "tool"}}, + {"attributes": {"traceloop.span.kind": "tool"}}, + ] + expected = {None: ["s1", "s1", "s2"], "": ["s1", "s1", "s2"], "t1": ["s1", "s1"], "missing": []} + assert tool_span_ids(spans, trace_id=trace_id) == expected[trace_id] + + +def test_trace_ids_skips_empty_and_preserves_first_appearance(): + spans = [{"traceId": value} for value in ["t2", None, "", "t1", "t2", "t3", "t1"]] + assert trace_ids(spans) == ["t2", "t1", "t3"] + + +SAMPLE_SPANS = [ + { + "scope": {"name": "agent"}, + "traceId": "trace-1", + "spanId": "span-1", + "name": "Agent.invoke", + "kind": "SPAN_KIND_SERVER", + "attributes": {"gen_ai.operation.name": "invoke_agent"}, + }, + { + "scope": {"name": "agent"}, + "traceId": "trace-1", + "spanId": "span-2", + "name": "Tool:search", + "kind": "SPAN_KIND_INTERNAL", + "attributes": {"gen_ai.operation.name": "execute_tool"}, + }, + { + "scope": {"name": "agent"}, + "traceId": "trace-1", + "spanId": "span-3", + "name": "Tool:calculator", + "kind": "SPAN_KIND_INTERNAL", + "attributes": {"gen_ai.operation.name": "execute_tool"}, + }, + { + "scope": {"name": "agent"}, + "traceId": "trace-2", + "spanId": "span-4", + "name": "Agent.invoke", + "kind": "SPAN_KIND_SERVER", + "attributes": {"gen_ai.operation.name": "invoke_agent"}, + }, + { + "scope": {"name": "agent"}, + "traceId": "trace-2", + "spanId": "span-5", + "name": "Tool:search", + "kind": "SPAN_KIND_INTERNAL", + "attributes": {"gen_ai.operation.name": "execute_tool"}, + }, +] + + +# --- Static helper tests --- + + +class TestExtractTraceIds: + def test_extracts_unique_ordered(self): + ids = trace_ids(SAMPLE_SPANS) + assert ids == ["trace-1", "trace-2"] + + def test_empty_spans(self): + assert trace_ids([]) == [] + + def test_skips_missing_trace_id(self): + spans = [{"spanId": "s1"}, {"traceId": "t1", "spanId": "s2"}] + assert trace_ids(spans) == ["t1"] + + +class TestExtractToolSpanIds: + def test_extracts_tool_spans(self): + ids = tool_span_ids(SAMPLE_SPANS) + assert ids == ["span-2", "span-3", "span-5"] + + def test_ignores_non_tool_spans(self): + spans = [ + {"name": "Agent.invoke", "kind": "SPAN_KIND_SERVER", "spanId": "s1"}, + {"name": "LLM.call", "kind": "SPAN_KIND_INTERNAL", "spanId": "s2"}, + ] + assert tool_span_ids(spans) == [] + + def test_empty_spans(self): + assert tool_span_ids([]) == [] + + def test_filters_by_trace_id(self): + ids = tool_span_ids(SAMPLE_SPANS, trace_id="trace-1") + assert ids == ["span-2", "span-3"] + + def test_filters_by_trace_id_no_match(self): + ids = tool_span_ids(SAMPLE_SPANS, trace_id="trace-999") + assert ids == [] + + def test_extracts_langgraph_tool_spans(self): + spans = [ + {"spanId": "s1", "traceId": "t1", "attributes": {"openinference.span.kind": "TOOL"}}, + {"spanId": "s2", "traceId": "t1", "attributes": {"openinference.span.kind": "LLM"}}, + ] + assert tool_span_ids(spans) == ["s1"] + + def test_extracts_traceloop_tool_spans(self): + spans = [ + {"spanId": "s1", "traceId": "t1", "attributes": {"traceloop.span.kind": "tool"}}, + {"spanId": "s2", "traceId": "t1", "attributes": {"traceloop.span.kind": "workflow"}}, + ] + assert tool_span_ids(spans) == ["s1"] + + def test_ignores_span_without_tool_attributes(self): + spans = [ + {"spanId": "s1", "traceId": "t1", "attributes": {"gen_ai.operation.name": "invoke_agent"}}, + {"spanId": "s2", "traceId": "t1", "attributes": {"some.other.attr": "value"}}, + {"spanId": "s3", "traceId": "t1"}, + ] + assert tool_span_ids(spans) == []