diff --git a/litellm/__init__.py b/litellm/__init__.py index 8e5dd68a186..0a007def7c8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -156,6 +156,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "compression_interception", "newrelic", "signoz", + "mixpanel", ] cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 38928f67f42..cb74c16cb82 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -309,6 +309,27 @@ }, "description": "Langsmith Logging Integration" }, + { + "id": "mixpanel", + "displayName": "Mixpanel", + "logo": "mixpanel.svg", + "supports_key_team_logging": true, + "dynamic_params": { + "mixpanel_project_token": { + "type": "password", + "ui_name": "Mixpanel Project Token", + "description": "Mixpanel project token for this team or key. Traces export to this project over OTLP.", + "required": false + }, + "mixpanel_region": { + "type": "text", + "ui_name": "Mixpanel Region (us, eu, or in)", + "description": "Data residency region for this project (us, eu, or in). Defaults to us.", + "required": false + } + }, + "description": "Mixpanel Agent Intelligence OpenTelemetry Integration. Setup: https://docs.mixpanel.com/docs/agent-intelligence" + }, { "id": "newrelic", "displayName": "New Relic", diff --git a/litellm/integrations/otel/README.md b/litellm/integrations/otel/README.md index 06c2da117b5..28ef489cfc5 100644 --- a/litellm/integrations/otel/README.md +++ b/litellm/integrations/otel/README.md @@ -322,7 +322,7 @@ nothing here imports outside it: global set with `excluded_services` in their own `callback_vars`; `"none"` means no datastore spans are excluded, while an omitted value follows the global set. These destination overrides require OTel v2 and are accepted only - on `langfuse_otel`, `arize`, `weave_otel` and `newrelic` callbacks. + on `langfuse_otel`, `arize`, `weave_otel`, `newrelic` and `mixpanel` callbacks. - [`baggage.py`](./model/baggage.py) — the single definition of which request-identity values are promoted into Baggage (so child spans inherit them) and under which attribute keys. @@ -341,7 +341,7 @@ nothing here imports outside it: - `genai` — the canonical OpenTelemetry GenAI vocabulary, always present. - `legacy` — an additional vocabulary using the older semconv-ai / Traceloop attribute key names, for backends that read those. - - `openinference`, `langfuse`, `weave`, `langtrace` — vendor vocabularies. + - `openinference`, `langfuse`, `weave`, `langtrace`, `mixpanel` — vendor vocabularies. - `resolve_mappers(names)` turns config names into mapper instances. ### Plumbing (`plumbing/`) diff --git a/litellm/integrations/otel/__init__.py b/litellm/integrations/otel/__init__.py index d7627d4d63d..424ea8a3a4a 100644 --- a/litellm/integrations/otel/__init__.py +++ b/litellm/integrations/otel/__init__.py @@ -60,6 +60,7 @@ from litellm.integrations.otel.model.semconv import ( NetworkTransport, RpcSystem, Server, + User, resolve_operation, resolve_output_type, resolve_provider, @@ -113,6 +114,7 @@ __all__ = [ "SpanError", "SpanRole", "SpanSpec", + "User", "db_system", "is_mcp_list_tools", "is_mcp_tool_call", diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index b5b65d299e0..d6546dbb541 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -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,6 +602,7 @@ class OpenTelemetryV2(CustomLogger): trace=call.trace, session_id=call.session_id, metadata_keys=tuple(self.config.baggage_metadata_keys), + named_user_id=call.caller_user_id or inbound_user_id(), operation=call.operation, ) end_time_ns: Final = to_ns(end_time) diff --git a/litellm/integrations/otel/mappers/__init__.py b/litellm/integrations/otel/mappers/__init__.py index 4e424e025c0..546f2aeb4c0 100644 --- a/litellm/integrations/otel/mappers/__init__.py +++ b/litellm/integrations/otel/mappers/__init__.py @@ -18,6 +18,7 @@ from litellm.integrations.otel.mappers.genai import GenAIMapper from litellm.integrations.otel.mappers.langfuse import LangfuseMapper from litellm.integrations.otel.mappers.langtrace import LangtraceMapper from litellm.integrations.otel.mappers.legacy import LegacyMapper +from litellm.integrations.otel.mappers.mixpanel import MixpanelMapper from litellm.integrations.otel.mappers.openinference import OpenInferenceMapper from litellm.integrations.otel.mappers.utils import tool_attr_budget from litellm.integrations.otel.mappers.weave import WeaveMapper @@ -35,6 +36,7 @@ _PLAIN_MAPPERS: Final[dict[str, Callable[[], AttributeMapper]]] = { "langfuse": LangfuseMapper, "weave": WeaveMapper, "langtrace": LangtraceMapper, + "mixpanel": MixpanelMapper, } @@ -63,6 +65,7 @@ __all__ = [ "LangfuseMapper", "LangtraceMapper", "LegacyMapper", + "MixpanelMapper", "OpenInferenceMapper", "WeaveMapper", "resolve_mappers", diff --git a/litellm/integrations/otel/mappers/mixpanel.py b/litellm/integrations/otel/mappers/mixpanel.py new file mode 100644 index 00000000000..801009e80ee --- /dev/null +++ b/litellm/integrations/otel/mappers/mixpanel.py @@ -0,0 +1,13 @@ +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 +from litellm.integrations.otel.model.semconv import User + + +class MixpanelMapper: + def map(self, data: SpanData) -> AttributeMap: + match data: + case LLMCallSpanData(): + return drop_none_pairs(((User.ID, data.identity.user_id or data.trace.user_id or None),)) + case _: + return drop_none_pairs(()) diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index b50e927f977..e5c2c00c752 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -44,6 +44,7 @@ class ExporterOwner(str, Enum): AGENTOPS = "agentops" NEWRELIC = "newrelic" SIGNOZ = "signoz" + MIXPANEL = "mixpanel" class _OTelV2Flag(BaseSettings): diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index 28d6ad54607..725aaccf17e 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -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,6 +73,9 @@ IDENTITY_METADATA_FIELDS: Final = ( REQUESTER_METADATA_KEY: Final = "requester_metadata" REQUESTER_METADATA_PATH: Final = f"{REQUESTER_METADATA_KEY}." +_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]: @@ -100,9 +105,14 @@ class RequestIdentity: provider_model: str | None = None request_route: str | None = None metadata: Mapping[str, str] = field(default_factory=dict) + user_id: str | None = None @classmethod - def from_payload(cls, payload: StandardLoggingPayload) -> RequestIdentity: + def from_payload( + cls, + payload: StandardLoggingPayload, + named_user_id: str | None = None, + ) -> RequestIdentity: """Parse caller identity out of a closed request's payload metadata. ``provider_model`` is resolved here too (see :func:`resolve_provider_model`) @@ -111,6 +121,8 @@ class RequestIdentity: """ raw_meta: Final = cast(Mapping[str, object], payload.get("metadata") or {}) 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``; @@ -120,10 +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")), + 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 @@ -188,7 +206,11 @@ class RequestContext: return self.identity.provider_model @classmethod - def from_standard_logging_payload(cls, payload: StandardLoggingPayload) -> RequestContext: + def from_standard_logging_payload( + cls, + payload: StandardLoggingPayload, + 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 {}) raw_response: Final = payload.get("response") @@ -203,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), + identity=RequestIdentity.from_payload(payload, named_user_id=named_user_id), ) @@ -271,6 +293,7 @@ class LLMCallEvent: completion_start_seconds: float | None trace: TraceControls session_id: str | None + caller_user_id: str | None = None @property def reached_a_provider(self) -> bool: @@ -300,9 +323,66 @@ class LLMCallEvent: completion_start_seconds=_epoch_seconds(kwargs.get("completion_start_time")), trace=trace, session_id=caller_session_id(kwargs, trace), + caller_user_id=caller_user_id_in_call(kwargs), ) +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 + ) + + +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: return to_seconds(value) if isinstance(value, (datetime, float, int, str)) and not isinstance(value, bool) else None diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 4238d72fc5e..31602c6c827 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -467,6 +467,7 @@ class LLMCallSpanData: session_id: str | None = None, *, metadata_keys: tuple[str, ...] = (), + named_user_id: str | None = None, operation: GenAIOperation | None = None, ) -> LLMCallSpanData: params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {}) @@ -474,7 +475,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) + 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) diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 8d186a2715a..5f9368d2582 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -255,6 +255,10 @@ class Server: PORT: Final = "server.port" +class User: + ID: Final = "user.id" + + class DB: """Database / cache client-span keys (OTel ``db.*`` semconv). diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index b328b281dde..439d65f17e6 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -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() diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index b5a5b85fc11..aec4b780324 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -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 @@ -152,7 +154,7 @@ def _otlp_traces_endpoint(endpoint: str | None) -> str | None: return endpoint endpoint = endpoint.rstrip("/") # Splunk Observability uses ``/v2/trace/otlp``; never rewrite it. - if endpoint.endswith("/v1/traces") or "/v2/trace/otlp" in endpoint: + if endpoint.endswith(("/v1/traces", "/v1/ai/otel/traces")) or "/v2/trace/otlp" in endpoint: return endpoint for other_signal in ("/v1/logs", "/v1/metrics"): if endpoint.endswith(other_signal): @@ -486,25 +488,46 @@ 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. + + That id replaces a ``user.id`` already on the span, such as what a mapper wrote for another backend, + so the whole request names one caller. 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() + if user_id is None: + return span + attributes: Final = span.attributes or _NO_ATTRIBUTES + 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 +713,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 +819,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 +1243,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 +1563,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 diff --git a/litellm/integrations/otel/presets/__init__.py b/litellm/integrations/otel/presets/__init__.py index 7c891c29409..b962366e547 100644 --- a/litellm/integrations/otel/presets/__init__.py +++ b/litellm/integrations/otel/presets/__init__.py @@ -21,6 +21,11 @@ from litellm.integrations.otel.presets.langfuse import ( ) from litellm.integrations.otel.presets.langtrace import langtrace_preset from litellm.integrations.otel.presets.levo import levo_preset +from litellm.integrations.otel.presets.mixpanel import ( + mixpanel_dynamic_endpoint, + mixpanel_dynamic_headers, + mixpanel_preset, +) from litellm.integrations.otel.presets.newrelic import ( newrelic_dynamic_endpoint, newrelic_dynamic_headers, @@ -48,6 +53,7 @@ PRESET_BY_CALLBACK: Final[Mapping[str, Preset]] = MappingProxyType( "langfuse_otel": langfuse_preset, "langtrace": langtrace_preset, "levo": levo_preset, + "mixpanel": mixpanel_preset, "newrelic": newrelic_preset, "signoz": signoz_preset, "weave_otel": weave_preset, @@ -63,6 +69,7 @@ DYNAMIC_HEADERS_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynami { "arize": arize_dynamic_headers, "langfuse_otel": langfuse_dynamic_headers, + "mixpanel": mixpanel_dynamic_headers, "newrelic": newrelic_dynamic_headers, "signoz": signoz_dynamic_headers, "weave_otel": weave_dynamic_headers, @@ -77,6 +84,7 @@ DYNAMIC_HEADERS_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynami DYNAMIC_ENDPOINT_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], str | None]]] = ( MappingProxyType( { + "mixpanel": mixpanel_dynamic_endpoint, "newrelic": newrelic_dynamic_endpoint, "signoz": signoz_dynamic_endpoint, } @@ -158,6 +166,7 @@ __all__ = [ "langfuse_preset", "langtrace_preset", "levo_preset", + "mixpanel_preset", "newrelic_preset", "phoenix_preset", "project_routing_headers", diff --git a/litellm/integrations/otel/presets/destinations.py b/litellm/integrations/otel/presets/destinations.py index 984b385a647..3b23827fe94 100644 --- a/litellm/integrations/otel/presets/destinations.py +++ b/litellm/integrations/otel/presets/destinations.py @@ -80,6 +80,13 @@ def _newrelic_destination(params: StandardCallbackDynamicParams) -> "_Destinatio return (endpoint, None) if endpoint else None +def _mixpanel_destination(params: StandardCallbackDynamicParams) -> "_Destination | None": + from litellm.integrations.otel.presets.mixpanel import mixpanel_dynamic_endpoint + + endpoint: Final = mixpanel_dynamic_endpoint(params) + return (endpoint, None) if endpoint else None + + #: Callback name -> destination resolver. A backend is destination-capable exactly #: when it appears here AND in ``DYNAMIC_HEADERS_BY_CALLBACK``: without a header #: builder the destination would carry no tenant credentials, and the exporter @@ -91,6 +98,7 @@ _DESTINATION_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicPa "arize": _arize_destination, "weave_otel": _weave_destination, "newrelic": _newrelic_destination, + "mixpanel": _mixpanel_destination, } ) ) @@ -105,6 +113,7 @@ _REQUIRED_HEADERS_BY_CALLBACK: Final[Mapping[str, frozenset[str]]] = MappingProx "arize": frozenset({"arize-space-id", "api_key"}), "weave_otel": frozenset({"Authorization", "project_id"}), "newrelic": frozenset({"api-key"}), + "mixpanel": frozenset({"Authorization"}), } ) diff --git a/litellm/integrations/otel/presets/mixpanel.py b/litellm/integrations/otel/presets/mixpanel.py new file mode 100644 index 00000000000..aad9178fd22 --- /dev/null +++ b/litellm/integrations/otel/presets/mixpanel.py @@ -0,0 +1,105 @@ +import base64 +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from litellm._logging import verbose_logger +from litellm.integrations.otel.model.config import ( + ExporterOwner, + ExporterSpec, + OpenTelemetryV2Config, +) +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" + +MIXPANEL_OTLP_ENDPOINT_BY_REGION: Final[Mapping[str, str]] = MappingProxyType( + { + "us": "https://ingestion-us.mixpanel.com/v1/ai/otel/traces", + "eu": "https://ingestion-eu.mixpanel.com/v1/ai/otel/traces", + "in": "https://ingestion-in.mixpanel.com/v1/ai/otel/traces", + } +) + +_DEFAULT_REGION: Final = "us" + + +class _MixpanelSettings(BaseSettings): + model_config = SettingsConfigDict(case_sensitive=False, extra="ignore", frozen=True) + + project_token: str | None = Field(default=None, validation_alias=MIXPANEL_PROJECT_TOKEN_ENV) + region: str | None = Field(default=None, validation_alias="MIXPANEL_REGION") + endpoint: str | None = Field(default=None, validation_alias="MIXPANEL_ENDPOINT") + + +def mixpanel_authorization_header(project_token: str) -> str: + encoded: Final = base64.b64encode(f"{project_token}:".encode()).decode() + return f"Basic {encoded}" + + +def _resolve_region_endpoint(region: str | None) -> str: + default_endpoint: Final = MIXPANEL_OTLP_ENDPOINT_BY_REGION[_DEFAULT_REGION] + if not region: + return default_endpoint + endpoint: Final = MIXPANEL_OTLP_ENDPOINT_BY_REGION.get(region.strip().lower()) + if endpoint is None: + verbose_logger.warning( + "Mixpanel: unknown mixpanel_region %r; supported regions: %s. Using the default (US) endpoint.", + region, + ", ".join(sorted(MIXPANEL_OTLP_ENDPOINT_BY_REGION)), + ) + return default_endpoint + return endpoint + + +def mixpanel_preset( + *, + config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, +) -> OpenTelemetryV2Config: + settings: Final = _MixpanelSettings() + 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) + 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": exporters, + "mapper_names": ensure_mappers(base.mapper_names, "genai", "mixpanel"), + } + ) + ) + + +def mixpanel_dynamic_headers( + params: StandardCallbackDynamicParams, +) -> dict[str, str]: # mutable-ok: DYNAMIC_HEADERS_BY_CALLBACK returns a dict + token: Final = (params.get("mixpanel_project_token") or "").strip() + if not token: + return {} + return {"Authorization": mixpanel_authorization_header(token)} + + +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")) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 74badc2c8c1..67ab0e95bb4 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -97,6 +97,7 @@ class CustomLoggerRegistry: "weave_otel": OpenTelemetry, "levo": OpenTelemetry, "signoz": OpenTelemetry, + "mixpanel": OpenTelemetry, "mlflow": MlflowLogger, "langfuse": LangfusePromptManagement, "otel": OpenTelemetry, diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index c3769686e45..f9e67f123c1 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -145,6 +145,8 @@ _supported_callback_params: Final[tuple[str, ...]] = ( "newrelic_region", "signoz_ingestion_endpoint", "signoz_ingestion_key", + "mixpanel_project_token", + "mixpanel_region", "turn_off_message_logging", ) @@ -160,6 +162,8 @@ _request_blocked_callback_params: Final = frozenset( "newrelic_region", "signoz_ingestion_endpoint", "signoz_ingestion_key", + "mixpanel_project_token", + "mixpanel_region", } ) @@ -174,6 +178,8 @@ _trusted_overlay_callback_params: Final = frozenset( "newrelic_region", "signoz_ingestion_endpoint", "signoz_ingestion_key", + "mixpanel_project_token", + "mixpanel_region", } ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9b6ec8d28d6..14d5422568d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6159,6 +6159,14 @@ def build_custom_logger_compatible_class( _in_memory_loggers.append(_signoz_logger) return _signoz_logger + elif logging_integration == "mixpanel": + from litellm.integrations.otel.model.config import is_otel_v2_enabled + + if not is_otel_v2_enabled(): + raise ValueError("The 'mixpanel' callback requires OpenTelemetry v2 (set LITELLM_OTEL_V2=true)") + + return _maybe_construct_otel_v2("mixpanel", _in_memory_loggers) + elif logging_integration == "mlflow": for callback in _in_memory_loggers: if isinstance(callback, MlflowLogger): @@ -6766,6 +6774,12 @@ def get_custom_logger_compatible_class( return callback if isinstance(callback, NewRelicLogger): return callback + elif logging_integration == "mixpanel": + from litellm.integrations.otel.logger import OpenTelemetryV2 + + for callback in _in_memory_loggers: + if isinstance(callback, OpenTelemetryV2) and callback.callback_name == "mixpanel": + return callback return None except Exception as e: diff --git a/litellm/proxy/_experimental/out/assets/logos/mixpanel.svg b/litellm/proxy/_experimental/out/assets/logos/mixpanel.svg new file mode 100644 index 00000000000..71f9e91877a --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/mixpanel.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e53b5d4ad52..72255950727 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3746,6 +3746,12 @@ class AllCallbacks(LiteLLMPydanticObjectBase): ], ) + mixpanel: CallbackOnUI = CallbackOnUI( + litellm_callback_name="mixpanel", + ui_callback_name="Mixpanel", + litellm_callback_params=("MIXPANEL_PROJECT_TOKEN", "MIXPANEL_REGION", "MIXPANEL_ENDPOINT"), + ) + signoz: CallbackOnUI = CallbackOnUI( litellm_callback_name="signoz", ui_callback_name="SigNoz", diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a52ad14e293..ff728a418e4 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -4403,6 +4403,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, @@ -4460,6 +4485,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 diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 302978d17a4..9cc449a6cb4 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -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 @@ -14,6 +14,12 @@ from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_V _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" @@ -26,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( @@ -61,6 +70,8 @@ def callback_config_error( ) if destination_setting_error is not None: return destination_setting_error + if callback_name == _MIXPANEL_CALLBACK: + return _mixpanel_config_error(callback_vars) if callback_name != _NEWRELIC_CALLBACK: return None return _newrelic_config_error(callback_vars) @@ -144,6 +155,7 @@ _VAR_FAMILIES: Final[Mapping[str, str]] = MappingProxyType( "humanloop_": "Humanloop", "langfuse_": "Langfuse", "langsmith_": "LangSmith", + "mixpanel_": "Mixpanel", "newrelic_": "New Relic", "posthog_": "PostHog", "wandb_": "Weights & Biases", @@ -306,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 @@ -445,3 +480,34 @@ def _newrelic_config_error(callback_vars: Mapping[str, str]) -> str | None: if region is not None and (not api_key or api_key == "None"): return "newrelic_region requires newrelic_api_key; the region rides the team's own key." 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.presets.mixpanel import MIXPANEL_OTLP_ENDPOINT_BY_REGION + + region: Final = callback_vars.get("mixpanel_region") + if region is not None and region.strip().lower() not in MIXPANEL_OTLP_ENDPOINT_BY_REGION: + return ( + f"Unknown mixpanel_region {region!r}. " + f"Supported regions: {', '.join(sorted(MIXPANEL_OTLP_ENDPOINT_BY_REGION))}." + ) + + token: Final = (callback_vars.get("mixpanel_project_token") or "").strip() + if not token 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 diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..e3c220a21e8 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1222,6 +1222,8 @@ def convert_key_logging_metadata_to_callback( continue if var.startswith("signoz_") and data.callback_name != "signoz": continue + if var.startswith("mixpanel_") and data.callback_name != "mixpanel": + continue if team_callback_settings_obj.callback_vars is None: team_callback_settings_obj.callback_vars = {} team_callback_settings_obj.callback_vars[var] = str(value) diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 5b778392664..cbafb63a066 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -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) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index cf2ae392677..d1e512a28ca 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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 @@ -1714,6 +1715,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, @@ -2141,6 +2145,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, @@ -2489,6 +2499,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, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 755a0a8cc7a..c1dbcfa3853 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4103,6 +4103,9 @@ class StandardCallbackDynamicParams(TypedDict, total=False): signoz_ingestion_endpoint: str | None # writable-ok: initialize_standard_callback_dynamic_params assigns it signoz_ingestion_key: str | None # writable-ok: initialize_standard_callback_dynamic_params assigns it + mixpanel_project_token: str | None # writable-ok: initialize_standard_callback_dynamic_params assigns it + mixpanel_region: str | None # writable-ok: initialize_standard_callback_dynamic_params assigns it + # Logging settings turn_off_message_logging: bool | None # when true will not log messages litellm_disabled_callbacks: list[str] | None diff --git a/tests/integration/_support/otlp_sink.py b/tests/integration/_support/otlp_sink.py index 561159fe197..b43626f9598 100644 --- a/tests/integration/_support/otlp_sink.py +++ b/tests/integration/_support/otlp_sink.py @@ -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: diff --git a/tests/integration/observability/test_mixpanel_user_id.py b/tests/integration/observability/test_mixpanel_user_id.py new file mode 100644 index 00000000000..e04cb57f663 --- /dev/null +++ b/tests/integration/observability/test_mixpanel_user_id.py @@ -0,0 +1,1046 @@ +"""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) + + at_project: Final = tuple((caller, _reached(mixpanel, spans, host)) for caller, spans in settled) + assert [caller for caller, spans in at_project if not spans] == [] + assert [(caller, _without_the_id(spans, caller)) for caller, spans in at_project] == [ + (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", " None ", "***REDACTED***", " ***REDACTED*** "], + ids=["empty", "whitespace", "null", "none-string", "padded-none-string", "masked", "padded-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))} + + +def test_the_operators_project_names_one_caller_when_a_request_also_sends_a_trace_user_id( + 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}" + traced_as: Final = f"traced-as-{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-two-ids-{uuid.uuid4().hex}") + trace_id: Final = _traced_call( + proxy, + key, + {"model": model, **_chat_call("hi"), "metadata": {"user.id": caller, "trace_user_id": traced_as}}, + ) + 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 _without_the_id(_outside_auth(at_mixpanel), caller) == [] + assert {span["attributes"].get(USER_ID) for span in at_other if OPERATION in span["attributes"]} == ( + {traced_as} if "langfuse_otel" in callbacks else {None} + ) + + +def test_a_team_project_names_one_caller_when_a_request_also_sends_a_trace_user_id( + combined_setup: tuple[OwnedProxy, tuple[str, ...]], provider: Wire, mixpanel: ConnectSink +) -> None: + owned, _ = 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, host, _token = _team_or_key(scenario, "team") + trace_id: Final = _traced_call( + proxy, + key, + { + "model": model, + **_chat_call("hi"), + "metadata": {"user.id": caller, "trace_user_id": f"traced-as-{uuid.uuid4().hex}"}, + }, + ) + spans: Final = _settled_trace(mixpanel.control_url, trace_id, _has_model_call) + + at_project: Final = _reached(mixpanel, spans, host) + assert [span for span in at_project if OPERATION in span["attributes"]], at_project + assert _without_the_id(at_project, caller) == [] + + +@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 diff --git a/tests/integration/spend/test_end_user_budget_enforcement.py b/tests/integration/spend/test_end_user_budget_enforcement.py index 92259d463b9..40abdf54d4f 100644 --- a/tests/integration/spend/test_end_user_budget_enforcement.py +++ b/tests/integration/spend/test_end_user_budget_enforcement.py @@ -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 diff --git a/tests/unit/integrations/otel/test_otel_v2_baggage.py b/tests/unit/integrations/otel/test_otel_v2_baggage.py index 5855cb60aca..4024df2b310 100644 --- a/tests/unit/integrations/otel/test_otel_v2_baggage.py +++ b/tests/unit/integrations/otel/test_otel_v2_baggage.py @@ -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, @@ -254,3 +265,117 @@ def test_baggage_processor_allowlist_can_be_widened(): assert span.attributes.get("custom.key") == "v" # team_alias not in this config's allowlist -> not promoted assert LiteLLM.TEAM_ALIAS not in span.attributes + + +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) + mp_cfg = mixpanel_preset(config_overrides=OpenTelemetryV2Config(exporter="in_memory")) + + 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(): + payload = _payload() + payload["end_user"] = "from-end-user" + payload["metadata"]["user_api_key_user_id"] = "from-key-owner" + payload["metadata"]["requester_metadata"] = {"user.id": "from-explicit-user-dot-id"} + data = LLMCallSpanData.from_standard_logging_payload(payload) + assert data.identity.user_id == "from-explicit-user-dot-id" + + key_owner_only_payload = _payload() + 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 diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index 6dd9ff35026..202e359abe5 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -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, @@ -5876,3 +5879,228 @@ class TestLangfuseRequestMetadataStaysWithLangfuse: assert [(span.dropped_attributes, span.dropped_events, span.dropped_links) for span in stripped] == [ (1, 1, 1) ] * 2 + + +def test_mixpanel_destination_resolves_endpoint_and_basic_auth(): + assert "mixpanel" in destination_capable_backends() + assert destination_for("mixpanel", {"mixpanel_region": "eu"}) is None + destination = destination_for("mixpanel", {"mixpanel_project_token": "team-tok", "mixpanel_region": "eu"}) + assert destination is not None + 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_read_at_auth_replaces_the_id_a_mapper_wrote(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") == "caller-1" + assert by_name(mixpanel)["POST /v1/chat/completions"].attributes.get("user.id") == "caller-1" + + def test_a_span_from_other_code_names_the_caller_only_for_mixpanel(self) -> None: + provider, mixpanel, other = self._fan_out() + + def run() -> None: + set_request_destinations((self.MIXPANEL, self.ARIZE)) + set_request_user_id("caller-1") + get_tracer(provider, "application").start_span("checkout", attributes={"user.id": "its-own"}).end() + + in_fresh_context(run) + + assert by_name(mixpanel)["checkout"].attributes.get("user.id") == "caller-1" + assert by_name(other)["checkout"].attributes.get("user.id") == "its-own" + + 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() diff --git a/tests/unit/integrations/otel/test_otel_v2_dynamic.py b/tests/unit/integrations/otel/test_otel_v2_dynamic.py index 8163ba06317..78e33ef092e 100644 --- a/tests/unit/integrations/otel/test_otel_v2_dynamic.py +++ b/tests/unit/integrations/otel/test_otel_v2_dynamic.py @@ -743,3 +743,34 @@ def test_signoz_keyless_team_endpoint_is_ignored_so_the_operator_key_never_reach owned = next(e for e in routed.exporters if e.owner == "signoz") assert owned.endpoint == "https://ingest.us.signoz.cloud:443" assert owned.headers == "signoz-ingestion-key=OPERATOR" + + +def test_mixpanel_dynamic_headers_and_endpoint_routing(monkeypatch, caplog): + expected_auth = "Basic " + base64.b64encode(b"team-token:").decode() + assert dynamic_otlp_headers("mixpanel", {"mixpanel_project_token": "team-token"}) == { + "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 ( + dynamic_otlp_endpoint("mixpanel", {"mixpanel_project_token": "team-token"}) + == "https://ingestion-us.mixpanel.com/v1/ai/otel/traces" + ) + assert ( + dynamic_otlp_endpoint("mixpanel", {"mixpanel_project_token": "team-token", "mixpanel_region": "eu"}) + == "https://ingestion-eu.mixpanel.com/v1/ai/otel/traces" + ) + assert ( + dynamic_otlp_endpoint("mixpanel", {"mixpanel_project_token": "team-token", "mixpanel_region": "IN"}) + == "https://ingestion-in.mixpanel.com/v1/ai/otel/traces" + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + fallback = dynamic_otlp_endpoint( + "mixpanel", {"mixpanel_project_token": "team-token", "mixpanel_region": "ap-south"} + ) + assert fallback == "https://ingestion-us.mixpanel.com/v1/ai/otel/traces" + assert "unknown mixpanel_region" in caplog.text diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index ac1f5557bf6..cd9673b5c96 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -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) diff --git a/tests/unit/integrations/otel/test_otel_v2_presets.py b/tests/unit/integrations/otel/test_otel_v2_presets.py index 570c27a24e1..3677b3eb732 100644 --- a/tests/unit/integrations/otel/test_otel_v2_presets.py +++ b/tests/unit/integrations/otel/test_otel_v2_presets.py @@ -278,3 +278,89 @@ def test_signoz_preset_accepts_the_factory_call_shape(monkeypatch): cfg = PRESET_BY_CALLBACK["signoz"](allow_missing_credentials=True) assert any(e.owner == ExporterOwner.SIGNOZ for e in cfg.exporters) + + +def test_mixpanel_preset_reads_token_and_defaults_to_us(monkeypatch): + import base64 + + monkeypatch.setenv("MIXPANEL_PROJECT_TOKEN", "mp-token-123") + monkeypatch.delenv("MIXPANEL_REGION", raising=False) + monkeypatch.delenv("MIXPANEL_ENDPOINT", raising=False) + from litellm.integrations.otel.model.config import ExporterOwner + from litellm.integrations.otel.plumbing.providers import _otlp_traces_endpoint + from litellm.integrations.otel.presets.mixpanel import mixpanel_preset + + cfg = mixpanel_preset() + spec = next(e for e in cfg.exporters if e.owner == ExporterOwner.MIXPANEL) + expected_basic = base64.b64encode(b"mp-token-123:").decode() + assert spec.kind == "otlp_http" + assert spec.endpoint == "https://ingestion-us.mixpanel.com/v1/ai/otel/traces" + assert _otlp_traces_endpoint(spec.endpoint) == "https://ingestion-us.mixpanel.com/v1/ai/otel/traces" + assert spec.headers == f"Authorization=Basic {expected_basic}" + assert spec.requires_headers is True + assert "genai" in cfg.mapper_names + assert "mixpanel" in cfg.mapper_names + assert "user.id" not in cfg.baggage_promoted_keys + + +def test_mixpanel_preset_region_and_endpoint_override(monkeypatch): + monkeypatch.setenv("MIXPANEL_PROJECT_TOKEN", "mp-token-123") + monkeypatch.setenv("MIXPANEL_REGION", "EU") + monkeypatch.delenv("MIXPANEL_ENDPOINT", raising=False) + from litellm.integrations.otel.model.config import ExporterOwner + from litellm.integrations.otel.presets.mixpanel import mixpanel_preset + + eu_cfg = mixpanel_preset() + eu_spec = next(e for e in eu_cfg.exporters if e.owner == ExporterOwner.MIXPANEL) + assert eu_spec.endpoint == "https://ingestion-eu.mixpanel.com/v1/ai/otel/traces" + + monkeypatch.setenv("MIXPANEL_REGION", "in") + in_cfg = mixpanel_preset() + in_spec = next(e for e in in_cfg.exporters if e.owner == ExporterOwner.MIXPANEL) + assert in_spec.endpoint == "https://ingestion-in.mixpanel.com/v1/ai/otel/traces" + + monkeypatch.setenv("MIXPANEL_ENDPOINT", "http://collector.internal:4318/v1/ai/otel/traces") + override_cfg = mixpanel_preset() + override_spec = next(e for e in override_cfg.exporters if e.owner == ExporterOwner.MIXPANEL) + assert override_spec.endpoint == "http://collector.internal:4318/v1/ai/otel/traces" + + +def test_mixpanel_preset_without_token_is_header_gated(monkeypatch): + monkeypatch.delenv("MIXPANEL_PROJECT_TOKEN", raising=False) + from litellm.integrations.otel.model.config import ExporterOwner + from litellm.integrations.otel.presets.mixpanel import mixpanel_preset + + cfg = mixpanel_preset() + 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 diff --git a/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py index 611e828cf84..92061469ca3 100644 --- a/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py @@ -17,6 +17,7 @@ from litellm.integrations.otel import GenAIOperation from litellm.integrations.otel.mappers import ( LangfuseMapper, LangtraceMapper, + MixpanelMapper, OpenInferenceMapper, WeaveMapper, resolve_mappers, @@ -1025,3 +1026,41 @@ def test_langfuse_root_trace_attributes_derive_the_name_only_when_the_caller_gav trace: TraceControls, call_type: str | None, expected: Mapping[str, object] ) -> None: assert LangfuseMapper.trace_attributes(trace, call_type) == expected + + +def test_mixpanel_mapper_stamps_user_id_and_omits_when_absent(): + from litellm.integrations.otel.model.payloads import MCPToolCallSpanData + + 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"} + + mcp_span = MCPToolCallSpanData( + operation=GenAIOperation.EXECUTE_TOOL, + method="tools/call", + tool_name="search", + server_name="srv", + server_address=None, + server_port=None, + session_id=None, + arguments_json=None, + result_json=None, + error=None, + response_cost=None, + identity=RequestIdentity(call_id="c2", user_id="mcp_user"), + ) + assert MixpanelMapper().map(mcp_span) == {} + assert MixpanelMapper().map(_llm_call()) == {} + + from litellm.integrations.otel.model.metadata import LLMCallEvent + + sdk_event = LLMCallEvent.from_dict( + { + "litellm_params": {"metadata": {"user.id": "sdk_user_789", "ignored_key": "x"}}, + "standard_logging_object": {"id": "c3", "model": "gpt-4o", "metadata": {}}, + } + ) + assert sdk_event.caller_user_id == "sdk_user_789" + assert LLMCallEvent.from_dict({}).caller_user_id is None diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index e8ea0c04f6f..34225feeafd 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -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 diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 8651d522ac8..b2722c960a7 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -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"}}, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 19f9d0dcc73..cedc58b7b51 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -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): diff --git a/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py index 34fa164fd6d..34347519eee 100644 --- a/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py @@ -481,3 +481,171 @@ class TestNewRelicKeyLoggingValidation: ) finally: is_otel_v2_enabled.cache_clear() + + +class TestMixpanelCallbackValidation: + def _team_data(self, callback_vars): + from litellm.proxy._types import AddTeamCallback + + return AddTeamCallback(callback_name="mixpanel", callback_type="success", callback_vars=callback_vars) + + def _key_metadata(self, callback_vars): + return {"logging": [{"callback_name": "mixpanel", "callback_type": "success", "callback_vars": callback_vars}]} + + def test_mixpanel_entry_in_callbacks_configs(self): + from litellm.proxy.common_utils.callback_utils import is_sensitive_callback_key + from litellm.proxy.management_endpoints.team_callback_endpoints import _redact_callback_secrets + + client = TestClient(app) + response = client.get("/callbacks/configs", headers={"Authorization": "Bearer sk-9876"}) + assert response.status_code == 200 + mixpanel = next((config for config in response.json() if config.get("id") == "mixpanel"), None) + assert mixpanel is not None + assert mixpanel["supports_key_team_logging"] is True + params = mixpanel["dynamic_params"] + assert params["mixpanel_project_token"]["type"] == "password" + assert "mixpanel_region" in params + assert is_sensitive_callback_key("mixpanel_project_token") is True + assert is_sensitive_callback_key("mixpanel_region") is False + + raw_metadata = { + "logging": [ + { + "callback_name": "mixpanel", + "callback_type": "success", + "callback_vars": {"mixpanel_project_token": "mp-secret-tok", "mixpanel_region": "eu"}, + } + ] + } + 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 + + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + try: + 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 "mixpanel_project_token must be" in str(exc.value.detail) + for unsendable in ("", " ", "None", " None ", "***REDACTED***", " ***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( + self._key_metadata({"mixpanel_project_token": "tok", "mixpanel_region": "EU"}) + ) + finally: + is_otel_v2_enabled.cache_clear() diff --git a/tests/unit/proxy/test_credential_slot_registry.py b/tests/unit/proxy/test_credential_slot_registry.py index 6332038feb8..07d102f5a52 100644 --- a/tests/unit/proxy/test_credential_slot_registry.py +++ b/tests/unit/proxy/test_credential_slot_registry.py @@ -102,6 +102,8 @@ CALLBACK_PARAM_CLASSIFICATION: Final[Mapping[str, Classification]] = MappingProx "newrelic_region": NotSecret("region name"), "signoz_ingestion_key": Unplanted(), "signoz_ingestion_endpoint": NotSecret("sink endpoint URL"), + "mixpanel_project_token": Unplanted(), + "mixpanel_region": NotSecret("region name"), "turn_off_message_logging": NotSecret("boolean logging switch"), "litellm_disabled_callbacks": NotSecret("list of callback names"), } diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 97c21fcc777..c69ffa06e0c 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -8970,6 +8970,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 diff --git a/ui/litellm-dashboard/public/assets/logos/mixpanel.svg b/ui/litellm-dashboard/public/assets/logos/mixpanel.svg new file mode 100644 index 00000000000..71f9e91877a --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/mixpanel.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx index 5fb9949538b..e9e3f4e0f1b 100644 --- a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx @@ -8,6 +8,7 @@ import lagoLogo from "../../public/assets/logos/lago.svg"; import langfuseLogo from "../../public/assets/logos/langfuse.png"; import langsmithLogo from "../../public/assets/logos/langsmith.png"; import newrelicLogo from "../../public/assets/logos/newrelic.png"; +import mixpanelLogo from "../../public/assets/logos/mixpanel.svg"; import openmeterLogo from "../../public/assets/logos/openmeter.png"; import otelLogo from "../../public/assets/logos/otel.png"; import pointfiveLogo from "../../public/assets/logos/pointfive.png"; @@ -91,6 +92,18 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [ }, descriptionKey: "callbackInfo.descriptions.datadog", }, + { + id: "mixpanel", + displayName: "Mixpanel", + logo: mixpanelLogo.src, + supports_key_team_logging: true, + dynamic_params: { + mixpanel_project_token: "password", + mixpanel_region: "text", + }, + description: + "Mixpanel Agent Intelligence OpenTelemetry Integration. Setup: https://docs.mixpanel.com/docs/agent-intelligence", + }, { id: "newrelic", displayName: "New Relic",