mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix(otel): give Mixpanel the caller's user.id on every span
The preset takes over from #48638. Spans after auth now carry the id the caller named, in the proxy's project and in a team's or key's own project, and a request's id never reaches another backend. `metadata["user.id"]` only labels spans, so customer limits and spend logs stay as they are on main. Mixpanel tokens that could never work get a 400 on every writer, and a request with no project token keeps the operator's region.
This commit is contained in:
parent
b79fc2dc9e
commit
0427b0e35f
28 changed files with 1885 additions and 122 deletions
|
|
@ -55,6 +55,7 @@ from litellm.integrations.otel.model.utils import to_ns
|
|||
from litellm.integrations.otel.plumbing.context import (
|
||||
active_phase,
|
||||
active_phase_span,
|
||||
inbound_user_id,
|
||||
is_recordable_span,
|
||||
mcp_message_transport_span,
|
||||
post_response_root,
|
||||
|
|
@ -601,7 +602,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
trace=call.trace,
|
||||
session_id=call.session_id,
|
||||
metadata_keys=tuple(self.config.baggage_metadata_keys),
|
||||
request_metadata=call.user_id_metadata,
|
||||
named_user_id=call.caller_user_id or inbound_user_id(),
|
||||
operation=call.operation,
|
||||
)
|
||||
end_time_ns: Final = to_ns(end_time)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,6 @@
|
|||
from litellm.integrations.otel.mappers.base import AttributeMap, SpanData
|
||||
from litellm.integrations.otel.mappers.utils import drop_none_pairs
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
LLMCallSpanData,
|
||||
MCPListToolsSpanData,
|
||||
MCPToolCallSpanData,
|
||||
)
|
||||
from litellm.integrations.otel.model.payloads import LLMCallSpanData
|
||||
from litellm.integrations.otel.model.semconv import User
|
||||
|
||||
|
||||
|
|
@ -13,7 +9,5 @@ class MixpanelMapper:
|
|||
match data:
|
||||
case LLMCallSpanData():
|
||||
return drop_none_pairs(((User.ID, data.identity.user_id or data.trace.user_id or None),))
|
||||
case MCPToolCallSpanData() | MCPListToolsSpanData():
|
||||
return drop_none_pairs(((User.ID, data.identity.user_id or None),))
|
||||
case _:
|
||||
return drop_none_pairs(())
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from types import MappingProxyType
|
|||
from typing import Final
|
||||
|
||||
from litellm.integrations.otel.model.metadata import RequestIdentity, allowlisted_metadata
|
||||
from litellm.integrations.otel.model.semconv import GenAI, LiteLLM, User
|
||||
from litellm.integrations.otel.model.semconv import GenAI, LiteLLM
|
||||
|
||||
# Attribute key -> value extractor over (identity, request_model,
|
||||
# team_metadata_keys). The single definition of what may be promoted and under
|
||||
|
|
@ -33,7 +33,6 @@ _PROMOTABLE: Final[dict[str, Callable[[RequestIdentity, str | None, tuple[str, .
|
|||
),
|
||||
LiteLLM.KEY_HASH: lambda identity, model, team_metadata_keys: identity.key_hash,
|
||||
LiteLLM.END_USER: lambda identity, model, team_metadata_keys: identity.end_user,
|
||||
User.ID: lambda identity, model, team_metadata_keys: identity.user_id,
|
||||
GenAI.REQUEST_MODEL: lambda identity, model, team_metadata_keys: model,
|
||||
LiteLLM.PROVIDER_MODEL: lambda identity, model, team_metadata_keys: identity.provider_model,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from __future__ import annotations
|
|||
from collections.abc import Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast, get_args
|
||||
|
||||
|
|
@ -52,6 +53,7 @@ from litellm.constants import (
|
|||
from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation, resolve_operation
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls
|
||||
from litellm.integrations.otel.model.utils import as_str, as_str_mapping, to_seconds
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.types.utils import InternalCallOrigin
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -71,14 +73,9 @@ IDENTITY_METADATA_FIELDS: Final = (
|
|||
|
||||
REQUESTER_METADATA_KEY: Final = "requester_metadata"
|
||||
REQUESTER_METADATA_PATH: Final = f"{REQUESTER_METADATA_KEY}."
|
||||
_USER_ID_METADATA_KEYS: Final[tuple[str, ...]] = (
|
||||
"requester_metadata.user.id",
|
||||
"user.id",
|
||||
"requester_metadata.user_id",
|
||||
"user_id",
|
||||
"requester_metadata.trace_user_id",
|
||||
"trace_user_id",
|
||||
)
|
||||
_USER_ID_METADATA_FIELDS: Final[tuple[str, ...]] = ("user.id", "user_id", "trace_user_id")
|
||||
_MAX_USER_ID_LENGTH: Final = 1024
|
||||
_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def allowlisted_metadata(metadata: Mapping[str, str], metadata_keys: tuple[str, ...]) -> Mapping[str, str]:
|
||||
|
|
@ -108,20 +105,13 @@ class RequestIdentity:
|
|||
provider_model: str | None = None
|
||||
request_route: str | None = None
|
||||
metadata: Mapping[str, str] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def user_id(self) -> str | None:
|
||||
return (
|
||||
next((value for key in _USER_ID_METADATA_KEYS if (value := self.metadata.get(key))), None)
|
||||
or self.end_user
|
||||
or self.metadata.get("user_api_key_end_user_id")
|
||||
)
|
||||
user_id: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_payload(
|
||||
cls,
|
||||
payload: StandardLoggingPayload,
|
||||
request_metadata: Mapping[str, str] = MappingProxyType({}),
|
||||
named_user_id: str | None = None,
|
||||
) -> RequestIdentity:
|
||||
"""Parse caller identity out of a closed request's payload metadata.
|
||||
|
||||
|
|
@ -130,9 +120,9 @@ class RequestIdentity:
|
|||
model, not just the user-facing one.
|
||||
"""
|
||||
raw_meta: Final = cast(Mapping[str, object], payload.get("metadata") or {})
|
||||
payload_view: Final = as_str_mapping(payload)
|
||||
params: Final = as_str_mapping(payload_view.get("model_parameters")) if payload_view is not None else None
|
||||
metadata: Final = MappingProxyType({**request_metadata, **dict(flatten_metadata(raw_meta))})
|
||||
metadata: Final = MappingProxyType(dict(flatten_metadata(raw_meta)))
|
||||
end_user: Final = as_str(payload.get("end_user")) or as_str(raw_meta.get("user_api_key_end_user_id"))
|
||||
params: Final = _as_mapping(cast(Mapping[str, object], payload).get("model_parameters")) # cast-ok: TypedDict
|
||||
return cls(
|
||||
call_id=as_str(payload.get("litellm_call_id")) or as_str(payload.get("id")),
|
||||
# StandardLoggingMetadata's canonical key is ``user_api_key_team_id``;
|
||||
|
|
@ -142,14 +132,16 @@ class RequestIdentity:
|
|||
team_alias=as_str(raw_meta.get("user_api_key_team_alias")) or as_str(raw_meta.get("team_alias")),
|
||||
team_metadata=_team_metadata_dict(raw_meta.get("user_api_key_team_metadata")),
|
||||
key_hash=as_str(raw_meta.get("user_api_key_hash")),
|
||||
end_user=(
|
||||
as_str(payload.get("end_user"))
|
||||
or as_str(raw_meta.get("user_api_key_end_user_id"))
|
||||
or (as_str(params.get("user")) if params is not None else None)
|
||||
),
|
||||
end_user=end_user,
|
||||
provider_model=resolve_provider_model(payload),
|
||||
request_route=as_str(raw_meta.get("user_api_key_request_route")),
|
||||
metadata=metadata,
|
||||
user_id=(
|
||||
named_user_id
|
||||
or caller_user_id(raw_meta)
|
||||
or usable_user_id(end_user)
|
||||
or usable_user_id(params.get("user"))
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
|
|
@ -217,7 +209,7 @@ class RequestContext:
|
|||
def from_standard_logging_payload(
|
||||
cls,
|
||||
payload: StandardLoggingPayload,
|
||||
request_metadata: Mapping[str, str] = MappingProxyType({}),
|
||||
named_user_id: str | None = None,
|
||||
) -> RequestContext:
|
||||
raw_meta: Final = cast(Mapping[str, object], payload.get("metadata") or {})
|
||||
hidden: Final = cast(Mapping[str, object], payload.get("hidden_params") or {})
|
||||
|
|
@ -233,7 +225,7 @@ class RequestContext:
|
|||
model_group=model_group,
|
||||
model_id=as_str(payload.get("model_id")) or _model_info_id(raw_meta.get("model_info")),
|
||||
api_base=as_str(payload.get("api_base")) or as_str(hidden.get("api_base")),
|
||||
identity=RequestIdentity.from_payload(payload, request_metadata=request_metadata),
|
||||
identity=RequestIdentity.from_payload(payload, named_user_id=named_user_id),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -301,7 +293,7 @@ class LLMCallEvent:
|
|||
completion_start_seconds: float | None
|
||||
trace: TraceControls
|
||||
session_id: str | None
|
||||
user_id_metadata: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
|
||||
caller_user_id: str | None = None
|
||||
|
||||
@property
|
||||
def reached_a_provider(self) -> bool:
|
||||
|
|
@ -331,19 +323,64 @@ class LLMCallEvent:
|
|||
completion_start_seconds=_epoch_seconds(kwargs.get("completion_start_time")),
|
||||
trace=trace,
|
||||
session_id=caller_session_id(kwargs, trace),
|
||||
user_id_metadata=caller_user_id_metadata(kwargs),
|
||||
caller_user_id=caller_user_id_in_call(kwargs),
|
||||
)
|
||||
|
||||
|
||||
def caller_user_id_metadata(kwargs: Mapping[str, object]) -> Mapping[str, str]:
|
||||
params: Final = as_str_mapping(kwargs.get("litellm_params"))
|
||||
if params is None:
|
||||
return MappingProxyType({})
|
||||
bodies: Final = (
|
||||
*(flatten_metadata(meta) if (meta := as_str_mapping(params.get("metadata"))) is not None else ()),
|
||||
*(flatten_metadata(lm_meta) if (lm_meta := as_str_mapping(params.get("litellm_metadata"))) is not None else ()),
|
||||
def _as_mapping(value: object) -> Mapping[str, object]:
|
||||
return (
|
||||
cast(Mapping[str, object], value) # cast-ok: isinstance cannot type the mapping's items
|
||||
if isinstance(value, Mapping)
|
||||
else _NO_METADATA
|
||||
)
|
||||
return MappingProxyType({key: value for key, value in bodies if key in _USER_ID_METADATA_KEYS})
|
||||
|
||||
|
||||
def caller_user_id_in_call(kwargs: Mapping[str, object]) -> str | None:
|
||||
params: Final = _as_mapping(kwargs.get("litellm_params"))
|
||||
return caller_user_id(params.get("metadata"), params.get("litellm_metadata"))
|
||||
|
||||
|
||||
def usable_user_id(value: object) -> str | None:
|
||||
"""``value`` as a ``user.id``, or ``None`` when it could not name a user."""
|
||||
if isinstance(value, bool) or not isinstance(value, (str, int)):
|
||||
return None
|
||||
text: Final = str(value).strip()
|
||||
return text if 0 < len(text) <= _MAX_USER_ID_LENGTH and text.isprintable() else None
|
||||
|
||||
|
||||
def _user_id_field(metadata: Mapping[str, object], field_name: str) -> str | None:
|
||||
named: Final = usable_user_id(metadata.get(field_name))
|
||||
if named is not None or field_name != "user.id":
|
||||
return named
|
||||
return usable_user_id(_as_mapping(metadata.get("user")).get("id"))
|
||||
|
||||
|
||||
def _named_user_ids(scopes: tuple[Mapping[str, object], ...]) -> Iterator[str]:
|
||||
for field_name in _USER_ID_METADATA_FIELDS:
|
||||
for scope in scopes:
|
||||
named = _user_id_field(scope, field_name)
|
||||
if named is not None:
|
||||
yield named
|
||||
|
||||
|
||||
def _caller_scopes(metadata: Mapping[str, object]) -> tuple[Mapping[str, object], Mapping[str, object]]:
|
||||
return (_as_mapping(metadata.get(REQUESTER_METADATA_KEY)), metadata)
|
||||
|
||||
|
||||
def caller_user_id(*sources: object) -> str | None:
|
||||
"""The id a caller named in its metadata, by key lookup so the cost never grows with the metadata's size."""
|
||||
scopes: Final = tuple(chain.from_iterable(_caller_scopes(_as_mapping(source)) for source in sources))
|
||||
return next(_named_user_ids(scopes), None)
|
||||
|
||||
|
||||
def request_user_id(end_user_id: object, body: Mapping[str, object], inbound_user_id: str | None = None) -> str | None:
|
||||
"""A request's ``user.id`` at the auth boundary, which reads the metadata as the caller sent it, JSON-encoded on a
|
||||
multipart or ``extra_body`` request. ``inbound_user_id`` is the id in the request's W3C Baggage."""
|
||||
metadata_sources: Final = tuple(
|
||||
safe_json_loads(value) if isinstance(value, str) else value
|
||||
for value in (body.get("metadata"), body.get("litellm_metadata"))
|
||||
)
|
||||
return caller_user_id(*metadata_sources) or inbound_user_id or usable_user_id(end_user_id)
|
||||
|
||||
|
||||
def _epoch_seconds(value: object) -> float | None:
|
||||
|
|
|
|||
|
|
@ -464,7 +464,7 @@ class LLMCallSpanData:
|
|||
session_id: str | None = None,
|
||||
*,
|
||||
metadata_keys: tuple[str, ...] = (),
|
||||
request_metadata: Mapping[str, str] = MappingProxyType({}),
|
||||
named_user_id: str | None = None,
|
||||
operation: GenAIOperation | None = None,
|
||||
) -> LLMCallSpanData:
|
||||
params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {})
|
||||
|
|
@ -472,7 +472,7 @@ class LLMCallSpanData:
|
|||
# The single parse of the request's metadata — the request-vs-provider
|
||||
# model split, the response model, api base, and identity all come from
|
||||
# here rather than being re-derived from the raw payload dicts.
|
||||
context: Final = RequestContext.from_standard_logging_payload(payload, request_metadata=request_metadata)
|
||||
context: Final = RequestContext.from_standard_logging_payload(payload, named_user_id=named_user_id)
|
||||
metadata: Final = payload.get("metadata")
|
||||
decision: Final = metadata["routing_decision"] if metadata and "routing_decision" in metadata else None
|
||||
promoted_metadata: Final = allowlisted_metadata(context.identity.metadata, metadata_keys)
|
||||
|
|
|
|||
|
|
@ -23,7 +23,8 @@ from opentelemetry.trace.propagation.tracecontext import (
|
|||
)
|
||||
|
||||
from litellm._internal_context import in_post_response_phase
|
||||
from litellm.integrations.otel.model.semconv import HTTP
|
||||
from litellm.integrations.otel.model.metadata import usable_user_id
|
||||
from litellm.integrations.otel.model.semconv import HTTP, User
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
|
|
@ -473,6 +474,25 @@ def reset_request_destinations(token: "Token[tuple[OtelDestination, ...]]") -> N
|
|||
_request_destinations.reset(token)
|
||||
|
||||
|
||||
_request_user_id: Final[ContextVar[str | None]] = ContextVar("litellm_otel_request_user_id", default=None)
|
||||
|
||||
|
||||
def set_request_user_id(user_id: str | None) -> None:
|
||||
"""Name the caller, inside the ``auth`` span so that span carries the id too. Each span reads the id as it ends,
|
||||
on the task that ended it, so a request's id never labels another request's span."""
|
||||
_request_user_id.set(user_id)
|
||||
|
||||
|
||||
def request_user_id() -> str | None:
|
||||
return _request_user_id.get()
|
||||
|
||||
|
||||
def inbound_user_id() -> str | None:
|
||||
"""The ``user.id`` in the W3C Baggage around this call, which an unauthenticated caller can write too, so it is
|
||||
read only once the request has passed auth."""
|
||||
return usable_user_id(baggage.get_baggage(User.ID))
|
||||
|
||||
|
||||
def request_destinations() -> 'tuple["OtelDestination", ...]':
|
||||
"""The destinations resolved for this request, empty outside a proxy request."""
|
||||
return _request_destinations.get()
|
||||
|
|
|
|||
|
|
@ -67,11 +67,13 @@ from litellm.integrations.otel.model.semconv import (
|
|||
LiteLLM,
|
||||
LiteLLMError,
|
||||
Server,
|
||||
User,
|
||||
)
|
||||
from litellm.integrations.otel.model.spans import LiteLLMSpanKind
|
||||
from litellm.integrations.otel.plumbing.context import (
|
||||
request_destinations,
|
||||
request_traces_excluded,
|
||||
request_user_id,
|
||||
suppressed_backends,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls
|
||||
|
|
@ -486,25 +488,45 @@ class _CountedSpanView(_SpanView):
|
|||
return self._dropped_counts[2]
|
||||
|
||||
|
||||
_LLM_CALL_ATTRIBUTES_BUT_USER_ID: Final = LangfuseMapper.LLM_CALL_ATTRIBUTES - {User.ID}
|
||||
|
||||
|
||||
def _is_litellm_span(span: ReadableSpan) -> bool:
|
||||
"""Whether LiteLLM's own tracer wrote ``span``, rather than application or instrumentation code."""
|
||||
scope: Final = span.instrumentation_scope
|
||||
return scope is not None and scope.name == LITELLM_TRACER_NAME
|
||||
|
||||
|
||||
def without_destination_langfuse_attributes(span: ReadableSpan, langfuse_mapped: bool) -> ReadableSpan:
|
||||
def with_request_user_id(span: ReadableSpan) -> ReadableSpan:
|
||||
"""``span`` carrying the id the request's caller named, for a backend that drops a span without one.
|
||||
|
||||
Called as the span ends, on the request's own task. A held span must already have
|
||||
been through here, since a tree can be released on another request's task.
|
||||
"""
|
||||
user_id: Final = request_user_id()
|
||||
attributes: Final = span.attributes or _NO_ATTRIBUTES
|
||||
if user_id is None or User.ID in attributes:
|
||||
return span
|
||||
return _CountedSpanView(span, MappingProxyType({**attributes, User.ID: user_id}))
|
||||
|
||||
|
||||
def without_destination_langfuse_attributes(
|
||||
span: ReadableSpan, langfuse_mapped: bool, keeps_user_id: bool = False
|
||||
) -> ReadableSpan:
|
||||
"""``span`` without the Langfuse attributes only a Langfuse destination asked for.
|
||||
|
||||
That is the request metadata the Langfuse mapper writes and, on a model call written by
|
||||
a logger whose own mappers leave the Langfuse vocabulary out (``langfuse_mapped`` off),
|
||||
every Langfuse attribute, since that logger wrote them for a key's or team's Langfuse
|
||||
destination alone. Only LiteLLM's own spans lose them: any other span keeps every key
|
||||
its own code set.
|
||||
its own code set. ``keeps_user_id`` spares ``user.id``, which Mixpanel reads too.
|
||||
"""
|
||||
attributes: Final = span.attributes or _NO_ATTRIBUTES
|
||||
dropped: Final = (
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES
|
||||
if langfuse_mapped or not is_llm_call_span(span)
|
||||
else _LLM_CALL_ATTRIBUTES_BUT_USER_ID
|
||||
if keeps_user_id
|
||||
else LangfuseMapper.LLM_CALL_ATTRIBUTES
|
||||
)
|
||||
if not (_is_litellm_span(span) and any(key in dropped for key in attributes)):
|
||||
|
|
@ -690,7 +712,9 @@ def _for_destination(
|
|||
source: Final = (
|
||||
span
|
||||
if destination.callback_name == ExporterOwner.LANGFUSE_OTEL
|
||||
else without_destination_langfuse_attributes(span, langfuse_mapped)
|
||||
else without_destination_langfuse_attributes(
|
||||
span, langfuse_mapped, keeps_user_id=destination.callback_name == ExporterOwner.MIXPANEL
|
||||
)
|
||||
)
|
||||
attributes: Final = source.attributes or _NO_ATTRIBUTES
|
||||
database: Final = _is_database_span(attributes)
|
||||
|
|
@ -794,10 +818,11 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
)
|
||||
):
|
||||
continue
|
||||
view = with_request_user_id(span) if destination.callback_name == ExporterOwner.MIXPANEL else span
|
||||
if not _is_sampled(destination):
|
||||
self._forward(span, destination)
|
||||
self._forward(view, destination)
|
||||
continue
|
||||
for held_span, held_destination in self._route(request, span, destination, failed):
|
||||
for held_span, held_destination in self._route(request, view, destination, failed):
|
||||
self._forward(held_span, held_destination)
|
||||
for held_span, held_destination in self._settle(request, span, failed):
|
||||
self._forward(held_span, held_destination)
|
||||
|
|
@ -1217,15 +1242,37 @@ class _DestinationLangfuseAttributeFilter(SpanProcessor):
|
|||
"""Strip the Langfuse attributes only a Langfuse destination asked for from LiteLLM's own spans before ``inner``
|
||||
exports them."""
|
||||
|
||||
def __init__(self, inner: SpanProcessor, langfuse_mapped: bool) -> None:
|
||||
def __init__(self, inner: SpanProcessor, langfuse_mapped: bool, keeps_user_id: bool = False) -> None:
|
||||
self._inner: Final = inner
|
||||
self._langfuse_mapped: Final = langfuse_mapped
|
||||
self._keeps_user_id: Final = keeps_user_id
|
||||
|
||||
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
|
||||
self._inner.on_start(span, parent_context)
|
||||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
self._inner.on_end(without_destination_langfuse_attributes(span, self._langfuse_mapped))
|
||||
self._inner.on_end(
|
||||
without_destination_langfuse_attributes(span, self._langfuse_mapped, keeps_user_id=self._keeps_user_id)
|
||||
)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._inner.shutdown()
|
||||
|
||||
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
||||
return self._inner.force_flush(timeout_millis)
|
||||
|
||||
|
||||
class _RequestUserIdStamp(SpanProcessor):
|
||||
"""Give every span the id the request's caller named before ``inner`` exports it."""
|
||||
|
||||
def __init__(self, inner: SpanProcessor) -> None:
|
||||
self._inner: Final = inner
|
||||
|
||||
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
|
||||
self._inner.on_start(span, parent_context)
|
||||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
self._inner.on_end(with_request_user_id(span))
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._inner.shutdown()
|
||||
|
|
@ -1515,10 +1562,12 @@ def _spec_processor(
|
|||
_exporter_from_spec(spec),
|
||||
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
|
||||
)
|
||||
labels_user_id: Final = spec.owner is ExporterOwner.MIXPANEL
|
||||
labeled: Final = _RequestUserIdStamp(exporting) if labels_user_id else exporting
|
||||
scoped: Final = (
|
||||
_DestinationLangfuseAttributeFilter(exporting, langfuse_mapped)
|
||||
_DestinationLangfuseAttributeFilter(labeled, langfuse_mapped, keeps_user_id=labels_user_id)
|
||||
if strip_destination_langfuse_attributes
|
||||
else exporting
|
||||
else labeled
|
||||
)
|
||||
return _MessageContentFilter(scoped) if strip_message_content else scoped
|
||||
|
||||
|
|
|
|||
|
|
@ -12,8 +12,7 @@ from litellm.integrations.otel.model.config import (
|
|||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import User
|
||||
from litellm.integrations.otel.presets.utils import ensure_mappers
|
||||
from litellm.integrations.otel.presets.utils import credential_gated_exporters, ensure_mappers
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
MIXPANEL_PROJECT_TOKEN_ENV: Final = "MIXPANEL_PROJECT_TOKEN"
|
||||
|
|
@ -66,19 +65,25 @@ def mixpanel_preset(
|
|||
base: Final = config_overrides or OpenTelemetryV2Config()
|
||||
token: Final = (settings.project_token or "").strip() or None
|
||||
endpoint: Final = (settings.endpoint or "").strip() or _resolve_region_endpoint(settings.region)
|
||||
spec: Final = ExporterSpec(
|
||||
kind="otlp_http",
|
||||
endpoint=endpoint,
|
||||
headers=(f"Authorization={mixpanel_authorization_header(token)}" if token else None),
|
||||
owner=ExporterOwner.MIXPANEL,
|
||||
requires_headers=True,
|
||||
exporters: Final = (
|
||||
(
|
||||
*base.exporters,
|
||||
ExporterSpec(
|
||||
kind="otlp_http",
|
||||
endpoint=endpoint,
|
||||
headers=f"Authorization={mixpanel_authorization_header(token)}",
|
||||
owner=ExporterOwner.MIXPANEL,
|
||||
requires_headers=True,
|
||||
),
|
||||
)
|
||||
if token
|
||||
else credential_gated_exporters(base.exporters, ExporterOwner.MIXPANEL)
|
||||
)
|
||||
return base.model_copy(
|
||||
update=MappingProxyType(
|
||||
{
|
||||
"exporters": (*base.exporters, spec),
|
||||
"exporters": exporters,
|
||||
"mapper_names": ensure_mappers(base.mapper_names, "genai", "mixpanel"),
|
||||
"baggage_promoted_keys": ensure_mappers(base.baggage_promoted_keys, User.ID),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
|
@ -93,5 +98,8 @@ def mixpanel_dynamic_headers(
|
|||
return {"Authorization": mixpanel_authorization_header(token)}
|
||||
|
||||
|
||||
def mixpanel_dynamic_endpoint(params: StandardCallbackDynamicParams) -> str:
|
||||
def mixpanel_dynamic_endpoint(params: StandardCallbackDynamicParams) -> str | None:
|
||||
"""The region endpoint of the project ``params`` name, or ``None`` to keep the operator's own endpoint."""
|
||||
if not (params.get("mixpanel_project_token") or "").strip():
|
||||
return None
|
||||
return _resolve_region_endpoint(params.get("mixpanel_region"))
|
||||
|
|
|
|||
|
|
@ -6188,16 +6188,11 @@ def build_custom_logger_compatible_class(
|
|||
|
||||
elif logging_integration == "mixpanel":
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.integrations.otel.presets.mixpanel import MIXPANEL_PROJECT_TOKEN_ENV
|
||||
|
||||
if not is_otel_v2_enabled():
|
||||
raise ValueError("The 'mixpanel' callback requires OpenTelemetry v2 (set LITELLM_OTEL_V2=true)")
|
||||
|
||||
_mixpanel_v2: Final = _maybe_construct_otel_v2("mixpanel", _in_memory_loggers)
|
||||
if _mixpanel_v2 is not None:
|
||||
return _mixpanel_v2
|
||||
|
||||
raise ValueError(f"{MIXPANEL_PROJECT_TOKEN_ENV} not found in environment variables")
|
||||
return _maybe_construct_otel_v2("mixpanel", _in_memory_loggers)
|
||||
|
||||
elif logging_integration == "mlflow":
|
||||
for callback in _in_memory_loggers:
|
||||
|
|
|
|||
|
|
@ -1920,11 +1920,9 @@ def get_end_user_id_from_request_body(
|
|||
if user_id_str:
|
||||
return user_id_str
|
||||
|
||||
# Check 5: 'metadata.user_id' or 'metadata.user.id' in request_body (another common pattern)
|
||||
# Check 5: 'metadata.user_id' in request_body (another common pattern)
|
||||
metadata_dict: Final = _as_dict(request_body.get("metadata"))
|
||||
user_id_str = _coerce_user_id_to_str(metadata_dict.get("user_id")) or _coerce_user_id_to_str(
|
||||
metadata_dict.get("user.id") # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType] # _as_dict returns untyped dict
|
||||
)
|
||||
user_id_str = _coerce_user_id_to_str(metadata_dict.get("user_id"))
|
||||
if user_id_str:
|
||||
return user_id_str
|
||||
|
||||
|
|
|
|||
|
|
@ -4311,6 +4311,31 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque
|
|||
verbose_proxy_logger.debug("OTel V2: tenant destination resolution failed: %s", exc)
|
||||
|
||||
|
||||
def _seed_request_user_id(
|
||||
user_api_key_dict: UserAPIKeyAuth, request_data: Mapping[str, object], request: Request
|
||||
) -> None:
|
||||
"""Name the caller for the spans a backend that drops an unlabeled span receives.
|
||||
|
||||
Called inside the ``auth`` phase span so that span carries the id too. Best-effort:
|
||||
trace labels must never fail auth.
|
||||
"""
|
||||
try:
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.integrations.otel.model.metadata import request_user_id
|
||||
from litellm.integrations.otel.plumbing.context import inbound_user_id, set_request_user_id
|
||||
|
||||
if is_otel_v2_enabled():
|
||||
headers: Final = cast( # cast-ok: the helper returns an unparameterized dict
|
||||
Mapping[str, object], safe_get_request_headers(request)
|
||||
)
|
||||
end_user_id: Final = user_api_key_dict.end_user_id or get_end_user_id_from_request_body(
|
||||
request_data, headers
|
||||
)
|
||||
set_request_user_id(request_user_id(end_user_id, request_data, inbound_user_id()))
|
||||
except Exception as exc: # noqa: BLE001 # telemetry labels are best-effort and must never break authentication
|
||||
verbose_proxy_logger.debug("OTel V2: request user id resolution failed: %s", exc)
|
||||
|
||||
|
||||
@tracer.wrap()
|
||||
async def user_api_key_auth(
|
||||
request: Request,
|
||||
|
|
@ -4368,6 +4393,11 @@ async def user_api_key_auth(
|
|||
safe_get_request_headers(request), user_api_key_auth_obj
|
||||
)
|
||||
_seed_request_destinations(user_api_key_auth_obj, request)
|
||||
_seed_request_user_id(
|
||||
user_api_key_auth_obj,
|
||||
cast(Mapping[str, object], request_data), # cast-ok: request bodies are dynamically typed
|
||||
request,
|
||||
)
|
||||
|
||||
# A body that never parsed is authenticated (so the trace carries identity
|
||||
# and this ``auth`` span) but not authorized: there is no model to check it
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ per-integration checks here.
|
|||
"""
|
||||
|
||||
import math
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings
|
||||
|
||||
|
|
@ -16,6 +16,10 @@ _NEWRELIC_CALLBACK: Final = "newrelic"
|
|||
_NEWRELIC_VAR_PREFIX: Final = "newrelic_"
|
||||
_MIXPANEL_CALLBACK: Final = "mixpanel"
|
||||
_MIXPANEL_VAR_PREFIX: Final = "mixpanel_"
|
||||
_MIXPANEL_VARS: Final = frozenset({"mixpanel_project_token", "mixpanel_region"})
|
||||
CALLBACK_VARS_REDACTED: Final = "***REDACTED***"
|
||||
_UNSENDABLE_MIXPANEL_TOKENS: Final = frozenset({"None", CALLBACK_VARS_REDACTED})
|
||||
_MIXPANEL_NEEDS_OTEL_V2: Final = "Per-team Mixpanel routing requires the proxy to run with LITELLM_OTEL_V2=true."
|
||||
_LANGFUSE_OTEL_CALLBACK: Final = "langfuse_otel"
|
||||
_LANGFUSE_SPAN_SCOPE_VAR: Final = "langfuse_span_scope"
|
||||
_ARIZE_CALLBACK: Final = "arize"
|
||||
|
|
@ -28,6 +32,9 @@ _ARIZE_SAMPLING_RATE_VARS: Final[frozenset[str]] = frozenset(
|
|||
def callback_config_error(
|
||||
callback_name: str | None, callback_vars: Mapping[str, str] | None, callback_type: str | None = None
|
||||
) -> str | None:
|
||||
mixpanel_v2_error: Final = _mixpanel_v2_error(callback_name)
|
||||
if mixpanel_v2_error is not None:
|
||||
return mixpanel_v2_error
|
||||
if not callback_vars:
|
||||
return None
|
||||
arize_error: Final = _arize_sampling_rate_error(callback_name, callback_vars) or _arize_otlp_protocol_error(
|
||||
|
|
@ -151,7 +158,6 @@ _VAR_FAMILIES: Final[Mapping[str, str]] = MappingProxyType(
|
|||
"mixpanel_": "Mixpanel",
|
||||
"newrelic_": "New Relic",
|
||||
"posthog_": "PostHog",
|
||||
"signoz_": "SigNoz",
|
||||
"wandb_": "Weights & Biases",
|
||||
"weave_": "Weights & Biases",
|
||||
}
|
||||
|
|
@ -312,6 +318,29 @@ def stored_callback_entries(logging_entries: object) -> tuple[tuple[str | None,
|
|||
return tuple((_entry_callback_name(entry), _entry_callback_vars(entry)) for entry in entries)
|
||||
|
||||
|
||||
def _mixpanel_callback_vars(metadata: Mapping[str, object]) -> Iterator[Mapping[str, str]]:
|
||||
yield from (
|
||||
_entry_callback_vars(entry)
|
||||
for entry in _object_sequence(metadata.get("logging")) or ()
|
||||
if _entry_callback_name(entry) == _MIXPANEL_CALLBACK
|
||||
)
|
||||
settings_vars: Final = _entry_callback_vars(metadata.get("callback_settings"))
|
||||
if any(key.startswith(_MIXPANEL_VAR_PREFIX) for key in settings_vars):
|
||||
yield settings_vars
|
||||
|
||||
|
||||
def mixpanel_logging_metadata_error(metadata: Mapping[str, object] | None) -> str | None:
|
||||
"""Validate only the Mixpanel settings of a team's metadata, for writers that leave other callbacks unchecked."""
|
||||
return next(
|
||||
(
|
||||
error
|
||||
for callback_vars in _mixpanel_callback_vars(metadata or MappingProxyType({}))
|
||||
if (error := callback_config_error(_MIXPANEL_CALLBACK, callback_vars)) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _object_sequence(value: object) -> Sequence[object] | None:
|
||||
return value if isinstance(value, Sequence) and not isinstance(value, (str, bytes)) else None
|
||||
|
||||
|
|
@ -453,16 +482,21 @@ def _newrelic_config_error(callback_vars: Mapping[str, str]) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def _mixpanel_v2_error(callback_name: str | None) -> str | None:
|
||||
if callback_name != _MIXPANEL_CALLBACK:
|
||||
return None
|
||||
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
|
||||
return None if is_otel_v2_enabled() else _MIXPANEL_NEEDS_OTEL_V2
|
||||
|
||||
|
||||
def _mixpanel_config_error(callback_vars: Mapping[str, str]) -> str | None:
|
||||
if not any(key.startswith(_MIXPANEL_VAR_PREFIX) for key in callback_vars):
|
||||
return None
|
||||
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.integrations.otel.presets.mixpanel import MIXPANEL_OTLP_ENDPOINT_BY_REGION
|
||||
|
||||
if not is_otel_v2_enabled():
|
||||
return "Per-team Mixpanel routing requires the proxy to run with LITELLM_OTEL_V2=true."
|
||||
|
||||
region: Final = callback_vars.get("mixpanel_region")
|
||||
if region is not None and region.strip().lower() not in MIXPANEL_OTLP_ENDPOINT_BY_REGION:
|
||||
return (
|
||||
|
|
@ -471,6 +505,9 @@ def _mixpanel_config_error(callback_vars: Mapping[str, str]) -> str | None:
|
|||
)
|
||||
|
||||
token: Final = callback_vars.get("mixpanel_project_token")
|
||||
if region is not None and (not token or not token.strip() or token == "None"):
|
||||
return "mixpanel_region requires mixpanel_project_token; the region rides the team's own project token."
|
||||
return None
|
||||
if not token or not token.strip() or token in _UNSENDABLE_MIXPANEL_TOKENS:
|
||||
return "mixpanel_project_token must be the project's own token, sent in full on every create or update."
|
||||
unknown: Final = tuple(
|
||||
key for key in callback_vars if key.startswith(_MIXPANEL_VAR_PREFIX) and key not in _MIXPANEL_VARS
|
||||
)
|
||||
return f"Unknown Mixpanel callback variable {unknown[0]!r}." if unknown else None
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.callback_config_validation import (
|
||||
CALLBACK_VARS_REDACTED,
|
||||
callback_config_error,
|
||||
conflicting_callback_var_error,
|
||||
conflicting_span_scope_error,
|
||||
|
|
@ -65,7 +66,7 @@ from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_V
|
|||
router: Final = APIRouter()
|
||||
|
||||
|
||||
_CALLBACK_VARS_REDACTED: Final = "***REDACTED***"
|
||||
_CALLBACK_VARS_REDACTED: Final = CALLBACK_VARS_REDACTED
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -131,6 +131,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
|||
evict_shared,
|
||||
publish_auth_cache_invalidation,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_config_validation import mixpanel_logging_metadata_error
|
||||
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
|
||||
from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
|
||||
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
|
||||
|
|
@ -1709,6 +1710,9 @@ async def new_team(
|
|||
```
|
||||
"""
|
||||
try:
|
||||
_raise_on_unsendable_mixpanel_callback(
|
||||
data.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
|
|
@ -2135,6 +2139,12 @@ async def _create_team_update_audit_log(
|
|||
)
|
||||
|
||||
|
||||
def _raise_on_unsendable_mixpanel_callback(metadata: Mapping[str, object] | None) -> None:
|
||||
error: Final = mixpanel_logging_metadata_error(metadata)
|
||||
if error is not None:
|
||||
raise HTTPException(status_code=400, detail={"error": error})
|
||||
|
||||
|
||||
async def _update_model_table(
|
||||
data: UpdateTeamRequest,
|
||||
model_id: int | None,
|
||||
|
|
@ -2483,6 +2493,9 @@ async def update_team(
|
|||
```
|
||||
"""
|
||||
try:
|
||||
_raise_on_unsendable_mixpanel_callback(
|
||||
data.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import is_audit_logging_enabled
|
||||
from litellm.proxy.proxy_server import (
|
||||
litellm_proxy_admin_name,
|
||||
|
|
|
|||
|
|
@ -250,7 +250,13 @@ class _ConnectHandler(_Handler):
|
|||
type(self)(wrapped, self.client_address, self.server)
|
||||
|
||||
|
||||
_MITM_HOSTS: Final = ("otlp.nr-data.net", "otlp.eu01.nr-data.net")
|
||||
_MITM_HOSTS: Final = (
|
||||
"otlp.nr-data.net",
|
||||
"otlp.eu01.nr-data.net",
|
||||
"ingestion-us.mixpanel.com",
|
||||
"ingestion-eu.mixpanel.com",
|
||||
"ingestion-in.mixpanel.com",
|
||||
)
|
||||
|
||||
|
||||
def _mitm_context(directory: Path) -> ssl.SSLContext:
|
||||
|
|
|
|||
995
tests/integration/observability/test_mixpanel_user_id.py
Normal file
995
tests/integration/observability/test_mixpanel_user_id.py
Normal file
|
|
@ -0,0 +1,995 @@
|
|||
"""Mixpanel joins agent traces to user events through ``user.id`` and drops a span without one, so every span a request
|
||||
produces once auth has named its caller reaches the project with that id, whichever callback setup the operator runs,
|
||||
and nothing else gets the id. Database and cache spans that close during auth, before the id is known, carry none."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, cast
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from _openinference_support import _spans as _wire_span_attributes
|
||||
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, held
|
||||
from integration._support.mcp import INITIALIZE, EntryPoint, McpCaller, echo_tool, register_mcp, scripted_peer
|
||||
from integration._support.otlp_sink import (
|
||||
AuditConfigWriter,
|
||||
ConnectSink,
|
||||
Span,
|
||||
SpanSinks,
|
||||
owned_connect_sink,
|
||||
recorded_requests,
|
||||
recorded_spans,
|
||||
spans_for_trace,
|
||||
)
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
OPERATION: Final = "gen_ai.operation.name"
|
||||
USER_ID: Final = "user.id"
|
||||
TEAM_TOKEN: Final = "team-project-token"
|
||||
KEY_TOKEN: Final = "key-project-token"
|
||||
OPERATOR_TOKEN: Final = "operator-project-token"
|
||||
OPERATOR_HOST: Final = "ingestion-us.mixpanel.com"
|
||||
PROVIDER_FAILS: Final = "provider-fails"
|
||||
CUSTOM_AUTH: Final = """
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
|
||||
async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
"""
|
||||
CHAT_COMPLETION: Final = {
|
||||
"id": "chatcmpl-mixpanel",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
CHAT_CHUNKS: Final = (
|
||||
b'data: {"id":"chatcmpl-mixpanel","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
|
||||
b'"choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}\n\n',
|
||||
b'data: {"id":"chatcmpl-mixpanel","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
|
||||
b'"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
)
|
||||
RESPONSE: Final = {
|
||||
"id": "resp_mixpanel",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_mixpanel",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
|
||||
|
||||
def _provider_reply(request: Request) -> Reply:
|
||||
if PROVIDER_FAILS.encode() in request.body:
|
||||
return Reply(status=400, body=json.dumps({"error": {"message": "rejected by the provider"}}).encode())
|
||||
if request.target.endswith("/responses"):
|
||||
return Reply(body=json.dumps(RESPONSE).encode())
|
||||
if json.loads(request.body).get("stream"):
|
||||
return Reply(content_type="text/event-stream", chunks=CHAT_CHUNKS)
|
||||
return Reply(body=json.dumps(CHAT_COMPLETION).encode())
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Call:
|
||||
path: str
|
||||
body: Callable[[str], dict[str, JsonValue]]
|
||||
|
||||
|
||||
def _chat_call(content: str) -> dict[str, JsonValue]:
|
||||
return {"messages": [{"role": "user", "content": content}]}
|
||||
|
||||
|
||||
CALLS: Final[Mapping[str, Call]] = {
|
||||
"chat": Call("/v1/chat/completions", _chat_call),
|
||||
"chat-stream": Call("/v1/chat/completions", lambda content: {**_chat_call(content), "stream": True}),
|
||||
"responses": Call("/v1/responses", lambda content: {"input": content}),
|
||||
"messages": Call("/v1/messages", lambda content: {**_chat_call(content), "max_tokens": 20}),
|
||||
}
|
||||
IDENTIFICATIONS: Final[Mapping[str, Callable[[str], dict[str, JsonValue]]]] = {
|
||||
"user": lambda caller: {"user": caller},
|
||||
"metadata-user-dot-id": lambda caller: {"metadata": {"user.id": caller}},
|
||||
"metadata-user-id": lambda caller: {"metadata": {"user_id": caller}},
|
||||
"metadata-trace-user-id": lambda caller: {"metadata": {"trace_user_id": caller}},
|
||||
"litellm-metadata-user-dot-id": lambda caller: {"litellm_metadata": {"user.id": caller}},
|
||||
"metadata-json-user-dot-id": lambda caller: {"metadata": json.dumps({"user.id": caller})},
|
||||
}
|
||||
|
||||
|
||||
def _authorization(token: str) -> str:
|
||||
return f"Basic {base64.b64encode(f'{token}:'.encode()).decode()}"
|
||||
|
||||
|
||||
def _team_or_key(scenario: Scenario, level: str, **access: JsonValue) -> tuple[str, str, str]:
|
||||
"""A key with a Mixpanel project of its own, or under a team with one, plus the host and token that project uses."""
|
||||
token, region, host = (
|
||||
(TEAM_TOKEN, "eu", "ingestion-eu.mixpanel.com")
|
||||
if level == "team"
|
||||
else (KEY_TOKEN, "in", "ingestion-in.mixpanel.com")
|
||||
)
|
||||
logging: Final[list[JsonValue]] = [
|
||||
{
|
||||
"callback_name": "mixpanel",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"mixpanel_project_token": token, "mixpanel_region": region},
|
||||
}
|
||||
]
|
||||
alias: Final = f"mixpanel-{uuid.uuid4().hex}"
|
||||
key: Final = (
|
||||
scenario.key(key_alias=alias, team_id=scenario.team(metadata={"logging": logging}, **access), **access)
|
||||
if level == "team"
|
||||
else scenario.key(key_alias=alias, metadata={"logging": logging}, **access)
|
||||
)
|
||||
return key, host, token
|
||||
|
||||
|
||||
def _call_on_a_new_connection(gateway: Gateway, key: str, body: Mapping[str, JsonValue]) -> str:
|
||||
"""A worker takes each connection, and only a new connection can land on one that has not served this key yet."""
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
with httpx.Client(base_url=gateway.client.base_url, trust_env=False, timeout=60) as client:
|
||||
response: Final = client.post(
|
||||
"/v1/chat/completions",
|
||||
json=body,
|
||||
headers={"Authorization": f"Bearer {key}", "traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return trace_id
|
||||
|
||||
|
||||
def _traced_call(gateway: Gateway, key: str, body: Mapping[str, JsonValue], path: str = "/v1/chat/completions") -> str:
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
response: Final = gateway.request(
|
||||
"POST", path, body, key=key, headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return trace_id
|
||||
|
||||
|
||||
def _request_finished(spans: tuple[Span, ...]) -> bool:
|
||||
return any(span["kind"] == 2 for span in spans) and any(OPERATION in span["attributes"] for span in spans)
|
||||
|
||||
|
||||
def _tool_call_finished(spans: tuple[Span, ...]) -> bool:
|
||||
return any(span["kind"] == 2 for span in spans) and any(_is_tool_call(span) for span in spans)
|
||||
|
||||
|
||||
def _is_tool_call(span: Span) -> bool:
|
||||
return str(span["name"]).startswith("tools/call ")
|
||||
|
||||
|
||||
def _has_model_call(spans: tuple[Span, ...]) -> bool:
|
||||
return any(OPERATION in span["attributes"] for span in spans)
|
||||
|
||||
|
||||
def _trace(url: str, trace_id: str) -> tuple[Span, ...]:
|
||||
return spans_for_trace(recorded_spans(url)[1], trace_id)
|
||||
|
||||
|
||||
def _settled_trace(
|
||||
url: str, trace_id: str, finished: Callable[[tuple[Span, ...]], bool] = _request_finished
|
||||
) -> tuple[Span, ...]:
|
||||
held(lambda: _trace(url, trace_id), finished, holding=3, seconds=40)
|
||||
return _trace(url, trace_id)
|
||||
|
||||
|
||||
def _without_the_id(spans: tuple[Span, ...], caller: str) -> list[tuple[str, JsonValue]]:
|
||||
return [
|
||||
(str(span["name"]), span["attributes"].get(USER_ID))
|
||||
for span in spans
|
||||
if span["attributes"].get(USER_ID) != caller
|
||||
]
|
||||
|
||||
|
||||
def _outside_auth(spans: tuple[Span, ...]) -> tuple[Span, ...]:
|
||||
auth_spans: Final = {span["span_id"] for span in spans if str(span["name"]).startswith("auth ")}
|
||||
return tuple(span for span in spans if span["parent_span_id"] not in auth_spans)
|
||||
|
||||
|
||||
def _inside_the_mcp_session(spans: tuple[Span, ...]) -> tuple[Span, ...]:
|
||||
"""Auth on ``/mcp`` runs on a task inside the transport, so the transport's ``POST /mcp`` span and the server
|
||||
lookups that come before auth never see the id. The model trace in the same list has no such span to leave out."""
|
||||
return tuple(
|
||||
span
|
||||
for span in _outside_auth(spans)
|
||||
if str(span["name"]) != "POST /mcp" and not str(span["name"]).startswith("postgres.select LiteLLM_")
|
||||
)
|
||||
|
||||
|
||||
def _ingest_accepted(_request: Request) -> Reply:
|
||||
return Reply(body=b"", content_type="application/x-protobuf")
|
||||
|
||||
|
||||
def _reached(sink: ConnectSink, spans: tuple[Span, ...], host: str) -> tuple[Span, ...]:
|
||||
exports: Final = tuple(
|
||||
cast("list[str]", export.get("span_ids", []))
|
||||
for export in recorded_requests(sink.control_url)
|
||||
if str(export.get("host", "")).partition(":")[0] == host
|
||||
)
|
||||
exported: Final = frozenset(chain.from_iterable(exports))
|
||||
return tuple(span for span in spans if span["span_id"] in exported)
|
||||
|
||||
|
||||
def _exports_of(sink: ConnectSink, spans: tuple[Span, ...]) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
span_ids: Final = {span["span_id"] for span in spans}
|
||||
return tuple(
|
||||
request
|
||||
for request in recorded_requests(sink.control_url)
|
||||
if span_ids & set(cast("list[str]", request.get("span_ids", [])))
|
||||
)
|
||||
|
||||
|
||||
def _sent_to(exports: tuple[Mapping[str, JsonValue], ...]) -> set[tuple[str, str]]:
|
||||
return {
|
||||
(
|
||||
str(export["host"]).partition(":")[0],
|
||||
next(
|
||||
str(value)
|
||||
for name, value in cast("dict[str, JsonValue]", export["headers"]).items()
|
||||
if name.lower() == "authorization"
|
||||
),
|
||||
)
|
||||
for export in exports
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def provider() -> Iterator[Wire]:
|
||||
with wire_server(_provider_reply) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def mixpanel(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ConnectSink]:
|
||||
with owned_connect_sink(tmp_path_factory.mktemp("mixpanel-ingest")) as sink:
|
||||
yield sink
|
||||
|
||||
|
||||
def _with_general_settings(config: Path, settings: Mapping[str, JsonValue]) -> Path:
|
||||
configured: Final = yaml.safe_load(config.read_text())
|
||||
config.write_text(
|
||||
yaml.safe_dump({**configured, "general_settings": {**configured["general_settings"], **settings}})
|
||||
)
|
||||
return config
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _booted(
|
||||
directory: Path,
|
||||
config: Path,
|
||||
mixpanel: ConnectSink,
|
||||
operator_token: str | None,
|
||||
collector: str | None = None,
|
||||
environment: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> Iterator[OwnedProxy]:
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
directory,
|
||||
{
|
||||
"LITELLM_OTEL_V2": "1",
|
||||
"OTEL_BSP_SCHEDULE_DELAY": "300",
|
||||
"HTTPS_PROXY": mixpanel.proxy_url,
|
||||
"NO_PROXY": "127.0.0.1,localhost",
|
||||
"REQUESTS_CA_BUNDLE": mixpanel.ca_pem,
|
||||
**({"MIXPANEL_PROJECT_TOKEN": operator_token} if operator_token else {}),
|
||||
**({"OTEL_EXPORTER": "otlp_http", "OTEL_ENDPOINT": collector} if collector else {}),
|
||||
**environment,
|
||||
},
|
||||
config=config,
|
||||
remove_environment=(
|
||||
*(() if operator_token else ("MIXPANEL_PROJECT_TOKEN",)),
|
||||
*(() if "MIXPANEL_REGION" in environment else ("MIXPANEL_REGION",)),
|
||||
*(() if "MIXPANEL_ENDPOINT" in environment else ("MIXPANEL_ENDPOINT",)),
|
||||
*(() if collector else ("OTEL_EXPORTER", "OTEL_ENDPOINT")),
|
||||
),
|
||||
workers=2,
|
||||
) as owned,
|
||||
):
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
scope="module",
|
||||
params=[(["mixpanel"], None), ([], None), (["mixpanel"], OPERATOR_TOKEN)],
|
||||
ids=["mixpanel-listed-no-token", "not-listed", "mixpanel-listed-with-token"],
|
||||
)
|
||||
def tenant_setup(
|
||||
request: pytest.FixtureRequest,
|
||||
mixpanel: ConnectSink,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
tmp_path_factory: pytest.TempPathFactory,
|
||||
) -> Iterator[OwnedProxy]:
|
||||
directory: Final = tmp_path_factory.mktemp("mixpanel-tenant")
|
||||
callbacks, operator_token = request.param
|
||||
with _booted(
|
||||
directory,
|
||||
otel_audit_config(directory, {"callbacks": cast("list[JsonValue]", callbacks)}),
|
||||
mixpanel,
|
||||
operator_token,
|
||||
) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def operator_setup(
|
||||
mixpanel: ConnectSink, otel_audit_config: AuditConfigWriter, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[OwnedProxy]:
|
||||
directory: Final = tmp_path_factory.mktemp("mixpanel-operator")
|
||||
with _booted(
|
||||
directory, otel_audit_config(directory, {"callbacks": ["mixpanel"]}), mixpanel, OPERATOR_TOKEN
|
||||
) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def eu_operator_setup(
|
||||
mixpanel: ConnectSink, otel_audit_config: AuditConfigWriter, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[OwnedProxy]:
|
||||
directory: Final = tmp_path_factory.mktemp("mixpanel-eu-operator")
|
||||
with _booted(
|
||||
directory,
|
||||
_with_general_settings(
|
||||
otel_audit_config(directory, {"callbacks": ["mixpanel"]}), {"allow_client_side_credentials": True}
|
||||
),
|
||||
mixpanel,
|
||||
OPERATOR_TOKEN,
|
||||
environment={"MIXPANEL_REGION": "eu"},
|
||||
) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def collector_beside_operator_setup(
|
||||
mixpanel: ConnectSink,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
tmp_path_factory: pytest.TempPathFactory,
|
||||
) -> Iterator[OwnedProxy]:
|
||||
directory: Final = tmp_path_factory.mktemp("mixpanel-beside-collector")
|
||||
with _booted(
|
||||
directory,
|
||||
otel_audit_config(directory, {"callbacks": ["mixpanel"]}),
|
||||
mixpanel,
|
||||
OPERATOR_TOKEN,
|
||||
collector=audit_sinks.operator,
|
||||
) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def custom_auth_setup(
|
||||
mixpanel: ConnectSink, otel_audit_config: AuditConfigWriter, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[OwnedProxy]:
|
||||
directory: Final = tmp_path_factory.mktemp("mixpanel-custom-auth")
|
||||
(directory / "mixpanel_custom_auth.py").write_text(CUSTOM_AUTH)
|
||||
with _booted(
|
||||
directory,
|
||||
_with_general_settings(
|
||||
otel_audit_config(directory, {"callbacks": ["mixpanel"]}),
|
||||
{"custom_auth": "mixpanel_custom_auth.authenticate"},
|
||||
),
|
||||
mixpanel,
|
||||
OPERATOR_TOKEN,
|
||||
) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=[{}, {"otel_tenant_destination_mode": "additive"}], ids=["default", "additive"])
|
||||
def collector_setup(
|
||||
request: pytest.FixtureRequest,
|
||||
mixpanel: ConnectSink,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
tmp_path_factory: pytest.TempPathFactory,
|
||||
) -> Iterator[OwnedProxy]:
|
||||
directory: Final = tmp_path_factory.mktemp("mixpanel-collector")
|
||||
with _booted(
|
||||
directory,
|
||||
otel_audit_config(directory, {"callbacks": ["otel"], **request.param}),
|
||||
mixpanel,
|
||||
None,
|
||||
collector=audit_sinks.operator,
|
||||
) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def additive_setup(
|
||||
mixpanel: ConnectSink, otel_audit_config: AuditConfigWriter, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[OwnedProxy]:
|
||||
directory: Final = tmp_path_factory.mktemp("mixpanel-additive")
|
||||
with _booted(
|
||||
directory,
|
||||
otel_audit_config(directory, {"callbacks": ["mixpanel"], "otel_tenant_destination_mode": "additive"}),
|
||||
mixpanel,
|
||||
OPERATOR_TOKEN,
|
||||
) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
scope="module",
|
||||
params=[("mixpanel", "otel"), ("otel", "mixpanel"), ("mixpanel", "langfuse_otel"), ("langfuse_otel", "mixpanel")],
|
||||
ids="+".join,
|
||||
)
|
||||
def combined_setup(
|
||||
request: pytest.FixtureRequest,
|
||||
mixpanel: ConnectSink,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
tmp_path_factory: pytest.TempPathFactory,
|
||||
) -> Iterator[tuple[OwnedProxy, tuple[str, ...]]]:
|
||||
callbacks: Final[tuple[str, ...]] = request.param
|
||||
directory: Final = tmp_path_factory.mktemp("mixpanel-combined")
|
||||
with _booted(
|
||||
directory,
|
||||
otel_audit_config(directory, {"callbacks": cast("list[JsonValue]", list(callbacks))}),
|
||||
mixpanel,
|
||||
OPERATOR_TOKEN,
|
||||
collector=audit_sinks.operator if "otel" in callbacks else None,
|
||||
environment={
|
||||
"LANGFUSE_HOST": audit_sinks.operator,
|
||||
"LANGFUSE_PUBLIC_KEY": "pk-lf-operator",
|
||||
"LANGFUSE_SECRET_KEY": "sk-lf-operator",
|
||||
}
|
||||
if "langfuse_otel" in callbacks
|
||||
else {},
|
||||
) as owned:
|
||||
yield owned, callbacks
|
||||
|
||||
|
||||
@pytest.mark.parametrize("identified", ["user", "metadata-user-dot-id"])
|
||||
@pytest.mark.parametrize("level", ["team", "key"])
|
||||
def test_a_team_or_key_project_gets_the_callers_id_on_every_span(
|
||||
level: str, identified: str, tenant_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = tenant_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key, host, token = _team_or_key(scenario, level)
|
||||
trace_id: Final = _traced_call(
|
||||
proxy, key, {"model": model, **_chat_call("hi"), **IDENTIFICATIONS[identified](caller)}
|
||||
)
|
||||
spans: Final = _settled_trace(mixpanel.control_url, trace_id)
|
||||
|
||||
at_project: Final = _reached(mixpanel, spans, host)
|
||||
assert _without_the_id(at_project, caller) == []
|
||||
names: Final = [str(span["name"]) for span in at_project]
|
||||
assert any(name.startswith("auth ") for name in names), names
|
||||
assert (host, _authorization(token)) in _sent_to(_exports_of(mixpanel, spans))
|
||||
assert OPERATION not in json.dumps([span["attributes"] for span in _reached(mixpanel, spans, OPERATOR_HOST)]), [
|
||||
span["name"] for span in spans
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call", sorted(CALLS))
|
||||
def test_every_endpoint_carries_the_id_to_a_team_project(
|
||||
call: str, tenant_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = tenant_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key, host, _ = _team_or_key(scenario, "team")
|
||||
trace_id: Final = _traced_call(
|
||||
proxy, key, {"model": model, **CALLS[call].body("hi"), "metadata": {"user.id": caller}}, CALLS[call].path
|
||||
)
|
||||
spans: Final = _settled_trace(mixpanel.control_url, trace_id)
|
||||
|
||||
at_project: Final = _reached(mixpanel, spans, host)
|
||||
assert at_project, [span["name"] for span in spans]
|
||||
assert _without_the_id(at_project, caller) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("level", ["team", "key"])
|
||||
def test_every_request_on_a_long_lived_key_names_its_own_caller(
|
||||
level: str, tenant_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = tenant_setup.gateway
|
||||
callers: Final = tuple(f"caller-{uuid.uuid4().hex}" for _ in range(12))
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key, host, _ = _team_or_key(scenario, level)
|
||||
traces: Final = tuple(
|
||||
(caller, _call_on_a_new_connection(proxy, key, {"model": model, **_chat_call("hi"), "user": caller}))
|
||||
for caller in callers
|
||||
)
|
||||
settled: Final = tuple((caller, _settled_trace(mixpanel.control_url, trace_id)) for caller, trace_id in traces)
|
||||
|
||||
assert [(caller, _without_the_id(_reached(mixpanel, spans, host), caller)) for caller, spans in settled] == [
|
||||
(caller, []) for caller in callers
|
||||
]
|
||||
|
||||
|
||||
def test_a_project_token_listed_without_one_stays_quiet_on_plain_traffic(
|
||||
tenant_setup: OwnedProxy, provider: Wire
|
||||
) -> None:
|
||||
proxy: Final = tenant_setup.gateway
|
||||
logged_before: Final = len(tenant_setup.log.read_text())
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
_traced_call(
|
||||
proxy, scenario.key(key_alias=f"mixpanel-plain-{uuid.uuid4().hex}"), {"model": model, **_chat_call("hi")}
|
||||
)
|
||||
held(
|
||||
lambda: tenant_setup.log.read_text()[logged_before:],
|
||||
lambda logged: "Error initializing custom logger" not in logged,
|
||||
holding=5,
|
||||
seconds=40,
|
||||
)
|
||||
|
||||
|
||||
def test_a_renamed_service_keeps_the_model_call_of_an_eu_proxy_in_its_region(
|
||||
eu_operator_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = eu_operator_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key: Final = scenario.key(
|
||||
key_alias=f"mixpanel-renamed-{uuid.uuid4().hex}", metadata={"otel_service_name": "renamed-service"}
|
||||
)
|
||||
_traced_call(
|
||||
proxy,
|
||||
key,
|
||||
{"model": model, **_chat_call("hi"), "metadata": {USER_ID: caller, "langsmith_project": "any-project"}},
|
||||
)
|
||||
model_spans: Final = eventually(
|
||||
lambda: tuple(
|
||||
span
|
||||
for span in recorded_spans(mixpanel.control_url)[1]
|
||||
if OPERATION in span["attributes"] and span["attributes"].get(USER_ID) == caller
|
||||
),
|
||||
bool,
|
||||
seconds=40,
|
||||
)
|
||||
|
||||
assert _reached(mixpanel, model_spans, "ingestion-eu.mixpanel.com") == model_spans
|
||||
assert _reached(mixpanel, model_spans, OPERATOR_HOST) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("identified", sorted(IDENTIFICATIONS))
|
||||
def test_the_operators_project_gets_the_callers_id_on_every_span_after_auth(
|
||||
identified: str, operator_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = operator_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key: Final = scenario.key(key_alias=f"mixpanel-operator-{uuid.uuid4().hex}")
|
||||
trace_id: Final = _traced_call(
|
||||
proxy, key, {"model": model, **_chat_call("hi"), **IDENTIFICATIONS[identified](caller)}
|
||||
)
|
||||
spans: Final = _settled_trace(mixpanel.control_url, trace_id)
|
||||
|
||||
assert any(str(span["name"]).startswith("auth ") for span in spans), [span["name"] for span in spans]
|
||||
assert _without_the_id(_outside_auth(spans), caller) == []
|
||||
assert _sent_to(_exports_of(mixpanel, spans)) == {(OPERATOR_HOST, _authorization(OPERATOR_TOKEN))}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad_id",
|
||||
[
|
||||
{"metadata": {"user.id": " "}},
|
||||
{"metadata": {"user.id": ""}},
|
||||
{"metadata": {"user.id": True}},
|
||||
{"metadata": {"user_id": False}},
|
||||
{"metadata": {"user.id": None}},
|
||||
{"metadata": {"user.id": ["list"]}},
|
||||
{"metadata": {"user.id": "x" * 2000}},
|
||||
],
|
||||
ids=["whitespace", "empty", "true", "false", "null", "list", "oversized"],
|
||||
)
|
||||
def test_an_unusable_id_gives_way_to_the_callers_real_one(
|
||||
bad_id: dict[str, JsonValue], operator_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = operator_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key: Final = scenario.key(key_alias=f"mixpanel-bad-id-{uuid.uuid4().hex}")
|
||||
trace_id: Final = _traced_call(proxy, key, {"model": model, **_chat_call("hi"), "user": caller, **bad_id})
|
||||
spans: Final = _settled_trace(mixpanel.control_url, trace_id)
|
||||
|
||||
model_call: Final = [span for span in spans if OPERATION in span["attributes"]]
|
||||
assert [span["attributes"].get(USER_ID) for span in model_call] == [caller]
|
||||
|
||||
|
||||
def test_the_baggage_header_names_the_caller_only_once_auth_has_passed(
|
||||
operator_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = operator_setup.gateway
|
||||
named: Final = f"named-{uuid.uuid4().hex}"
|
||||
rejected_trace: Final = uuid.uuid4().hex
|
||||
accepted_trace: Final = uuid.uuid4().hex
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key: Final = scenario.key(key_alias=f"mixpanel-baggage-{uuid.uuid4().hex}")
|
||||
rejected: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, **_chat_call("hi")},
|
||||
key="sk-not-a-key",
|
||||
headers={
|
||||
"baggage": f"{USER_ID}={named}",
|
||||
"traceparent": f"00-{rejected_trace}-{uuid.uuid4().hex[:16]}-01",
|
||||
},
|
||||
)
|
||||
assert rejected.status_code == 401, rejected.text
|
||||
accepted: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, **_chat_call("hi")},
|
||||
key=key,
|
||||
headers={
|
||||
"baggage": f"{USER_ID}={named}",
|
||||
"traceparent": f"00-{accepted_trace}-{uuid.uuid4().hex[:16]}-01",
|
||||
},
|
||||
)
|
||||
assert accepted.status_code == 200, accepted.text
|
||||
accepted_spans: Final = _settled_trace(mixpanel.control_url, accepted_trace)
|
||||
rejected_spans: Final = _settled_trace(
|
||||
mixpanel.control_url, rejected_trace, lambda spans: any(span["kind"] == 2 for span in spans)
|
||||
)
|
||||
|
||||
assert named not in json.dumps([span["attributes"] for span in rejected_spans])
|
||||
assert _without_the_id(_outside_auth(accepted_spans), named) == []
|
||||
|
||||
|
||||
def test_a_collector_never_gets_the_id_only_a_team_project_asked_for(
|
||||
collector_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink, audit_sinks: SpanSinks
|
||||
) -> None:
|
||||
proxy: Final = collector_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key, _, _ = _team_or_key(scenario, "team")
|
||||
trace_id: Final = _traced_call(
|
||||
proxy, key, {"model": model, **_chat_call("hi"), "metadata": {"user.id": caller}}
|
||||
)
|
||||
at_mixpanel: Final = _settled_trace(mixpanel.control_url, trace_id)
|
||||
at_collector: Final = _settled_trace(audit_sinks.operator, trace_id)
|
||||
|
||||
assert _without_the_id(at_mixpanel, caller) == []
|
||||
assert caller not in json.dumps([span["attributes"] for span in at_collector])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"token",
|
||||
["", " ", None, "None", "***REDACTED***"],
|
||||
ids=["empty", "whitespace", "null", "none-string", "masked"],
|
||||
)
|
||||
def test_a_team_or_key_callback_refuses_a_token_it_could_never_send(
|
||||
token: JsonValue, operator_setup: OwnedProxy, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = operator_setup.gateway
|
||||
callback_vars: Final[dict[str, JsonValue]] = {"mixpanel_project_token": token}
|
||||
logging: Final[list[JsonValue]] = [
|
||||
{"callback_name": "mixpanel", "callback_type": "success", "callback_vars": callback_vars}
|
||||
]
|
||||
with proxy.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
team: Final = proxy.request(
|
||||
"POST",
|
||||
f"/team/{team_id}/callback",
|
||||
{"callback_name": "mixpanel", "callback_type": "success", "callback_vars": callback_vars},
|
||||
)
|
||||
new_team: Final = proxy.request(
|
||||
"POST", "/team/new", {"team_alias": f"mixpanel-{uuid.uuid4().hex}", "metadata": {"logging": logging}}
|
||||
)
|
||||
updated_team: Final = proxy.request(
|
||||
"POST", "/team/update", {"team_id": team_id, "metadata": {"logging": logging}}
|
||||
)
|
||||
generated: Final = proxy.request(
|
||||
"POST", "/key/generate", {"key_alias": f"mixpanel-{uuid.uuid4().hex}", "metadata": {"logging": logging}}
|
||||
)
|
||||
created: Final = scenario.key(key_alias=f"mixpanel-{uuid.uuid4().hex}")
|
||||
updated: Final = proxy.request("POST", "/key/update", {"key": created, "metadata": {"logging": logging}})
|
||||
|
||||
for refused in (team, new_team, updated_team, generated, updated):
|
||||
assert refused.status_code in (400, 422), refused.text
|
||||
assert "mixpanel_project_token" in refused.text, refused.text
|
||||
|
||||
|
||||
def test_each_project_of_an_additive_proxy_gets_the_callers_id_on_every_span(
|
||||
additive_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = additive_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key, host, token = _team_or_key(scenario, "team")
|
||||
trace_id: Final = _traced_call(proxy, key, {"model": model, **_chat_call("hi"), "user": caller})
|
||||
spans: Final = _settled_trace(
|
||||
mixpanel.control_url,
|
||||
trace_id,
|
||||
lambda found: len([span for span in found if OPERATION in span["attributes"]]) == 2,
|
||||
)
|
||||
|
||||
assert _without_the_id(_outside_auth(spans), caller) == []
|
||||
assert [str(span["name"]) for span in _reached(mixpanel, spans, host) if OPERATION in span["attributes"]]
|
||||
assert [str(span["name"]) for span in _reached(mixpanel, spans, OPERATOR_HOST) if OPERATION in span["attributes"]]
|
||||
assert _sent_to(_exports_of(mixpanel, spans)) == {
|
||||
(host, _authorization(token)),
|
||||
(OPERATOR_HOST, _authorization(OPERATOR_TOKEN)),
|
||||
}
|
||||
|
||||
|
||||
def test_mixpanel_listed_with_another_callback_gets_one_labeled_copy_of_the_model_call(
|
||||
combined_setup: tuple[OwnedProxy, tuple[str, ...]], provider: Wire, mixpanel: ConnectSink, audit_sinks: SpanSinks
|
||||
) -> None:
|
||||
owned, callbacks = combined_setup
|
||||
proxy: Final = owned.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key: Final = scenario.key(key_alias=f"mixpanel-combined-{uuid.uuid4().hex}")
|
||||
trace_id: Final = _traced_call(proxy, key, {"model": model, **_chat_call("hi"), "user": caller})
|
||||
at_mixpanel: Final = _settled_trace(mixpanel.control_url, trace_id, _has_model_call)
|
||||
at_other: Final = _settled_trace(audit_sinks.operator, trace_id, _has_model_call)
|
||||
|
||||
assert len([span for span in at_mixpanel if OPERATION in span["attributes"]]) == 1, at_mixpanel
|
||||
assert len([span for span in at_other if OPERATION in span["attributes"]]) == 1, at_other
|
||||
assert _without_the_id(_outside_auth(at_mixpanel), caller) == []
|
||||
assert any(span["kind"] == 2 for span in at_mixpanel) == (callbacks[0] == "mixpanel"), [
|
||||
span["name"] for span in at_mixpanel
|
||||
]
|
||||
assert _sent_to(_exports_of(mixpanel, at_mixpanel)) == {(OPERATOR_HOST, _authorization(OPERATOR_TOKEN))}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry", ["mcp", "rest"])
|
||||
def test_a_team_project_gets_the_callers_id_on_an_mcp_tool_call(
|
||||
entry: EntryPoint, tenant_setup: OwnedProxy, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = tenant_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
alias: Final = f"mixpanelmcp{uuid.uuid4().hex[:8]}"
|
||||
with scripted_peer(echo_tool("echo")) as peer, proxy.scenario() as scenario:
|
||||
server_id: Final = register_mcp(scenario, peer, alias)
|
||||
key, host, _ = _team_or_key(scenario, "team", object_permission={"mcp_servers": [server_id]})
|
||||
outcome: Final = McpCaller(
|
||||
proxy,
|
||||
key,
|
||||
entry,
|
||||
alias,
|
||||
headers={
|
||||
"baggage": f"{USER_ID}={caller}",
|
||||
"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01",
|
||||
},
|
||||
).call(f"{alias}-echo", {"note": "hi"}, server_id=server_id)
|
||||
assert outcome.error is None, outcome.raw
|
||||
spans: Final = _settled_trace(mixpanel.control_url, trace_id, _tool_call_finished)
|
||||
|
||||
at_project: Final = _reached(mixpanel, spans, host)
|
||||
assert [span for span in at_project if _is_tool_call(span)], [span["name"] for span in spans]
|
||||
assert _without_the_id(at_project, caller) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("later_message_names", [True, False], ids=["its-own-caller", "nobody"])
|
||||
def test_a_later_message_of_an_mcp_session_carries_its_own_caller_and_never_the_first(
|
||||
later_message_names: bool, tenant_setup: OwnedProxy, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = tenant_setup.gateway
|
||||
first_caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
later_caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
alias: Final = f"mixpanelmcp{uuid.uuid4().hex[:8]}"
|
||||
with scripted_peer(echo_tool("echo")) as peer, proxy.scenario() as scenario:
|
||||
server_id: Final = register_mcp(scenario, peer, alias)
|
||||
key, host, _ = _team_or_key(scenario, "team", object_permission={"mcp_servers": [server_id]})
|
||||
opened: Final = McpCaller(proxy, key, "mcp", alias, headers={"baggage": f"{USER_ID}={first_caller}"}).rpc(
|
||||
"initialize", INITIALIZE
|
||||
)
|
||||
assert opened.status_code == 200, opened.text
|
||||
session: Final = {"Mcp-Session-Id": opened.headers["mcp-session-id"]}
|
||||
initialized: Final = proxy.client.post(
|
||||
"/mcp",
|
||||
headers=McpCaller(proxy, key, "mcp", alias, headers=session)._headers(),
|
||||
json={"jsonrpc": "2.0", "method": "notifications/initialized"},
|
||||
)
|
||||
assert initialized.status_code == 202, initialized.text
|
||||
outcome: Final = McpCaller(
|
||||
proxy,
|
||||
key,
|
||||
"mcp",
|
||||
alias,
|
||||
headers={
|
||||
**session,
|
||||
**({"baggage": f"{USER_ID}={later_caller}"} if later_message_names else {}),
|
||||
"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01",
|
||||
},
|
||||
).call(f"{alias}-echo", {"note": "hi"}, server_id=server_id)
|
||||
assert outcome.error is None, outcome.raw
|
||||
spans: Final = _settled_trace(mixpanel.control_url, trace_id, _tool_call_finished)
|
||||
|
||||
at_project: Final = _reached(mixpanel, spans, host)
|
||||
assert [span for span in at_project if _is_tool_call(span)], [span["name"] for span in spans]
|
||||
assert first_caller not in json.dumps([span["attributes"] for span in spans])
|
||||
assert {span["attributes"].get(USER_ID) for span in at_project} == (
|
||||
{later_caller} if later_message_names else {None}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry", ["mcp", "rest"])
|
||||
def test_a_collector_beside_the_operators_project_gets_no_id_on_a_model_or_tool_call(
|
||||
entry: EntryPoint,
|
||||
collector_beside_operator_setup: OwnedProxy,
|
||||
provider: Wire,
|
||||
mixpanel: ConnectSink,
|
||||
audit_sinks: SpanSinks,
|
||||
) -> None:
|
||||
proxy: Final = collector_beside_operator_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
tool_trace: Final = uuid.uuid4().hex
|
||||
alias: Final = f"mixpanelmcp{uuid.uuid4().hex[:8]}"
|
||||
with scripted_peer(echo_tool("echo")) as peer, proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
server_id: Final = register_mcp(scenario, peer, alias)
|
||||
key: Final = scenario.key(
|
||||
key_alias=f"mixpanel-collector-{uuid.uuid4().hex}", object_permission={"mcp_servers": [server_id]}
|
||||
)
|
||||
model_trace: Final = _traced_call(proxy, key, {"model": model, **_chat_call("hi"), "user": caller})
|
||||
outcome: Final = McpCaller(
|
||||
proxy,
|
||||
key,
|
||||
entry,
|
||||
alias,
|
||||
headers={
|
||||
"x-litellm-customer-id": caller,
|
||||
"traceparent": f"00-{tool_trace}-{uuid.uuid4().hex[:16]}-01",
|
||||
},
|
||||
).call(f"{alias}-echo", {"note": "hi"}, server_id=server_id)
|
||||
assert outcome.error is None, outcome.raw
|
||||
at_mixpanel: Final = (
|
||||
*_settled_trace(mixpanel.control_url, model_trace),
|
||||
*_settled_trace(mixpanel.control_url, tool_trace, _tool_call_finished),
|
||||
)
|
||||
at_collector: Final = (
|
||||
*_settled_trace(audit_sinks.operator, model_trace),
|
||||
*_settled_trace(audit_sinks.operator, tool_trace, _tool_call_finished),
|
||||
)
|
||||
|
||||
assert [span for span in at_mixpanel if _is_tool_call(span)], [span["name"] for span in at_mixpanel]
|
||||
labeled: Final = _outside_auth(at_mixpanel) if entry == "rest" else _inside_the_mcp_session(at_mixpanel)
|
||||
assert _without_the_id(labeled, caller) == []
|
||||
assert [span["name"] for span in at_collector if USER_ID in span["attributes"]] == []
|
||||
|
||||
|
||||
def test_a_custom_auth_proxy_names_the_caller_from_the_request(
|
||||
custom_auth_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = custom_auth_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
trace_id: Final = _traced_call(proxy, "sk-any-caller", {"model": model, **_chat_call("hi"), "user": caller})
|
||||
spans: Final = _settled_trace(mixpanel.control_url, trace_id)
|
||||
|
||||
assert _without_the_id(_outside_auth(spans), caller) == []
|
||||
|
||||
|
||||
def test_a_failed_model_call_reaches_a_team_project_with_the_callers_id(
|
||||
tenant_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = tenant_setup.gateway
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key, host, _ = _team_or_key(scenario, "team")
|
||||
response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, **_chat_call(PROVIDER_FAILS), "user": caller},
|
||||
key=key,
|
||||
headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
spans: Final = _settled_trace(mixpanel.control_url, trace_id)
|
||||
|
||||
at_project: Final = _reached(mixpanel, spans, host)
|
||||
assert at_project, [span["name"] for span in spans]
|
||||
assert _without_the_id(at_project, caller) == []
|
||||
|
||||
|
||||
def test_concurrent_callers_never_get_each_others_id(
|
||||
operator_setup: OwnedProxy, provider: Wire, mixpanel: ConnectSink
|
||||
) -> None:
|
||||
proxy: Final = operator_setup.gateway
|
||||
callers: Final = tuple(
|
||||
(f"caller-{uuid.uuid4().hex}", sorted(CALLS)[index % len(CALLS)], uuid.uuid4().hex) for index in range(32)
|
||||
)
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key: Final = scenario.key(key_alias=f"mixpanel-burst-{uuid.uuid4().hex}")
|
||||
team_key, _, _ = _team_or_key(scenario, "team")
|
||||
with ThreadPoolExecutor(max_workers=len(callers)) as executor:
|
||||
statuses: Final = tuple(
|
||||
executor.map(
|
||||
lambda indexed: (
|
||||
proxy.request(
|
||||
"POST",
|
||||
CALLS[indexed[1][1]].path,
|
||||
{"model": model, **CALLS[indexed[1][1]].body("hi"), "metadata": {"user.id": indexed[1][0]}},
|
||||
key=team_key if indexed[0] % 2 else key,
|
||||
headers={"traceparent": f"00-{indexed[1][2]}-{uuid.uuid4().hex[:16]}-01"},
|
||||
).status_code
|
||||
),
|
||||
enumerate(callers),
|
||||
)
|
||||
)
|
||||
assert statuses == (200,) * len(callers), statuses
|
||||
traces: Final = tuple(
|
||||
(caller, _settled_trace(mixpanel.control_url, trace_id)) for caller, _, trace_id in callers
|
||||
)
|
||||
|
||||
assert [(caller, _without_the_id(_outside_auth(spans), caller)) for caller, spans in traces] == [
|
||||
(caller, []) for caller, _ in traces
|
||||
]
|
||||
|
||||
|
||||
def test_a_project_outage_costs_spans_and_never_a_request(
|
||||
mixpanel: ConnectSink, provider: Wire, otel_audit_config: AuditConfigWriter, tmp_path: Path
|
||||
) -> None:
|
||||
with wire_server(_ingest_accepted) as first_listener:
|
||||
port: Final = urlsplit(first_listener.url).port
|
||||
caller: Final = f"caller-{uuid.uuid4().hex}"
|
||||
with _booted(
|
||||
tmp_path,
|
||||
otel_audit_config(tmp_path, {"callbacks": ["mixpanel"]}),
|
||||
mixpanel,
|
||||
OPERATOR_TOKEN,
|
||||
environment={"MIXPANEL_ENDPOINT": f"http://127.0.0.1:{port}/v1/ai/otel/traces"},
|
||||
) as owned:
|
||||
proxy: Final = owned.gateway
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key: Final = scenario.key(key_alias=f"mixpanel-outage-{uuid.uuid4().hex}")
|
||||
during: Final = tuple(
|
||||
proxy.request("POST", "/v1/chat/completions", {"model": model, **_chat_call("hi")}, key=key).status_code
|
||||
for _ in range(10)
|
||||
)
|
||||
assert during == (200,) * len(during), during
|
||||
with wire_server(_ingest_accepted, port=port) as recovered:
|
||||
assert (
|
||||
proxy.request(
|
||||
"POST", "/v1/chat/completions", {"model": model, **_chat_call("hi"), "user": caller}, key=key
|
||||
).status_code
|
||||
== 200
|
||||
)
|
||||
landed: Final = eventually(
|
||||
lambda: tuple(_wire_span_attributes(recovered.drain())),
|
||||
lambda spans: any(span.get(USER_ID) == caller and OPERATION in span for span in spans),
|
||||
seconds=40,
|
||||
)
|
||||
assert any(span.get(USER_ID) == caller for span in landed), landed
|
||||
|
|
@ -649,3 +649,35 @@ def test_streamed_call_is_refused_as_json_once_over_budget(gateway: Gateway) ->
|
|||
assert "data:" not in denied.text, denied.text
|
||||
_assert_budget_refusal(denied, end_user, 0.05)
|
||||
assert _observed_requests(upstream) == []
|
||||
|
||||
|
||||
def test_metadata_user_dot_id_is_not_an_end_user(gateway: Gateway) -> None:
|
||||
labeled: Final = f"end-user-label-{uuid.uuid4().hex}"
|
||||
limited: Final = f"end-user-limited-{uuid.uuid4().hex}"
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
key: Final = scenario.key(key_alias=f"end-user-label-{uuid.uuid4().hex}")
|
||||
for customer in (labeled, limited):
|
||||
gateway.post("/customer/new", {"user_id": customer, "rpm_limit": 1})
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [customer]})
|
||||
|
||||
def sent(metadata: Mapping[str, JsonValue]) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": "hi"}], "metadata": dict(metadata)},
|
||||
key=key,
|
||||
)
|
||||
|
||||
labeled_replies: Final = tuple(sent({"user.id": labeled}) for _ in range(4))
|
||||
assert [reply.status_code for reply in labeled_replies] == [200, 200, 200, 200], labeled_replies[-1].text
|
||||
response_id: Final = string_value(object_value(labeled_replies[-1].json())["id"])
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT end_user FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (response_id,)),
|
||||
lambda found: len(found) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert rows == [{"end_user": ""}]
|
||||
limited_statuses: Final = [sent({"user_id": limited}).status_code for _ in range(4)]
|
||||
|
||||
assert 429 in limited_statuses, limited_statuses
|
||||
|
|
|
|||
|
|
@ -2,10 +2,15 @@
|
|||
and the two antipattern boundaries: http.* is never promoted, and the full
|
||||
metadata blob is never promoted (only the bounded allowlist)."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("opentelemetry")
|
||||
|
||||
from opentelemetry import baggage # noqa: E402
|
||||
from opentelemetry.context import attach, detach # noqa: E402
|
||||
|
||||
from litellm.integrations.otel import ( # noqa: E402
|
||||
HTTP,
|
||||
GenAI,
|
||||
|
|
@ -15,6 +20,12 @@ from litellm.integrations.otel import ( # noqa: E402
|
|||
)
|
||||
from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402
|
||||
from litellm.integrations.otel.model.baggage import BAGGAGE_PROMOTED_KEYS # noqa: E402
|
||||
from litellm.integrations.otel.model.metadata import ( # noqa: E402
|
||||
caller_user_id,
|
||||
caller_user_id_in_call,
|
||||
request_user_id,
|
||||
usable_user_id,
|
||||
)
|
||||
from litellm.integrations.otel.model.payloads import ( # noqa: E402
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
|
|
@ -256,34 +267,19 @@ def test_baggage_processor_allowlist_can_be_widened():
|
|||
assert LiteLLM.TEAM_ALIAS not in span.attributes
|
||||
|
||||
|
||||
def test_user_id_not_promoted_by_default_and_promoted_for_mixpanel(monkeypatch):
|
||||
def test_the_mixpanel_preset_never_promotes_user_id_through_baggage(monkeypatch):
|
||||
monkeypatch.setenv("MIXPANEL_PROJECT_TOKEN", "mp-token")
|
||||
from litellm.integrations.otel.presets.mixpanel import mixpanel_preset
|
||||
|
||||
payload = _payload()
|
||||
payload["end_user"] = "user_123"
|
||||
data = LLMCallSpanData.from_standard_logging_payload(payload)
|
||||
|
||||
default_bag = promoted_baggage(data.identity, data.request_model, BAGGAGE_PROMOTED_KEYS)
|
||||
assert "user.id" not in default_bag
|
||||
|
||||
mp_cfg = mixpanel_preset(config_overrides=OpenTelemetryV2Config(exporter="in_memory"))
|
||||
engine, exporter = _engine_and_exporter(mp_cfg)
|
||||
mp_bag = promoted_baggage(data.identity, data.request_model, tuple(mp_cfg.baggage_promoted_keys))
|
||||
assert mp_bag["user.id"] == "user_123"
|
||||
|
||||
ctx = ctx_mod.set_request_baggage(mp_bag)
|
||||
root = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions", ctx)
|
||||
root_ctx = ctx_mod.context_from_span(root, ctx)
|
||||
engine.emit(SpanRole.LLM_CALL, data, parent_context=root_ctx)
|
||||
engine.emit(SpanRole.GUARDRAIL, GuardrailSpanData("presidio", status="success"), root_ctx)
|
||||
engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), root_ctx)
|
||||
root.end()
|
||||
|
||||
spans = exporter.get_finished_spans()
|
||||
assert len(spans) == 4
|
||||
for span in spans:
|
||||
assert span.attributes.get("user.id") == "user_123"
|
||||
assert "user.id" not in mp_cfg.baggage_promoted_keys
|
||||
assert "user.id" not in promoted_baggage(
|
||||
data.identity, data.request_model, (*mp_cfg.baggage_promoted_keys, "user.id")
|
||||
)
|
||||
|
||||
|
||||
def test_user_id_precedence_and_key_owner_exclusion():
|
||||
|
|
@ -298,3 +294,88 @@ def test_user_id_precedence_and_key_owner_exclusion():
|
|||
key_owner_only_payload["metadata"]["user_api_key_user_id"] = "from-key-owner"
|
||||
key_owner_only_data = LLMCallSpanData.from_standard_logging_payload(key_owner_only_payload)
|
||||
assert key_owner_only_data.identity.user_id is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"unusable",
|
||||
["", " ", True, False, None, ["list"], {"a": 1}, "x" * 1025],
|
||||
ids=["empty", "whitespace", "true", "false", "null", "list", "mapping", "oversized"],
|
||||
)
|
||||
def test_an_unusable_named_id_gives_way_to_the_next_source(unusable):
|
||||
payload = _payload()
|
||||
payload["end_user"] = "from-end-user"
|
||||
payload["metadata"]["requester_metadata"] = {"user.id": unusable}
|
||||
assert LLMCallSpanData.from_standard_logging_payload(payload).identity.user_id == "from-end-user"
|
||||
assert caller_user_id({"user.id": unusable, "user_id": "the-next-field"}) == "the-next-field"
|
||||
|
||||
|
||||
def test_a_control_character_is_not_part_of_an_id():
|
||||
assert usable_user_id("a\nb=c") is None
|
||||
assert usable_user_id("usuario né") == "usuario né"
|
||||
|
||||
|
||||
def test_the_id_in_baggage_is_read_from_the_current_context():
|
||||
token = attach(baggage.set_baggage("user.id", "in-baggage"))
|
||||
try:
|
||||
assert ctx_mod.inbound_user_id() == "in-baggage"
|
||||
finally:
|
||||
detach(token)
|
||||
assert ctx_mod.inbound_user_id() is None
|
||||
unusable = attach(baggage.set_baggage("user.id", "a\nb"))
|
||||
try:
|
||||
assert ctx_mod.inbound_user_id() is None
|
||||
finally:
|
||||
detach(unusable)
|
||||
|
||||
|
||||
def test_a_numeric_id_and_surrounding_space_are_kept_as_the_callers_id():
|
||||
assert caller_user_id({"user.id": 12345}) == "12345"
|
||||
assert caller_user_id({"user.id": " padded "}) == "padded"
|
||||
|
||||
|
||||
def test_named_id_precedence_between_fields_and_sources():
|
||||
assert caller_user_id({"trace_user_id": "t", "user_id": "u", "user.id": "dot"}) == "dot"
|
||||
assert caller_user_id({"trace_user_id": "t", "user_id": "u"}) == "u"
|
||||
assert caller_user_id({"trace_user_id": "t"}) == "t"
|
||||
assert caller_user_id({"user": {"id": "nested"}}) == "nested"
|
||||
assert caller_user_id({"user": "not-a-mapping"}) is None
|
||||
assert caller_user_id({"user.id": "top", "requester_metadata": {"user.id": "requester"}}) == "requester"
|
||||
assert caller_user_id({"user_id": "top", "requester_metadata": {"user.id": "requester"}}) == "requester"
|
||||
assert caller_user_id({"user.id": "top", "requester_metadata": {"user_id": "requester"}}) == "top"
|
||||
assert caller_user_id({"user_id": "metadata"}, {"user.id": "litellm-metadata"}) == "litellm-metadata"
|
||||
assert caller_user_id("not-a-mapping", None, {"user_id": "found"}) == "found"
|
||||
|
||||
|
||||
def test_the_id_at_the_auth_boundary_reads_encoded_metadata_and_falls_back_to_the_end_user():
|
||||
assert request_user_id(None, {"metadata": {"user.id": "plain"}}) == "plain"
|
||||
assert request_user_id(None, {"metadata": json.dumps({"user.id": "encoded"})}) == "encoded"
|
||||
assert request_user_id(None, {"litellm_metadata": {"user.id": "litellm"}}) == "litellm"
|
||||
assert request_user_id("end-user", {"metadata": {"user.id": "named"}}) == "named"
|
||||
assert request_user_id("end-user", {"metadata": "not json"}) == "end-user"
|
||||
assert request_user_id("end-user", {}) == "end-user"
|
||||
assert request_user_id(None, {"metadata": {"session": "s"}}) is None
|
||||
assert request_user_id(" ", {}) is None
|
||||
assert request_user_id("end-user", {}, "in-baggage") == "in-baggage"
|
||||
assert request_user_id("end-user", {"metadata": {"user.id": "named"}}, "in-baggage") == "named"
|
||||
|
||||
|
||||
def test_the_call_event_reads_the_named_id_from_either_metadata_and_the_payload_uses_it():
|
||||
def event_id(params):
|
||||
return caller_user_id_in_call({"litellm_params": params})
|
||||
|
||||
assert event_id({"metadata": {"user.id": "from-metadata"}}) == "from-metadata"
|
||||
assert event_id({"litellm_metadata": {"user.id": "from-litellm-metadata"}}) == "from-litellm-metadata"
|
||||
assert event_id({"metadata": None}) is None
|
||||
assert caller_user_id_in_call({"litellm_params": None}) is None
|
||||
|
||||
payload = _payload()
|
||||
payload["end_user"] = "from-end-user"
|
||||
assert LLMCallSpanData.from_standard_logging_payload(payload, named_user_id="named").identity.user_id == "named"
|
||||
|
||||
|
||||
def test_the_openai_user_param_is_the_last_resort_and_leaves_the_end_user_alone():
|
||||
payload = _payload()
|
||||
payload["model_parameters"] = {"user": "from-param"}
|
||||
identity = LLMCallSpanData.from_standard_logging_payload(payload).identity
|
||||
assert identity.user_id == "from-param"
|
||||
assert not identity.end_user
|
||||
|
|
|
|||
|
|
@ -14,8 +14,9 @@ from types import MappingProxyType
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from opentelemetry import baggage
|
||||
from opentelemetry import trace as trace_api
|
||||
from opentelemetry.context import Context
|
||||
from opentelemetry.context import Context, attach
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
|
|
@ -60,7 +61,9 @@ from litellm.integrations.otel.plumbing import providers as otel_providers
|
|||
from litellm.integrations.otel.plumbing.context import (
|
||||
destination_backends,
|
||||
request_destinations,
|
||||
request_user_id,
|
||||
set_request_destinations,
|
||||
set_request_user_id,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing.providers import (
|
||||
TenantFanOutSpanProcessor,
|
||||
|
|
@ -5886,3 +5889,205 @@ def test_mixpanel_destination_resolves_endpoint_and_basic_auth():
|
|||
assert destination.endpoint == "https://ingestion-eu.mixpanel.com/v1/ai/otel/traces"
|
||||
assert destination.headers["Authorization"] == f"Basic {b64encode(b'team-tok:').decode()}"
|
||||
assert destination.callback_name == "mixpanel"
|
||||
|
||||
|
||||
class TestMixpanelGetsTheCallersId:
|
||||
"""Mixpanel drops a span without ``user.id``, and no other backend asked for it."""
|
||||
|
||||
MIXPANEL: Final = OtelDestination(
|
||||
endpoint="https://ingestion-eu.mixpanel.com/v1/ai/otel/traces",
|
||||
headers={"Authorization": "Basic dGVhbTo="},
|
||||
callback_name="mixpanel",
|
||||
)
|
||||
ARIZE: Final = OtelDestination(endpoint="https://arize.local/v1/traces", callback_name="arize")
|
||||
|
||||
@staticmethod
|
||||
def _fan_out(
|
||||
langfuse_mapped: bool = False,
|
||||
) -> tuple[TracerProvider, InMemorySpanExporter, InMemorySpanExporter]:
|
||||
mixpanel_exporter, other_exporter = InMemorySpanExporter(), InMemorySpanExporter()
|
||||
provider: Final = TracerProvider()
|
||||
provider.add_span_processor(
|
||||
TenantFanOutSpanProcessor(
|
||||
processor_factory=lambda destination: SimpleSpanProcessor(
|
||||
mixpanel_exporter if destination.callback_name == "mixpanel" else other_exporter
|
||||
),
|
||||
langfuse_mapped=langfuse_mapped,
|
||||
)
|
||||
)
|
||||
return provider, mixpanel_exporter, other_exporter
|
||||
|
||||
@staticmethod
|
||||
def _operator(monkeypatch: pytest.MonkeyPatch) -> tuple[TracerProvider, dict[object, InMemorySpanExporter]]:
|
||||
exporters: dict[object, InMemorySpanExporter] = {} # mutable-ok: filled as specs build
|
||||
kind: Final = "lit10367_mixpanel"
|
||||
monkeypatch.setitem(
|
||||
otel_providers._EXPORTER_FACTORIES,
|
||||
kind,
|
||||
lambda spec: exporters.setdefault(spec.owner, InMemorySpanExporter()),
|
||||
)
|
||||
config: Final = OpenTelemetryV2Config(
|
||||
mapper_names=["genai", "mixpanel"],
|
||||
exporters=[ExporterSpec(kind=kind, owner=owner) for owner in (ExporterOwner.MIXPANEL, None)],
|
||||
)
|
||||
return build_tracer_provider(config, use_simple_processor=True), exporters
|
||||
|
||||
@staticmethod
|
||||
def _request(
|
||||
provider: TracerProvider,
|
||||
destinations: tuple[OtelDestination, ...] = (),
|
||||
user_id: str | None = "caller-1",
|
||||
model_call_attributes: Mapping[str, str] = MappingProxyType({"gen_ai.operation.name": "chat"}),
|
||||
) -> None:
|
||||
def run() -> None:
|
||||
set_request_destinations(destinations)
|
||||
set_request_user_id(user_id)
|
||||
tracer = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions"):
|
||||
with tracer.start_as_current_span("chat gpt-4o") as span:
|
||||
span.set_attributes(dict(model_call_attributes))
|
||||
get_tracer(provider, "instrumentation").start_span("redis").end()
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
def test_a_mixpanel_destination_gets_the_id_on_every_span_and_a_second_destination_on_none(self) -> None:
|
||||
provider, mixpanel, other = self._fan_out()
|
||||
|
||||
self._request(provider, (self.MIXPANEL, self.ARIZE))
|
||||
|
||||
assert {name: span.attributes.get("user.id") for name, span in by_name(mixpanel).items()} == {
|
||||
"POST /v1/chat/completions": "caller-1",
|
||||
"chat gpt-4o": "caller-1",
|
||||
"redis": "caller-1",
|
||||
}
|
||||
assert [span.name for span in other.get_finished_spans() if "user.id" in span.attributes] == []
|
||||
assert len(other.get_finished_spans()) == 3
|
||||
|
||||
def test_the_id_the_mapper_wrote_survives_the_langfuse_strip_only_for_mixpanel(self) -> None:
|
||||
provider, mixpanel, other = self._fan_out()
|
||||
|
||||
self._request(
|
||||
provider,
|
||||
(self.MIXPANEL, self.ARIZE),
|
||||
user_id=None,
|
||||
model_call_attributes={"gen_ai.operation.name": "chat", "user.id": "from-mapper", "session.id": "s"},
|
||||
)
|
||||
|
||||
assert by_name(mixpanel)["chat gpt-4o"].attributes.get("user.id") == "from-mapper"
|
||||
assert "session.id" not in by_name(mixpanel)["chat gpt-4o"].attributes
|
||||
assert "user.id" not in by_name(other)["chat gpt-4o"].attributes
|
||||
|
||||
def test_the_id_the_mapper_wrote_wins_over_the_id_read_at_auth(self) -> None:
|
||||
provider, mixpanel, _ = self._fan_out()
|
||||
|
||||
self._request(
|
||||
provider,
|
||||
(self.MIXPANEL,),
|
||||
model_call_attributes={"gen_ai.operation.name": "chat", "user.id": "from-mapper"},
|
||||
)
|
||||
|
||||
assert by_name(mixpanel)["chat gpt-4o"].attributes.get("user.id") == "from-mapper"
|
||||
assert by_name(mixpanel)["POST /v1/chat/completions"].attributes.get("user.id") == "caller-1"
|
||||
|
||||
def test_a_request_without_an_id_adds_nothing(self) -> None:
|
||||
provider, mixpanel, _ = self._fan_out()
|
||||
|
||||
self._request(provider, (self.MIXPANEL,), user_id=None)
|
||||
|
||||
assert [span.name for span in mixpanel.get_finished_spans() if "user.id" in span.attributes] == []
|
||||
assert len(mixpanel.get_finished_spans()) == 3
|
||||
|
||||
def test_a_held_tree_released_on_another_requests_task_keeps_its_own_id(self) -> None:
|
||||
provider, mixpanel, _ = self._fan_out()
|
||||
sampled: Final = self.MIXPANEL.model_copy(update={"success_sampling_rate": 1.0})
|
||||
|
||||
def run() -> None:
|
||||
set_request_destinations((sampled,))
|
||||
tracer = get_tracer(provider, "litellm")
|
||||
for index in range(otel_providers._MAX_PENDING_TREES + 1):
|
||||
set_request_user_id(f"caller-{index}")
|
||||
root = tracer.start_span(f"POST {index}", context=trace_api.Context())
|
||||
tracer.start_span(f"chat {index}", context=trace_api.set_span_in_context(root)).end()
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
assert [(span.name, span.attributes.get("user.id")) for span in mixpanel.get_finished_spans()] == [
|
||||
("chat 0", "caller-0")
|
||||
]
|
||||
|
||||
def test_the_operators_project_gets_the_id_on_every_span_and_a_collector_beside_it_on_none(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
provider, exporters = self._operator(monkeypatch)
|
||||
|
||||
self._request(
|
||||
provider, model_call_attributes={"gen_ai.operation.name": "chat", "user.id": "caller-1", "session.id": "s"}
|
||||
)
|
||||
|
||||
assert {
|
||||
name: span.attributes.get("user.id") for name, span in by_name(exporters[ExporterOwner.MIXPANEL]).items()
|
||||
} == {
|
||||
"POST /v1/chat/completions": "caller-1",
|
||||
"chat gpt-4o": "caller-1",
|
||||
"redis": "caller-1",
|
||||
}
|
||||
assert "session.id" not in by_name(exporters[ExporterOwner.MIXPANEL])["chat gpt-4o"].attributes
|
||||
assert [span.name for span in exporters[None].get_finished_spans() if "user.id" in span.attributes] == []
|
||||
assert len(exporters[None].get_finished_spans()) == 3
|
||||
|
||||
def test_the_operators_project_keeps_the_mapped_id_when_auth_never_read_one(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
provider, exporters = self._operator(monkeypatch)
|
||||
|
||||
self._request(
|
||||
provider, user_id=None, model_call_attributes={"gen_ai.operation.name": "chat", "user.id": "from-mapper"}
|
||||
)
|
||||
|
||||
assert by_name(exporters[ExporterOwner.MIXPANEL])["chat gpt-4o"].attributes.get("user.id") == "from-mapper"
|
||||
assert "user.id" not in by_name(exporters[None])["chat gpt-4o"].attributes
|
||||
|
||||
def test_auth_names_the_caller_for_the_request_and_a_failure_never_breaks_auth(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.auth.user_api_key_auth import _seed_request_user_id
|
||||
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
def request_with(headers: tuple[tuple[bytes, bytes], ...] = ()) -> Request:
|
||||
return Request(scope={"type": "http", "headers": list(headers)})
|
||||
|
||||
def seeded(end_user_id: str | None, body: object) -> str | None:
|
||||
def run() -> str | None:
|
||||
_seed_request_user_id(UserAPIKeyAuth(end_user_id=end_user_id), body, request_with())
|
||||
return request_user_id()
|
||||
|
||||
return in_fresh_context(run)
|
||||
|
||||
def seeded_from_header() -> str | None:
|
||||
_seed_request_user_id(
|
||||
UserAPIKeyAuth(), {}, request_with(((b"x-litellm-customer-id", b"from-header"),))
|
||||
)
|
||||
return request_user_id()
|
||||
|
||||
assert seeded(None, {"metadata": {"user.id": "named"}}) == "named"
|
||||
assert seeded("end-user", {}) == "end-user"
|
||||
assert seeded(None, {"user": "from-body"}) == "from-body"
|
||||
assert in_fresh_context(seeded_from_header) == "from-header"
|
||||
assert seeded(None, {}) is None
|
||||
assert seeded("end-user", None) is None
|
||||
|
||||
def seeded_under_baggage() -> str | None:
|
||||
attach(baggage.set_baggage("user.id", "in-baggage"))
|
||||
_seed_request_user_id(UserAPIKeyAuth(end_user_id="end-user"), {}, request_with())
|
||||
return request_user_id()
|
||||
|
||||
assert in_fresh_context(seeded_under_baggage) == "in-baggage"
|
||||
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "false")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
assert seeded("end-user", {}) is None
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
|
|
|||
|
|
@ -751,6 +751,8 @@ def test_mixpanel_dynamic_headers_and_endpoint_routing(monkeypatch, caplog):
|
|||
"Authorization": expected_auth
|
||||
}
|
||||
assert dynamic_otlp_headers("mixpanel", {}) is None
|
||||
assert dynamic_otlp_endpoint("mixpanel", {"langsmith_project": "another-callback"}) is None
|
||||
assert dynamic_otlp_endpoint("mixpanel", {"mixpanel_project_token": " ", "mixpanel_region": "eu"}) is None
|
||||
|
||||
monkeypatch.setenv("MIXPANEL_REGION", "eu")
|
||||
assert (
|
||||
|
|
|
|||
|
|
@ -23,7 +23,8 @@ import pytest
|
|||
|
||||
pytest.importorskip("opentelemetry")
|
||||
|
||||
from opentelemetry import trace # noqa: E402
|
||||
from opentelemetry import baggage, trace # noqa: E402
|
||||
from opentelemetry.context import attach # noqa: E402
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402
|
||||
InMemorySpanExporter,
|
||||
)
|
||||
|
|
@ -45,7 +46,7 @@ from litellm.integrations.otel import ( # noqa: E402
|
|||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
|
||||
from litellm.integrations.otel.model.config import CaptureMessageContent, ExporterSpec # noqa: E402
|
||||
from litellm.integrations.otel.model.config import CaptureMessageContent, ExporterOwner, ExporterSpec # noqa: E402
|
||||
from litellm.integrations.otel.model.destination import OtelDestination # noqa: E402
|
||||
from litellm.integrations.otel.model.spans import ( # noqa: E402
|
||||
LITELLM_PROXY_REQUEST_SPAN_NAME,
|
||||
|
|
@ -156,6 +157,37 @@ def _emit_llm(logger, kwargs=None, *, ambient=None, fail=False):
|
|||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("call_metadata", "in_baggage", "expected"),
|
||||
[
|
||||
({"user.id": "named-in-call"}, "in-baggage", "named-in-call"),
|
||||
({}, "in-baggage", "in-baggage"),
|
||||
({}, None, None),
|
||||
],
|
||||
ids=["the-call-wins", "baggage-when-the-call-names-nobody", "neither"],
|
||||
)
|
||||
def test_an_sdk_call_names_its_user_in_the_call_or_in_the_baggage_around_it(
|
||||
monkeypatch, call_metadata, in_baggage, expected
|
||||
):
|
||||
exporter = InMemorySpanExporter()
|
||||
monkeypatch.setitem(providers._EXPORTER_FACTORIES, "sdk_mixpanel", lambda _spec: exporter)
|
||||
cfg = OpenTelemetryV2Config(
|
||||
mapper_names=["genai", "mixpanel"],
|
||||
exporters=[ExporterSpec(kind="sdk_mixpanel", owner=ExporterOwner.MIXPANEL)],
|
||||
)
|
||||
logger = OpenTelemetryV2(config=cfg, tracer_provider=providers.build_tracer_provider(cfg, use_simple_processor=True))
|
||||
|
||||
def emit():
|
||||
if in_baggage is not None:
|
||||
attach(baggage.set_baggage("user.id", in_baggage))
|
||||
_emit_llm(logger, {**_kwargs(), "litellm_params": {"metadata": call_metadata}})
|
||||
|
||||
contextvars.copy_context().run(emit)
|
||||
|
||||
model_calls = [span for span in exporter.get_finished_spans() if "gen_ai.operation.name" in span.attributes]
|
||||
assert [span.attributes.get("user.id") for span in model_calls] == [expected]
|
||||
|
||||
|
||||
def test_to_ns_handles_datetime_and_float():
|
||||
dt = datetime(2026, 5, 26, 12, 0, 0, tzinfo=timezone.utc)
|
||||
assert to_ns(dt) == int(dt.timestamp() * 1e9)
|
||||
|
|
|
|||
|
|
@ -300,7 +300,7 @@ def test_mixpanel_preset_reads_token_and_defaults_to_us(monkeypatch):
|
|||
assert spec.requires_headers is True
|
||||
assert "genai" in cfg.mapper_names
|
||||
assert "mixpanel" in cfg.mapper_names
|
||||
assert "user.id" in cfg.baggage_promoted_keys
|
||||
assert "user.id" not in cfg.baggage_promoted_keys
|
||||
|
||||
|
||||
def test_mixpanel_preset_region_and_endpoint_override(monkeypatch):
|
||||
|
|
@ -334,3 +334,33 @@ def test_mixpanel_preset_without_token_is_header_gated(monkeypatch):
|
|||
spec = next(e for e in cfg.exporters if e.owner == ExporterOwner.MIXPANEL)
|
||||
assert spec.headers is None
|
||||
assert spec.requires_headers is True
|
||||
|
||||
|
||||
def test_mixpanel_without_a_token_or_collector_stays_on_v2_and_exports_nowhere(monkeypatch):
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2
|
||||
from litellm.litellm_core_utils.litellm_logging import _exports_nowhere, _maybe_construct_otel_v2
|
||||
|
||||
monkeypatch.delenv("MIXPANEL_PROJECT_TOKEN", raising=False)
|
||||
monkeypatch.delenv("OTEL_EXPORTER", raising=False)
|
||||
monkeypatch.delenv("OTEL_ENDPOINT", raising=False)
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
try:
|
||||
logger = _maybe_construct_otel_v2("mixpanel", [])
|
||||
finally:
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
assert isinstance(logger, OpenTelemetryV2)
|
||||
assert logger.callback_name == "mixpanel"
|
||||
assert _exports_nowhere(logger.config)
|
||||
|
||||
|
||||
def test_mixpanel_with_a_token_keeps_the_operators_collector_beside_its_own_exporter(monkeypatch):
|
||||
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.presets.mixpanel import mixpanel_preset
|
||||
|
||||
monkeypatch.setenv("MIXPANEL_PROJECT_TOKEN", "mp-token")
|
||||
collector = ExporterSpec(kind="otlp_http", endpoint="http://collector.local")
|
||||
cfg = mixpanel_preset(config_overrides=OpenTelemetryV2Config(exporters=[collector]))
|
||||
assert [spec.owner for spec in cfg.exporters] == [None, ExporterOwner.MIXPANEL]
|
||||
assert "user.id" not in cfg.baggage_promoted_keys
|
||||
|
|
|
|||
|
|
@ -992,8 +992,8 @@ def test_langfuse_root_trace_attributes_derive_the_name_only_when_the_caller_gav
|
|||
def test_mixpanel_mapper_stamps_user_id_and_omits_when_absent():
|
||||
from litellm.integrations.otel.model.payloads import MCPToolCallSpanData
|
||||
|
||||
with_end_user = _llm_call(identity=RequestIdentity(call_id="c1", end_user="user_123"))
|
||||
assert MixpanelMapper().map(with_end_user) == {"user.id": "user_123"}
|
||||
with_named_user = _llm_call(identity=RequestIdentity(call_id="c1", user_id="user_123"))
|
||||
assert MixpanelMapper().map(with_named_user) == {"user.id": "user_123"}
|
||||
|
||||
with_trace_user = _llm_call(trace=TraceControls(user_id="trace_user_456"))
|
||||
assert MixpanelMapper().map(with_trace_user) == {"user.id": "trace_user_456"}
|
||||
|
|
@ -1010,9 +1010,9 @@ def test_mixpanel_mapper_stamps_user_id_and_omits_when_absent():
|
|||
result_json=None,
|
||||
error=None,
|
||||
response_cost=None,
|
||||
identity=RequestIdentity(call_id="c2", end_user="mcp_user"),
|
||||
identity=RequestIdentity(call_id="c2", user_id="mcp_user"),
|
||||
)
|
||||
assert MixpanelMapper().map(mcp_span) == {"user.id": "mcp_user"}
|
||||
assert MixpanelMapper().map(mcp_span) == {}
|
||||
assert MixpanelMapper().map(_llm_call()) == {}
|
||||
|
||||
from litellm.integrations.otel.model.metadata import LLMCallEvent
|
||||
|
|
@ -1023,5 +1023,5 @@ def test_mixpanel_mapper_stamps_user_id_and_omits_when_absent():
|
|||
"standard_logging_object": {"id": "c3", "model": "gpt-4o", "metadata": {}},
|
||||
}
|
||||
)
|
||||
assert sdk_event.user_id_metadata == {"user.id": "sdk_user_789"}
|
||||
assert LLMCallEvent.from_dict({}).user_id_metadata == {}
|
||||
assert sdk_event.caller_user_id == "sdk_user_789"
|
||||
assert LLMCallEvent.from_dict({}).caller_user_id is None
|
||||
|
|
|
|||
|
|
@ -9622,6 +9622,38 @@ def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch):
|
|||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
|
||||
def test_mixpanel_dispatch_needs_otel_v2_and_builds_one_logger(monkeypatch):
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.litellm_core_utils import litellm_logging as logging_module
|
||||
|
||||
def built() -> object:
|
||||
return logging_module.build_custom_logger_compatible_class(
|
||||
"mixpanel", internal_usage_cache=None, llm_router=None
|
||||
)
|
||||
|
||||
logging_module._in_memory_loggers.clear()
|
||||
monkeypatch.setenv("MIXPANEL_PROJECT_TOKEN", "test-token")
|
||||
monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
try:
|
||||
with pytest.raises(ValueError, match="LITELLM_OTEL_V2"):
|
||||
built()
|
||||
assert logging_module.get_custom_logger_compatible_class("mixpanel") is None
|
||||
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
logger = built()
|
||||
assert isinstance(logger, OpenTelemetryV2)
|
||||
assert logger.callback_name == "mixpanel"
|
||||
assert built() is logger
|
||||
assert logging_module.get_custom_logger_compatible_class("mixpanel") is logger
|
||||
finally:
|
||||
logging_module._in_memory_loggers.clear()
|
||||
monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
|
||||
def test_newrelic_dispatch_keeps_legacy_agent_when_flag_off(monkeypatch):
|
||||
from litellm.integrations.newrelic import NewRelicLogger
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
|
|
|
|||
|
|
@ -2619,6 +2619,12 @@ class TestGetEndUserIdDropsMalformedBodyValues:
|
|||
|
||||
assert result is None
|
||||
|
||||
def test_metadata_user_dot_id_is_only_a_span_label(self):
|
||||
with patch("litellm.proxy.proxy_server.general_settings", {}):
|
||||
result = get_end_user_id_from_request_body(request_body={"metadata": {"user.id": "label"}}, request_headers={})
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_dict_metadata_user_id_returns_none(self):
|
||||
request_body = {
|
||||
"metadata": {"user_id": {"device_id": "abc"}},
|
||||
|
|
|
|||
|
|
@ -281,6 +281,39 @@ async def test_aaauser_personal_budgets(key_ownership):
|
|||
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("body", "headers"),
|
||||
[
|
||||
({"metadata": {"user.id": "named-caller"}}, []),
|
||||
({"user": "named-caller"}, []),
|
||||
({}, [(b"x-litellm-customer-id", b"named-caller")]),
|
||||
],
|
||||
ids=["metadata", "user-field", "customer-header"],
|
||||
)
|
||||
async def test_auth_names_the_caller_for_the_spans_of_the_request(
|
||||
monkeypatch: pytest.MonkeyPatch, body: dict[str, object], headers: list[tuple[bytes, bytes]]
|
||||
) -> None:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.integrations.otel.plumbing.context import request_user_id
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", None)
|
||||
request: Final = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": headers})
|
||||
request._url = URL(url="/chat/completions")
|
||||
request._body = json.dumps(body).encode()
|
||||
|
||||
await user_api_key_auth(request=request, api_key="Bearer " + MASTER_KEY)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
assert request_user_id() == "named-caller"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("prohibited_param", ["api_base", "base_url"])
|
||||
async def test_user_api_key_auth_fails_with_prohibited_params(prohibited_param):
|
||||
|
|
|
|||
|
|
@ -520,10 +520,89 @@ class TestMixpanelCallbackValidation:
|
|||
redacted = _redact_callback_secrets(raw_metadata)
|
||||
assert redacted["logging"][0]["callback_vars"]["mixpanel_project_token"] == "***REDACTED***"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"metadata",
|
||||
[
|
||||
{"logging": [{"callback_name": "mixpanel", "callback_vars": {"mixpanel_project_token": "***REDACTED***"}}]},
|
||||
{"logging": [{"callback_name": "mixpanel", "callback_vars": {"mixpanel_region": "eu"}}]},
|
||||
{"callback_settings": {"callback_vars": {"mixpanel_project_token": " "}}},
|
||||
],
|
||||
ids=["masked-logging-entry", "region-without-token", "blank-callback-settings"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_writers_refuse_a_mixpanel_token_they_could_never_send(self, metadata, monkeypatch):
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.proxy._types import NewTeamRequest, ProxyException, UpdateTeamRequest, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.team_endpoints import new_team, update_team
|
||||
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
admin = UserAPIKeyAuth(user_role="proxy_admin")
|
||||
request = MagicMock(spec=Request)
|
||||
try:
|
||||
with pytest.raises(ProxyException) as created:
|
||||
await new_team(data=NewTeamRequest(metadata=metadata), http_request=request, user_api_key_dict=admin)
|
||||
with pytest.raises(ProxyException) as updated:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="team", metadata=metadata),
|
||||
http_request=request,
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
finally:
|
||||
monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
for refusal in (created.value, updated.value):
|
||||
assert str(refusal.code) == "400"
|
||||
assert "mixpanel_project_token must be" in refusal.message
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"metadata",
|
||||
[
|
||||
None,
|
||||
{},
|
||||
{"logging": "not-a-list"},
|
||||
{"logging": [{"callback_name": "newrelic", "callback_vars": {"newrelic_license_key": ""}}]},
|
||||
{"callback_settings": {"callback_vars": {"langfuse_public_key": "pk"}}},
|
||||
],
|
||||
ids=["none", "empty", "not-a-list", "another-callback", "another-callback-settings"],
|
||||
)
|
||||
@pytest.mark.parametrize("otel_v2", ["true", "false"])
|
||||
def test_the_team_check_leaves_other_callbacks_alone(self, metadata, otel_v2, monkeypatch):
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.proxy.common_utils.callback_config_validation import mixpanel_logging_metadata_error
|
||||
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", otel_v2)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
try:
|
||||
assert mixpanel_logging_metadata_error(metadata) is None
|
||||
finally:
|
||||
monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
def test_the_team_check_accepts_a_real_mixpanel_token(self, monkeypatch):
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.proxy.common_utils.callback_config_validation import mixpanel_logging_metadata_error
|
||||
|
||||
real_token = {
|
||||
"logging": [{"callback_name": "mixpanel", "callback_vars": {"mixpanel_project_token": "real"}}],
|
||||
"callback_settings": {"callback_vars": {"mixpanel_project_token": "real"}},
|
||||
}
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
try:
|
||||
assert mixpanel_logging_metadata_error(real_token) is None
|
||||
finally:
|
||||
monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
def test_team_and_key_validation(self, monkeypatch):
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.proxy.common_utils.callback_config_validation import callback_config_error
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import raise_on_invalid_key_logging_config
|
||||
from litellm.proxy.management_endpoints.team_callback_endpoints import _validate_team_callback
|
||||
|
||||
|
|
@ -533,15 +612,36 @@ class TestMixpanelCallbackValidation:
|
|||
with pytest.raises(HTTPException) as exc:
|
||||
_validate_team_callback(self._team_data({"mixpanel_project_token": "tok"}))
|
||||
assert "LITELLM_OTEL_V2" in str(exc.value.detail)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_validate_team_callback(self._team_data({}))
|
||||
assert "LITELLM_OTEL_V2" in str(exc.value.detail)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
raise_on_invalid_key_logging_config(self._key_metadata({}))
|
||||
assert "LITELLM_OTEL_V2" in str(exc.value.detail)
|
||||
assert callback_config_error("langfuse", {}) is None
|
||||
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
_validate_team_callback(self._team_data({}))
|
||||
raise_on_invalid_key_logging_config(self._key_metadata({}))
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_validate_team_callback(self._team_data({"mixpanel_project_token": "tok", "mixpanel_region": "mars"}))
|
||||
assert "Unknown mixpanel_region" in str(exc.value.detail)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
raise_on_invalid_key_logging_config(self._key_metadata({"mixpanel_region": "eu"}))
|
||||
assert "requires mixpanel_project_token" in str(exc.value.detail)
|
||||
assert "mixpanel_project_token must be" in str(exc.value.detail)
|
||||
for unsendable in ("", " ", "None", "***REDACTED***"):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_validate_team_callback(self._team_data({"mixpanel_project_token": unsendable}))
|
||||
assert "mixpanel_project_token must be" in str(exc.value.detail)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
raise_on_invalid_key_logging_config(self._key_metadata({"mixpanel_project_token": unsendable}))
|
||||
assert "mixpanel_project_token must be" in str(exc.value.detail)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
raise_on_invalid_key_logging_config(
|
||||
self._key_metadata({"mixpanel_project_token": "tok", "mixpanel_foo": "bar"})
|
||||
)
|
||||
assert "mixpanel_foo" in str(exc.value.detail)
|
||||
_validate_team_callback(self._team_data({"capture_message_content": "span_only"}))
|
||||
_validate_team_callback(self._team_data({"mixpanel_project_token": "tok", "mixpanel_region": "in"}))
|
||||
raise_on_invalid_key_logging_config(
|
||||
|
|
|
|||
|
|
@ -8969,6 +8969,33 @@ def test_vertex_sends_exactly_one_authorization_header():
|
|||
assert _authorization_values(vertex_request_headers) == [GOOGLE_ACCESS_TOKEN]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixpanel_vars_under_another_callback_never_reach_the_trusted_field():
|
||||
key_with_langfuse_callback = UserAPIKeyAuth(
|
||||
api_key="hashed-key",
|
||||
metadata={
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "langfuse",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"langfuse_public_key": "pk", "mixpanel_project_token": "tok"},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}]},
|
||||
request=_callback_credential_request_mock(),
|
||||
user_api_key_dict=key_with_langfuse_callback,
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
|
||||
assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"langfuse_public_key": "pk"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_newrelic_team_callback_vars_reach_trusted_field():
|
||||
"""A key with a newrelic team callback stamps its vars into the proxy-owned
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue