From aebf9d5398c520e0d8cc3a240755a460f973b7ea Mon Sep 17 00:00:00 2001 From: mrinal Date: Fri, 2 Oct 2026 22:44:00 +0000 Subject: [PATCH] feat(otel): let a team strip captured message content from its own OTel v2 destination Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integrations/otel/model/destination.py | 19 +- .../integrations/otel/plumbing/providers.py | 90 +++- litellm/integrations/otel/plumbing/routing.py | 16 +- .../integrations/otel/presets/destinations.py | 1 + .../initialize_dynamic_callback_params.py | 17 +- litellm/proxy/_types.py | 6 + .../callback_config_validation.py | 45 +- .../team_callback_endpoints.py | 5 +- litellm/types/utils.py | 4 + .../otel/test_otel_v2_destinations.py | 403 +++++++++++++++++- ...test_initialize_dynamic_callback_params.py | 14 + .../test_callback_config_validation.py | 54 ++- .../test_team_callback_endpoints.py | 45 ++ 13 files changed, 689 insertions(+), 30 deletions(-) diff --git a/litellm/integrations/otel/model/destination.py b/litellm/integrations/otel/model/destination.py index c9c035f24a1..5887059e4b3 100644 --- a/litellm/integrations/otel/model/destination.py +++ b/litellm/integrations/otel/model/destination.py @@ -10,7 +10,7 @@ from urllib.parse import quote from pydantic import BaseModel, ConfigDict, Field -from litellm.types.utils import OtelSpanScope +from litellm.types.utils import OtelSpanScope, TeamCaptureMessageContent class OtelDestination(BaseModel): @@ -31,6 +31,17 @@ class OtelDestination(BaseModel): default="full", description="``llm_only`` keeps just the model-call spans; the rest of the request tree is not forwarded.", ) + capture_message_content: TeamCaptureMessageContent | None = Field( + default=None, + description=( + "``no_content`` strips prompt and response content from this destination's copy. ``span_only`` " + "and ``None`` (omitted) forward whatever the global capture policy collected." + ), + ) + + @property + def redacts_message_content(self) -> bool: + return self.capture_message_content == "no_content" def header_string(self) -> str: """Render headers as the ``k=v,k2=v2`` form an ``ExporterSpec`` expects. @@ -45,9 +56,9 @@ class OtelDestination(BaseModel): def cache_key(self) -> tuple[str, tuple[tuple[str, str], ...], tuple[tuple[str, str], ...], str | None]: """Identity for processor reuse, so one destination means one exporter. - ``span_scope`` is left out on purpose: the scope decides which spans reach the - processor, not how the processor exports them, so a full and an ``llm_only`` - view of the same account share one exporter. + ``span_scope`` and ``capture_message_content`` are left out on purpose: they + decide what reaches the processor, not how the processor exports it, so two + views of the same account share one exporter. """ return ( self.endpoint, diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 8b01750b8f2..c4d0b5e933e 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -39,7 +39,11 @@ from opentelemetry.util.types import Attributes, AttributeValue from litellm._logging import verbose_logger from litellm._version import version as litellm_version -from litellm.integrations.otel.mappers.langfuse import LANGFUSE_TRACE_NAME +from litellm.integrations.otel.mappers.langfuse import ( + LANGFUSE_OBSERVATION_INPUT, + LANGFUSE_OBSERVATION_OUTPUT, + LANGFUSE_TRACE_NAME, +) from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec, OpenTelemetryV2Config from litellm.integrations.otel.model.semconv import ( DB, @@ -384,6 +388,26 @@ _CAPTURED_HEADER_PREFIXES: Final = ("http.request.header.", "http.response.heade # ``?key=`` query parameter. _URL_KEYS: Final = frozenset({"http.url", "http.target", "url.full"}) _URL_QUERY_KEY: Final = "url.query" +# Prompt and response content, in every vocabulary the mappers write it in once the +# global capture policy collects it: GenAI semconv (MCP tool-call arguments and results +# included), Langfuse, OpenInference's blobs and per-index messages, Langtrace, Weave. +_MESSAGE_CONTENT_KEYS: Final = frozenset( + { + GenAI.INPUT_MESSAGES, + GenAI.OUTPUT_MESSAGES, + GenAI.SYSTEM_INSTRUCTIONS, + GenAI.TOOL_CALL_ARGUMENTS, + GenAI.TOOL_CALL_RESULT, + LANGFUSE_OBSERVATION_INPUT, + LANGFUSE_OBSERVATION_OUTPUT, + "input.value", + "output.value", + "llm.prompts", + "llm.completions", + "weave.output", + } +) +_INDEXED_MESSAGE_PREFIXES: Final = ("llm.input_messages.", "llm.output_messages.") class _SpanView(ReadableSpan): @@ -414,6 +438,19 @@ class _SpanView(ReadableSpan): ) +def _is_message_content(key: str) -> bool: + return key in _MESSAGE_CONTENT_KEYS or key.startswith(_INDEXED_MESSAGE_PREFIXES) + + +def without_message_content(span: ReadableSpan) -> ReadableSpan: + """``span`` without its prompt and response content, everything else left as it was.""" + attributes: Final = span.attributes or _NO_ATTRIBUTES + if not any(_is_message_content(key) for key in attributes): + return span + kept: Final = MappingProxyType({key: value for key, value in attributes.items() if not _is_message_content(key)}) + return _SpanView(span, span.resource, kept, span.events, span.status, parent=span.parent) + + def _is_database_span(attributes: Mapping[str, AttributeValue]) -> bool: return any(key in attributes for key in _DB_SYSTEM_KEYS) @@ -510,19 +547,21 @@ def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> Read exception it raised and names the operator's guardrail endpoint. Stack traces walk the operator's install and come off every span, as do the headers the operator captures on the server span, whose request side holds the caller's bearer token, - and the query string of the request URL, which can hold the same key. The span - itself stays, so the tenant still gets the whole trace tree. + and the query string of the request URL, which can hold the same key. A + destination set to ``no_content`` loses the prompt and response content too. The + span itself stays, so the tenant still gets the whole trace tree. """ extra: Final = destination.resource_attributes attributes: Final = span.attributes or _NO_ATTRIBUTES database: Final = _is_database_span(attributes) owned: Final = _is_tenant_owned_span(attributes) unreachable: Final = _guardrail_unreachable(attributes) + redacted: Final = destination.redacts_message_content kept: Final = MappingProxyType( { key: _without_query(key, value) for key, value in attributes.items() - if _tenant_visible(key, database, owned, unreachable) + if _tenant_visible(key, database, owned, unreachable) and not (redacted and _is_message_content(key)) } ) recorded: Final = span.events @@ -816,7 +855,8 @@ class _OverriddenBackendFilter(SpanProcessor): ``llm_only`` the model-call spans go through as trace roots and the rest of the tree is held back, unless a destination of the request names ``sink``, the account this exporter writes to, with a wider scope: the fan-out then delivers the rest of - the tree there and the model call keeps its place in it. + the tree there and the model call keeps its place in it. A destination naming + ``sink`` with ``no_content`` strips the content from this exporter's copy. """ def __init__( @@ -837,7 +877,19 @@ class _OverriddenBackendFilter(SpanProcessor): def on_end(self, span: ReadableSpan) -> None: if self._owner in suppressed_backends() or not _in_scope(span, self._scope): return - self._inner.on_end(_scoped(span, self._account_scope())) + scoped: Final = _scoped(span, self._account_scope()) + self._inner.on_end(without_message_content(scoped) if self._account_redacts_content() else scoped) + + def _account_redacts_content(self) -> bool: + """Whether a destination of this request that names ``sink`` asked for no content. + + The fan-out skips a destination this exporter already writes to, so this copy is + the one that account gets and has to honor the destination's restriction. + """ + return self._sink is not None and any( + destination.redacts_message_content and _sink_key(destination.endpoint, destination.headers) == self._sink + for destination in request_destinations() + ) def _account_scope(self) -> "OtelSpanScope": if self._scope == "full" or self._sink is None: @@ -856,6 +908,25 @@ class _OverriddenBackendFilter(SpanProcessor): return self._inner.force_flush(timeout_millis) +class _MessageContentFilter(SpanProcessor): + """Strip prompt and response content from every span 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(without_message_content(span)) + + def shutdown(self) -> None: + self._inner.shutdown() + + def force_flush(self, timeout_millis: int = 30000) -> bool: + return self._inner.force_flush(timeout_millis) + + def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter: """Build a single exporter from the top-level config fields. @@ -1113,6 +1184,7 @@ def build_tracer_provider( baggage_processor: SpanProcessor | None = None, use_simple_processor: bool | None = None, tenant_overrides: bool = False, + content_redacted_owner: str | None = None, ) -> TracerProvider: """Build the shared :class:`TracerProvider`. @@ -1131,6 +1203,8 @@ def build_tracer_provider( ``config.langfuse_span_scope`` narrows the exporter owned by ``langfuse_otel`` alone; a collector or any other backend in the same config keeps the full tree. + ``content_redacted_owner`` strips the prompt and response content from the exporter + that backend owns, again leaving every other exporter in the config alone. """ provider: Final = TracerProvider(resource=build_resource(config)) if baggage_processor is None: @@ -1151,12 +1225,14 @@ def build_tracer_provider( exp, (spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor), ) + if content_redacted_owner is not None and spec.owner is not None and spec.owner.value == content_redacted_owner: + processor = _MessageContentFilter(processor) owner = spec.owner.value if tenant_overrides and spec.owner is not None else None scope = _operator_scope(config, spec) sink = _sink_key(spec.endpoint, parse_headers(spec.headers)) if _exports_to_the_wire(spec) else None provider.add_span_processor( _OverriddenBackendFilter(processor, owner, scope, sink) - if owner is not None or scope != "full" + if owner is not None or scope != "full" or sink is not None else processor ) return provider diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py index d7b1cadfc92..cbd1be92adb 100644 --- a/litellm/integrations/otel/plumbing/routing.py +++ b/litellm/integrations/otel/plumbing/routing.py @@ -69,7 +69,7 @@ _MAX_RETIRED_PROVIDERS: Final = 64 _HeaderItems: TypeAlias = tuple[tuple[str, str], ...] -_RouteKey: TypeAlias = tuple[_HeaderItems, _HeaderItems, str | None, str | None] +_RouteKey: TypeAlias = tuple[_HeaderItems, _HeaderItems, str | None, str | None, bool] _NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) @@ -249,15 +249,23 @@ class TenantTracerCache: # A fixed per-integration region endpoint (New Relic us/eu), never a # caller-supplied host; ``None`` keeps the preset's own endpoint. endpoint: Final = dynamic_otlp_endpoint(self._callback_name, dynamic_params) + # The tenant's own account gets the content restriction its callback vars name; the + # operator's account, reached with a project header alone, keeps the global policy. + redacted: Final = ( + bool(credential_headers) + and dynamic_params is not None + and dynamic_params.get("capture_message_content") == "no_content" + ) cache_key: Final = ( tuple(sorted(credential_headers.items())), tuple(sorted(project_headers.items())), endpoint, service_name, + redacted, ) with self._lock: provider: Final = self._cached_provider_locked( - cache_key, credential_headers, project_headers, endpoint, service_name + cache_key, credential_headers, project_headers, endpoint, service_name, redacted ) self._open_span_counts[provider] = self._open_span_counts.get(provider, 0) + 1 evicted: Final = self._evicted_on_overflow_locked() @@ -276,13 +284,15 @@ class TenantTracerCache: project_headers: Mapping[str, str], endpoint: str | None, service_name: str | None, + redacted: bool, ) -> TracerProvider: cached: Final = self._providers.get(cache_key) if cached is not None: self._providers.move_to_end(cache_key) return cached built: Final = build_tracer_provider( - self._routed_config(credential_headers, project_headers, endpoint, service_name) + self._routed_config(credential_headers, project_headers, endpoint, service_name), + content_redacted_owner=self._callback_name if redacted else None, ) self._providers[cache_key] = built return built diff --git a/litellm/integrations/otel/presets/destinations.py b/litellm/integrations/otel/presets/destinations.py index e6cb775af1d..275b178050b 100644 --- a/litellm/integrations/otel/presets/destinations.py +++ b/litellm/integrations/otel/presets/destinations.py @@ -156,4 +156,5 @@ def destination_for( callback_name=callback_name, protocol=protocol, span_scope=_span_scope(callback_name, params), + capture_message_content=params.get("capture_message_content"), ) diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index 3100ca6fba1..642e09fb72c 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -4,7 +4,12 @@ from contextlib import contextmanager from contextvars import ContextVar from typing import Any, Final -from litellm.types.utils import OTEL_SPAN_SCOPES, TRUSTED_CALLBACK_VARS_FIELD, StandardCallbackDynamicParams +from litellm.types.utils import ( + OTEL_SPAN_SCOPES, + TEAM_CAPTURE_MESSAGE_CONTENT_VALUES, + TRUSTED_CALLBACK_VARS_FIELD, + StandardCallbackDynamicParams, +) _CLIENT_CALLBACK_METADATA_SLOTS: Final[tuple[str, ...]] = ("litellm_metadata", "metadata") _inherited_message_logging_disabled: Final[ContextVar[bool]] = ContextVar( @@ -81,6 +86,13 @@ def validate_langfuse_span_scope_value(value: str) -> None: raise ValueError(f"Invalid langfuse_span_scope {value!r}: must be one of {sorted(OTEL_SPAN_SCOPES)}") +def validate_capture_message_content_value(value: str) -> None: + if value not in TEAM_CAPTURE_MESSAGE_CONTENT_VALUES: + raise ValueError( + f"Invalid capture_message_content {value!r}: must be one of {sorted(TEAM_CAPTURE_MESSAGE_CONTENT_VALUES)}" + ) + + # Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict _supported_callback_params: Final[tuple[str, ...]] = ( "langfuse_public_key", @@ -116,6 +128,7 @@ _supported_callback_params: Final[tuple[str, ...]] = ( "signoz_ingestion_endpoint", "signoz_ingestion_key", "turn_off_message_logging", + "capture_message_content", ) _request_blocked_callback_params: Final = frozenset( @@ -130,6 +143,7 @@ _request_blocked_callback_params: Final = frozenset( "newrelic_region", "signoz_ingestion_endpoint", "signoz_ingestion_key", + "capture_message_content", } ) @@ -144,6 +158,7 @@ _trusted_overlay_callback_params: Final = frozenset( "newrelic_region", "signoz_ingestion_endpoint", "signoz_ingestion_key", + "capture_message_content", } ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index da69b77286e..adfb89acb4a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -24,6 +24,7 @@ from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._uuid import uuid from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_COMMANDS from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + validate_capture_message_content_value, validate_langfuse_environment_value, validate_langfuse_span_scope_value, validate_no_callback_env_reference, @@ -2347,6 +2348,8 @@ class AddTeamCallback(LiteLLMPydanticObjectBase): validate_langfuse_environment_value(callback_vars[key]) if key == "langfuse_span_scope": validate_langfuse_span_scope_value(callback_vars[key]) + if key == "capture_message_content": + validate_capture_message_content_value(callback_vars[key]) return values @@ -2397,6 +2400,9 @@ class TeamCallbackMetadata(LiteLLMPydanticObjectBase): for key in callback_vars: if key not in valid_keys: raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}") + capture_message_content: Final = callback_vars.get("capture_message_content") + if capture_message_content is not None: + validate_capture_message_content_value(str(capture_message_content)) return values diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 6a0fcbe0bb3..5931a8a8e01 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -14,6 +14,10 @@ _NEWRELIC_CALLBACK: Final = "newrelic" _NEWRELIC_VAR_PREFIX: Final = "newrelic_" _LANGFUSE_OTEL_CALLBACK: Final = "langfuse_otel" _LANGFUSE_SPAN_SCOPE_VAR: Final = "langfuse_span_scope" +_CAPTURE_MESSAGE_CONTENT_VAR: Final = "capture_message_content" +#: Options every entry of one team or key shares, because the entries are flattened +#: into one set of callback vars before a request reads them. +_SHARED_OPTION_VARS: Final = (_LANGFUSE_SPAN_SCOPE_VAR, _CAPTURE_MESSAGE_CONTENT_VAR) _ARIZE_CALLBACK: Final = "arize" _ARIZE_SAMPLING_RATE_VARS: Final[frozenset[str]] = frozenset( {"arize_success_sampling_rate", "arize_error_sampling_rate"} @@ -31,6 +35,9 @@ def callback_config_error(callback_name: str | None, callback_vars: Mapping[str, ) if langfuse_error is not None: return langfuse_error + capture_error: Final = _capture_message_content_error(callback_name, callback_vars) + if capture_error is not None: + return capture_error if callback_name != _NEWRELIC_CALLBACK: return None return _newrelic_config_error(callback_vars) @@ -75,6 +82,27 @@ def _langfuse_span_scope_error(callback_name: str | None, callback_vars: Mapping return None +def _capture_message_content_error(callback_name: str | None, callback_vars: Mapping[str, str]) -> str | None: + value: Final = callback_vars.get(_CAPTURE_MESSAGE_CONTENT_VAR) + if value is None: + return None + from litellm.integrations.otel.presets.destinations import destination_capable_backends + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + validate_capture_message_content_value, + ) + + supported: Final = sorted(destination_capable_backends()) + if callback_name not in supported: + return ( + f"{_CAPTURE_MESSAGE_CONTENT_VAR} applies to the OTel v2 callbacks {supported} only, not {callback_name!r}" + ) + try: + validate_capture_message_content_value(value) + except ValueError as e: + return str(e) + return None + + # Which credential family a dynamic variable belongs to. The families are the # integrations that share one account: every langfuse_* variable configures the # same Langfuse project whether it rides the classic callback or the OTel one, @@ -94,7 +122,7 @@ _VAR_FAMILIES: Final[Mapping[str, str]] = MappingProxyType( } ) -_FAMILY_OPTION_VARS: Final[frozenset[str]] = frozenset({_LANGFUSE_SPAN_SCOPE_VAR, *_ARIZE_SAMPLING_RATE_VARS}) +_FAMILY_OPTION_VARS: Final[frozenset[str]] = frozenset({*_SHARED_OPTION_VARS, *_ARIZE_SAMPLING_RATE_VARS}) def _family_of(var: str) -> str | None: @@ -165,19 +193,20 @@ def cross_entry_family_error( ) -def conflicting_span_scope_error( +def conflicting_shared_option_error( callback_vars: Mapping[str, str] | None, stored_vars_by_entry: Sequence[Mapping[str, str]], ) -> str | None: - incoming: Final = None if callback_vars is None else callback_vars.get(_LANGFUSE_SPAN_SCOPE_VAR) - if incoming is None: + if not callback_vars: return None return next( ( - f"{_LANGFUSE_SPAN_SCOPE_VAR} is already set to {stored!r} by another callback entry. " - f"Every entry shares one scope: remove that entry or send the same value." + f"{var} is already set to {stored!r} by another callback entry. " + f"Every entry shares one value: remove that entry or send the same value." + for var in _SHARED_OPTION_VARS + if (incoming := callback_vars.get(var)) is not None for entry in stored_vars_by_entry - if (stored := entry.get(_LANGFUSE_SPAN_SCOPE_VAR)) not in (None, incoming) + if (stored := entry.get(var)) not in (None, incoming) ), None, ) @@ -196,7 +225,7 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str error for error in ( *(_logging_entry_error(entry) for entry in entries), - *(conflicting_span_scope_error(entry_vars[i], entry_vars[:i]) for i in range(len(entry_vars))), + *(conflicting_shared_option_error(entry_vars[i], entry_vars[:i]) for i in range(len(entry_vars))), ) if error is not None ), diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index ac6169d25bd..7567a16136d 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -31,7 +31,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_config_error, - conflicting_span_scope_error, + conflicting_shared_option_error, cross_entry_family_error, ) from litellm.proxy.common_utils.callback_utils import ( @@ -283,6 +283,7 @@ async def add_team_callbacks( - langfuse_host: The host for the Langfuse callback - langfuse_environment: The tracing environment for the Langfuse callback (lowercase; falls back to LANGFUSE_TRACING_ENVIRONMENT) - langfuse_span_scope: For langfuse_otel, "full" (default) sends the whole request trace, "llm_only" sends only the model-call spans + - capture_message_content: For the OTel v2 callbacks (langfuse_otel, arize, weave_otel, newrelic), "no_content" strips prompt and response content from this team's exported spans, "span_only" keeps whatever the global capture policy collects. Omitted behaves like "span_only". It only narrows: a team gets content only when the proxy's global OTel v2 capture is enabled - gcs_bucket_name: The name of the GCS bucket - gcs_path_service_account: The path to the GCS service account - langsmith_api_key: The API key for the Langsmith callback @@ -348,7 +349,7 @@ async def add_team_callbacks( decrypted_logging: Final = decrypt_callback_vars(team_metadata).get("logging") stored_entries: Final = decrypted_logging if isinstance(decrypted_logging, list) else () stored_entry_vars: Final = [entry.get("callback_vars") or {} for entry in stored_entries] - scope_error: Final = conflicting_span_scope_error(data.callback_vars, stored_entry_vars) + scope_error: Final = conflicting_shared_option_error(data.callback_vars, stored_entry_vars) if scope_error is not None: raise _callback_config_error(scope_error) # One entry has to own a credential family end to end. The entries are diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8c10b9e3497..02eb3ab440c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3634,6 +3634,8 @@ OPENAI_RESPONSE_HEADERS: Final = [ OtelSpanScope = Literal["full", "llm_only"] OTEL_SPAN_SCOPES: Final[frozenset[str]] = frozenset(get_args(OtelSpanScope)) +TeamCaptureMessageContent = Literal["no_content", "span_only"] +TEAM_CAPTURE_MESSAGE_CONTENT_VALUES: Final[frozenset[str]] = frozenset(get_args(TeamCaptureMessageContent)) class StandardCallbackDynamicParams(TypedDict, total=False): @@ -3645,6 +3647,8 @@ class StandardCallbackDynamicParams(TypedDict, total=False): langfuse_environment: ReadOnly[str | None] langfuse_span_scope: ReadOnly[OtelSpanScope | None] + capture_message_content: ReadOnly[TeamCaptureMessageContent | None] + # Langfuse prompt version langfuse_prompt_version: int | None diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index cb986e61229..991f1639349 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -4,6 +4,7 @@ import contextvars import time from base64 import b64encode from collections.abc import Mapping +from dataclasses import replace from functools import reduce from types import MappingProxyType @@ -16,6 +17,7 @@ from opentelemetry.trace import Status, StatusCode import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.otel import GenAIOperation from litellm.integrations.otel import logger as otel_logger from litellm.integrations.otel.logger import ( OpenTelemetryV2, @@ -23,6 +25,7 @@ from litellm.integrations.otel.logger import ( fan_out_provider, publish_global_otel_v2_provider, ) +from litellm.integrations.otel.mappers import resolve_mappers from litellm.integrations.otel.model.config import ( ExporterOwner, ExporterSpec, @@ -30,6 +33,15 @@ from litellm.integrations.otel.model.config import ( is_otel_v2_enabled, ) from litellm.integrations.otel.model.destination import OtelDestination +from litellm.integrations.otel.model.payloads import ( + LLMCallSpanData, + LLMRequestParams, + LLMUsage, + MCPToolCallSpanData, + RequestIdentity, + ServerInfo, + ToolDefinition, +) from litellm.integrations.otel.plumbing import providers as otel_providers from litellm.integrations.otel.plumbing.context import ( destination_backends, @@ -51,7 +63,7 @@ from litellm.integrations.otel.presets.destinations import ( destination_for, ) from litellm.integrations.otel.presets.langfuse import langfuse_preset -from litellm.proxy._types import AddTeamCallback, UserAPIKeyAuth +from litellm.proxy._types import AddTeamCallback, TeamCallbackMetadata, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import ( convert_key_logging_metadata_to_callback, resolve_tenant_otel_destinations, @@ -1116,7 +1128,9 @@ class TestProviderWiring: assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"redis"}) - @pytest.mark.parametrize("otel", [None, True, "on", "", []], ids=["null", "true", "on", "empty_string", "empty_list"]) + @pytest.mark.parametrize( + "otel", [None, True, "on", "", []], ids=["null", "true", "on", "empty_string", "empty_list"] + ) def test_a_non_mapping_otel_block_falls_back_to_the_published_logger_config(self, monkeypatch, otel): monkeypatch.setattr(litellm, "callback_settings", {"otel": otel}, raising=False) preset = OpenTelemetryV2( @@ -2856,8 +2870,8 @@ class TestEvictionSafety: import threading from litellm.integrations.otel.plumbing.providers import ( - _DrainPool, _MAX_CACHED_DESTINATION_PROCESSORS, + _DrainPool, ) class GatedDrain(_DrainPool): @@ -3384,3 +3398,386 @@ class TestTenantHostSsrfGuard: destination_for("langfuse_otel", self._langfuse("http://10.0.0.5:3000")) assert sum("provider_url_destination_allowed_hosts" in record.message for record in caplog.records) == 1 + + +_ALL_MAPPERS = ("genai", "legacy", "openinference", "langfuse", "weave", "langtrace") +_SECRET = "SECRET-MARKER" +_MODEL_CALL_WITH_CONTENT = LLMCallSpanData( + operation=GenAIOperation.CHAT, + provider="openai", + request_model="gpt-4o", + response_model="gpt-4o-2024", + response_id="resp_1", + request_params=LLMRequestParams(temperature=0.5, max_tokens=128), + usage=LLMUsage(input_tokens=12, output_tokens=8, total_tokens=20), + finish_reasons=("tool_calls",), + error=None, + response_cost=0.001, + server=ServerInfo("api.openai.com", 443), + identity=RequestIdentity(call_id="c1", team_id="t1"), + tools=(ToolDefinition(name="lookup", description="Find a city", parameters_json='{"type":"object"}'),), + messages_in=( + {"role": "system", "content": f"system {_SECRET}"}, + {"role": "user", "content": f"prompt {_SECRET}"}, + ), + choices_out=( + { + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": f"answer {_SECRET}", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": f'{{"city": "{_SECRET}"}}'}, + } + ], + }, + }, + ), +) +_MODEL_CALL_WITHOUT_CONTENT = replace(_MODEL_CALL_WITH_CONTENT, messages_in=(), choices_out=()) +_TOOL_CALL_WITH_CONTENT = MCPToolCallSpanData( + operation=GenAIOperation.EXECUTE_TOOL, + method="tools/call", + tool_name="lookup", + server_name="maps", + server_address="maps.local", + server_port=443, + session_id="s1", + arguments_json=f'{{"city": "{_SECRET}"}}', + result_json=f'{{"weather": "{_SECRET}"}}', + error=None, + response_cost=None, + identity=RequestIdentity(call_id="c2", team_id="t1"), +) +_TOOL_CALL_WITHOUT_CONTENT = replace(_TOOL_CALL_WITH_CONTENT, arguments_json=None, result_json=None) + +NO_CONTENT_DEST = OtelDestination( + endpoint="http://team-a.local/api/public/otel", + headers=MappingProxyType({"Authorization": "Basic YQ=="}), + callback_name="langfuse_otel", + capture_message_content="no_content", +) +INHERITING_DEST = OtelDestination( + endpoint="http://team-b.local/v1", + headers=MappingProxyType({"api_key": "k", "arize-space-id": "s"}), + callback_name="arize", +) + + +def mapped(data) -> Mapping[str, object]: + """What the configured mappers write on the span, the way the logger stamps it.""" + return MappingProxyType(reduce(lambda acc, mapper: {**acc, **mapper.map(data)}, resolve_mappers(_ALL_MAPPERS), {})) + + +def recorded(attributes: Mapping[str, object]) -> dict[str, object]: + """``attributes`` the way the SDK records them, sequences frozen into tuples.""" + return {key: tuple(value) if isinstance(value, list) else value for key, value in attributes.items()} + + +def model_call_tree(provider: TracerProvider, attributes: Mapping[str, object]) -> None: + 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 llm: + llm.set_attributes(attributes) + llm.add_event("gen_ai.content.first_chunk", {"gen_ai.response.model": "gpt-4o-2024"}) + + +def by_name(exporter: InMemorySpanExporter): + return {span.name: span for span in exporter.get_finished_spans()} + + +def carries_content(span) -> bool: + return any(_SECRET in str(value) for value in span.attributes.values()) + + +class TestCaptureMessageContent: + """A team narrows the globally captured content for its own destination only.""" + + @staticmethod + def _fan_out(operator: InMemorySpanExporter, exporters: Mapping[str, InMemorySpanExporter]) -> TracerProvider: + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(operator)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda d: SimpleSpanProcessor(exporters[d.endpoint])) + ) + return provider + + @staticmethod + def _run(provider, destinations, attributes): + def run(): + set_request_destinations(destinations) + model_call_tree(provider, attributes) + + in_fresh_context(run) + + @pytest.mark.parametrize( + ("globally_captured", "setting", "content_exported"), + [ + (True, None, True), + (True, "span_only", True), + (True, "no_content", False), + (False, None, False), + (False, "span_only", False), + (False, "no_content", False), + ], + ) + def test_the_team_setting_only_ever_narrows_the_global_capture(self, globally_captured, setting, content_exported): + attributes = mapped(_MODEL_CALL_WITH_CONTENT if globally_captured else _MODEL_CALL_WITHOUT_CONTENT) + destination = OtelDestination( + endpoint="http://team.local/api/public/otel", + headers=MappingProxyType({"Authorization": "Basic dA=="}), + callback_name="langfuse_otel", + capture_message_content=setting, + ) + operator, tenant = InMemorySpanExporter(), InMemorySpanExporter() + + self._run(self._fan_out(operator, {destination.endpoint: tenant}), (destination,), attributes) + + exported = by_name(tenant)["chat gpt-4o"] + assert carries_content(exported) is content_exported + assert carries_content(by_name(operator)["chat gpt-4o"]) is globally_captured, ( + "the operator's copy is untouched" + ) + assert exported.attributes["gen_ai.request.model"] == "gpt-4o" + assert exported.attributes["gen_ai.usage.input_tokens"] == 12 + assert exported.attributes["gen_ai.usage.output_tokens"] == 8 + + @pytest.mark.parametrize( + ("with_content", "without_content"), + [ + (_MODEL_CALL_WITH_CONTENT, _MODEL_CALL_WITHOUT_CONTENT), + (_TOOL_CALL_WITH_CONTENT, _TOOL_CALL_WITHOUT_CONTENT), + ], + ids=["model_call", "mcp_tool_call"], + ) + def test_no_content_removes_exactly_what_global_capture_adds_in_every_mapper_vocabulary( + self, with_content, without_content + ): + captured, uncaptured = mapped(with_content), mapped(without_content) + assert frozenset(captured) - frozenset(uncaptured), "the fixture must exercise captured content" + operator, tenant = InMemorySpanExporter(), InMemorySpanExporter() + + self._run(self._fan_out(operator, {NO_CONTENT_DEST.endpoint: tenant}), (NO_CONTENT_DEST,), captured) + + exported = by_name(tenant)["chat gpt-4o"] + assert dict(exported.attributes) == recorded(uncaptured), "only content goes, every other attribute stays" + assert not carries_content(exported) + + def test_restricting_one_destination_leaves_the_original_span_and_the_other_destination_alone(self): + attributes = mapped(_MODEL_CALL_WITH_CONTENT) + operator, restricted, inheriting = InMemorySpanExporter(), InMemorySpanExporter(), InMemorySpanExporter() + exporters = {NO_CONTENT_DEST.endpoint: restricted, INHERITING_DEST.endpoint: inheriting} + + self._run(self._fan_out(operator, exporters), (NO_CONTENT_DEST, INHERITING_DEST), attributes) + + original = by_name(operator) + kept = by_name(inheriting) + stripped = by_name(restricted) + assert dict(original["chat gpt-4o"].attributes) == recorded(attributes) + assert dict(kept["chat gpt-4o"].attributes) == recorded(attributes) + assert not carries_content(stripped["chat gpt-4o"]) + for copy in (kept, stripped): + for name, span in copy.items(): + assert span.context == original[name].context, "same trace id and span id" + assert span.parent == original[name].parent, "same place in the tree" + assert [e.name for e in copy["chat gpt-4o"].events] == [e.name for e in original["chat gpt-4o"].events] + + def test_an_omitted_setting_and_span_only_export_the_same_span(self): + attributes = mapped(_MODEL_CALL_WITH_CONTENT) + span_only = INHERITING_DEST.model_copy(update={"capture_message_content": "span_only"}) + omitted_exporter, span_only_exporter = InMemorySpanExporter(), InMemorySpanExporter() + + for destination, exporter in ((INHERITING_DEST, omitted_exporter), (span_only, span_only_exporter)): + self._run( + self._fan_out(InMemorySpanExporter(), {destination.endpoint: exporter}), (destination,), attributes + ) + + assert {n: dict(s.attributes) for n, s in by_name(omitted_exporter).items()} == { + n: dict(s.attributes) for n, s in by_name(span_only_exporter).items() + } + + @pytest.mark.parametrize(("setting", "content_exported"), [("no_content", False), (None, True)]) + def test_the_operators_exporter_on_the_same_account_honors_the_teams_setting_once( + self, monkeypatch, setting, content_exported + ): + """The fan-out skips a destination the operator already writes to, so the operator's + copy is the one that account receives and must not bypass the team's restriction.""" + monkeypatch.setattr(litellm, "otel_tenant_destination_mode", "additive", raising=False) + shared = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor( + _OverriddenBackendFilter( + SimpleSpanProcessor(shared), "langfuse_otel", "full", TestRoutingMode.OPERATOR_SINK + ) + ) + provider.add_span_processor( + TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(shared), + operator_sinks=MappingProxyType({TestRoutingMode.OPERATOR_SINK: "full"}), + ) + ) + destination = OtelDestination( + endpoint=TestRoutingMode.SAME_ACCOUNT_ENDPOINT, + headers=MappingProxyType({"Authorization": "Basic op"}), + callback_name="langfuse_otel", + capture_message_content=setting, + ) + + self._run(provider, (destination,), mapped(_MODEL_CALL_WITH_CONTENT)) + + model_calls = [s for s in shared.get_finished_spans() if s.name == "chat gpt-4o"] + assert len(model_calls) == 1, "the same account received the span twice" + assert carries_content(model_calls[0]) is content_exported + + @pytest.mark.parametrize("mode", ["override", "additive"]) + def test_an_operator_collector_on_the_teams_account_does_not_bypass_no_content(self, monkeypatch, mode): + monkeypatch.setattr(litellm, "otel_tenant_destination_mode", mode, raising=False) + shared = InMemorySpanExporter() + monkeypatch.setattr(otel_providers, "_exporter_from_spec", lambda _spec: shared) + config = OpenTelemetryV2Config( + exporters=[ + ExporterSpec( + kind="otlp_http", + endpoint=TestRoutingMode.OPERATOR_SINK[0], + headers="authorization=Basic op", + ) + ] + ) + provider = build_tracer_provider(config, use_simple_processor=True, tenant_overrides=True) + provider.add_span_processor( + TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(shared), + operator_sinks=operator_sink_scopes(config), + ) + ) + destination = OtelDestination( + endpoint=TestRoutingMode.SAME_ACCOUNT_ENDPOINT, + headers=MappingProxyType({"Authorization": "Basic op"}), + callback_name="langfuse_otel", + capture_message_content="no_content", + ) + + self._run(provider, (destination,), mapped(_MODEL_CALL_WITH_CONTENT)) + + model_calls = [s for s in shared.get_finished_spans() if s.name == "chat gpt-4o"] + assert model_calls, "the account must receive the model call before its content can be judged absent" + assert not any(carries_content(span) for span in model_calls) + + @pytest.mark.parametrize( + ("setting", "content_exported"), [("no_content", False), ("span_only", True), (None, True)] + ) + def test_a_request_routed_to_the_teams_credentials_honors_the_setting(self, monkeypatch, setting, content_exported): + """A failure-only entry, or a destination the fan-out could not deliver, leaves the + model call on the per-request tracer route with the team's credentials.""" + exporters = {} + monkeypatch.setattr( + otel_providers, + "_exporter_from_spec", + lambda spec: exporters.setdefault(spec.owner, InMemorySpanExporter()), + ) + config = OpenTelemetryV2Config( + exporters=[ + ExporterSpec( + kind="otlp_http", + endpoint="http://op.local", + owner=ExporterOwner.LANGFUSE_OTEL, + use_simple_processor=True, + ), + ExporterSpec(kind="in_memory"), + ] + ) + cache = TenantTracerCache(config, "langfuse_otel", "litellm") + params = {"langfuse_public_key": "pk-team", "langfuse_secret_key": "sk-team"} + if setting is not None: + params["capture_message_content"] = setting + + route = cache.route_for(get_tracer(TracerProvider(), "litellm"), params) + with route.tracer.start_as_current_span("chat gpt-4o") as llm: + llm.set_attributes(mapped(_MODEL_CALL_WITH_CONTENT)) + cache.release(route.provider) + + assert route.detached is True + team_copy = by_name(exporters[ExporterOwner.LANGFUSE_OTEL])["chat gpt-4o"] + assert carries_content(team_copy) is content_exported + assert team_copy.attributes["gen_ai.usage.input_tokens"] == 12 + assert carries_content(by_name(exporters[None])["chat gpt-4o"]), "the operator's collector keeps the content" + + def test_a_restricted_and_an_unrestricted_team_never_share_a_routed_provider(self): + cache = TenantTracerCache( + OpenTelemetryV2Config( + exporters=[ + ExporterSpec(kind="otlp_http", endpoint="http://op.local", owner=ExporterOwner.LANGFUSE_OTEL) + ] + ), + "langfuse_otel", + "litellm", + ) + default = get_tracer(TracerProvider(), "litellm") + creds = {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"} + + restricted = cache.route_for(default, {**creds, "capture_message_content": "no_content"}) + unrestricted = cache.route_for(default, creds) + + assert restricted.provider is not unrestricted.provider + cache.release(restricted.provider) + cache.release(unrestricted.provider) + + @pytest.mark.usefixtures("allow_test_hosts") + @pytest.mark.parametrize("setting", ["no_content", "span_only", None]) + def test_the_setting_rides_the_teams_destination_and_omission_stays_omitted(self, monkeypatch, setting): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + callback_vars = { + "langfuse_public_key": "pk-team", + "langfuse_secret_key": "sk-team", + "langfuse_host": "http://team.local", + } + if setting is not None: + callback_vars["capture_message_content"] = setting + auth = UserAPIKeyAuth( + team_metadata={ + "logging": [ + {"callback_name": "langfuse_otel", "callback_type": "success", "callback_vars": callback_vars}, + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {"arize_api_key": "k", "arize_space_id": "s"}, + }, + ] + } + ) + + destinations = {d.callback_name: d for d in resolve_tenant_otel_destinations(auth)} + + assert destinations["langfuse_otel"].capture_message_content == setting + assert destinations["arize"].capture_message_content is None, "the setting stays on its own callback" + + @pytest.mark.parametrize("value", ["full", "NO_CONTENT", "span_and_event", "", "true"]) + def test_an_unsupported_value_fails_registration(self, value): + with pytest.raises(ValueError, match=r"Invalid capture_message_content .*\['no_content', 'span_only'\]"): + AddTeamCallback( + callback_name="langfuse_otel", + callback_type="success", + callback_vars={ + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + "capture_message_content": value, + }, + ) + + @pytest.mark.parametrize("value", ["no_content", "span_only"]) + def test_a_supported_value_is_stored_as_given(self, value): + saved = AddTeamCallback( + callback_name="langfuse_otel", + callback_type="success", + callback_vars={"langfuse_public_key": "pk", "langfuse_secret_key": "sk", "capture_message_content": value}, + ) + + assert saved.callback_vars["capture_message_content"] == value + + def test_the_legacy_team_callback_metadata_rejects_an_unsupported_value(self): + with pytest.raises(ValueError, match="Invalid capture_message_content"): + TeamCallbackMetadata(success_callback=["langfuse_otel"], callback_vars={"capture_message_content": "all"}) diff --git a/tests/unit/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/unit/litellm_core_utils/test_initialize_dynamic_callback_params.py index a9a6509f6dd..94970fd89d8 100644 --- a/tests/unit/litellm_core_utils/test_initialize_dynamic_callback_params.py +++ b/tests/unit/litellm_core_utils/test_initialize_dynamic_callback_params.py @@ -279,3 +279,17 @@ def test_arize_sampling_rates_are_picked_up_from_metadata(): assert params.get("arize_success_sampling_rate") == "0.5" assert params.get("arize_error_sampling_rate") == "0.1" + + +def test_capture_message_content_comes_only_from_the_trusted_team_vars(): + from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD + + caller_only = { + "capture_message_content": "span_only", + "metadata": {"capture_message_content": "span_only"}, + "litellm_params": {"metadata": {"capture_message_content": "span_only"}}, + } + assert initialize_standard_callback_dynamic_params(caller_only).get("capture_message_content") is None + + stamped = {**caller_only, TRUSTED_CALLBACK_VARS_FIELD: {"capture_message_content": "no_content"}} + assert initialize_standard_callback_dynamic_params(stamped).get("capture_message_content") == "no_content" diff --git a/tests/unit/proxy/common_utils/test_callback_config_validation.py b/tests/unit/proxy/common_utils/test_callback_config_validation.py index 33ae986a511..cac60dff0b4 100644 --- a/tests/unit/proxy/common_utils/test_callback_config_validation.py +++ b/tests/unit/proxy/common_utils/test_callback_config_validation.py @@ -2,7 +2,7 @@ import pytest from litellm.proxy.common_utils.callback_config_validation import ( callback_config_error, - conflicting_span_scope_error, + conflicting_shared_option_error, cross_entry_family_error, logging_metadata_config_error, ) @@ -59,7 +59,7 @@ def test_a_bad_span_scope_is_reported_even_when_the_environment_is_fine(): def test_one_span_scope_per_team(new_vars, stored, rejected): """The entries flatten last-wins, so a second scope would export whichever entry was stored last. An entry that names no scope leaves the stored one in charge.""" - error = conflicting_span_scope_error(new_vars, stored) + error = conflicting_shared_option_error(new_vars, stored) assert (error is not None) is rejected if rejected: assert "langfuse_span_scope" in error and stored[-1]["langfuse_span_scope"] in error @@ -128,3 +128,53 @@ def test_arize_sampling_rates_are_not_family_credentials(): stored = [{"arize_api_key": "k1", "arize_success_sampling_rate": "0.5"}] assert cross_entry_family_error({"arize_success_sampling_rate": "0.1"}, stored) is None assert cross_entry_family_error({"arize_error_sampling_rate": "0.5"}, stored) is None + + +@pytest.mark.parametrize("callback_name", ["langfuse_otel", "arize", "weave_otel", "newrelic"]) +@pytest.mark.parametrize("value", ["no_content", "span_only"]) +def test_capture_message_content_is_accepted_on_every_otel_v2_destination(callback_name, value): + assert callback_config_error(callback_name, {"capture_message_content": value}) is None + + +@pytest.mark.parametrize("callback_name", ["langfuse", "datadog", "otel", "arize_phoenix", None]) +def test_capture_message_content_is_rejected_where_it_would_never_take_effect(callback_name): + error = callback_config_error(callback_name, {"capture_message_content": "no_content"}) + assert error is not None and "capture_message_content" in error and "langfuse_otel" in error + + +def test_an_unsupported_capture_message_content_is_rejected_on_key_logging_metadata(): + error = logging_metadata_config_error( + { + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": {"capture_message_content": "full"}, + } + ] + } + ) + assert error is not None and "Invalid capture_message_content" in error + + +@pytest.mark.parametrize( + "new_vars, stored, rejected", + [ + ({"capture_message_content": "no_content"}, [{"capture_message_content": "span_only"}], True), + ( + {"capture_message_content": "span_only"}, + [{"langfuse_public_key": "pk"}, {"capture_message_content": "no_content"}], + True, + ), + ({"capture_message_content": "no_content"}, [{"capture_message_content": "no_content"}], False), + ({"capture_message_content": "no_content"}, [{"langfuse_span_scope": "llm_only"}], False), + ({"langfuse_public_key": "pk"}, [{"capture_message_content": "no_content"}], False), + ], +) +def test_one_capture_message_content_per_team(new_vars, stored, rejected): + """The entries flatten last-wins on the routed path, so a second value would apply + whichever entry was stored last to the team's traffic.""" + error = conflicting_shared_option_error(new_vars, stored) + assert (error is not None) is rejected + if rejected: + assert "capture_message_content" in error diff --git a/tests/unit/proxy/management_endpoints/test_team_callback_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_callback_endpoints.py index e368a26155b..b7c4bdf7cfd 100644 --- a/tests/unit/proxy/management_endpoints/test_team_callback_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_callback_endpoints.py @@ -1712,3 +1712,48 @@ def test_add_team_callback_accepts_arize_sampling_rate_vars(): ) assert data.callback_vars["arize_success_sampling_rate"] == "0.5" assert data.callback_vars["arize_error_sampling_rate"] == "1.0" + + +@pytest.mark.asyncio +async def test_add_team_callbacks_rejects_capture_message_content_on_a_non_otel_v2_callback(patched_prisma): + data = AddTeamCallback( + callback_name="langfuse", + callback_type="success", + callback_vars={ + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + "capture_message_content": "no_content", + }, + ) + with pytest.raises(HTTPException) as exc: + await add_team_callbacks( + data=data, + http_request=Mock(spec=Request), + team_id="team-victim", + user_api_key_dict=_admin_auth(), + ) + assert exc.value.status_code == 400 + assert "capture_message_content" in str(exc.value.detail) + patched_prisma.db.litellm_teamtable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_team_callbacks_stores_capture_message_content_in_the_existing_callback_metadata(patched_prisma): + data = AddTeamCallback( + callback_name="langfuse_otel", + callback_type="success", + callback_vars={ + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + "capture_message_content": "no_content", + }, + ) + await add_team_callbacks( + data=data, + http_request=Mock(spec=Request), + team_id="team-victim", + user_api_key_dict=_admin_auth(), + ) + patched_prisma.db.litellm_teamtable.update.assert_awaited_once() + stored = json.loads(patched_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) + assert [entry["callback_vars"].get("capture_message_content") for entry in stored["logging"]] == ["no_content"]