diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 0579eb15baf..385522c2435 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -898,6 +898,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def redacts_messages_itself(self) -> bool: return False + def exempt_from_message_redaction(self) -> bool: + """Router pre-call checks that key routing off the request and never log it return True, so + `turn_off_message_logging` leaves the success event they receive unredacted.""" + return False + def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict: """ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9c49efc529f..9b6ec8d28d6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4405,7 +4405,11 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(callback, CustomLogger): # custom logger class from litellm.integrations.shadow_eval_logger import ShadowEvalLogger - model_call_details: dict = self.model_call_details + model_call_details: dict = ( + self._model_call_details_with_request_messages() + if callback.exempt_from_message_redaction() + else self.model_call_details + ) ################################## # call redaction hook for custom logger model_call_details = callback.redact_standard_logging_payload_from_model_call_details( @@ -4969,6 +4973,21 @@ class Logging(LiteLLMLoggingBaseClass): ) -> None: executor.submit(self.success_handler, result, start_time, end_time, cache_hit) + def _model_call_details_with_request_messages(self) -> dict[str, object]: # mutable-ok: loggers take a dict + """The success-event payload with the request messages put back where `turn_off_message_logging` + redacted them in place, for a callback that routes on the request and never logs it.""" + request_messages: Final = cast(object, self.messages) # cast-ok: the request messages are only put back as-is + standard_logging_object: Final = self.model_call_details.get("standard_logging_object") + return { + **self.model_call_details, + "messages": request_messages, + **( + {"standard_logging_object": {**standard_logging_object, "messages": request_messages}} + if isinstance(standard_logging_object, dict) + else {} + ), + } + def _get_completed_success_callbacks( self, *, asynchronous: bool ) -> tuple[str | Callable[..., object] | CustomLogger, ...]: diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 81012d38b65..e5d944b3aef 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -163,6 +163,9 @@ class PromptCachingDeploymentCheck(CustomLogger): return healthy_deployments + def exempt_from_message_redaction(self) -> bool: + return True + async def async_log_success_event( self, kwargs: _PromptCachingSuccessEventKwargs, diff --git a/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py new file mode 100644 index 00000000000..13bf67582b3 --- /dev/null +++ b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py @@ -0,0 +1,95 @@ +import collections +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import eventually, gateway_from_environment, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.wire import Reply, Wire, wire_server +from integration.routing._prompt_cache_pin_support import ( + DEFAULT_MARK, + anthropic_deployment, + ask_turn, + chat_body, + claude_peer, + new_marker, + one_new_pin, + pin_config, + post_chat, + read_pins, + system_turn, +) +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +from litellm.constants import REDACTED_BY_LITELLM + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +_GROUP: Final = "pin-logging-off" +_LEFT: Final = "pin-logging-off-left" +_RIGHT: Final = "pin-logging-off-right" +_FOLLOW_UPS: Final = 20 +_REDACTED_MESSAGES: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}] +_LOGGED_EVENT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + redis: "Redis[bytes]" + logger_sink: Wire + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("prompt-cache-pin-message-logging-off") + with ( + gateway_from_environment() as gateway, + owned_redis(directory) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + wire_server(lambda request: Reply()) as logger_sink, + ): + config: Final = pin_config( + directory, + [anthropic_deployment(_GROUP, _LEFT, left.url), anthropic_deployment(_GROUP, _RIGHT, right.url)], + litellm_settings={ + "turn_off_message_logging": True, + "callbacks": ["integration._support.redaction_probe.probe"], + "cache": False, + "enable_redis_auth_cache": False, + }, + ) + environment: Final = { + "REDIS_HOST": cache.host, + "REDIS_PORT": str(cache.port), + "REDACTION_PROBE_URL": logger_sink.url, + } + with ( + owned_proxy_process(gateway, directory, environment, config=config, workers=1) as owned, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + yield _Rig(owned, redis, logger_sink) + + +def test_follow_ups_stay_on_the_first_deployment_while_the_logger_sees_redacted_messages(rig: _Rig) -> None: + before: Final = frozenset(read_pins(rig.redis)) + body: Final = chat_body(_GROUP, [system_turn(new_marker(), DEFAULT_MARK), ask_turn(new_marker())]) + + _, first_served = post_chat(rig.owned.gateway, body) + _, pin = one_new_pin(rig.redis, before) + assert pin.model_id == first_served, (pin, first_served) + + served: Final = tuple(post_chat(rig.owned.gateway, body)[1] for _ in range(_FOLLOW_UPS)) + assert served == (first_served,) * _FOLLOW_UPS, collections.Counter(served) + + eventually(rig.logger_sink.received.qsize, lambda count: count >= 1 + _FOLLOW_UPS, seconds=30) + logged: Final = tuple( + object_value(_LOGGED_EVENT.validate_json(request.body)["standard"])["messages"] + for request in rig.logger_sink.drain() + ) + assert len(logged) == 1 + _FOLLOW_UPS, len(logged) + assert all(messages == _REDACTED_MESSAGES for messages in logged), logged diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 9487e4d6321..5169642bc43 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1480,3 +1480,58 @@ async def test_a_pin_written_without_tools_is_read_when_tools_is_an_empty_list() ) assert filtered == [deployments[0]] + + +class _LoggedMessagesProbe(CustomLogger): + def __init__(self) -> None: + self.logged_messages: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_messages.append((kwargs["standard_logging_object"]["messages"], kwargs["messages"])) + + +@pytest.mark.asyncio +async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_redacted_messages( + monkeypatch, local_model_cost_map +): + """ + `turn_off_message_logging` redacted the shared success-event payload before the check read it, so the + pin was keyed off `redacted-by-litellm`, never written, and a cacheable conversation bounced across the + group. The pin must land whatever the flag says, while a real logger alongside still sees only the + redacted messages. + """ + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + probe = _LoggedMessagesProbe() + monkeypatch.setattr(litellm, "callbacks", [probe]) + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-1", "dep-2") + ], + optional_pre_call_checks=["prompt_caching"], + ) + messages = _messages(word_count=5000) + + first = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok") + served_by = first._hidden_params["model_id"] + + affinity_key = PromptCachingCache.get_prompt_caching_cache_key(messages, None) + pin = await _eventually(lambda: router.prompt_caching_pin_cache.get_cache(key=affinity_key)) + assert pin == {"model_id": served_by}, pin + + subsequent = [ + (await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok"))._hidden_params[ + "model_id" + ] + for _ in range(4) + ] + assert subsequent == [served_by] * 4 + + await _eventually(lambda: len(probe.logged_messages) == 5) + redacted = [{"role": "user", "content": "redacted-by-litellm"}] + assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages