scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(caching): keep whitespace replies and text completion fills cached

Browse source
This commit is contained in:
mateo-berri 2026-11-06 18:43:00 -08:00
parent d21e29981b
commit 905a9c4ece
2 changed files with 99 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
litellm/caching/caching_handler.py
View file

@ -154,13 +154,9 @@ def _content_block_payload(event: Mapping[str, object]) -> object:
return None
def _is_blank(value: object) -> bool:
return value in ("", [], {}, None) or (isinstance(value, str) and not value.strip())
def _content_block_carries_output(block: object) -> bool:
fields: Final = as_str_mapping(block)
return fields is not None and any(key != "type" and not _is_blank(value) for key, value in fields.items())
return fields is not None and any(key != "type" and value not in ("", [], None) for key, value in fields.items())
def _sse_line_carries_output(line: str) -> bool:
@ -182,26 +178,29 @@ def _message_field_carries_output(key: str, value: object) -> bool:
return False
case "provider_specific_fields":
fields: Final = as_str_mapping(value) or {}
return any(name not in _REPLAY_METADATA_FIELDS and not _is_blank(field) for name, field in fields.items())
return any(
name not in _REPLAY_METADATA_FIELDS and field not in ("", [], {}, None)
for name, field in fields.items()
)
case _:
return not _is_blank(value)
return value not in ("", [], {}, None)
def _choice_output_fields(choice: object) -> Mapping[str, object]:
def _chat_message_fields(choice: object) -> Mapping[str, object] | None:
match choice:
case Choices():
return choice.message.model_dump(exclude_none=True)
case StreamingChoices():
return choice.delta.model_dump(exclude_none=True)
case TextChoices():
return {"text": choice.text}
return None
case _:
cached: Final = as_str_mapping(choice) or {}
return as_str_mapping(cached.get("message")) or {"text": cached.get("text")}
return as_str_mapping((as_str_mapping(choice) or {}).get("message"))
def _choice_carries_output(choice: object) -> bool:
return any(_message_field_carries_output(key, value) for key, value in _choice_output_fields(choice).items())
message: Final = _chat_message_fields(choice)
return message is None or any(_message_field_carries_output(key, value) for key, value in message.items())
def _choices_carry_output(choices: object) -> bool:
@ -228,6 +227,17 @@ def is_response_without_output(result: object) -> bool:
return False
def _text_choice_carries_output(choice: TextChoices) -> bool:
text: Final[object] = choice.text
return text != ""
def _assembled_stream_without_output(response: ModelResponse | TextCompletionResponse) -> bool:
if isinstance(response, TextCompletionResponse):
return not any(_text_choice_carries_output(choice) for choice in response.choices)
return is_response_without_output(response)
def _stream_replay_requested(kwargs: Mapping[str, object]) -> bool:
if kwargs.get("stream", False) is True:
return True
@ -1391,7 +1401,7 @@ class LLMCachingHandler:
is_async=True,
)
)
if complete_streaming_response is None or is_response_without_output(complete_streaming_response):
if complete_streaming_response is None or _assembled_stream_without_output(complete_streaming_response):
return
await self.async_set_cache(
result=complete_streaming_response,
@ -1421,7 +1431,7 @@ class LLMCachingHandler:
)
)
if complete_streaming_response is None or is_response_without_output(complete_streaming_response):
if complete_streaming_response is None or _assembled_stream_without_output(complete_streaming_response):
return
self.sync_set_cache(
result=complete_streaming_response,

108
tests/unit/caching/test_caching_handler.py
View file

@ -2850,18 +2850,6 @@ def _empty_response_cases() -> list[EmptyResponseCase]:
CallTypes.aanthropic_messages.value,
{"messages": _unique_messages(), "max_tokens": 16},
),
(
_anthropic_message_with_only_a_whitespace_text_block(),
aanthropic_messages,
CallTypes.aanthropic_messages.value,
{"messages": _unique_messages(), "max_tokens": 16},
),
(
litellm.ModelResponse(choices=[litellm.Choices(message=litellm.Message(content="\n\n"))]),
litellm.acompletion,
CallTypes.acompletion.value,
{"messages": _unique_messages()},
),
(
_empty_reply_with_provider_metadata({"thought_signatures": ["c2lnbmF0dXJl"]}),
litellm.acompletion,
@ -2913,15 +2901,6 @@ def _empty_response_cases() -> list[EmptyResponseCase]:
CallTypes.aanthropic_messages.value,
{"messages": _unique_messages(), "max_tokens": 16, "stream": True},
),
(
_anthropic_stream_with_only_empty_blocks(
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "\n\n"}},
),
aanthropic_messages,
CallTypes.aanthropic_messages.value,
{"messages": _unique_messages(), "max_tokens": 16, "stream": True},
),
]
@ -2976,26 +2955,72 @@ async def test_async_get_cache_treats_stored_response_without_output_as_miss(
assert lookup.cached_result is None
@pytest.mark.asyncio
async def test_async_set_cache_keeps_reply_whose_text_is_padded_with_whitespace():
setup_cache()
handler: Final = LLMCachingHandler(original_function=litellm.acompletion, request_kwargs={}, start_time=_FIXED_START)
kwargs: Final = {"messages": _unique_messages()}
reply: Final = litellm.ModelResponse(choices=[litellm.Choices(message=litellm.Message(content="\n\nParis\n"))])
def _chat_reply(content: str) -> ModelResponse:
return litellm.ModelResponse(choices=[litellm.Choices(message=litellm.Message(content=content))])
await handler.async_set_cache(result=reply, original_function=litellm.acompletion, kwargs=kwargs)
def _text_completion_fill(text: str) -> TextCompletionResponse:
return litellm.TextCompletionResponse(choices=[litellm.utils.TextChoices(text=text)])
def _replies_that_look_empty_but_carry_output() -> list[EmptyResponseCase]:
return [
(_chat_reply("\n\nParis\n"), litellm.acompletion, CallTypes.acompletion.value, {"messages": _unique_messages()}),
(_chat_reply("\n\n"), litellm.acompletion, CallTypes.acompletion.value, {"messages": _unique_messages()}),
(
_anthropic_message_with_only_a_whitespace_text_block(),
aanthropic_messages,
CallTypes.aanthropic_messages.value,
{"messages": _unique_messages(), "max_tokens": 16},
),
(
_anthropic_stream_with_only_empty_blocks(
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "\n\n"}},
),
aanthropic_messages,
CallTypes.aanthropic_messages.value,
{"messages": _unique_messages(), "max_tokens": 16, "stream": True},
),
(_text_completion_fill(" "), litellm.atext_completion, CallTypes.atext_completion.value, {"prompt": str(uuid.uuid4())}),
(_text_completion_fill(""), litellm.atext_completion, CallTypes.atext_completion.value, {"prompt": str(uuid.uuid4())}),
]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"result, original_function, call_type, kwargs",
_replies_that_look_empty_but_carry_output(),
ids=[
"chat-padded",
"chat-whitespace",
"messages-whitespace",
"messages-stream-whitespace",
"text-completion-indent",
"text-completion-empty-fill",
],
)
async def test_async_set_cache_keeps_whitespace_replies_and_text_completion_fills(
result: EmptyResult,
original_function: Callable[..., Awaitable[object]],
call_type: str,
kwargs: dict[str, object],
):
setup_cache()
handler: Final = LLMCachingHandler(original_function=original_function, request_kwargs={}, start_time=_FIXED_START)
await handler.async_set_cache(result=result, original_function=original_function, kwargs=kwargs)
await asyncio.gather(*_PENDING_CACHE_WRITES)
lookup: Final = await handler.async_get_cache(
model="gpt-3.5-turbo",
original_function=litellm.acompletion,
logging_obj=_completion_logging_obj(CallTypes.acompletion.value),
original_function=original_function,
logging_obj=_completion_logging_obj(call_type),
start_time=_FIXED_START,
call_type=CallTypes.acompletion.value,
call_type=call_type,
kwargs=kwargs,
)
assert isinstance(lookup.cached_result, ModelResponse)
assert lookup.cached_result.choices[0].message.content == "\n\nParis\n"
assert lookup.cached_result is not None
@pytest.mark.asyncio
@ -3196,6 +3221,23 @@ async def test_async_streamed_answer_with_content_is_cached():
assert _cached_content(litellm.cache.get_cache(**kwargs)) == "hi"
def _streamed_text_completion_closing_chunk(text: str) -> TextCompletionResponse:
return litellm.TextCompletionResponse(choices=[litellm.utils.TextChoices(text=text, finish_reason="stop")])
@pytest.mark.asyncio
@pytest.mark.parametrize("text, cached", [("", False), (" ", True)], ids=["empty", "indent"])
async def test_async_streamed_text_completion_is_cached_only_with_text(text: str, cached: bool):
setup_cache()
kwargs = {"model": "gpt-3.5-turbo-instruct", "prompt": str(uuid.uuid4())}
handler = LLMCachingHandler(original_function=litellm.atext_completion, request_kwargs=kwargs, start_time=_FIXED_START)
await handler.add_streaming_response_to_cache(_streamed_text_completion_closing_chunk(text))
await asyncio.gather(*_PENDING_CACHE_WRITES)
assert (litellm.cache.get_cache(**kwargs) is not None) is cached
def test_sync_streamed_answer_without_output_is_never_cached():
setup_cache()
kwargs = {"model": "gpt-4o", "messages": _unique_messages()}