Merge 6b15cf58e9 into 08f2a98250

litellm/caching/caching_handler.py

@@ -36,6 +36,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
@@ -67,6 +68,7 @@ _StreamResultT = TypeVar("_StreamResultT")
 
 
 from litellm.litellm_core_utils.core_helpers import (
+    as_str_mapping,
     get_parent_otel_span_from_kwargs,
 )
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
@@ -128,9 +130,37 @@ 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 _sse_line_carries_output(line: str) -> bool:
+    event: Final = _sse_event_data(line)
+    payload: Final = as_str_mapping(_content_block_payload(event)) if event is not None else None
+    return payload is not None and any(key != "type" and value not in ("", [], None) for key, value in payload.items())
+
+
+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))
 
 
 def is_response_without_output(result: object) -> bool:
@@ -142,7 +172,7 @@ def is_response_without_output(result: object) -> bool:
         return False
     cached_stream_events: Final = result.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)
+        return not any(_sse_frame_carries_output(event) for event in cached_stream_events)
     if "choices" in result:
         return not result["choices"]
     if result.get("object") == "response":

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
 

tests/unit/caching/test_caching_handler.py

@@ -2765,6 +2765,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]]
 
@@ -2814,6 +2830,39 @@ 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": ""}},
+            ),
+            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},
+        ),
     ]
 
 

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"])]
     )