diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 6db26c09c32..1444e735fed 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -63,10 +63,10 @@ class LangfuseOtelLogger(OpenTelemetry): Sets OpenTelemetry span attributes for Langfuse observability. Uses the same attribute setting logic as Arize Phoenix for consistency. """ - from opentelemetry.sdk.trace import ReadableSpan + from opentelemetry import trace - # A span that already carries attributes (a caller's active span) does not get the caller identity. - fresh: Final = not (isinstance(span, ReadableSpan) and span.attributes) + # The caller's active span, which LiteLLM writes onto instead of starting its own, gets no caller identity. + own_span: Final = span is not trace.get_current_span() _utils.set_attributes(span, kwargs, response_obj, LangfuseLLMObsOTELAttributes) span.set_attribute("langfuse.observation.type", "generation") @@ -74,7 +74,7 @@ class LangfuseOtelLogger(OpenTelemetry): # Set Langfuse specific attributes ######################################################### LangfuseOtelLogger._set_langfuse_specific_attributes(span=span, kwargs=kwargs, response_obj=response_obj) - if fresh: + if own_span: LangfuseOtelLogger._set_request_metadata_attributes( span=span, kwargs=_MAPPING_ADAPTER.validate_python(kwargs) ) diff --git a/tests/integration/sdk/test_langfuse_otel_caller_span_sdk.py b/tests/integration/sdk/test_langfuse_otel_caller_span_sdk.py new file mode 100644 index 00000000000..2e0f59974bd --- /dev/null +++ b/tests/integration/sdk/test_langfuse_otel_caller_span_sdk.py @@ -0,0 +1,113 @@ +import asyncio +import json +import time +from collections.abc import Mapping +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +import litellm +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + +_IDENTITY: Final = { + "user_api_key_alias": "caller-span-alias", + "user_api_key_user_id": "caller-span-user", + "user_api_key_team_id": "caller-span-team", +} +_CHAT_REPLY: Final = Reply( + body=json.dumps( + { + "id": "chatcmpl-caller-span", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ).encode() +) + + +def _is_identity(key: str) -> bool: + return key == "langfuse.observation.metadata" or key.startswith("langfuse.trace.metadata.user_api_key_") + + +def _exported(langfuse: Wire, seen: list[tuple[str, dict[str, object]]]) -> list[tuple[str, dict[str, object]]]: + for batch in langfuse.drain(): + request: ExportTraceServiceRequest = ExportTraceServiceRequest.FromString(batch.body) + seen.extend( + ( + span.name, + { + attribute.key: getattr(attribute.value, attribute.value.WhichOneof("value")) + for attribute in span.attributes + }, + ) + for resource in request.resource_spans + for scope in resource.scope_spans + for span in scope.spans + ) + return seen + + +async def _exported_until( + langfuse: Wire, seen: list[tuple[str, dict[str, object]]], name: str, prompt: str +) -> list[tuple[str, dict[str, object]]]: + """The spans the Langfuse sink got, once a ``name`` span for ``prompt`` is among them; polls without blocking the loop.""" + deadline: Final = time.monotonic() + 25 + while not any(span == name and prompt in json.dumps(attributes) for span, attributes in _exported(langfuse, seen)): + assert time.monotonic() < deadline, [span for span, _ in seen] + await asyncio.sleep(0.2) + return seen + + +async def _complete(upstream: Wire, prompt: str) -> None: + await litellm.acompletion( + model="openai/gpt-4o-mini", + api_base=upstream.url + "/v1", + api_key="synthetic-caller-span-key", + messages=[{"role": "user", "content": prompt}], + metadata=dict(_IDENTITY), + ) + await GLOBAL_LOGGING_WORKER.flush() + + +@pytest.mark.asyncio +async def test_langfuse_otel_keeps_the_caller_identity_off_the_callers_active_span( + monkeypatch: pytest.MonkeyPatch, +) -> None: + app_exporter: Final = InMemorySpanExporter() + app_provider: Final = TracerProvider() + app_provider.add_span_processor(SimpleSpanProcessor(app_exporter)) + + def chat(_request: Request) -> Reply: + return _CHAT_REPLY + + with wire_server(chat) as upstream, wire_server(lambda _request: Reply()) as langfuse: + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-caller-span") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-caller-span") + monkeypatch.setenv("LANGFUSE_OTEL_HOST", langfuse.url) + monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", "100") + monkeypatch.setattr(litellm, "callbacks", ["langfuse_otel"]) + + seen: Final[list[tuple[str, dict[str, object]]]] = [] # mutable-ok: drain() consumes the sink queue + with app_provider.get_tracer("app").start_as_current_span("app.operation"): + await _complete(upstream, "inside the caller span") + await _exported_until(langfuse, seen, "raw_gen_ai_request", "inside the caller span") + await _complete(upstream, "no caller span") + exported: Final = await _exported_until(langfuse, seen, "litellm_request", "no caller span") + + (app_span,) = app_exporter.get_finished_spans() + app_attributes: Final[Mapping[str, object]] = app_span.attributes or {} + assert app_attributes["langfuse.observation.type"] == "generation", sorted(app_attributes) + assert not [key for key in app_attributes if _is_identity(key)], sorted(app_attributes) + (own_span,) = [attributes for name, attributes in exported if name == "litellm_request"] + assert json.loads(str(own_span["langfuse.observation.metadata"])) == _IDENTITY, own_span + assert {key: value for key, value in own_span.items() if key.startswith("langfuse.trace.metadata.")} == { + f"langfuse.trace.metadata.{field}": value for field, value in _IDENTITY.items() + }, sorted(own_span) diff --git a/tests/unit/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py index fec14efa477..5b72f7e191e 100644 --- a/tests/unit/integrations/test_langfuse_otel.py +++ b/tests/unit/integrations/test_langfuse_otel.py @@ -2,12 +2,14 @@ import base64 import json import os from collections.abc import Mapping +from contextlib import nullcontext from datetime import datetime, timezone from types import MappingProxyType from typing import Final from unittest.mock import MagicMock, patch import pytest +from opentelemetry import trace from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -720,13 +722,15 @@ class TestLangfuseOtelIntegration: max_span_attributes: int | None, carried: dict[str, str] | None = None, added_after: Mapping[str, str | int] = MappingProxyType({}), + active: bool = False, ) -> ReadableSpan: exporter: Final = InMemorySpanExporter() provider: Final = TracerProvider(span_limits=SpanLimits(max_span_attributes=max_span_attributes)) provider.add_span_processor(SimpleSpanProcessor(exporter)) span: Final = provider.get_tracer("langfuse-otel-test").start_span("litellm_request", attributes=carried) response_obj: Final = {"id": "chatcmpl-1", "choices": [{"message": {"role": "assistant", "content": "hello"}}]} - LangfuseOtelLogger.set_langfuse_otel_attributes(span, kwargs, response_obj) + with trace.use_span(span, end_on_exit=False) if active else nullcontext(): + LangfuseOtelLogger.set_langfuse_otel_attributes(span, kwargs, response_obj) span.set_attributes(added_after) span.end() (finished,) = exporter.get_finished_spans() @@ -786,13 +790,13 @@ class TestLangfuseOtelIntegration: kept_identity.append(len(set(failed_span.attributes or {}) & identity_keys)) assert kept_identity == list(range(len(identity_keys) + 1)) - def test_request_metadata_skips_a_span_that_already_carries_attributes(self) -> None: + @pytest.mark.parametrize("carried", [None, {"caller.attribute": "kept"}], ids=["bare", "with_attributes"]) + def test_request_metadata_skips_the_callers_active_span(self, carried: dict[str, str] | None) -> None: attributes: Final = dict( - self._exported_span(self._tool_call_kwargs(), SpanLimits.UNSET, {"caller.attribute": "kept"}).attributes - or {} + self._exported_span(self._tool_call_kwargs(), SpanLimits.UNSET, carried, active=True).attributes or {} ) - assert attributes["caller.attribute"] == "kept" + assert {key: attributes[key] for key in carried or {}} == (carried or {}) assert attributes["langfuse.observation.type"] == "generation" assert "langfuse.observation.metadata" not in attributes assert not [key for key in attributes if key.startswith("langfuse.trace.metadata.")] @@ -1376,3 +1380,38 @@ def test_langfuse_otel_does_not_start_proxy_request_span() -> None: finally: langfuse_provider.shutdown() generic_provider.shutdown() + + +def test_langfuse_otel_success_writes_identity_onto_its_own_span_and_never_onto_the_callers() -> None: + langfuse_exporter: Final = InMemorySpanExporter() + langfuse_provider: Final = TracerProvider() + langfuse_provider.add_span_processor(SimpleSpanProcessor(langfuse_exporter)) + app_exporter: Final = InMemorySpanExporter() + app_provider: Final = TracerProvider() + app_provider.add_span_processor(SimpleSpanProcessor(app_exporter)) + with patch.dict(os.environ, LANGFUSE_ENV_ONLY, clear=True): + logger: Final = LangfuseOtelLogger(callback_name="langfuse_otel", tracer_provider=langfuse_provider) + response_obj: Final = {"id": "chatcmpl-1", "choices": [{"message": {"role": "assistant", "content": "hello"}}]} + started_at: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + try: + with app_provider.get_tracer("app").start_as_current_span("app.operation"): + logger._handle_success( + TestLangfuseOtelIntegration._tool_call_kwargs(), response_obj, started_at, started_at + ) + logger._handle_success(TestLangfuseOtelIntegration._tool_call_kwargs(), response_obj, started_at, started_at) + + (app_span,) = app_exporter.get_finished_spans() + (own_span,) = [span for span in langfuse_exporter.get_finished_spans() if span.name == "litellm_request"] + app_attributes: Final = dict(app_span.attributes or {}) + own_attributes: Final = dict(own_span.attributes or {}) + assert app_attributes["langfuse.observation.type"] == "generation" + identity_prefixes: Final = ("langfuse.observation.metadata", "langfuse.trace.metadata.") + assert not [key for key in app_attributes if key.startswith(identity_prefixes)] + assert json.loads(str(own_attributes["langfuse.observation.metadata"])) == { + "user_api_key_alias": "prod-key", + "user_api_key_team_id": "team-1", + } + assert own_attributes["langfuse.trace.metadata.user_api_key_team_id"] == "team-1" + finally: + langfuse_provider.shutdown() + app_provider.shutdown()