diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 82804d2d707..ab5d2332c21 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( from litellm.litellm_core_utils.logging_utils import ( assemble_complete_response_from_streaming_chunks, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.caching import CACHED_STREAM_EVENTS_KEY, EMBEDDING_CACHE_FORMAT_VERSION, CachedEmbedding from litellm.types.integrations.custom_logger import converted_stream_requested from litellm.types.llms.base import LiteLLMBaseModel @@ -68,6 +69,8 @@ _StreamResultT = TypeVar("_StreamResultT") from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + as_str_mapping, get_parent_otel_span_from_kwargs, ) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -133,41 +136,108 @@ def _is_chat_completion_cached_dict(cached_result: dict) -> bool: return "choices" in cached_result -def _starts_content_block(event: object) -> bool: - text: Final = event.decode("utf-8", errors="replace") if isinstance(event, bytes) else event - return isinstance(text, str) and "event: content_block_start" in text.splitlines() +_SSE_EVENT_DATA: Final = TypeAdapter(Mapping[str, object]) + + +def _sse_event_data(line: str) -> Mapping[str, object] | None: + if not line.startswith("data:"): + return None + try: + return _SSE_EVENT_DATA.validate_json(line.removeprefix("data:")) + except ValidationError: + return None + + +def _content_block_payload(event: Mapping[str, object]) -> object: + match event.get("type"): + case "content_block_start": + return event.get("content_block") + case "content_block_delta": + return event.get("delta") + case _: + return None + + +def _content_block_carries_output(block: object) -> bool: + fields: Final = as_str_mapping(block) + 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: + event: Final = _sse_event_data(line) + return event is not None and _content_block_carries_output(_content_block_payload(event)) + + +def _sse_frame_carries_output(frame: object) -> bool: + text: Final = frame.decode("utf-8", errors="replace") if isinstance(frame, bytes) else frame + return isinstance(text, str) and any(_sse_line_carries_output(line) for line in split_sse_lines(text)) + + +_REPLAY_METADATA_FIELDS: Final = frozenset({"thought_signatures", "continuation_token"}) + + +def _message_field_carries_output(key: str, value: object) -> bool: + match key: + case "role": + return False + case "provider_specific_fields": + fields: Final = as_str_mapping(value) or {} + return any( + name not in _REPLAY_METADATA_FIELDS and field not in ("", [], {}, None) + for name, field in fields.items() + ) + case _: + return value not in ("", [], {}, None) + + +def _chat_message_fields(choice: object) -> Mapping[str, object] | None: + match choice: + case Choices(): + return choice.message.model_dump(exclude_none=True) + case TextChoices(): + return None + case _: + return as_str_mapping((as_str_mapping(choice) or {}).get("message")) + + +def _choice_carries_output(choice: object) -> bool: + 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: + return any(_choice_carries_output(choice) for choice in as_object_list(choices) or ()) def is_response_without_output(result: object) -> bool: if isinstance(result, (ModelResponse, TextCompletionResponse)): - return not result.choices + return not _choices_carry_output(result.choices) if isinstance(result, ResponsesAPIResponse): return not getattr(result, "output", None) - if not isinstance(result, dict): + cached: Final = as_str_mapping(result) + if cached is None: return False - cached_stream_events: Final = result.get(CACHED_STREAM_EVENTS_KEY) + cached_stream_events: Final = cached.get(CACHED_STREAM_EVENTS_KEY) if isinstance(cached_stream_events, (list, tuple)): - return not any(_starts_content_block(event) for event in cached_stream_events) - if "choices" in result: - return not result["choices"] - if result.get("object") == "response": - return not result.get("output") - if result.get("type") == "message": - return not result.get("content") + return not any(_sse_frame_carries_output(event) for event in cached_stream_events) + if "choices" in cached: + return not _choices_carry_output(cached["choices"]) + if cached.get("object") == "response": + return not cached.get("output") + if cached.get("type") == "message": + return not any(_content_block_carries_output(block) for block in as_object_list(cached.get("content")) or ()) return False -def _choice_carries_output(choice: Choices | TextChoices) -> bool: - payload: Final[Mapping[str, object]] = ( - choice.message.model_dump(exclude_none=True) if isinstance(choice, Choices) else {"text": choice.text} - ) - return any(key != "role" and value not in ("", [], {}) for key, value in payload.items()) +def _text_choice_carries_output(choice: TextChoices) -> bool: + text: Final[object] = choice.text + return text != "" def _assembled_stream_without_output(response: ModelResponse | TextCompletionResponse) -> bool: - """The stream wrapper closes a stream whose chunks carried no choice with an empty choice of its own, so the - assembled response has choices and the no-output check on it alone would store that empty answer.""" - return not any(_choice_carries_output(choice) for choice in response.choices) + 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: diff --git a/litellm/llms/anthropic/messages/native/response_cache.py b/litellm/llms/anthropic/messages/native/response_cache.py index d2425c3f484..11c0b501966 100644 --- a/litellm/llms/anthropic/messages/native/response_cache.py +++ b/litellm/llms/anthropic/messages/native/response_cache.py @@ -115,7 +115,7 @@ class AnthropicMessagesStreamCacheWriter: return cached_payload: Final = {CACHED_STREAM_EVENTS_KEY: events} if is_response_without_output(cached_payload): - verbose_logger.debug("LiteLLM Cache: not caching a stream with no content blocks") + verbose_logger.debug("LiteLLM Cache: not caching a stream with no output") return dual_cache: Final = self.caching_handler.dual_cache diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 51754c80dd6..024f9de2fb3 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -2402,7 +2402,7 @@ def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatc async def _short_lived_script(): await handler.async_set_cache( - result=litellm.ModelResponse(), + result=chat_completion_response, original_function=acompletion, kwargs={}, ) @@ -2434,7 +2434,7 @@ def test_async_cache_write_runs_in_the_post_response_phase_without_leaking_it(mo monkeypatch.setattr(litellm, "cache", _PhaseRecordingCache()) async def _request(): - await handler.async_set_cache(result=litellm.ModelResponse(), original_function=acompletion, kwargs={}) + await handler.async_set_cache(result=chat_completion_response, original_function=acompletion, kwargs={}) leaked = in_post_response_phase() await asyncio.gather(*_PENDING_CACHE_WRITES) return leaked @@ -2689,7 +2689,7 @@ async def test_response_cache_lookup_and_write_declare_the_llm_response_target(m call_type=CallTypes.acompletion.value, kwargs={"messages": [{"role": "user", "content": "hi"}]}, ) - await handler.async_set_cache(result=litellm.ModelResponse(), original_function=acompletion, kwargs={}) + await handler.async_set_cache(result=chat_completion_response, original_function=acompletion, kwargs={}) await asyncio.gather(*_PENDING_CACHE_WRITES) assert seen == {"get": "llm_response", "set": "llm_response"} @@ -2755,6 +2755,26 @@ def _anthropic_message_missing_content() -> dict[str, object]: return {"id": f"msg_{uuid.uuid4()}", "type": "message", "role": "assistant", "model": "claude-sonnet-5"} +def _anthropic_message_with_only_an_empty_text_block() -> dict[str, object]: + return { + "id": f"msg_{uuid.uuid4()}", + "type": "message", + "role": "assistant", + "model": "gemini/gemini-3.8-flash", + "content": [{"type": "text", "text": ""}], + } + + +def _anthropic_message_with_only_a_whitespace_text_block() -> dict[str, object]: + return {**_anthropic_message_with_only_an_empty_text_block(), "content": [{"type": "text", "text": "\n\n"}]} + + +def _empty_reply_with_provider_metadata(provider_specific_fields: dict[str, object]) -> ModelResponse: + return litellm.ModelResponse( + choices=[litellm.Choices(message=litellm.Message(content="", provider_specific_fields=provider_specific_fields))] + ) + + def _anthropic_stream_without_content_blocks() -> dict[str, object]: return { CACHED_STREAM_EVENTS_KEY: [ @@ -2765,6 +2785,22 @@ def _anthropic_stream_without_content_blocks() -> dict[str, object]: } +def _anthropic_sse_event(payload: dict[str, object]) -> str: + return f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n" + + +def _anthropic_stream_with_only_empty_blocks(*blocks_and_deltas: dict[str, object]) -> dict[str, object]: + return { + CACHED_STREAM_EVENTS_KEY: [ + 'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_1", "content": []}}\n\n', + *(_anthropic_sse_event(event) for event in blocks_and_deltas), + 'event: content_block_stop\ndata: {"type": "content_block_stop", "index": 0}\n\n', + 'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}}\n\n', + 'event: message_stop\ndata: {"type": "message_stop"}\n\n', + ] + } + + EmptyResult = ModelResponse | TextCompletionResponse | ResponsesAPIResponse | dict[str, object] EmptyResponseCase = tuple[EmptyResult, Callable[..., Awaitable[object]], str, dict[str, object]] @@ -2808,12 +2844,63 @@ def _empty_response_cases() -> list[EmptyResponseCase]: CallTypes.aanthropic_messages.value, {"messages": _unique_messages(), "max_tokens": 16}, ), + ( + _anthropic_message_with_only_an_empty_text_block(), + aanthropic_messages, + CallTypes.aanthropic_messages.value, + {"messages": _unique_messages(), "max_tokens": 16}, + ), + ( + _empty_reply_with_provider_metadata({"thought_signatures": ["c2lnbmF0dXJl"]}), + litellm.acompletion, + CallTypes.acompletion.value, + {"messages": _unique_messages()}, + ), + ( + _empty_reply_with_provider_metadata({"continuation_token": "dG9rZW4="}), + litellm.acompletion, + CallTypes.acompletion.value, + {"messages": _unique_messages()}, + ), ( _anthropic_stream_without_content_blocks(), aanthropic_messages, 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": ""}}, + ), + aanthropic_messages, + 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": "", "citations": None}, + }, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": ""}}, + ), + aanthropic_messages, + 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": "", "citations": []}, + }, + ), + aanthropic_messages, + CallTypes.aanthropic_messages.value, + {"messages": _unique_messages(), "max_tokens": 16, "stream": True}, + ), ] @@ -2868,6 +2955,74 @@ async def test_async_get_cache_treats_stored_response_without_output_as_miss( assert lookup.cached_result is None +def _chat_reply(content: str) -> ModelResponse: + return litellm.ModelResponse(choices=[litellm.Choices(message=litellm.Message(content=content))]) + + +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=original_function, + logging_obj=_completion_logging_obj(call_type), + start_time=_FIXED_START, + call_type=call_type, + kwargs=kwargs, + ) + + assert lookup.cached_result is not None + + @pytest.mark.asyncio async def test_async_get_cache_heals_stored_completion_without_choices(): setup_cache() @@ -3014,6 +3169,18 @@ def _closing_chunk_without_output() -> ModelResponseStream: return _stream_chunk([StreamingChoices(finish_reason="stop", index=0, delta=Delta())]) +def _closing_chunk_with_only_a_thought_signature() -> ModelResponseStream: + return _stream_chunk( + [ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="", provider_specific_fields={"thought_signatures": ["c2lnbmF0dXJl"]}), + ) + ] + ) + + def _cached_content(cached: object) -> object: assert cached is not None stored = cached if isinstance(cached, dict) else json.loads(cached) @@ -3025,12 +3192,17 @@ def _content_chunk(content: str) -> ModelResponseStream: @pytest.mark.asyncio -async def test_async_streamed_answer_without_output_is_never_cached(): +@pytest.mark.parametrize( + "closing_chunk", + [_closing_chunk_without_output(), _closing_chunk_with_only_a_thought_signature()], + ids=["no-delta", "only-a-thought-signature"], +) +async def test_async_streamed_answer_without_output_is_never_cached(closing_chunk: ModelResponseStream): setup_cache() kwargs = {"model": "gpt-4o", "messages": _unique_messages()} handler = LLMCachingHandler(original_function=litellm.acompletion, request_kwargs=kwargs, start_time=_FIXED_START) - await handler.add_streaming_response_to_cache(_closing_chunk_without_output()) + await handler.add_streaming_response_to_cache(closing_chunk) await asyncio.gather(*_PENDING_CACHE_WRITES) assert litellm.cache.get_cache(**kwargs) is None @@ -3049,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()} diff --git a/tests/unit/llms/anthropic/messages/native/test_response_cache.py b/tests/unit/llms/anthropic/messages/native/test_response_cache.py index bc3b1f6ac76..ac104b55bde 100644 --- a/tests/unit/llms/anthropic/messages/native/test_response_cache.py +++ b/tests/unit/llms/anthropic/messages/native/test_response_cache.py @@ -1,4 +1,5 @@ import asyncio +import json from typing import Any, AsyncIterator, Dict, Final, List import pytest @@ -136,9 +137,48 @@ async def test_cache_key_separates_anthropic_native_params(local_cache, request_ assert len(fake_handler.calls) == 2 +TOOL_USE_WITHOUT_ARGUMENTS_EVENTS: Final[List[bytes]] = [ + STREAM_EVENTS[0], + b'event: content_block_start\ndata: {"type": "content_block_start", "index": 0, ' + b'"content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_time", "input": {}}}\n\n', + b'event: content_block_delta\ndata: {"type": "content_block_delta", "index": 0, ' + b'"delta": {"type": "input_json_delta", "partial_json": ""}}\n\n', + STREAM_EVENTS[3], + b'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, ' + b'"usage": {"output_tokens": 3}}\n\n', + STREAM_EVENTS[5], +] + +CRLF_STREAM_IN_ONE_CHUNK: Final[List[bytes]] = [b"".join(STREAM_EVENTS).replace(b"\n", b"\r\n")] + +LINE_SEPARATOR_TEXT_EVENTS: Final[List[bytes]] = [ + *STREAM_EVENTS[:2], + b"event: content_block_delta\ndata: " + + json.dumps( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "A
B"}}, + ensure_ascii=False, + ).encode() + + b"\n\n", + *STREAM_EVENTS[3:], +] + +NON_JSON_DATA_LINE_EVENTS: Final[List[bytes]] = [STREAM_EVENTS[0], b"data: [DONE]\n\n", *STREAM_EVENTS[1:]] + + @pytest.mark.asyncio -async def test_streaming_request_is_replayed_from_cache(local_cache, request_kwargs, monkeypatch): - fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _byte_stream([b"event: never_used\n\n"])]) +@pytest.mark.parametrize( + "events", + [ + STREAM_EVENTS, + TOOL_USE_WITHOUT_ARGUMENTS_EVENTS, + CRLF_STREAM_IN_ONE_CHUNK, + LINE_SEPARATOR_TEXT_EVENTS, + NON_JSON_DATA_LINE_EVENTS, + ], + ids=["text", "tool-use-without-arguments", "crlf-in-one-chunk", "line-separator-in-text", "non-json-data-line"], +) +async def test_streaming_request_is_replayed_from_cache(local_cache, request_kwargs, monkeypatch, events): + fake_handler = _CountingHandler([_byte_stream(events), _byte_stream([b"event: never_used\n\n"])]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) @@ -147,8 +187,8 @@ async def test_streaming_request_is_replayed_from_cache(local_cache, request_kwa second = await _collect(second_stream) assert len(fake_handler.calls) == 1 - assert first == STREAM_EVENTS - assert second == STREAM_EVENTS + assert first == events + assert second == events assert second_stream._hidden_params["cache_hit"] is True @@ -165,8 +205,15 @@ async def test_streaming_cache_is_not_shared_with_non_streaming(local_cache, req @pytest.mark.asyncio -async def test_stream_without_content_blocks_is_not_cached(local_cache, request_kwargs, monkeypatch): - empty_events = [STREAM_EVENTS[0], STREAM_EVENTS[4], STREAM_EVENTS[5]] +@pytest.mark.parametrize( + "empty_events", + [ + [STREAM_EVENTS[0], STREAM_EVENTS[4], STREAM_EVENTS[5]], + [STREAM_EVENTS[0], STREAM_EVENTS[1], STREAM_EVENTS[3], STREAM_EVENTS[4], STREAM_EVENTS[5]], + ], + ids=["no-blocks", "only-an-empty-text-block"], +) +async def test_stream_without_output_is_not_cached(local_cache, request_kwargs, monkeypatch, empty_events): fake_handler = _CountingHandler( [_byte_stream(empty_events), _byte_stream(STREAM_EVENTS), _byte_stream([b"event: never_used\n\n"])] )