From 905a9c4ece9d2b4bd0f4f9cea859d4e3bfa8eaa7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 18:43:00 -0800 Subject: [PATCH] fix(caching): keep whitespace replies and text completion fills cached --- litellm/caching/caching_handler.py | 38 +++++--- tests/unit/caching/test_caching_handler.py | 108 ++++++++++++++------- 2 files changed, 99 insertions(+), 47 deletions(-) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 1d06b6804fc..a4fbfa69191 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -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, diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 2b7d7fa0e14..024f9de2fb3 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -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()}