feat: correlate telemetry spans with the conversation

livekit-agents/livekit/agents/telemetry/gen_ai.py

@@ -333,8 +333,6 @@ def set_request_attributes(
     if stream:
         # "if and only if the request is streaming; if unset, assumed non-streaming"
         attrs[trace_types.ATTR_GEN_AI_REQUEST_STREAM] = True
-    if (conv := _conversation_id()) is not None:
-        attrs[trace_types.ATTR_GEN_AI_CONVERSATION_ID] = conv
     if output_type:
         attrs[trace_types.ATTR_GEN_AI_OUTPUT_TYPE] = output_type
     span.set_attributes(attrs)
@@ -419,11 +417,6 @@ def set_tool_attributes(
     if agent_name:
         # "the human-readable name of the agent executing the tool", conditionally required
         attrs[trace_types.ATTR_GEN_AI_AGENT_NAME] = agent_name
-    # not in the convention's execute_tool table, but Datadog groups a session by this
-    # attribute rather than by trace membership, so a tool span without it drops out of
-    # the session view
-    if (conv := _conversation_id()) is not None:
-        attrs[trace_types.ATTR_GEN_AI_CONVERSATION_ID] = conv
     if _capture_content:
         if description:
             attrs[trace_types.ATTR_GEN_AI_TOOL_DESCRIPTION] = description
@@ -464,8 +457,6 @@ def set_agent_attributes(
         attrs[trace_types.ATTR_GEN_AI_PROVIDER_NAME] = normalized
     if model:
         attrs[trace_types.ATTR_GEN_AI_REQUEST_MODEL] = model
-    if (conv := _conversation_id()) is not None:
-        attrs[trace_types.ATTR_GEN_AI_CONVERSATION_ID] = conv
     span.set_attributes(attrs)
 
 
@@ -478,8 +469,6 @@ def set_workflow_attributes(span: trace.Span, *, name: str) -> None:
         trace_types.ATTR_GEN_AI_OPERATION_NAME: trace_types.GenAIOperationName.INVOKE_WORKFLOW,
         trace_types.ATTR_GEN_AI_WORKFLOW_NAME: name,
     }
-    if (conv := _conversation_id()) is not None:
-        attrs[trace_types.ATTR_GEN_AI_CONVERSATION_ID] = conv
     span.set_attributes(attrs)
 
 

livekit-agents/livekit/agents/telemetry/traces.py

@@ -73,7 +73,7 @@
     recording_enabled,
 )
 from ..utils import is_given
-from . import pii, trace_types, utils as telemetry_utils
+from . import gen_ai, pii, trace_types, utils as telemetry_utils
 
 if TYPE_CHECKING:
     from ..llm import ChatItem
@@ -165,9 +165,18 @@ def _serialize_session_options(options: AgentSessionOptions) -> dict[str, Any]:
 
 
 _USE_SPAN_SIGNATURE = inspect.signature(trace_api.use_span)
+_START_SPAN_SIGNATURE = inspect.signature(Tracer.start_span)
 _START_AS_CURRENT_SPAN_SIGNATURE = inspect.signature(Tracer.start_as_current_span)
 
 
+def _with_conversation_id(
+    attributes: Mapping[str, AttributeValue] | None,
+) -> Mapping[str, AttributeValue] | None:
+    if (conversation_id := gen_ai._conversation_id()) is None:
+        return attributes
+    return {trace_types.ATTR_GEN_AI_CONVERSATION_ID: conversation_id, **(attributes or {})}
+
+
 class _DynamicTracer(Tracer):
     def __init__(self, instrumenting_module_name: str) -> None:
         self._instrumenting_module_name = instrumenting_module_name
@@ -182,7 +191,9 @@ def set_provider(self, tracer_provider: trace_api.TracerProvider) -> None:
         )
 
     def start_span(self, *args: Any, **kwargs: Any) -> Span:
-        return self._tracer.start_span(*args, **kwargs)
+        bound = _START_SPAN_SIGNATURE.bind(self._tracer, *args, **kwargs)
+        bound.arguments["attributes"] = _with_conversation_id(bound.arguments.get("attributes"))
+        return self._tracer.start_span(*bound.args[1:], **bound.kwargs)
 
     @_agnosticcontextmanager
     def use_span(self, *args: Any, **kwargs: Any) -> Iterator[Span]:
@@ -214,7 +225,7 @@ def detached_span(
         it and becomes the accidental parent of unrelated spans those tasks emit for the rest
         of the session. The parent is ``context`` when given, else the ambient context; the
         exception, if any, is recorded redaction-aware and the span is ended."""
-        span = self._tracer.start_span(name, context=context, attributes=attributes)
+        span = self.start_span(name, context=context, attributes=attributes)
         try:
             yield span
         except Exception as e:
@@ -226,6 +237,7 @@ def detached_span(
     @_agnosticcontextmanager
     def start_as_current_span(self, *args: Any, **kwargs: Any) -> Iterator[Span]:
         bound = _START_AS_CURRENT_SPAN_SIGNATURE.bind(self._tracer, *args, **kwargs)
+        bound.arguments["attributes"] = _with_conversation_id(bound.arguments.get("attributes"))
         record_exception = bound.arguments.get("record_exception", True)
         set_status_on_exception = bound.arguments.get("set_status_on_exception", True)
         bound.arguments.update(record_exception=False, set_status_on_exception=False)

tests/test_telemetry_conversation_id.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from contextlib import AbstractContextManager, nullcontext
+from functools import partial
+from types import MappingProxyType
+
+import pytest
+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.trace import Span, SpanKind
+
+from livekit.agents.telemetry import gen_ai, set_tracer_provider, tracer
+
+pytestmark = [pytest.mark.unit, pytest.mark.no_concurrent]
+
+
+@pytest.fixture
+def span_exporter() -> Iterator[InMemorySpanExporter]:
+    original_provider = tracer._tracer_provider
+    provider = TracerProvider()
+    exporter = InMemorySpanExporter()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    set_tracer_provider(provider)
+    try:
+        yield exporter
+    finally:
+        set_tracer_provider(original_provider)
+        provider.shutdown()
+
+
+@pytest.mark.parametrize("conversation_id", [None, "RM_test"])
+@pytest.mark.parametrize("explicit_id", [None, "explicit"])
+@pytest.mark.parametrize("api", ["start_span", "start_as_current_span", "detached_span"])
+def test_conversation_id_preserves_caller_attributes(
+    span_exporter: InMemorySpanExporter,
+    monkeypatch: pytest.MonkeyPatch,
+    conversation_id: str | None,
+    explicit_id: str | None,
+    api: str,
+) -> None:
+    monkeypatch.setattr(gen_ai, "_conversation_id", lambda: conversation_id)
+    attributes = {"test.attribute": "value"}
+    if explicit_id:
+        attributes["gen_ai.conversation.id"] = explicit_id
+    original = attributes.copy()
+    # The standard tracer APIs also allow positional, read-only attribute mappings.
+    if api == "start_span":
+        tracer.start_span("test", None, SpanKind.INTERNAL, MappingProxyType(attributes)).end()
+    elif api == "start_as_current_span":
+        with tracer.start_as_current_span(
+            "test", None, SpanKind.INTERNAL, MappingProxyType(attributes)
+        ):
+            pass
+    else:
+        with tracer.detached_span("test", attributes=attributes):
+            pass
+    assert attributes == original
+    expected = dict(original)
+    if explicit_id or conversation_id:
+        expected["gen_ai.conversation.id"] = explicit_id or conversation_id
+    assert dict(span_exporter.get_finished_spans()[0].attributes) == expected
+
+
+def test_spans_without_attributes_receive_conversation_id(
+    span_exporter: InMemorySpanExporter, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(gen_ai, "_conversation_id", lambda: "RM_test")
+    tracer.start_span("test").end()
+    with tracer.start_as_current_span("test"):
+        pass
+    with tracer.detached_span("test"):
+        pass
+    assert len(span_exporter.get_finished_spans()) == 3
+    for span in span_exporter.get_finished_spans():
+        assert span.attributes["gen_ai.conversation.id"] == "RM_test"
+
+
+@pytest.mark.parametrize(
+    "set_attributes",
+    [
+        pytest.param(partial(gen_ai.set_request_attributes, operation="chat"), id="request"),
+        pytest.param(partial(gen_ai.set_tool_attributes, name="get_weather"), id="tool"),
+        pytest.param(
+            partial(gen_ai.set_agent_attributes, operation="invoke_agent", agent_name="agent"),
+            id="agent",
+        ),
+        pytest.param(partial(gen_ai.set_workflow_attributes, name="agent_session"), id="workflow"),
+    ],
+)
+@pytest.mark.parametrize("explicit_id", [None, "explicit"])
+@pytest.mark.parametrize("conversation_id", [None, "RM_test"])
+@pytest.mark.parametrize("api", ["start_span", "start_as_current_span", "detached_span"])
+def test_gen_ai_setters_preserve_conversation_id(
+    span_exporter: InMemorySpanExporter,
+    monkeypatch: pytest.MonkeyPatch,
+    set_attributes: Callable[[Span], None],
+    explicit_id: str | None,
+    conversation_id: str | None,
+    api: str,
+) -> None:
+    monkeypatch.setattr(gen_ai, "_conversation_id", lambda: conversation_id)
+    attributes = {"gen_ai.conversation.id": explicit_id} if explicit_id is not None else {}
+    span_context: AbstractContextManager[Span]
+    if api == "start_span":
+        span_context = nullcontext(tracer.start_span("test", attributes=attributes))
+    elif api == "start_as_current_span":
+        span_context = tracer.start_as_current_span("test", attributes=attributes)
+    else:
+        span_context = tracer.detached_span("test", attributes=attributes)
+
+    with span_context as span:
+        set_attributes(span)
+        monkeypatch.setattr(gen_ai, "_conversation_id", lambda: "RM_other")
+        set_attributes(span)
+        if api == "start_span":
+            span.end()
+
+    exported = span_exporter.get_finished_spans()[0]
+    assert exported.attributes.get("gen_ai.conversation.id") == (explicit_id or conversation_id)

tests/test_telemetry_gen_ai.py

@@ -263,15 +263,3 @@ def test_the_provider_tables_only_target_registry_values() -> None:
         *(v for _, v in trace_types._PROVIDER_BY_HOST_SUFFIX),
     }
     assert targets <= trace_types.GEN_AI_PROVIDER_NAMES
-
-
-def test_tool_spans_carry_the_conversation_id(monkeypatch: pytest.MonkeyPatch) -> None:
-    # execute_tool is a first-class GenAI operation; a backend that groups by
-    # gen_ai.conversation.id would otherwise drop tool spans out of the session view
-    monkeypatch.setattr(gen_ai, "_conversation_id", lambda: "room-42")
-
-    span, exporter = _exporting_span("function_tool")
-    gen_ai.set_tool_attributes(span, name="get_weather", call_id="call_1")
-    span.end()
-
-    assert _attributes(exporter)["gen_ai.conversation.id"] == "room-42"

tests/test_telemetry_metadata.py

@@ -82,6 +82,7 @@ def test_exported_spans_keep_provider_metadata(
             "langfuse.session.id": "customer-session",
             "job_id": "provider-job" if state == "outside" else "job-a",
             "room_id": "provider-room" if state == "outside" else "room-a",
+            **({"gen_ai.conversation.id": "room-a"} if state != "outside" else {}),
         }
 
 
@@ -107,6 +108,7 @@ def test_job_fallback_metadata_does_not_cross_jobs(
         "langfuse.session.id": "customer-session",
         "job_id": "job-a",
         "room_id": "room-a",
+        "gen_ai.conversation.id": "room-a",
     }
 
     with tracer.start_as_current_span("worker"):