From bbd9c17a2096ba154e1e486e26036f205fd33fde Mon Sep 17 00:00:00 2001 From: Gaurav Mittal Date: Mon, 6 Jul 2026 07:33:11 -0700 Subject: [PATCH] Add CAUSED_BY_GENERATION attribution links to tool call spans Extend TraceloopCallbackHandler to record which LLM generation triggered each tool invocation using OTel span links with a structured link_type attribute. On on_llm_end the handler saves the generation SpanContext and maps each structured tool_call_id to the LLM run id; on on_tool_start it looks up the saved context by tool_call_id (or parent run id fallback) and attaches a Link(gen_ctx, attributes={gen_ai.attribution.link_type: CAUSED_BY_GENERATION}) to the tool span. Changes: - Import Link from opentelemetry.trace - Add _generation_contexts and _tool_call_id_to_generation state dicts - Thread optional links: List through _create_span and _create_task_span - Capture SpanContext + map tool_call_ids in on_llm_end before ending span - Build attribution link in on_tool_start and pass to _create_task_span - Add test_attribution_links.py with two unit tests (2 passed) --- .../langchain/callback_handler.py | 43 ++++++- .../tests/test_attribution_links.py | 118 ++++++++++++++++++ 2 files changed, 159 insertions(+), 2 deletions(-) create mode 100644 packages/opentelemetry-instrumentation-langchain/tests/test_attribution_links.py diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py index 15ed817e15..efc8d21828 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py +++ b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py @@ -69,7 +69,7 @@ SpanAttributes, TraceloopSpanKindValues, ) -from opentelemetry.trace import SpanKind, Tracer, set_span_in_context +from opentelemetry.trace import Link, SpanKind, Tracer, set_span_in_context from opentelemetry.instrumentation.langchain.patch import ( LANGGRAPH_FLOW_KEY, LANGGRAPH_GRAPH_SPAN_KEY, @@ -176,6 +176,11 @@ def __init__( self.spans: dict[UUID, SpanHolder] = {} self.run_inline = True self._callback_manager: CallbackManager | AsyncCallbackManager = None + # Attribution tracking: maps str(llm_run_id) → SpanContext, saved just + # before the LLM span ends so tool-call spans can link back to it. + self._generation_contexts: dict[str, Any] = {} + # Maps structured tool_call_id → str(llm_run_id) for direct attribution. + self._tool_call_id_to_generation: dict[str, str] = {} @staticmethod def _get_name_from_callback( @@ -285,6 +290,7 @@ def _create_span( entity_name: str = "", entity_path: str = "", metadata: Optional[dict[str, Any]] = None, + links: Optional[List[Any]] = None, ) -> Span: """Create and register a span holder, replacing any existing holder for the run id.""" old_holder = self.spans.get(run_id) @@ -314,11 +320,13 @@ def _create_span( # This doesn't affect the core span functionality pass + span_links = links or [] if parent_run_id is not None and parent_run_id in self.spans: span = self.tracer.start_span( span_name, context=set_span_in_context(self.spans[parent_run_id].span), kind=kind, + links=span_links, ) else: # Check if we're in a LangGraph flow and this is the first child @@ -332,6 +340,7 @@ def _create_span( span_name, context=graph_span_holder.context, kind=kind, + links=span_links, ) # Flip the flag so subsequent spans use normal parenting. # Note: This mutable list pattern is intentional. LangGraph callbacks @@ -339,7 +348,7 @@ def _create_span( # The mutable list is used because OTel context values are immutable. first_child_pending[0] = False else: - span = self.tracer.start_span(span_name, kind=kind) + span = self.tracer.start_span(span_name, kind=kind, links=span_links) token = self._safe_attach_context(span) @@ -382,6 +391,7 @@ def _create_task_span( entity_path: str = "", metadata: Optional[dict[str, Any]] = None, serialized: Optional[dict[str, Any]] = None, + links: Optional[List[Any]] = None, ) -> Span: """Create a workflow, task, or tool span with LangChain and GenAI attributes.""" # Determine span type @@ -409,6 +419,7 @@ def _create_task_span( entity_name=entity_name, entity_path=entity_path, metadata=metadata, + links=links, ) # Set traceloop attributes for backwards compatibility @@ -831,6 +842,18 @@ def on_llm_end( }, ) + # Save span context for attribution: on_tool_start will link back here. + self._generation_contexts[str(run_id)] = span.get_span_context() + for gen_list in response.generations: + for gen in gen_list: + msg = getattr(gen, "message", None) + if msg is None: + continue + for tc in getattr(msg, "tool_calls", None) or []: + tc_id = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) + if tc_id: + self._tool_call_id_to_generation[tc_id] = str(run_id) + self._end_span(span, run_id) @dont_throw @@ -854,6 +877,21 @@ def on_tool_start( workflow_name = self.get_workflow_name(parent_run_id) entity_path = self.get_entity_path(parent_run_id) + # Build CAUSED_BY_GENERATION attribution link. + # Prefer structured tool_call_id lookup; fall back to parent LLM run id. + attribution_links = [] + tool_call_id = kwargs.get("tool_call_id") + gen_ctx = None + if tool_call_id and tool_call_id in self._tool_call_id_to_generation: + llm_run_key = self._tool_call_id_to_generation[tool_call_id] + gen_ctx = self._generation_contexts.get(llm_run_key) + if gen_ctx is None and parent_run_id is not None: + gen_ctx = self._generation_contexts.get(str(parent_run_id)) + if gen_ctx is not None and gen_ctx.is_valid: + attribution_links = [ + Link(gen_ctx, attributes={"gen_ai.attribution.link_type": "CAUSED_BY_GENERATION"}) + ] + span = self._create_task_span( run_id, parent_run_id, @@ -864,6 +902,7 @@ def on_tool_start( entity_path, metadata=metadata, serialized=serialized, + links=attribution_links, ) if not should_emit_events() and should_send_prompts(): input_json = json.dumps( diff --git a/packages/opentelemetry-instrumentation-langchain/tests/test_attribution_links.py b/packages/opentelemetry-instrumentation-langchain/tests/test_attribution_links.py new file mode 100644 index 0000000000..89ca2898f3 --- /dev/null +++ b/packages/opentelemetry-instrumentation-langchain/tests/test_attribution_links.py @@ -0,0 +1,118 @@ +"""Unit tests for CAUSED_BY_GENERATION attribution links on tool call spans.""" + +from uuid import uuid4 + +import pytest +from langchain_core.messages import AIMessage +from langchain_core.outputs import ChatGeneration, LLMResult +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from opentelemetry.instrumentation.langchain.callback_handler import ( + TraceloopCallbackHandler, +) + + +@pytest.fixture +def _exporter(): + return InMemorySpanExporter() + + +@pytest.fixture +def _handler(_exporter): + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(_exporter)) + tracer = provider.get_tracer("test") + + reader = InMemoryMetricReader() + meter = MeterProvider(metric_readers=[reader]).get_meter("test") + duration_hist = meter.create_histogram("gen_ai.client.operation.duration") + token_hist = meter.create_histogram("gen_ai.client.token.usage") + return TraceloopCallbackHandler(tracer, duration_hist, token_hist) + + +def _llm_result_with_tool_call(tool_call_id: str, tool_name: str) -> LLMResult: + msg = AIMessage( + content="", + tool_calls=[{"id": tool_call_id, "name": tool_name, "args": {}}], + ) + return LLMResult(generations=[[ChatGeneration(message=msg)]]) + + +def test_caused_by_generation_link_structured_tool_call(_handler, _exporter): + """Tool span must carry a CAUSED_BY_GENERATION link to the LLM span that triggered it.""" + root_id = uuid4() + llm_id = uuid4() + tool_id = uuid4() + tc_id = "call_test_abc123" + + _handler.on_chain_start( + {"id": ["FakeChain"], "name": "FakeChain"}, + {}, + run_id=root_id, + ) + _handler.on_chat_model_start( + {"id": ["FakeLLM"], "name": "FakeLLM"}, + [[]], + run_id=llm_id, + parent_run_id=root_id, + ) + _handler.on_llm_end( + _llm_result_with_tool_call(tc_id, "search"), + run_id=llm_id, + parent_run_id=root_id, + ) + _handler.on_tool_start( + {"id": ["SearchTool"], "name": "search"}, + "test query", + run_id=tool_id, + parent_run_id=root_id, + tool_call_id=tc_id, + ) + _handler.on_tool_end("result", run_id=tool_id, parent_run_id=root_id) + _handler.on_chain_end({}, run_id=root_id) + + finished = _exporter.get_finished_spans() + tool_spans = [s for s in finished if "execute_tool" in s.name] + assert len(tool_spans) == 1, f"Unexpected spans: {[s.name for s in finished]}" + + tool_span = tool_spans[0] + assert len(tool_span.links) == 1, "Tool span must have exactly one attribution link" + + link = tool_span.links[0] + assert link.attributes.get("gen_ai.attribution.link_type") == "CAUSED_BY_GENERATION" + + # Link must point to the LLM span's context + llm_spans = [s for s in finished if ".chat" in s.name] + assert len(llm_spans) == 1 + llm_span = llm_spans[0] + assert link.context.trace_id == llm_span.context.trace_id + assert link.context.span_id == llm_span.context.span_id + + +def test_no_attribution_link_without_tool_call_id(_handler, _exporter): + """Tool span without a structured tool_call_id must carry no attribution link.""" + root_id = uuid4() + tool_id = uuid4() + + _handler.on_chain_start( + {"id": ["FakeChain"], "name": "FakeChain"}, + {}, + run_id=root_id, + ) + _handler.on_tool_start( + {"id": ["SomeTool"], "name": "mytool"}, + "input", + run_id=tool_id, + parent_run_id=root_id, + ) + _handler.on_tool_end("output", run_id=tool_id, parent_run_id=root_id) + _handler.on_chain_end({}, run_id=root_id) + + finished = _exporter.get_finished_spans() + tool_spans = [s for s in finished if "execute_tool" in s.name] + assert len(tool_spans) == 1 + assert len(tool_spans[0].links) == 0, "Tool span with no LLM predecessor must have no links"