diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index d41647fe5f1..b61bb6cdc90 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1702,6 +1702,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) return { **user_api_key_metadata, + **LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_auth), **parent_correlation.as_search_metadata(), "model_group": search_tool_name, "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_auth), diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 4d99dc5db34..4dd6496812e 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -658,6 +658,7 @@ class LiteLLMExecutedBatchRunner: def _row_metadata(self, run: _BatchRun) -> dict[str, object]: # mutable-ok: router updates metadata in place return { **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(run.user_api_key_dict), + **LiteLLMProxyRequestSetup.get_tenant_logging_metadata(run.user_api_key_dict), "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(run.user_api_key_dict), "user_api_end_user_max_budget": run.user_api_key_dict.end_user_max_budget, "tags": list(run.request_tags), diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index c2829f28e61..00e9d56428e 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -83,6 +83,7 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj return { **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict), + **LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict), "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), } diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index e1fc725ae63..68a6177cabc 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -67,7 +67,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # deployment_params_provider, get_request_model_access_groups, sanitize_error_information_for_spend_logs, # pyright: ignore[reportUnusedImport] # backwards-compatible package export - should_store_prompts_and_responses_in_spend_logs, + should_store_prompts_for_call, ) from litellm.proxy.utils import ( REJECTED_BEFORE_REQUEST_SETUP, @@ -425,7 +425,7 @@ class _ProxyDBLogger(CustomLogger): response_obj, start_time, end_time, - store_bodies=should_store_prompts_and_responses_in_spend_logs(), + store_bodies=should_store_prompts_for_call(kwargs), ) if isinstance(event, SpendEventBuildError): verbose_proxy_logger.warning("collector: tracking cost in-process, event not buildable: %s", event.reason) @@ -489,6 +489,7 @@ class _ProxyDBLogger(CustomLogger): LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) _metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) + _metadata.update(LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict)) _metadata["status"] = "failure" _error_information = StandardLoggingPayloadSetup.get_error_information( original_exception=original_exception, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 546521f9048..8a23262acc1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -23,6 +23,7 @@ from fastapi import HTTPException, Request, status from pydantic import TypeAdapter from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger, verbose_proxy_logger @@ -1643,7 +1644,19 @@ def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: } +class TenantLoggingMetadata(TypedDict): + user_api_key_metadata: ReadOnly[Mapping[str, object] | None] + user_api_key_team_metadata: ReadOnly[Mapping[str, object] | None] + + class LiteLLMProxyRequestSetup: + @staticmethod + def get_tenant_logging_metadata(user_api_key_dict: UserAPIKeyAuth) -> TenantLoggingMetadata: + return { + "user_api_key_metadata": strip_callback_config(user_api_key_dict.metadata), + "user_api_key_team_metadata": strip_callback_config(user_api_key_dict.team_metadata), + } + @staticmethod def _get_timeout_from_request(headers: dict) -> float | None: """ diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 627164abb93..1eed00cc72c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -301,7 +301,11 @@ def get_passthrough_router_request_metadata( ) if request is not None: LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_data) - return _METADATA_ADAPTER.validate_python(request_data["litellm_metadata"]) + metadata: Final = _METADATA_ADAPTER.validate_python(request_data["litellm_metadata"]) + return { + **metadata, + **LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict), + } async def llm_passthrough_factory_proxy_route( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 8111e46c0b7..53317010af4 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -114,6 +114,7 @@ from litellm.proxy.common_request_processing import ( open_sse_before_first_byte, resolve_litellm_call_id, ) +from litellm.proxy.common_utils.callback_utils import strip_callback_config from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -730,6 +731,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation _metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups + _metadata["user_api_key_metadata"] = _metadata["user_api_key_auth_metadata"] + _metadata["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata) # The per-model budget counters are keyed off these. get_sanitized_user_information_from_key # returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this # the post-call increment finds nothing and every passthrough request goes untracked and diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3297a298af4..848bc246404 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -173,7 +173,6 @@ from litellm.proxy.common_utils.callback_utils import ( is_sensitive_callback_key, normalize_callback_names, process_callback, - strip_callback_config, ) from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 # legacy module exports _realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -19265,7 +19264,7 @@ async def async_queue_request( logged_api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) data["metadata"]["user_api_key"] = logged_api_key data["metadata"]["user_api_key_hash"] = logged_api_key - data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) + data["metadata"].update(LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict)) _headers: Final = safe_get_request_headers(request).copy() _headers.pop("authorization", None) # do not store the original `sk-..` api key in the db data["metadata"]["headers"] = _headers diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index e4492ee8992..15b2d18eeae 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -50,6 +50,7 @@ from litellm.litellm_core_utils.classifier_logging import ( without_classifier_audit, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, get_litellm_metadata_from_kwargs, proxy_stamped_kubernetes_pod_routing, proxy_stamped_used_client_oauth_token, @@ -87,6 +88,7 @@ from litellm.responses.litellm_completion_transformation.compaction_items import summary_only_compaction_content, with_compaction_summary, ) +from litellm.secret_managers.main import str_to_bool from litellm.types.router import DeploymentTypedDict, LiteLLM_Params, SilentExperimentFlag from litellm.types.utils import ( PROMPT_CARRYING_GUARDRAIL_FIELDS, @@ -113,6 +115,12 @@ if TYPE_CHECKING: from litellm.router import Router +_TENANT_PROMPT_STORAGE_FIELDS: Final = ("user_api_key_metadata", "user_api_key_team_metadata") +_SPEND_LOGS_PRIVACY_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter( + Mapping[str, object] | None +) + + def _get_max_string_length_prompt_in_db() -> int: """ Resolve prompt truncation threshold at runtime so values loaded later via @@ -336,9 +344,11 @@ def _get_spend_logs_metadata( clean_metadata["batch_failed_requests"] = batch_failed_requests clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload( - vector_store_request_metadata + vector_store_request_metadata, metadata=_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata) + ) + clean_metadata["guardrail_information"] = _sanitize_guardrail_information_for_spend_logs( + guardrail_information, metadata=_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata) ) - clean_metadata["guardrail_information"] = _sanitize_guardrail_information_for_spend_logs(guardrail_information) clean_metadata["usage_object"] = usage_object clean_metadata["model_map_information"] = model_map_information clean_metadata["cold_storage_object_key"] = cold_storage_object_key @@ -906,7 +916,9 @@ def get_logging_payload( if drop_requester_ip_address: clean_metadata["requester_ip_address"] = None clean_metadata["error_information"] = sanitize_error_information_for_spend_logs( - clean_metadata["error_information"], original_exception=failure_exception + clean_metadata["error_information"], + original_exception=failure_exception, + metadata=_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata), ) if litellm.cache is None: @@ -1248,7 +1260,7 @@ def _get_messages_for_spend_logs_payload( standard_logging_payload: StandardLoggingPayload | None, metadata: dict | None = None, ) -> str: - if should_store_prompts_and_responses_in_spend_logs(): + if should_store_prompts_and_responses_in_spend_logs(_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata)): if standard_logging_payload is not None: call_type: Final = standard_logging_payload.get("call_type", "") if call_type == "_arealtime": @@ -1291,8 +1303,8 @@ def stores_latest_turn_only() -> bool: return general_settings_view().get("spend_logs_prompt_storage") == _LATEST_TURN_STORAGE -def stores_full_prompts_in_spend_logs() -> bool: - return should_store_prompts_and_responses_in_spend_logs() and not stores_latest_turn_only() +def stores_full_prompts_in_spend_logs(metadata: Mapping[str, object] | None = None) -> bool: + return should_store_prompts_and_responses_in_spend_logs(metadata) and not stores_latest_turn_only() def _block_type(block: object) -> str | None: @@ -1651,6 +1663,7 @@ def _redact_prompt_leaks_in_error_string(text: str) -> str: def _sanitize_guardrail_information_for_spend_logs( guardrail_information: list[StandardLoggingGuardrailInformation] | None, + metadata: Mapping[str, object] | None = None, ) -> list[StandardLoggingGuardrailInformation] | None: """ When ``store_prompts_in_spend_logs`` is False, redact prompt-carrying fields @@ -1671,7 +1684,7 @@ def _sanitize_guardrail_information_for_spend_logs( here to match OTEL's defensive read pattern; otherwise iteration would yield the dict's keys and crash the whole spend-log write. """ - if guardrail_information is None or stores_full_prompts_in_spend_logs(): + if guardrail_information is None or stores_full_prompts_in_spend_logs(metadata): return guardrail_information entries: Final = [guardrail_information] if isinstance(guardrail_information, dict) else guardrail_information return [_redact_prompt_fields_in_guardrail_entry(entry) for entry in entries if isinstance(entry, dict)] @@ -1721,6 +1734,7 @@ def _redact_prompt_fields_in_guardrail_entry( def sanitize_error_information_for_spend_logs( error_information: StandardLoggingPayloadErrorInformation | None, original_exception: BaseException | None = None, + metadata: Mapping[str, object] | None = None, ) -> StandardLoggingPayloadErrorInformation | None: """ Sanitize ``error_information`` before it lands in ``LiteLLM_SpendLogs.metadata``. @@ -1749,7 +1763,7 @@ def sanitize_error_information_for_spend_logs( ) sanitized = cast(dict, {**persisted}) - if not stores_full_prompts_in_spend_logs(): + if not stores_full_prompts_in_spend_logs(metadata): for field in ("error_message", "traceback"): value = sanitized.get(field) if isinstance(value, str): @@ -2057,7 +2071,7 @@ def _get_proxy_server_request_for_spend_logs_payload( If turn_off_message_logging is enabled, redact messages in the request body. """ - if should_store_prompts_and_responses_in_spend_logs(): + if should_store_prompts_and_responses_in_spend_logs(_SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(metadata)): _proxy_server_request: Final = cast(dict | None, litellm_params.get("proxy_server_request", EMPTY_MAPPING)) if _proxy_server_request is not None: _request_body = _proxy_server_request.get("body", EMPTY_MAPPING) or EMPTY_MAPPING @@ -2098,11 +2112,12 @@ def _get_proxy_server_request_for_spend_logs_payload( def _get_vector_store_request_for_spend_logs_payload( vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None, + metadata: Mapping[str, object] | None = None, ) -> list[StandardLoggingVectorStoreRequest] | None: """ If user does not want to store prompts and responses, then remove the content from the vector store request metadata """ - if stores_full_prompts_in_spend_logs(): + if stores_full_prompts_in_spend_logs(metadata): return vector_store_request_metadata if vector_store_request_metadata is None: @@ -2136,7 +2151,9 @@ def _get_response_for_spend_logs_payload( ) -> str: if payload is None: return "{}" - if should_store_prompts_and_responses_in_spend_logs(): + if should_store_prompts_and_responses_in_spend_logs( + _SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(get_litellm_metadata_from_kwargs(kwargs or {})) + ): response_obj: object = payload.get("response") if response_obj is None: return "{}" @@ -2186,10 +2203,27 @@ def _get_response_for_spend_logs_payload( return "{}" -def should_store_prompts_and_responses_in_spend_logs() -> bool: +def _disables_prompt_storage(tenant_metadata: object) -> bool: + settings: Final = as_str_mapping(tenant_metadata) + value: Final = settings.get("store_prompts_in_spend_logs") if settings is not None else None + return value is False or (isinstance(value, str) and str_to_bool(value) is False) + + +def should_store_prompts_for_call(kwargs: Mapping[str, object]) -> bool: + return should_store_prompts_and_responses_in_spend_logs( + _SPEND_LOGS_PRIVACY_METADATA_ADAPTER.validate_python(get_litellm_metadata_from_kwargs(dict(kwargs))) + ) + + +def should_store_prompts_and_responses_in_spend_logs(metadata: Mapping[str, object] | None = None) -> bool: from litellm.proxy.proxy_server import general_settings from litellm.secret_managers.main import get_secret_bool + if metadata is not None and any( + _disables_prompt_storage(metadata.get(field)) for field in _TENANT_PROMPT_STORAGE_FIELDS + ): + return False + # Check general_settings (from DB or proxy_config.yaml) store_prompts_value: Final = general_settings.get("store_prompts_in_spend_logs") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a83e7d62616..503dfb69793 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4258,7 +4258,12 @@ class ProxyLogging: ) request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP}) - normalized_metadata_slot(request_data, "metadata").update(user_api_key_logged_metadata) + normalized_metadata_slot(request_data, "metadata").update( + { + **user_api_key_logged_metadata, + **LiteLLMProxyRequestSetup.get_tenant_logging_metadata(user_api_key_dict), + } + ) if litellm_logging_obj is not None: ## UPDATE LOGGING INPUT diff --git a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py index 93292a3a864..f9918d5a1ee 100644 --- a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -726,6 +726,33 @@ async def test_create_dispatches_each_row_with_the_key_team_and_project_tags_as_ assert inherited == [("key-tag", "shared", "®ion:eu", "project-tag")] * 2 +@pytest.mark.parametrize( + ("key_metadata", "team_metadata", "stored"), + [ + ({"store_prompts_in_spend_logs": False}, {}, False), + ({}, {"store_prompts_in_spend_logs": False}, False), + ({}, {}, True), + ], +) +async def test_create_dispatches_rows_that_keep_the_key_and_team_prompt_storage_policy( + monkeypatch: pytest.MonkeyPatch, key_metadata: dict[str, object], team_metadata: dict[str, object], stored: bool +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.spend_tracking_utils import should_store_prompts_and_responses_in_spend_logs + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True}) + harness: Final = make_runner( + user=UserAPIKeyAuth(api_key="sk-batch-key", metadata=key_metadata, team_metadata=team_metadata) + ) + await harness.create_and_finish() + + decisions: Final = [ + should_store_prompts_and_responses_in_spend_logs(call.kwargs["metadata"]) + for call in harness.router.acompletion.await_args_list + ] + assert decisions == [stored] * 2 + + async def test_create_dispatches_cli_session_rows_under_the_per_user_alias_not_the_login_token() -> None: session = UserAPIKeyAuth( api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa", diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index f036fb322cb..d06f95111aa 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -498,6 +498,50 @@ async def test_async_post_call_failure_hook_honors_the_redaction_header_on_a_lit assert stored_request["litellm_metadata"]["headers"] == request_headers +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("key_metadata", "team_metadata", "code_word_is_stored"), + [ + ({"store_prompts_in_spend_logs": False}, {}, False), + ({}, {"store_prompts_in_spend_logs": False}, False), + ({"store_prompts_in_spend_logs": True}, {"store_prompts_in_spend_logs": False}, False), + ({}, {}, True), + ], +) +async def test_async_post_call_failure_hook_honors_key_and_team_prompt_storage( + monkeypatch: pytest.MonkeyPatch, + key_metadata: dict[str, object], + team_metadata: dict[str, object], + code_word_is_stored: bool, +): + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + writer: Final = MagicMock(spec=DBSpendUpdateWriter) + writer.update_database = AsyncMock() + request_input: Final = "The code word is VERMILION-42. Reply OK." + + await ProxyDBLogger(spend_writer=lambda: writer).async_post_call_failure_hook( + request_data={ + "model": "gpt-5.6", + "input": request_input, + "litellm_metadata": {}, + "proxy_server_request": { + "url": "http://localhost/v1/responses", + "body": {"model": "gpt-5.6", "input": request_input}, + }, + }, + original_exception=Exception("temperature must be at most 2"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key", metadata=key_metadata, team_metadata=team_metadata), + ) + + row: Final = get_logging_payload( + kwargs=writer.update_database.call_args.kwargs["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert ("VERMILION-42" in row["proxy_server_request"]) is code_word_is_stored + + @pytest.mark.asyncio @pytest.mark.parametrize( "used_client_oauth_token, custom_llm_provider, expected", @@ -3398,6 +3442,43 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski assert event.response_cost == 0.0125 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tenant_metadata", "stored"), + [ + ({"user_api_key_metadata": {"store_prompts_in_spend_logs": False}}, False), + ({"user_api_key_team_metadata": {"store_prompts_in_spend_logs": False}}, False), + ({"user_api_key_team_metadata": {"store_prompts_in_spend_logs": "false"}}, False), + ({"user_api_key_metadata": {}, "user_api_key_team_metadata": {}}, True), + ], +) +async def test_spend_event_leaves_out_bodies_for_tenants_that_turned_off_prompt_storage( + socket_path: Path, monkeypatch: pytest.MonkeyPatch, tenant_metadata: dict[str, object], stored: bool +): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True}) + handler: Final = _RecordingHandler() + consumer: Final = SpendEventConsumer(handler) + address: Final = UnixAddress(path=str(socket_path)) + server: Final = await consumer.serve(address) + producer: Final = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback + ) + kwargs: Final[dict[str, object]] = _offload_kwargs() + kwargs["litellm_params"]["metadata"].update(tenant_metadata) + + await ProxyDBLogger(producer).async_log_success_event(kwargs, _offload_response(), datetime.now(), datetime.now()) + await producer.close(drain_timeout=5.0) + server.close() + assert await consumer.drain(timeout=5.0) == 0 + + event: Final = decode_spend_event(handler.lines[0]) + assert not isinstance(event, SpendEventDecodeError) + assert (event.litellm_params["proxy_server_request"] is not None) is stored + assert ("messages" in event.standard_logging_object) is stored + + def _completed_batch() -> LiteLLMBatch: return LiteLLMBatch( id="batch_abc", diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index db99818ebc4..e11e30dfa76 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/tests/unit/proxy/proxy_server/test_routes_queue.py @@ -8,14 +8,19 @@ from __future__ import annotations from collections.abc import Callable from contextlib import AbstractContextManager +from datetime import datetime, timezone +from queue import SimpleQueue from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import litellm import pytest from fastapi.testclient import TestClient from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy import proxy_server +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from .conftest import normalize # type: ignore[import-not-found] @@ -133,6 +138,65 @@ def _queue_router( return router +class _QueueRequestLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.requests: Final = SimpleQueue[dict[str, object]]() + + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + self.requests.put(kwargs) + + +@pytest.mark.parametrize("storage_mode", [None, "latest_turn"]) +@pytest.mark.parametrize( + ("key_setting", "team_setting", "stored"), + [(None, False, False), (True, False, False), (False, None, False), (None, None, True), (True, True, True)], +) +def test_queued_requests_apply_authenticated_prompt_storage_policy( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + storage_mode: str | None, + key_setting: bool | None, + team_setting: bool | None, + stored: bool, +) -> None: + _queue_router(monkeypatch, default_priority=0) + logger: Final = _QueueRequestLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr( + proxy_server, + "general_settings", + {"store_prompts_in_spend_logs": True, "spend_logs_prompt_storage": storage_mode}, + ) + with auth_as( + metadata={} if key_setting is None else {"store_prompts_in_spend_logs": key_setting}, + team_metadata={} if team_setting is None else {"store_prompts_in_spend_logs": team_setting}, + ): + response: Final = client.post( + "/queue/chat/completions", + json={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "queue-private-prompt"}], + "metadata": { + "user_api_key_metadata": {"store_prompts_in_spend_logs": not stored}, + "user_api_key_team_metadata": {"store_prompts_in_spend_logs": not stored}, + }, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "queued" + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + payload: Final = get_logging_payload( + kwargs=logger.requests.get_nowait(), + response_obj=litellm.ModelResponse.model_validate(response.json()), + start_time=now, + end_time=now, + ) + assert ("queue-private-prompt" in payload["proxy_server_request"]) is stored + assert ("queued" in payload["response"]) is stored + + @pytest.mark.parametrize( ("default_priority", "priority_fields", "expected_priority"), [ diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index b9411eacc8f..af8bc254e2c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -8,6 +8,7 @@ from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import Request from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -32,10 +33,14 @@ from litellm.constants import ( MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import SpendLogsMetadata, SpendLogsMetadataFields, SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import get_passthrough_router_request_metadata +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import HttpPassThroughEndpointHelpers +from litellm.types.passthrough_endpoints.pass_through_endpoints import PassthroughStandardLoggingPayload from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import ( _extract_usage_for_ocr_call, @@ -211,6 +216,178 @@ def test_classifier_audit_spend_storage_obeys_privacy_and_truncation(monkeypatch assert audit["classifier_input"]["system"] == "rubric" * 1000 +_TENANT_PROMPT: Final = "tenant-private-prompt" +_TENANT_RESPONSE: Final = "tenant-private-response" + + +def _tenant_settings(setting: bool | str | None) -> dict[str, object]: + return {} if setting is None else {"store_prompts_in_spend_logs": setting} + + +def _spend_log_for(litellm_params: Mapping[str, object]) -> tuple[SpendLogsPayload, dict[str, object]]: + now: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc) + messages: Final = [{"role": "user", "content": _TENANT_PROMPT}] + kwargs: Final[dict[str, object]] = { + "model": "gpt-4o-mini", + "litellm_params": {**litellm_params, "proxy_server_request": {"body": {"messages": messages}}}, + "standard_logging_object": { + "call_type": "_arealtime", + "model_map_information": None, + "messages": messages, + "response": {"choices": [{"message": {"content": _TENANT_RESPONSE}}]}, + "metadata": {}, + }, + } + response: Final = litellm.ModelResponse( + id="tenant-storage", + choices=[{"message": {"role": "assistant", "content": _TENANT_RESPONSE}}], + usage=litellm.Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + return get_logging_payload(kwargs=kwargs, response_obj=response, start_time=now, end_time=now), kwargs + + +def _stores_tenant_content(payload: SpendLogsPayload) -> tuple[bool, bool, bool]: + return ( + _TENANT_PROMPT in payload["messages"], + _TENANT_PROMPT in payload["proxy_server_request"], + _TENANT_RESPONSE in payload["response"], + ) + + +@pytest.mark.parametrize("storage_mode", [None, "latest_turn"]) +@pytest.mark.parametrize( + ("key_setting", "team_setting", "global_store", "stored"), + [ + (False, None, True, False), + (None, False, True, False), + (True, False, True, False), + (False, True, True, False), + (None, "false", True, False), + (None, None, True, True), + (True, True, True, True), + (True, True, False, False), + ], +) +def test_tenant_prompt_storage_policy_decides_what_spend_logs_keep( + monkeypatch: pytest.MonkeyPatch, + key_setting: bool | str | None, + team_setting: bool | str | None, + global_store: bool, + stored: bool, + storage_mode: str | None, +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr( + proxy_server, + "general_settings", + {"store_prompts_in_spend_logs": global_store, "spend_logs_prompt_storage": storage_mode}, + ) + monkeypatch.delenv("STORE_PROMPTS_IN_SPEND_LOGS", raising=False) + payload, kwargs = _spend_log_for( + { + "metadata": { + "user_api_key_metadata": _tenant_settings(key_setting), + "user_api_key_team_metadata": _tenant_settings(team_setting), + "store_prompts_in_spend_logs": True, + "headers": {"litellm-disable-message-redaction": "true"}, + } + } + ) + assert _stores_tenant_content(payload) == (stored, stored, stored) + assert kwargs["standard_logging_object"]["messages"] == [{"role": "user", "content": _TENANT_PROMPT}] + assert kwargs["standard_logging_object"]["response"] == {"choices": [{"message": {"content": _TENANT_RESPONSE}}]} + + +def _passthrough_logging_params(auth: UserAPIKeyAuth, request: Request) -> Mapping[str, object]: + context: Final = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( + request=request, + user_api_key_dict=auth, + passthrough_logging_payload=PassthroughStandardLoggingPayload(url="https://api.openai.com/v1/chat/completions"), + logging_obj=LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="tenant-storage", + function_id="tenant-storage", + ), + _parsed_body={ + "messages": [{"role": "user", "content": _TENANT_PROMPT}], + "metadata": { + "user_api_key_metadata": {"store_prompts_in_spend_logs": True}, + "user_api_key_team_metadata": {"store_prompts_in_spend_logs": True}, + }, + }, + ) + return TypeAdapter[Mapping[str, object]](Mapping[str, object]).validate_python(context["litellm_params"]) + + +def _router_passthrough_logging_params(auth: UserAPIKeyAuth, request: Request) -> Mapping[str, object]: + return {"litellm_metadata": dict(get_passthrough_router_request_metadata(auth, request))} + + +@pytest.mark.parametrize("logging_params_for", [_passthrough_logging_params, _router_passthrough_logging_params]) +@pytest.mark.parametrize( + ("key_setting", "team_setting", "stored"), + [(False, None, False), (None, False, False), (True, False, False), (None, None, True)], +) +def test_passthrough_spend_logs_follow_the_authenticated_tenant_storage_policy( + monkeypatch: pytest.MonkeyPatch, + logging_params_for: Callable[[UserAPIKeyAuth, Request], Mapping[str, object]], + key_setting: bool | None, + team_setting: bool | None, + stored: bool, +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True}) + monkeypatch.delenv("STORE_PROMPTS_IN_SPEND_LOGS", raising=False) + auth: Final = UserAPIKeyAuth( + api_key="test-storage-key", metadata=_tenant_settings(key_setting), team_metadata=_tenant_settings(team_setting) + ) + request: Final = Request( + {"type": "http", "method": "POST", "path": "/openai/v1/chat/completions", "headers": [], "query_string": b""} + ) + payload, _ = _spend_log_for(logging_params_for(auth, request)) + assert _stores_tenant_content(payload) == (stored, stored, stored) + + +@pytest.mark.parametrize("tenant_field", ["user_api_key_metadata", "user_api_key_team_metadata"]) +def test_tenant_prompt_restriction_redacts_ancillary_fields_without_mutating_callbacks( + monkeypatch: pytest.MonkeyPatch, tenant_field: str +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True}) + metadata: Final = {tenant_field: {"store_prompts_in_spend_logs": False}} + guardrails: Final = [{"guardrail_name": "test", "guardrail_request": "private-guardrail"}] + vectors: Final = [ + {"vector_store_search_response": {"data": [{"content": [{"text": "private-vector", "type": "text"}]}]}} + ] + error: Final = { + "error_message": 'bad request {"messages": [{"role": "user", "content": "private-error"}]}', + "error_class": "BadRequestError", + } + cleaned: Final = _get_spend_logs_metadata( + metadata=metadata, guardrail_information=guardrails, vector_store_request_metadata=vectors + ) + assert "private-guardrail" not in json.dumps(cleaned["guardrail_information"]) + assert "private-vector" not in json.dumps(cleaned["vector_store_request_metadata"]) + assert guardrails == [{"guardrail_name": "test", "guardrail_request": "private-guardrail"}] + assert vectors[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == "private-vector" + assert "private-error" not in json.dumps(sanitize_error_information_for_spend_logs(error, metadata=metadata)) + assert "private-error" in error["error_message"] + + +def test_request_body_metadata_cannot_disable_spend_prompt_storage(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True}) + assert should_store_prompts_and_responses_in_spend_logs({"store_prompts_in_spend_logs": False}) is True + + def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_tokens(): additional_usage_values = _get_additional_usage_values_for_usage( litellm.Usage( diff --git a/ui/litellm-dashboard/src/components/Teams.integration.test.tsx b/ui/litellm-dashboard/src/components/Teams.integration.test.tsx index a6801e3c996..85c871a4c59 100644 --- a/ui/litellm-dashboard/src/components/Teams.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.integration.test.tsx @@ -47,6 +47,20 @@ afterEach(() => { }); describe("Teams creation integration", () => { + it("sends the prompt storage restriction in team metadata", async () => { + renderWithProviders(); + fireEvent.click(screen.getAllByRole("button", { name: /create team/i })[0]); + fireEvent.change(await screen.findByLabelText(/team name/i), { target: { value: "Private Team" } }); + fireEvent.click(screen.getByRole("switch", { name: "Store prompts in spend logs" })); + fireEvent.click(screen.getAllByRole("button", { name: /create team/i }).at(-1)!); + await waitFor(() => + expect(vi.mocked(fetch).mock.calls.some(([url]) => String(url).includes("/team/new"))).toBe(true), + ); + const request = vi.mocked(fetch).mock.calls.find(([url]) => String(url).includes("/team/new")); + const payload = JSON.parse(String(request?.[1]?.body)) as { metadata: Record }; + expect(payload.metadata.store_prompts_in_spend_logs).toBe(false); + }); + it("sends a configured default member per-model budget", async () => { const user = userEvent.setup({ delay: null }); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 9191b2bd4f7..734b768a63c 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -71,6 +71,7 @@ import DeleteResourceModal from "./common_components/DeleteResourceModal"; import { teamCreateCall } from "./networking"; import { normalizeTeamModelSelection } from "./team/teamModelAccess"; import { ModelSelect } from "./ModelSelect/ModelSelect"; +import { StorePromptsToggle } from "./common_components/StorePromptsToggle"; const SUPPRESSED_BY_DESCRIPTION = ""; @@ -877,6 +878,19 @@ const Teams: React.FC = ({ accessToken, userID, userRole, isViewOnly {t("teams.page.create.fields.metadata.label")} + + {({ value, onChange }) => ( + + onChange([ + ...(value ?? []).filter((pair) => pair.key !== "store_prompts_in_spend_logs"), + { key: "store_prompts_in_spend_logs", value: JSON.stringify(store) }, + ]) + } + /> + )} + { + it("leaves an unset default untouched and writes false while preserving metadata", () => { + const onChange = vi.fn(); + render(); + expect(onChange).not.toHaveBeenCalled(); + fireEvent.click(screen.getByRole("switch", { name: "Store prompts in spend logs" })); + expect(JSON.parse(onChange.mock.calls[0][0])).toEqual({ note: "keep", store_prompts_in_spend_logs: false }); + }); + + it("can remove an existing restriction", () => { + const onChange = vi.fn(); + render(); + expect(screen.getByRole("switch")).not.toBeChecked(); + fireEvent.click(screen.getByRole("switch")); + expect(JSON.parse(onChange.mock.calls[0][0])).toEqual({ store_prompts_in_spend_logs: true }); + }); + + it("shows a saved string false restriction as off", () => { + render(); + expect(screen.getByRole("switch", { name: "Store prompts in spend logs" })).not.toBeChecked(); + }); + + it("does not overwrite invalid metadata", () => { + const onChange = vi.fn(); + render(); + expect(screen.getByRole("switch")).toHaveAttribute("aria-disabled", "true"); + fireEvent.click(screen.getByRole("switch")); + expect(onChange).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/StorePromptsToggle.tsx b/ui/litellm-dashboard/src/components/common_components/StorePromptsToggle.tsx new file mode 100644 index 00000000000..c163edb0bc9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/StorePromptsToggle.tsx @@ -0,0 +1,54 @@ +import { useId } from "react"; +import { useTranslation } from "react-i18next"; +import { Field, FieldDescription, FieldLabel } from "@/components/ui/field"; +import { Switch } from "@/components/ui/switch"; + +interface StorePromptsToggleProps { + value?: unknown; + onChange: (value: boolean) => void; + disabled?: boolean; +} + +function disablesPromptStorage(value: unknown): boolean { + return value === false || (typeof value === "string" && value.trim().toLowerCase() === "false"); +} + +export function StorePromptsToggle({ value, onChange, disabled }: StorePromptsToggleProps) { + const id = useId(); + const { t } = useTranslation(); + return ( + + {t("storePromptsToggle.label")} + + {t("storePromptsToggle.description")} + + ); +} + +export function StorePromptsMetadataToggle({ + value, + onChange, +}: { + value: string | undefined; + onChange: (value: string) => void; +}) { + const metadata = readMetadata(value); + return ( + onChange(JSON.stringify({ ...metadata, store_prompts_in_spend_logs: store }, null, 2))} + /> + ); +} + +function readMetadata(value: string | undefined): Record | null { + try { + const parsed: unknown = JSON.parse(value || "{}"); + return parsed !== null && typeof parsed === "object" && !Array.isArray(parsed) + ? (parsed as Record) + : null; + } catch { + return null; + } +} diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts index 7c6d5def8da..c544a89930d 100644 --- a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts @@ -22,6 +22,13 @@ const baseInput: KeyCreateInput = { const build = (formValues: Record, overrides: Partial = {}): KeyPayloadResult => buildKeyCreatePayload({ ...baseInput, ...overrides, formValues }); +it.each([false, true])("preserves the prompt storage restriction %s in key metadata", (store) => { + const result = build({ metadata: JSON.stringify({ note: "keep", store_prompts_in_spend_logs: store }) }); + expect(result.kind).toBe("ok"); + if (result.kind !== "ok") throw new Error(result.kind); + expect(JSON.parse(String(result.payload.metadata))).toEqual({ note: "keep", store_prompts_in_spend_logs: store }); +}); + const payloadOf = (result: KeyPayloadResult): Record => { expect(result.kind).toBe("ok"); if (result.kind !== "ok") throw new Error("unreachable"); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index ab0efe08a64..ce780a221fd 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -37,6 +37,7 @@ import BudgetDurationDropdown from "../common_components/budget_duration_dropdow import SchemaFormFields from "../common_components/check_openapi_schema"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; import ModelAliasManager from "../common_components/ModelAliasManager"; +import { StorePromptsMetadataToggle } from "../common_components/StorePromptsToggle"; import { MountedFormField, MountedFormProvider, @@ -1571,12 +1572,18 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp className="mt-4" > {(control) => ( -