mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(otel v2): keep the proxy end user off a request root another v2 callback started
With another v2 callback listed before langfuse_otel (otel, arize_phoenix), that callback's provider starts the proxy root and exports it to its own backend, so the end-user fallback wrote user.id onto the generic or Phoenix root. Fill the end user only on a root this logger's provider started, and keep the caller's trace controls everywhere as on main
This commit is contained in:
parent
5b2d7cfc2a
commit
44634cadc9
4 changed files with 144 additions and 3 deletions
|
|
@ -1,6 +1,9 @@
|
|||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2
|
||||
from litellm.integrations.otel.mappers.langfuse import (
|
||||
|
|
@ -9,7 +12,7 @@ from litellm.integrations.otel.mappers.langfuse import (
|
|||
LangfuseMapper,
|
||||
)
|
||||
from litellm.integrations.otel.model.request_io import request_input, response_output, stream_output
|
||||
from litellm.integrations.otel.model.trace_controls import langfuse_trace_controls
|
||||
from litellm.integrations.otel.model.trace_controls import caller_trace_controls, langfuse_trace_controls
|
||||
from litellm.integrations.otel.plumbing.context import request_root_span
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -19,14 +22,22 @@ if TYPE_CHECKING:
|
|||
|
||||
class LangfuseOpenTelemetryV2(OpenTelemetryV2):
|
||||
"""Stamps the caller's trace controls (name, user, session, tags) on the request. Langfuse reads them off
|
||||
the root observation, and the proxy's root span is still recording when the LLM call starts."""
|
||||
the root observation, and the proxy's root span is still recording when the LLM call starts.
|
||||
|
||||
The proxy end user fills a missing user only on a root this logger's own provider started. With another v2
|
||||
callback listed first, that callback's provider starts the root and exports it to its own backend.
|
||||
"""
|
||||
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
|
||||
root: Final = request_root_span()
|
||||
if root is not None and root.is_recording():
|
||||
root.set_attributes(LangfuseMapper.trace_attributes(langfuse_trace_controls(kwargs)))
|
||||
controls: Final = langfuse_trace_controls(kwargs) if self._started(root) else caller_trace_controls(kwargs)
|
||||
root.set_attributes(LangfuseMapper.trace_attributes(controls))
|
||||
super().log_pre_api_call(model, messages, kwargs)
|
||||
|
||||
def _started(self, span: Span) -> bool:
|
||||
return isinstance(span, ReadableSpan) and span.resource is self.tracer_provider.resource
|
||||
|
||||
|
||||
class LangfuseContentOpenTelemetryV2(LangfuseOpenTelemetryV2):
|
||||
"""Stamps the request's input and output on the root observation while it is still recording.
|
||||
|
|
|
|||
|
|
@ -271,6 +271,34 @@ def _marker_trace_span_attributes(sink: Wire, batches: list[bytes], marker: str)
|
|||
return tuple(attributes for trace_id, attributes in spans if trace_id == marker_trace)
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
marker_trace: Final = next(
|
||||
(
|
||||
trace_id
|
||||
for trace_id, _has_parent, attributes in spans
|
||||
if any(isinstance(value, str) and marker in value for value in attributes.values())
|
||||
),
|
||||
None,
|
||||
)
|
||||
return tuple(
|
||||
attributes for trace_id, has_parent, attributes in spans if trace_id == marker_trace and not has_parent
|
||||
)
|
||||
|
||||
|
||||
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")}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,79 @@
|
|||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from _langfuse_otel import (
|
||||
_generation_span_attributes,
|
||||
_marker_trace_root_attributes,
|
||||
_observability_proxy,
|
||||
_sink,
|
||||
_upstream_reply_for,
|
||||
)
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.wire import wire_server
|
||||
|
||||
|
||||
def _generic_otel_env(collector_url: str) -> Mapping[str, str]:
|
||||
return {"OTEL_EXPORTER": "otlp_http", "OTEL_ENDPOINT": collector_url + "/v1/traces"}
|
||||
|
||||
|
||||
def _phoenix_env(collector_url: str) -> Mapping[str, str]:
|
||||
return {"PHOENIX_COLLECTOR_ENDPOINT": collector_url + "/v1/traces", "PHOENIX_API_KEY": "phoenix-integration"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("first_callback", "first_env"),
|
||||
[("otel", _generic_otel_env), ("arize_phoenix", _phoenix_env)],
|
||||
ids=["generic-otel-first", "arize-phoenix-first"],
|
||||
)
|
||||
def test_langfuse_otel_v2_keeps_the_end_user_off_another_callbacks_root(
|
||||
gateway: Gateway, tmp_path: Path, first_callback: str, first_env: Callable[[str], Mapping[str, str]]
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with (
|
||||
wire_server(lambda request: _upstream_reply_for(request, marker)) as provider,
|
||||
wire_server(_sink) as first_collector,
|
||||
wire_server(_sink) as langfuse_collector,
|
||||
_observability_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{
|
||||
"LITELLM_OTEL_V2": "1",
|
||||
"LANGFUSE_PUBLIC_KEY": "pk-integration",
|
||||
"LANGFUSE_SECRET_KEY": "sk-integration",
|
||||
"LANGFUSE_HOST": langfuse_collector.url,
|
||||
**first_env(first_collector.url),
|
||||
},
|
||||
callbacks=(first_callback, "langfuse_otel"),
|
||||
config_name="shared_root.yaml",
|
||||
) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
owner: Final = scenario.user()
|
||||
key: Final = scenario.key(user_id=owner)
|
||||
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=provider.url + "/v1")
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": marker}], "cache": {"no-cache": True}},
|
||||
key=key,
|
||||
headers={"x-litellm-end-user-id": f"end-user-{marker}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
first_batches: Final[list[bytes]] = []
|
||||
roots: Final = eventually(
|
||||
lambda: _marker_trace_root_attributes(first_collector, first_batches, marker),
|
||||
lambda spans: len(spans) == 1,
|
||||
seconds=30,
|
||||
)
|
||||
assert roots[0].get("http.route") == "/v1/chat/completions", roots
|
||||
assert "user.id" not in roots[0], roots
|
||||
langfuse_batches: Final[list[bytes]] = []
|
||||
generation: Final = eventually(
|
||||
lambda: _generation_span_attributes(langfuse_collector, langfuse_batches, marker),
|
||||
lambda spans: len(spans) == 1,
|
||||
seconds=30,
|
||||
)[0]
|
||||
assert generation.get("user.id") == f"end-user-{marker}", generation
|
||||
|
|
@ -466,6 +466,29 @@ def test_a_callers_trace_user_id_still_wins_over_the_proxy_end_user():
|
|||
assert generation_attrs["user.id"] == "caller-1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("metadata", "expected_root_user"),
|
||||
[
|
||||
({"user_api_key_end_user_id": "end-1"}, None),
|
||||
({"trace_user_id": "caller-1", "user_api_key_end_user_id": "end-1"}, "caller-1"),
|
||||
],
|
||||
ids=["end-user-stays-off", "caller-trace-user-as-before"],
|
||||
)
|
||||
def test_a_root_another_v2_logger_started_gets_only_the_callers_trace_user(metadata, expected_root_user):
|
||||
first, first_exporter = _logger(mappers=("genai",))
|
||||
langfuse, _ = _logger()
|
||||
root: Final = _start_root(first)
|
||||
|
||||
langfuse.log_pre_api_call(
|
||||
model="gpt-5.4-mini",
|
||||
messages=[],
|
||||
kwargs={"litellm_call_id": "call_1", "litellm_params": {"metadata": metadata}},
|
||||
)
|
||||
root.end()
|
||||
|
||||
assert _root_attrs(first_exporter).get("user.id") == expected_root_user
|
||||
|
||||
|
||||
def test_caller_metadata_cannot_override_the_proxy_team_identity():
|
||||
logger, exporter = _logger()
|
||||
response: Final = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue