diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index 9add77ff424..6040bcafbcc 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -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 diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 11733ce4cee..8dd0b2da0ea 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -80,6 +80,7 @@ class Cache: "rerank", "responses", "aresponses", + "anthropic_messages", ], # s3 Bucket, boto3 configuration azure_account_url: Optional[str] = None, diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 3cf1d911d7f..55645b3bb68 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -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 diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index b4fa5cb60aa..92d04a6d78e 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -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]: """ diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index d0780c82d06..ba74ebdfdb7 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -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', '')})" diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py index f704ed2c9d1..dce9ed92a56 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py @@ -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 diff --git a/litellm/types/caching.py b/litellm/types/caching.py index f8050b292c7..7158fabe66f 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -29,6 +29,7 @@ CachingSupportedCallTypes = Literal[ "rerank", "responses", "aresponses", + "anthropic_messages", ] diff --git a/tests/test_litellm/test_caching/__init__.py b/tests/test_litellm/test_caching/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/test_caching/test_anthropic_messages_cache_key.py b/tests/test_litellm/test_caching/test_anthropic_messages_cache_key.py new file mode 100644 index 00000000000..df5f809948c --- /dev/null +++ b/tests/test_litellm/test_caching/test_anthropic_messages_cache_key.py @@ -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 diff --git a/tests/test_litellm/test_caching/test_anthropic_messages_cache_read.py b/tests/test_litellm/test_caching/test_anthropic_messages_cache_read.py new file mode 100644 index 00000000000..29a94c4f39d --- /dev/null +++ b/tests/test_litellm/test_caching/test_anthropic_messages_cache_read.py @@ -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 diff --git a/tests/test_litellm/test_caching/test_anthropic_messages_caching_e2e.py b/tests/test_litellm/test_caching/test_anthropic_messages_caching_e2e.py new file mode 100644 index 00000000000..12fde086e1f --- /dev/null +++ b/tests/test_litellm/test_caching/test_anthropic_messages_caching_e2e.py @@ -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!" diff --git a/tests/test_litellm/test_caching/test_anthropic_messages_stream_cache_write.py b/tests/test_litellm/test_caching/test_anthropic_messages_stream_cache_write.py new file mode 100644 index 00000000000..cf26d0ecff7 --- /dev/null +++ b/tests/test_litellm/test_caching/test_anthropic_messages_stream_cache_write.py @@ -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