diff --git a/tests/integration/observability/_langfuse_otel.py b/tests/integration/observability/_langfuse_otel.py index 72092dc7e7a..5bdfe3398bf 100644 --- a/tests/integration/observability/_langfuse_otel.py +++ b/tests/integration/observability/_langfuse_otel.py @@ -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")} diff --git a/tests/integration/observability/test_langfuse_otel_identity_chaos.py b/tests/integration/observability/test_langfuse_otel_identity_chaos.py index 72364e6e1da..ee72f1a8fe7 100644 --- a/tests/integration/observability/test_langfuse_otel_identity_chaos.py +++ b/tests/integration/observability/test_langfuse_otel_identity_chaos.py @@ -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} diff --git a/tests/integration/observability/test_langfuse_otel_identity_surfaces.py b/tests/integration/observability/test_langfuse_otel_identity_surfaces.py index 469242e628b..43fd06ec60b 100644 --- a/tests/integration/observability/test_langfuse_otel_identity_surfaces.py +++ b/tests/integration/observability/test_langfuse_otel_identity_surfaces.py @@ -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 diff --git a/tests/unit/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py index 2d7241c0fc8..948a8c70e61 100644 --- a/tests/unit/integrations/test_langfuse_otel.py +++ b/tests/unit/integrations/test_langfuse_otel.py @@ -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)