test: count unique spans by span id and split stacked comprehensions in the Langfuse identity helpers
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled

The sink-outage chaos cell deduplicated by trace id and attributes, which could hide a second emission of the same call; an exporter retry re-sends the same span id. Also type the unit helpers and keep each comprehension to one for clause
This commit is contained in:
yucheng-berri 2026-10-16 20:14:37 -07:00
parent 44634cadc9
commit a7eb8b146b
4 changed files with 88 additions and 176 deletions

View file

@ -1,5 +1,6 @@
import itertools
import json
from collections.abc import Iterator, Mapping
from collections.abc import Iterable, Iterator, Mapping
from contextlib import contextmanager
from pathlib import Path
from typing import Final
@ -10,6 +11,7 @@ from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
from opentelemetry.proto.collector.trace.v1 import trace_service_pb2
from opentelemetry.proto.common.v1.common_pb2 import AnyValue
from opentelemetry.proto.trace.v1.trace_pb2 import Span
MARKER_JSON_KEYS: Final = ("llm.response.id", "gen_ai.response.id")
@ -19,18 +21,20 @@ def _attribute_value(value: AnyValue) -> object:
return getattr(value, kind) if kind is not None else None
def _spans(body: bytes) -> tuple[tuple[str, dict[str, object]], ...]:
def _span_messages(body: bytes) -> Iterator[Span]:
export: Final = trace_service_pb2.ExportTraceServiceRequest()
export.ParseFromString(body)
return tuple(
(
span.trace_id.hex(),
{attribute.key: _attribute_value(attribute.value) for attribute in span.attributes},
)
for resource in export.resource_spans
for scope in resource.scope_spans
for span in scope.spans
)
for resource in export.resource_spans:
for scope in resource.scope_spans:
yield from scope.spans
def _batch_span_messages(batches: Iterable[bytes]) -> Iterator[Span]:
return itertools.chain.from_iterable(_span_messages(body) for body in batches)
def _attributes(span: Span) -> dict[str, object]:
return {attribute.key: _attribute_value(attribute.value) for attribute in span.attributes}
def _sse_frame(value: Mapping[str, object]) -> bytes:
@ -184,7 +188,13 @@ def _sink(_request: Request) -> Reply:
def _drained_spans(sink: Wire, batches: list[bytes]) -> tuple[tuple[str, dict[str, object]], ...]:
batches.extend(request.body for request in sink.drain())
return tuple(span for body in batches for span in _spans(body))
return tuple((span.trace_id.hex(), _attributes(span)) for span in _batch_span_messages(batches))
def _drained_unique_spans(sink: Wire, batches: list[bytes]) -> tuple[tuple[str, dict[str, object]], ...]:
batches.extend(request.body for request in sink.drain())
by_span_id: Final = {span.span_id: span for span in _batch_span_messages(batches)}
return tuple((span.trace_id.hex(), _attributes(span)) for span in by_span_id.values())
def _generation_span_attributes(sink: Wire, batches: list[bytes], marker: str) -> tuple[dict[str, object], ...]:
@ -212,17 +222,6 @@ def _generation_marker_span_attributes(sink: Wire, batches: list[bytes], marker:
)
def _dedupe_spans(spans: tuple[tuple[str, dict[str, object]], ...]) -> tuple[tuple[str, dict[str, object]], ...]:
seen: Final = set()
unique: Final = []
for trace_id, attributes in spans:
fingerprint: Final = (trace_id, frozenset(attributes.items()))
if fingerprint not in seen:
seen.add(fingerprint)
unique.append((trace_id, attributes))
return tuple(unique)
def _span_containing_marker(
spans: tuple[tuple[str, dict[str, object]], ...], marker: str
) -> tuple[tuple[str, dict[str, object]], ...]:
@ -274,11 +273,7 @@ def _marker_trace_span_attributes(sink: Wire, batches: list[bytes], marker: str)
def _marker_trace_root_attributes(sink: Wire, batches: list[bytes], marker: str) -> tuple[dict[str, object], ...]:
batches.extend(request.body for request in sink.drain())
spans: Final = tuple(
(span.trace_id.hex(), bool(span.parent_span_id), {a.key: _attribute_value(a.value) for a in span.attributes})
for body in batches
for resource in _export(body).resource_spans
for scope in resource.scope_spans
for span in scope.spans
(span.trace_id.hex(), bool(span.parent_span_id), _attributes(span)) for span in _batch_span_messages(batches)
)
marker_trace: Final = next(
(
@ -293,12 +288,6 @@ def _marker_trace_root_attributes(sink: Wire, batches: list[bytes], marker: str)
)
def _export(body: bytes) -> trace_service_pb2.ExportTraceServiceRequest:
export: Final = trace_service_pb2.ExportTraceServiceRequest()
export.ParseFromString(body)
return export
def _user_id_session_id(attributes: Mapping[str, object]) -> dict[str, object]:
return {key: attributes.get(key) for key in ("user.id", "session.id")}

View file

@ -7,8 +7,7 @@ from typing import Final
import httpx
from _langfuse_otel import (
_dedupe_spans,
_drained_spans,
_drained_unique_spans,
_langfuse_proxy,
_marker_from_body,
_proxy_config,
@ -98,7 +97,7 @@ def _send_to_live_gateway(holder: dict[str, Gateway], model: str, endpoint: str,
def _landed_counts(sink: Wire, batches: list[bytes], markers: tuple[str, ...]) -> dict[str, int]:
spans: Final = _dedupe_spans(_drained_spans(sink, batches))
spans: Final = _drained_unique_spans(sink, batches)
generations: Final = tuple(span for span in spans if span[1].get("langfuse.observation.type") == "generation")
return {marker: len(_span_containing_marker(generations, marker)) for marker in markers}

View file

@ -1,3 +1,4 @@
import itertools
import uuid
from pathlib import Path
from typing import Final
@ -230,9 +231,8 @@ def test_langfuse_otel_openai_sdk_streaming_end_user(gateway: Gateway, tmp_path:
with stream:
chunks: Final = tuple(stream)
assert {chunk.id for chunk in chunks} == {marker}
assert f"reply {marker}" in "".join(
choice.delta.content or "" for chunk in chunks for choice in chunk.choices
)
choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
assert f"reply {marker}" in "".join(choice.delta.content or "" for choice in choices)
batches: Final[list[bytes]] = []
attributes: Final = _await_generation(collector, batches, marker)
assert _identity(attributes) == {"user.id": f"end-user-{marker}", "session.id": None}, attributes

View file

@ -2,6 +2,7 @@ import base64
import json
import logging
import os
from collections.abc import Mapping
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Final
@ -9,6 +10,7 @@ from unittest.mock import MagicMock, patch
import pytest
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import Span
from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
@ -111,19 +113,13 @@ class TestLangfuseOtelIntegration:
mock_kwargs = {"test": "kwargs"}
mock_response = {"test": "response"}
with patch(
"litellm.integrations.arize._utils.set_attributes"
) as mock_set_attributes:
LangfuseOtelLogger.set_langfuse_otel_attributes(
mock_span, mock_kwargs, mock_response
)
with patch("litellm.integrations.arize._utils.set_attributes") as mock_set_attributes:
LangfuseOtelLogger.set_langfuse_otel_attributes(mock_span, mock_kwargs, mock_response)
mock_set_attributes.assert_called_once_with(
mock_span, mock_kwargs, mock_response, LangfuseLLMObsOTELAttributes, emit_session_and_user=False
)
mock_span.set_attribute.assert_any_call(
"langfuse.observation.type", "generation"
)
mock_span.set_attribute.assert_any_call("langfuse.observation.type", "generation")
def test_set_langfuse_environment_attribute(self):
"""Test that Langfuse environment is set correctly when environment variable is present."""
@ -132,17 +128,11 @@ class TestLangfuseOtelIntegration:
test_env = "staging"
with patch.dict(os.environ, {"LANGFUSE_TRACING_ENVIRONMENT": test_env}):
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(
mock_span, mock_kwargs, {}
)
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, mock_kwargs, {})
# safe_set_attribute(span, key, value) → positional args
mock_safe_set_attribute.assert_called_once_with(
mock_span, "langfuse.environment", test_env
)
mock_safe_set_attribute.assert_called_once_with(mock_span, "langfuse.environment", test_env)
def test_set_langfuse_environment_attribute_prefers_dynamic_param(self):
"""Per-key/team langfuse_environment beats the deployment env var."""
@ -155,18 +145,10 @@ class TestLangfuseOtelIntegration:
self.attributes[key] = value
span = _RecordingSpan()
mock_kwargs = {
"standard_callback_dynamic_params": {
"langfuse_environment": "team-a-env"
}
}
mock_kwargs = {"standard_callback_dynamic_params": {"langfuse_environment": "team-a-env"}}
with patch.dict(
os.environ, {"LANGFUSE_TRACING_ENVIRONMENT": "deployment-wide"}
):
LangfuseOtelLogger._set_langfuse_specific_attributes(
span, mock_kwargs, {}
)
with patch.dict(os.environ, {"LANGFUSE_TRACING_ENVIRONMENT": "deployment-wide"}):
LangfuseOtelLogger._set_langfuse_specific_attributes(span, mock_kwargs, {})
assert span.attributes["langfuse.environment"] == "team-a-env"
@ -230,12 +212,8 @@ class TestLangfuseOtelIntegration:
kwargs = {"litellm_params": {"metadata": metadata}}
# Capture calls to safe_set_attribute
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(
MagicMock(), kwargs, None
)
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), kwargs, None)
# Build expected calls manually for clarity
from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes
@ -256,9 +234,7 @@ class TestLangfuseOtelIntegration:
LangfuseSpanAttributes.TRACE_METADATA.value: json.dumps({"k": "v"}),
LangfuseSpanAttributes.RELEASE.value: "rel-1",
LangfuseSpanAttributes.EXISTING_TRACE_ID.value: "existing-id",
LangfuseSpanAttributes.UPDATE_TRACE_KEYS.value: json.dumps(
["key1", "key2"]
),
LangfuseSpanAttributes.UPDATE_TRACE_KEYS.value: json.dumps(["key1", "key2"]),
LangfuseSpanAttributes.DEBUG_LANGFUSE.value: True,
}
@ -268,9 +244,7 @@ class TestLangfuseOtelIntegration:
for call in mock_safe_set_attribute.call_args_list
}
assert (
actual == expected
), "Mismatch between expected and actual OTEL attribute mapping."
assert actual == expected, "Mismatch between expected and actual OTEL attribute mapping."
@pytest.mark.parametrize(
"metadata, expected_version",
@ -295,16 +269,10 @@ class TestLangfuseOtelIntegration:
def test_version_emitted_on_langfuse_v4_key(self, metadata, expected_version):
kwargs = {"litellm_params": {"metadata": {"trace_release": "rel-9", **metadata}}}
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(
MagicMock(), kwargs, None
)
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), kwargs, None)
emitted = {
call.args[1]: call.args[2] for call in mock_safe_set_attribute.call_args_list
}
emitted = {call.args[1]: call.args[2] for call in mock_safe_set_attribute.call_args_list}
if expected_version is None:
assert "langfuse.version" not in emitted
@ -342,12 +310,8 @@ class TestLangfuseOtelIntegration:
"messages": [{"role": "user", "content": "What's the weather in Tokyo?"}],
}
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(
MagicMock(), kwargs, response_obj
)
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), kwargs, response_obj)
expect_output = {
LangfuseSpanAttributes.OBSERVATION_INPUT.value: [
@ -360,14 +324,9 @@ class TestLangfuseOtelIntegration:
}
# Flatten the actual calls into {key: value}
actual = {
call.args[1]: json.loads(call.args[2])
for call in mock_safe_set_attribute.call_args_list
}
actual = {call.args[1]: json.loads(call.args[2]) for call in mock_safe_set_attribute.call_args_list}
assert (
actual == expect_output
), "Mismatch in observation input/output OTEL attributes."
assert actual == expect_output, "Mismatch in observation input/output OTEL attributes."
def test_set_langfuse_specific_attributes_with_tool_calls(self):
"""Test that _set_langfuse_specific_attributes correctly sets observation.output with tool calls in Langfuse format."""
@ -391,9 +350,7 @@ class TestLangfuseOtelIntegration:
"content": None,
"tool_calls": [
ChatCompletionMessageToolCall(
function=Function(
arguments='{"location":"Tokyo"}', name="get_weather"
),
function=Function(arguments='{"location":"Tokyo"}', name="get_weather"),
id="call_123",
type="function",
)
@ -403,12 +360,8 @@ class TestLangfuseOtelIntegration:
],
)
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(
MagicMock(), {}, response_obj
)
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), {}, response_obj)
expected = {
LangfuseSpanAttributes.OBSERVATION_OUTPUT.value: [
@ -423,13 +376,8 @@ class TestLangfuseOtelIntegration:
}
# Flatten the actual calls into {key: value}
actual = {
call.args[1]: json.loads(call.args[2])
for call in mock_safe_set_attribute.call_args_list
}
assert (
actual == expected
), "Mismatch in observation output OTEL attribute for tool calls."
actual = {call.args[1]: json.loads(call.args[2]) for call in mock_safe_set_attribute.call_args_list}
assert actual == expected, "Mismatch in observation output OTEL attribute for tool calls."
def test_construct_dynamic_otel_headers_with_langfuse_keys(self):
"""Test that construct_dynamic_otel_headers creates proper auth headers when langfuse keys are provided."""
@ -580,9 +528,7 @@ class TestLangfuseOtelKeyDynamicConfig:
logger = LangfuseOtelLogger()
assert logger.OTEL_EXPORTER == "console"
tracer = logger.get_tracer_to_use_for_request(
{"standard_callback_dynamic_params": self._dynamic_params()}
)
tracer = logger.get_tracer_to_use_for_request({"standard_callback_dynamic_params": self._dynamic_params()})
assert tracer is not logger.tracer
assert len(logger._tracer_provider_cache) == 1
@ -643,9 +589,7 @@ class TestLangfuseOtelKeyDynamicConfig:
with self._clean_env():
logger = LangfuseOtelLogger()
with patch.object(otel_module.verbose_logger, "debug", side_effect=_spy):
logger.get_tracer_to_use_for_request(
{"standard_callback_dynamic_params": self._dynamic_params()}
)
logger.get_tracer_to_use_for_request({"standard_callback_dynamic_params": self._dynamic_params()})
logged = "\n".join(recorded_arguments)
assert "initializing span processor" in logged
@ -705,12 +649,8 @@ class TestLangfuseOtelResponsesAPI:
LangfuseLLMObsOTELAttributes,
)
with patch(
"litellm.integrations.arize._utils.set_attributes"
) as mock_set_attributes:
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
with patch("litellm.integrations.arize._utils.set_attributes") as mock_set_attributes:
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
logger = LangfuseOtelLogger()
logger.set_langfuse_otel_attributes(mock_span, kwargs, mock_response)
@ -723,9 +663,7 @@ class TestLangfuseOtelResponsesAPI:
mock_safe_set_attribute.assert_any_call(
mock_span, "langfuse.generation.name", "responses_test_generation"
)
mock_safe_set_attribute.assert_any_call(
mock_span, "langfuse.trace.name", "responses_api_trace"
)
mock_safe_set_attribute.assert_any_call(mock_span, "langfuse.trace.name", "responses_api_trace")
def test_responses_api_metadata_extraction(self):
"""Test that metadata is correctly extracted from ResponsesAPI kwargs."""
@ -775,9 +713,7 @@ class TestLangfuseOtelResponsesAPI:
mock_span = MagicMock()
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, {})
# Verify specific attributes were set
@ -859,21 +795,15 @@ class TestLangfuseOtelResponsesAPI:
kwargs = {
"call_type": "responses",
"messages": [
{"role": "user", "content": "What's the weather in San Francisco?"}
],
"messages": [{"role": "user", "content": "What's the weather in San Francisco?"}],
"model": "gpt-4o",
"optional_params": {},
}
mock_span = MagicMock()
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(
mock_span, kwargs, response_obj
)
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj)
# Verify observation output was set
output_calls = [
@ -892,17 +822,11 @@ class TestLangfuseOtelResponsesAPI:
# Verify reasoning summary
assert output_data[0]["role"] == "reasoning_summary"
assert (
output_data[0]["content"]
== "Let me analyze this problem step by step..."
)
assert output_data[0]["content"] == "Let me analyze this problem step by step..."
# Verify message
assert output_data[1]["role"] == "assistant"
assert (
output_data[1]["content"]
== "The weather in San Francisco is sunny, 20°C."
)
assert output_data[1]["content"] == "The weather in San Francisco is sunny, 20°C."
def test_responses_api_with_function_calls(self):
"""Test Langfuse OTEL logger with Responses API function_call output."""
@ -927,21 +851,15 @@ class TestLangfuseOtelResponsesAPI:
kwargs = {
"call_type": "responses",
"messages": [
{"role": "user", "content": "What's the weather in San Francisco?"}
],
"messages": [{"role": "user", "content": "What's the weather in San Francisco?"}],
"model": "gpt-4o",
"optional_params": {},
}
mock_span = MagicMock()
with patch(
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(
mock_span, kwargs, response_obj
)
with patch("litellm.integrations.arize._utils.safe_set_attribute") as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj)
# Verify observation output was set
output_calls = [
@ -996,9 +914,11 @@ class TestLangfuseOtelResponsesAPI:
mock_span = MagicMock()
with patch( # test-quality-ok: the span attribute sink is the observable boundary; sibling tests in this class stub the same seam
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute:
with (
patch( # test-quality-ok: the span attribute sink is the observable boundary; sibling tests in this class stub the same seam
"litellm.integrations.arize._utils.safe_set_attribute"
) as mock_safe_set_attribute
):
LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj)
output_calls = [
@ -1014,12 +934,16 @@ class TestLangfuseOtelResponsesAPI:
class TestLangfuseOtelTraceIdentity:
def _recording_span(self):
from opentelemetry.sdk.trace import TracerProvider
def _recording_span(self) -> Span:
return TracerProvider().get_tracer("test").start_span("generation")
def _kwargs(self, slp_metadata=None, request_metadata=None, model_parameters=None, slp_extra=None):
def _kwargs(
self,
slp_metadata: Mapping[str, object] | None = None,
request_metadata: Mapping[str, object] | None = None,
model_parameters: Mapping[str, object] | None = None,
slp_extra: Mapping[str, object] | None = None,
) -> dict[str, object]:
return {
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hello"}],
@ -1033,7 +957,7 @@ class TestLangfuseOtelTraceIdentity:
},
}
def _response_obj(self):
def _response_obj(self) -> dict[str, object]:
return {
"id": "chatcmpl-1",
"model": "gpt-4o",
@ -1041,7 +965,7 @@ class TestLangfuseOtelTraceIdentity:
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
def _identity(self, kwargs):
def _identity(self, kwargs: Mapping[str, object]) -> tuple[dict[str, object], dict[str, object]]:
span = self._recording_span()
LangfuseOtelLogger.set_langfuse_otel_attributes(span, kwargs, self._response_obj())
attributes = dict(span.attributes or {})
@ -1178,9 +1102,7 @@ class TestLangfuseOtelTraceIdentity:
LANGFUSE_ENV_ONLY: Final = {
key: value
for key, value in os.environ.items()
if not key.startswith(("LANGFUSE_", "OTEL_"))
key: value for key, value in os.environ.items() if not key.startswith(("LANGFUSE_", "OTEL_"))
}
@ -1191,7 +1113,9 @@ def _decoded_basic_auth(header: str) -> str:
def test_langfuse_otel_env_config_headers_carry_v4_ingestion_and_basic_auth() -> None:
with patch.dict(
os.environ, {**LANGFUSE_ENV_ONLY, "LANGFUSE_PUBLIC_KEY": "pk-lf-123", "LANGFUSE_SECRET_KEY": "sk-lf-123"}, clear=True
os.environ,
{**LANGFUSE_ENV_ONLY, "LANGFUSE_PUBLIC_KEY": "pk-lf-123", "LANGFUSE_SECRET_KEY": "sk-lf-123"},
clear=True,
):
logger: Final = LangfuseOtelLogger()
headers: Final = OpenTelemetry._get_headers_dictionary(logger.config.headers)