mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(router): write the prompt caching pin when turn_off_message_logging is on (#50276)
* fix(router): write the prompt caching pin when turn_off_message_logging is on The success event PromptCachingDeploymentCheck reads had its messages redacted in place by turn_off_message_logging before any callback ran, so the pin was keyed off redacted-by-litellm, never written, and every follow-up bounced across the group. Callbacks that route on the request and never log it now opt out of that redaction and get the request messages back, while every other callback and the spend log keep the redacted payload. * fix(logging): type the request-messages payload helper as dict[str, object] The bare dict return grew the LIT001 count past the merge base and left the returned fields unchecked * test(integration): type the pin rig's Redis client as Redis[bytes] The untyped generic grew the basedpyright reportUnknownMemberType count past the merge base * test(integration): quote the Redis[bytes] annotation so the rig collects at runtime redis.Redis is not generic at runtime, so the unquoted subscript raised TypeError on import * fix(logging): read the request messages once as object in the exempt payload helper Two partially unknown reads of self.messages grew the basedpyright reportUnknownMemberType count past the merge base --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
05e4190918
commit
3f13c1ca34
5 changed files with 178 additions and 1 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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, ...]:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue