scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-22 00:31:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 51

feat(caching): add response caching support for anthropic_messages call type

Browse source 
Adds caching support for the anthropic_messages call type (/v1/messages endpoint),
enabling response caching for both streaming and non-streaming requests via the
existing litellm.cache infrastructure.

Changes:
- Add 'anthropic_messages' to CachingSupportedCallTypes
- Include Anthropic-specific params (system, stop_sequences, top_k) in cache key
- Add cache read/replay branch in _convert_cached_result_to_model_response
- Add streaming cache write in AgenticAnthropicStreamingIterator._persist_to_cache()
- Comprehensive test coverage (24 tests)

Fixes #25653
This commit is contained in:
Maxx Rodriguez 2026-05-08 23:25:44 -07:00
parent 144279eb57
commit 1ec4b9e3ad
12 changed files with 996 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit-misc.yml vendored
View file

@ -31,6 +31,7 @@ jobs:
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/passthrough
tests/test_litellm/test_caching
tests/test_litellm/vector_stores
tests/test_litellm/test_*.py
workers: 2

1
litellm/caching/caching.py
View file

@ -80,6 +80,7 @@ class Cache:
"rerank",
"responses",
"aresponses",
"anthropic_messages",
],
# s3 Bucket, boto3 configuration
azure_account_url: Optional[str] = None,

19
litellm/caching/caching_handler.py
View file

@ -29,6 +29,7 @@ from typing import (
Optional,
Tuple,
Union,
cast,
)
from pydantic import BaseModel
@ -883,6 +884,24 @@ class LLMCachingHandler:
else:
cached_result = response_obj
elif (call_type == CallTypes.anthropic_messages.value) and isinstance(
cached_result, dict
):
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
if kwargs.get("stream", False) is True:
cached_result = FakeAnthropicMessagesStreamIterator(
response=cast(AnthropicMessagesResponse, cached_result),
logging_obj=logging_obj,
cache_hit=True,
)
# else: return dict as-is (it's already AnthropicMessagesResponse)
if (
hasattr(cached_result, "_hidden_params")
and cached_result._hidden_params is not None

27
litellm/litellm_core_utils/model_param_helper.py
View file

@ -74,12 +74,17 @@ class ModelParamHelper:
)
exclude_kwargs = ModelParamHelper._get_exclude_kwargs()
anthropic_messages_kwargs = (
ModelParamHelper._get_litellm_supported_anthropic_messages_kwargs()
)
combined_kwargs = chat_completion_kwargs.union(
text_completion_kwargs,
embedding_kwargs,
transcription_kwargs,
rerank_kwargs,
responses_api_kwargs,
anthropic_messages_kwargs,
)
combined_kwargs = combined_kwargs.difference(exclude_kwargs)
return combined_kwargs
@ -190,6 +195,28 @@ class ModelParamHelper:
)
return non_streaming_params.union(streaming_params)
@staticmethod
def _get_litellm_supported_anthropic_messages_kwargs() -> Set[str]:
"""
Get the Anthropic Messages API-specific kwargs that should be included
in cache key generation.
These params are not part of the OpenAI API spec but are critical for
cache key uniqueness when using the Anthropic Messages API directly.
"""
return {
"system",
"stop_sequences",
"top_k",
"betas",
"context_management",
"output_format",
"output_config",
"inference_geo",
"speed",
"reasoning_effort",
}
@staticmethod
def _get_exclude_kwargs() -> Set[str]:
"""

69
litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py
View file

@ -8,9 +8,11 @@ to run through agentic completion hooks. If an agentic hook fires, the
follow-up response is chained as Phase 2 of the same iterator.
"""
import asyncio
import json
from typing import Any, AsyncIterator, Dict, List, Optional, cast
import litellm
from litellm._logging import verbose_logger
@ -175,6 +177,7 @@ class AgenticAnthropicStreamingIterator:
self._stream_exhausted = False
self._hook_processing_done = False
self._follow_up_iterator: Optional[AsyncIterator] = None
self._response_cached = False
def __aiter__(self):
return self
@ -198,6 +201,70 @@ class AgenticAnthropicStreamingIterator:
raise StopAsyncIteration
def _persist_to_cache(self, rebuilt_response: Dict[str, Any]) -> None:
"""Persist the rebuilt streaming response to the LiteLLM cache."""
if self._response_cached:
return
try:
caching_handler = getattr(self._logging_obj, "_llm_caching_handler", None)
if caching_handler is None:
return
request_kwargs = getattr(caching_handler, "request_kwargs", None)
if (
not isinstance(request_kwargs, dict)
or request_kwargs.get("stream") is not True
):
return
request_kwargs = request_kwargs.copy()
preset_cache_key = getattr(caching_handler, "preset_cache_key", None)
request_cache_key = request_kwargs.pop("cache_key", None)
if preset_cache_key is None:
preset_cache_key = request_cache_key
if request_kwargs.get("metadata") is None:
request_kwargs.pop("metadata", None)
request_kwargs.pop("custom_llm_provider", None)
if preset_cache_key is not None:
request_kwargs["cache_key"] = preset_cache_key
if not caching_handler._should_store_result_in_cache(
original_function=caching_handler.original_function,
kwargs=request_kwargs,
):
return
if litellm.cache is None:
return
cached_response = json.dumps(rebuilt_response)
self._cache_write_task = asyncio.create_task(
litellm.cache.async_add_cache(
cached_response,
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
**request_kwargs,
)
)
self._cache_write_task.add_done_callback(
lambda task: (
verbose_logger.warning(
"AgenticStreamingIterator: Cache write failed: %s",
task.exception(),
)
if not task.cancelled() and task.exception()
else None
)
)
self._response_cached = True
except Exception as e:
verbose_logger.warning(
"AgenticStreamingIterator: Failed to persist response to cache: %s",
str(e),
)
async def _process_agentic_hooks(self) -> None:
"""Rebuild the Anthropic response from collected SSE bytes and call hooks."""
if self._hook_processing_done:
@ -215,6 +282,8 @@ class AgenticAnthropicStreamingIterator:
)
return
self._persist_to_cache(rebuilt)
[
(
f"{b.get('type')}({b.get('name', '')})"

33
litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py
View file

@ -8,9 +8,11 @@ Used when WebSearch interception converts stream=True to stream=False but
the LLM doesn't make a tool call, and we need to return a stream to the user.
"""
import asyncio
import json
from typing import Any, Dict, List, cast
from typing import Any, Dict, List, Optional, cast
from litellm._logging import verbose_logger
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -33,10 +35,18 @@ class FakeAnthropicMessagesStreamIterator:
- message_stop
"""
def __init__(self, response: AnthropicMessagesResponse):
def __init__(
self,
response: AnthropicMessagesResponse,
logging_obj: Optional[Any] = None,
cache_hit: bool = False,
):
self.response = response
self.chunks = self._create_streaming_chunks()
self.current_index = 0
self._logging_obj = logging_obj
self._cache_hit = cache_hit
self._logging_done = False
def _create_content_block_chunks(
self, block_dict: Dict[str, Any], index: int
@ -214,12 +224,31 @@ class FakeAnthropicMessagesStreamIterator:
async def __anext__(self):
if self.current_index >= len(self.chunks):
self._log_stream_completion()
raise StopAsyncIteration
chunk = self.chunks[self.current_index]
self.current_index += 1
return chunk
def _log_stream_completion(self) -> None:
"""Log the cache hit to the success handler when stream is exhausted."""
if self._logging_done or self._logging_obj is None or not self._cache_hit:
return
self._logging_done = True
try:
asyncio.create_task(
self._logging_obj.async_success_handler(
self.response, None, None, self._cache_hit
)
)
except Exception as e:
verbose_logger.debug(
"FakeAnthropicMessagesStreamIterator: Failed to log cache hit: %s",
str(e),
)
def __iter__(self):
return self

1
litellm/types/caching.py
View file

@ -29,6 +29,7 @@ CachingSupportedCallTypes = Literal[
"rerank",
"responses",
"aresponses",
"anthropic_messages",
]

0
tests/test_litellm/test_caching/__init__.py Normal file
View file

104
tests/test_litellm/test_caching/test_anthropic_messages_cache_key.py Normal file
View file

@ -0,0 +1,104 @@
"""
Tests that Anthropic Messages API-specific params (system, stop_sequences, top_k)
are included in cache key generation, preventing false cache hits.
"""
import pytest
from litellm.caching.caching import Cache
from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
class TestAnthropicMessagesCacheKey:
"""Tests for cache key correctness with Anthropic-specific params."""
def setup_method(self):
self.cache = Cache()
def test_anthropic_params_in_all_llm_api_params(self):
"""Verify system, stop_sequences, and top_k are in the combined param set."""
all_params = ModelParamHelper._get_all_llm_api_params()
assert "system" in all_params
assert "stop_sequences" in all_params
assert "top_k" in all_params
def test_different_system_prompts_produce_different_cache_keys(self):
"""Two requests with different system prompts must NOT share a cache key."""
key1 = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
system="You are a helpful assistant.",
)
key2 = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
system="You are a pirate.",
)
assert key1 != key2
def test_different_stop_sequences_produce_different_cache_keys(self):
"""Two requests with different stop_sequences must NOT share a cache key."""
key1 = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
stop_sequences=["END"],
)
key2 = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
stop_sequences=["STOP", "DONE"],
)
assert key1 != key2
def test_different_top_k_values_produce_different_cache_keys(self):
"""Two requests with different top_k values must NOT share a cache key."""
key1 = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
top_k=10,
)
key2 = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
top_k=50,
)
assert key1 != key2
def test_same_params_produce_same_cache_key(self):
"""Identical requests must produce the same cache key."""
kwargs = dict(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
system="You are a helpful assistant.",
stop_sequences=["END"],
top_k=10,
)
key1 = self.cache.get_cache_key(**kwargs)
key2 = self.cache.get_cache_key(**kwargs)
assert key1 == key2
def test_system_absent_vs_present_produces_different_cache_keys(self):
"""A request with system param vs without must produce different keys."""
key_with_system = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
system="You are a helpful assistant.",
)
key_without_system = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
)
assert key_with_system != key_without_system
def test_top_k_absent_vs_present_produces_different_cache_keys(self):
"""A request with top_k vs without must produce different keys."""
key_with_top_k = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
top_k=40,
)
key_without_top_k = self.cache.get_cache_key(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello"}],
)
assert key_with_top_k != key_without_top_k

140
tests/test_litellm/test_caching/test_anthropic_messages_cache_read.py Normal file
View file

@ -0,0 +1,140 @@
"""
Tests for _convert_cached_result_to_model_response handling of
call_type="anthropic_messages".
Verifies:
1. Non-streaming: returns the cached dict as-is
2. Streaming: returns a FakeAnthropicMessagesStreamIterator
"""
import datetime
import sys
from unittest.mock import MagicMock, patch
import pytest
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.types.utils import CallTypes
# Patch update_response_metadata where it's actually used in caching_handler
_CACHING_HANDLER_MODULE = sys.modules["litellm.caching.caching_handler"]
def _make_caching_handler() -> LLMCachingHandler:
"""Create a LLMCachingHandler with minimal mocked dependencies."""
handler = LLMCachingHandler.__new__(LLMCachingHandler)
handler.start_time = datetime.datetime.now()
return handler
def _sample_anthropic_response() -> dict:
"""Return a minimal valid AnthropicMessagesResponse dict."""
return {
"id": "msg_123",
"type": "message",
"role": "assistant",
"model": "claude-3-5-sonnet-20241022",
"content": [{"type": "text", "text": "Hello!"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {
"input_tokens": 10,
"output_tokens": 5,
},
}
class TestAnthropicMessagesCacheRead:
"""Tests for cache read conversion with anthropic_messages call type."""
@patch.object(_CACHING_HANDLER_MODULE, "update_response_metadata")
def test_non_streaming_returns_dict_as_is(self, mock_update_metadata):
"""Non-streaming anthropic_messages cache hit returns the dict unchanged."""
handler = _make_caching_handler()
cached_dict = _sample_anthropic_response()
logging_obj = MagicMock()
result = handler._convert_cached_result_to_model_response(
cached_result=cached_dict,
call_type=CallTypes.anthropic_messages.value,
kwargs={"stream": False},
logging_obj=logging_obj,
model="claude-3-5-sonnet-20241022",
args=(),
custom_llm_provider="anthropic",
)
# Should return the same dict object, not wrapped
assert result is cached_dict
assert isinstance(result, dict)
assert result["id"] == "msg_123"
assert result["content"][0]["text"] == "Hello!"
@patch.object(_CACHING_HANDLER_MODULE, "update_response_metadata")
def test_streaming_returns_fake_stream_iterator(self, mock_update_metadata):
"""Streaming anthropic_messages cache hit returns FakeAnthropicMessagesStreamIterator."""
handler = _make_caching_handler()
cached_dict = _sample_anthropic_response()
logging_obj = MagicMock()
result = handler._convert_cached_result_to_model_response(
cached_result=cached_dict,
call_type=CallTypes.anthropic_messages.value,
kwargs={"stream": True},
logging_obj=logging_obj,
model="claude-3-5-sonnet-20241022",
args=(),
custom_llm_provider="anthropic",
)
assert isinstance(result, FakeAnthropicMessagesStreamIterator)
# Verify the iterator was constructed with the response dict
assert result.response is cached_dict
@patch.object(_CACHING_HANDLER_MODULE, "update_response_metadata")
def test_non_streaming_without_stream_key(self, mock_update_metadata):
"""When stream key is absent from kwargs, should treat as non-streaming."""
handler = _make_caching_handler()
cached_dict = _sample_anthropic_response()
logging_obj = MagicMock()
result = handler._convert_cached_result_to_model_response(
cached_result=cached_dict,
call_type=CallTypes.anthropic_messages.value,
kwargs={},
logging_obj=logging_obj,
model="claude-3-5-sonnet-20241022",
args=(),
custom_llm_provider="anthropic",
)
# Without stream key, should return dict as-is
assert result is cached_dict
assert isinstance(result, dict)
@patch.object(_CACHING_HANDLER_MODULE, "update_response_metadata")
def test_does_not_match_non_dict_cached_result(self, mock_update_metadata):
"""If cached_result is not a dict, the anthropic_messages branch is skipped."""
handler = _make_caching_handler()
# Simulate a non-dict cached result (e.g., a ModelResponse object)
cached_result = MagicMock()
cached_result._hidden_params = {"cache_hit": False}
logging_obj = MagicMock()
result = handler._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=CallTypes.anthropic_messages.value,
kwargs={"stream": False},
logging_obj=logging_obj,
model="claude-3-5-sonnet-20241022",
args=(),
custom_llm_provider="anthropic",
)
# Should not be converted to FakeAnthropicMessagesStreamIterator
assert not isinstance(result, FakeAnthropicMessagesStreamIterator)
# The _hidden_params["cache_hit"] should be set to True by the final block
assert cached_result._hidden_params["cache_hit"] is True

214
tests/test_litellm/test_caching/test_anthropic_messages_caching_e2e.py Normal file
View file

@ -0,0 +1,214 @@
"""
End-to-end integration tests for anthropic_messages caching flow.
Verifies the full cache lifecycle:
1. Cache miss → provider call → store in cache
2. Cache hit → return from cache (no provider call)
3. Cache bypass via cache={"no-cache": True}
4. Different parameters produce separate cache entries (no false hits)
"""
import asyncio
from unittest.mock import AsyncMock, patch
import pytest
import litellm
from litellm.caching.caching import Cache
def _mock_anthropic_response(text: str = "Hello!", msg_id: str = "msg_test_123"):
"""Return a minimal valid AnthropicMessagesResponse dict."""
return {
"id": msg_id,
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [{"type": "text", "text": text}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 5},
}
@pytest.fixture(autouse=True)
def setup_and_teardown_cache():
"""Set up a local in-memory cache before each test and clean up after."""
original_cache = litellm.cache
litellm.cache = Cache(type="local")
yield
litellm.cache = original_cache
class TestAnthropicMessagesCachingE2E:
"""End-to-end tests for anthropic_messages caching."""
@pytest.mark.asyncio
async def test_non_streaming_cache_miss_then_hit(self):
"""
First call should hit the provider (cache miss).
Second identical call should return from cache (no provider call).
"""
mock_response = _mock_anthropic_response()
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_anthropic_messages_handler",
new_callable=AsyncMock,
return_value=mock_response,
) as mock_handler:
# First call - cache miss, should call provider
result1 = await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Hello"}],
max_tokens=100,
api_key="fake-key",
)
assert mock_handler.call_count == 1
assert result1["id"] == "msg_test_123"
assert result1["content"][0]["text"] == "Hello!"
# Allow the async cache write task to complete
await asyncio.sleep(0.1)
# Second call - identical params, should hit cache
result2 = await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Hello"}],
max_tokens=100,
api_key="fake-key",
)
# Provider should NOT have been called again
assert mock_handler.call_count == 1
# Result should match the cached response
assert result2["id"] == "msg_test_123"
assert result2["content"][0]["text"] == "Hello!"
@pytest.mark.asyncio
async def test_cache_bypass_with_no_cache(self):
"""
After a cached response exists, passing cache={"no-cache": True}
should force a fresh provider call.
"""
mock_response = _mock_anthropic_response(
text="First response", msg_id="msg_first"
)
mock_response_fresh = _mock_anthropic_response(
text="Fresh response", msg_id="msg_fresh"
)
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_anthropic_messages_handler",
new_callable=AsyncMock,
) as mock_handler:
mock_handler.return_value = mock_response
# First call - populates cache
result1 = await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Hello"}],
max_tokens=100,
api_key="fake-key",
)
assert mock_handler.call_count == 1
assert result1["content"][0]["text"] == "First response"
# Second call with no-cache - should bypass cache
mock_handler.return_value = mock_response_fresh
result2 = await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Hello"}],
max_tokens=100,
api_key="fake-key",
cache={"no-cache": True},
)
# Provider should have been called again
assert mock_handler.call_count == 2
assert result2["content"][0]["text"] == "Fresh response"
@pytest.mark.asyncio
async def test_non_streaming_cache_hit_returns_dict(self):
"""
Verify that a cache hit for non-streaming returns the dict directly
through _convert_cached_result_to_model_response.
"""
mock_response = _mock_anthropic_response(
text="Cached dict response", msg_id="msg_dict_cache"
)
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_anthropic_messages_handler",
new_callable=AsyncMock,
return_value=mock_response,
) as mock_handler:
# First call - cache miss
await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Cache dict test"}],
max_tokens=50,
api_key="fake-key",
)
assert mock_handler.call_count == 1
# Allow cache write
await asyncio.sleep(0.1)
# Second call - cache hit, exercises _convert_cached_result_to_model_response
result = await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Cache dict test"}],
max_tokens=50,
api_key="fake-key",
)
# Should NOT call provider again
assert mock_handler.call_count == 1
# Should return the cached dict
assert result["id"] == "msg_dict_cache"
assert result["content"][0]["text"] == "Cached dict response"
@pytest.mark.asyncio
async def test_different_system_prompts_no_false_hit(self):
"""
Two calls with different system prompts should both call the provider.
The cache should not return a false hit for different system prompts.
"""
mock_response_assistant = _mock_anthropic_response(
text="I am an assistant", msg_id="msg_assistant"
)
mock_response_pirate = _mock_anthropic_response(
text="Arrr matey!", msg_id="msg_pirate"
)
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_anthropic_messages_handler",
new_callable=AsyncMock,
) as mock_handler:
mock_handler.return_value = mock_response_assistant
# First call with system="You are a helpful assistant."
result1 = await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Who are you?"}],
max_tokens=100,
system="You are a helpful assistant.",
api_key="fake-key",
)
assert mock_handler.call_count == 1
assert result1["content"][0]["text"] == "I am an assistant"
# Second call with different system prompt
mock_handler.return_value = mock_response_pirate
result2 = await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Who are you?"}],
max_tokens=100,
system="You are a pirate.",
api_key="fake-key",
)
# Provider should have been called again (different cache key)
assert mock_handler.call_count == 2
assert result2["content"][0]["text"] == "Arrr matey!"

389
tests/test_litellm/test_caching/test_anthropic_messages_stream_cache_write.py Normal file
View file

@ -0,0 +1,389 @@
"""
Tests for streaming cache write in AgenticAnthropicStreamingIterator.
Verifies that when a stream completes and the response is rebuilt from SSE events,
the complete response is persisted to the LiteLLM cache.
"""
import asyncio
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
AgenticAnthropicStreamingIterator,
)
def _make_sse_bytes() -> bytes:
"""Create valid SSE bytes representing a complete Anthropic response."""
events = [
(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_test_123",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 10, "output_tokens": 0},
},
},
),
(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
),
(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "Hello!"},
},
),
(
"content_block_stop",
{"type": "content_block_stop", "index": 0},
),
(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 5},
},
),
(
"message_stop",
{"type": "message_stop"},
),
]
lines = []
for event_type, data in events:
lines.append(f"event: {event_type}")
lines.append(f"data: {json.dumps(data)}")
lines.append("")
return "\n".join(lines).encode()
async def _make_async_iter(data: bytes):
"""Create an async iterator that yields the data in one chunk."""
yield data
def _make_mock_logging_obj(
should_store: bool = True,
preset_cache_key: str = "test-cache-key",
):
"""Create a mock logging_obj with _llm_caching_handler attached."""
caching_handler = MagicMock()
caching_handler.request_kwargs = {
"stream": True,
"model": "claude-sonnet-4-20250514",
"messages": [{"role": "user", "content": "Hi"}],
}
caching_handler.preset_cache_key = preset_cache_key
caching_handler.original_function = MagicMock()
caching_handler._should_store_result_in_cache = MagicMock(return_value=should_store)
caching_handler.dual_cache = MagicMock()
logging_obj = MagicMock()
logging_obj._llm_caching_handler = caching_handler
logging_obj.litellm_call_id = "test-call-id"
return logging_obj
def _make_mock_http_handler():
"""Create a mock http_handler whose agentic hook returns None (no follow-up)."""
handler = MagicMock()
handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
return handler
def _create_iterator(
sse_bytes: bytes,
logging_obj=None,
http_handler=None,
):
"""Helper to create an AgenticAnthropicStreamingIterator for testing."""
if logging_obj is None:
logging_obj = _make_mock_logging_obj()
if http_handler is None:
http_handler = _make_mock_http_handler()
return AgenticAnthropicStreamingIterator(
completion_stream=_make_async_iter(sse_bytes),
http_handler=http_handler,
model="claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Hi"}],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=logging_obj,
custom_llm_provider="anthropic",
kwargs={},
)
@pytest.mark.asyncio
async def test_persist_to_cache_called_on_stream_completion():
"""Should persist rebuilt response to cache when stream completes."""
mock_cache = MagicMock()
mock_cache.async_add_cache = AsyncMock(return_value=None)
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj()
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", mock_cache):
# Consume the entire stream
chunks = []
async for chunk in iterator:
chunks.append(chunk)
# Verify we got the SSE bytes
assert len(chunks) >= 1
# Verify async_add_cache was called
mock_cache.async_add_cache.assert_called_once()
# Verify the cached response is the rebuilt dict as JSON
call_args = mock_cache.async_add_cache.call_args
cached_json = call_args[0][0]
cached_response = json.loads(cached_json)
assert cached_response["id"] == "msg_test_123"
assert cached_response["model"] == "claude-sonnet-4-20250514"
assert cached_response["role"] == "assistant"
assert cached_response["stop_reason"] == "end_turn"
assert cached_response["content"][0]["type"] == "text"
assert cached_response["content"][0]["text"] == "Hello!"
assert cached_response["usage"]["input_tokens"] == 10
assert cached_response["usage"]["output_tokens"] == 5
# Verify dual_cache was passed
assert call_args[1]["dynamic_cache_object"] is not None
# Verify cache_key was passed
assert call_args[1]["cache_key"] == "test-cache-key"
@pytest.mark.asyncio
async def test_persist_to_cache_uses_request_cache_key_as_fallback():
"""Should use request_kwargs cache_key when preset_cache_key is None."""
mock_cache = MagicMock()
mock_cache.async_add_cache = AsyncMock(return_value=None)
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj(preset_cache_key=None)
# Set a cache_key in request_kwargs
logging_obj._llm_caching_handler.request_kwargs["cache_key"] = "fallback-key"
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", mock_cache):
async for _ in iterator:
pass
mock_cache.async_add_cache.assert_called_once()
call_kwargs = mock_cache.async_add_cache.call_args[1]
assert call_kwargs["cache_key"] == "fallback-key"
@pytest.mark.asyncio
async def test_persist_to_cache_skipped_when_no_caching_handler():
"""Should not attempt cache write when logging_obj has no _llm_caching_handler."""
mock_cache = MagicMock()
mock_cache.async_add_cache = AsyncMock(return_value=None)
sse_bytes = _make_sse_bytes()
logging_obj = MagicMock()
logging_obj._llm_caching_handler = None
logging_obj.litellm_call_id = "test-call-id"
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", mock_cache):
async for _ in iterator:
pass
mock_cache.async_add_cache.assert_not_called()
@pytest.mark.asyncio
async def test_persist_to_cache_skipped_when_should_store_returns_false():
"""Should not cache when _should_store_result_in_cache returns False."""
mock_cache = MagicMock()
mock_cache.async_add_cache = AsyncMock(return_value=None)
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj(should_store=False)
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", mock_cache):
async for _ in iterator:
pass
mock_cache.async_add_cache.assert_not_called()
@pytest.mark.asyncio
async def test_persist_to_cache_skipped_when_stream_not_in_request_kwargs():
"""Should not cache when request_kwargs.stream is not True."""
mock_cache = MagicMock()
mock_cache.async_add_cache = AsyncMock(return_value=None)
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj()
logging_obj._llm_caching_handler.request_kwargs["stream"] = False
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", mock_cache):
async for _ in iterator:
pass
mock_cache.async_add_cache.assert_not_called()
@pytest.mark.asyncio
async def test_persist_to_cache_skipped_when_litellm_cache_is_none():
"""Should not attempt cache write when litellm.cache is None."""
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj()
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", None):
async for _ in iterator:
pass
# No assertion needed - just verify no exception was raised
assert iterator._response_cached is False
@pytest.mark.asyncio
async def test_persist_to_cache_not_called_twice():
"""Should only persist to cache once even if called multiple times."""
mock_cache = MagicMock()
mock_cache.async_add_cache = AsyncMock(return_value=None)
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj()
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", mock_cache):
async for _ in iterator:
pass
# Call persist again directly - should be a no-op
rebuilt = (
AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
[sse_bytes]
)
)
iterator._persist_to_cache(rebuilt)
# Should only have been called once
mock_cache.async_add_cache.assert_called_once()
@pytest.mark.asyncio
async def test_persist_to_cache_removes_metadata_if_none():
"""Should remove metadata from request_kwargs if it's None."""
mock_cache = MagicMock()
mock_cache.async_add_cache = AsyncMock(return_value=None)
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj()
logging_obj._llm_caching_handler.request_kwargs["metadata"] = None
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", mock_cache):
async for _ in iterator:
pass
call_kwargs = mock_cache.async_add_cache.call_args[1]
assert "metadata" not in call_kwargs
@pytest.mark.asyncio
async def test_persist_to_cache_removes_custom_llm_provider():
"""Should remove custom_llm_provider from request_kwargs."""
mock_cache = MagicMock()
mock_cache.async_add_cache = AsyncMock(return_value=None)
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj()
logging_obj._llm_caching_handler.request_kwargs["custom_llm_provider"] = "anthropic"
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", mock_cache):
async for _ in iterator:
pass
call_kwargs = mock_cache.async_add_cache.call_args[1]
assert "custom_llm_provider" not in call_kwargs
@pytest.mark.asyncio
async def test_persist_to_cache_with_real_local_cache():
"""Should persist to a real in-memory cache and be retrievable."""
from litellm.caching.caching import Cache
real_cache = Cache(type="local")
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj(preset_cache_key="real-cache-test-key")
# Set dual_cache to None so real cache doesn't try to await a mock
logging_obj._llm_caching_handler.dual_cache = None
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", real_cache):
# Consume the entire stream — triggers _persist_to_cache
async for _ in iterator:
pass
# Allow the async cache write task to complete
await asyncio.sleep(0.1)
# Verify the response was actually stored in the real cache
cached = await real_cache.async_get_cache(cache_key="real-cache-test-key")
assert cached is not None
# The cache may return a string or dict depending on the backend
if isinstance(cached, str):
cached_response = json.loads(cached)
else:
cached_response = cached
assert cached_response["id"] == "msg_test_123"
assert cached_response["content"][0]["text"] == "Hello!"
assert cached_response["stop_reason"] == "end_turn"
@pytest.mark.asyncio
async def test_persist_to_cache_handles_exceptions_gracefully():
"""Should not raise when cache write setup fails."""
sse_bytes = _make_sse_bytes()
logging_obj = _make_mock_logging_obj()
# Make _should_store_result_in_cache raise an exception
logging_obj._llm_caching_handler._should_store_result_in_cache.side_effect = (
RuntimeError("unexpected error")
)
iterator = _create_iterator(sse_bytes, logging_obj=logging_obj)
with patch("litellm.cache", MagicMock()):
# Should not raise
async for _ in iterator:
pass
assert iterator._response_cached is False