mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(caching): log cached anthropic stream replay only once
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
d2a5de2e04
commit
b6cf4066f4
2 changed files with 36 additions and 1 deletions
|
|
@ -96,6 +96,7 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat
|
|||
super().__init__(litellm_logging_obj=litellm_logging_obj, request_body=request_body)
|
||||
self.chunks: list[bytes] = [event.encode("utf-8") for event in events]
|
||||
self.current_index = 0
|
||||
self.logged = False
|
||||
self._hidden_params: dict[str, Any] = {"cache_hit": True}
|
||||
litellm_logging_obj.model_call_details["cache_hit"] = True
|
||||
|
||||
|
|
@ -104,7 +105,9 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat
|
|||
|
||||
async def __anext__(self) -> bytes:
|
||||
if self.current_index >= len(self.chunks):
|
||||
await self._handle_streaming_logging(self.chunks)
|
||||
if not self.logged:
|
||||
self.logged = True
|
||||
await self._handle_streaming_logging(self.chunks)
|
||||
raise StopAsyncIteration
|
||||
chunk = self.chunks[self.current_index]
|
||||
self.current_index += 1
|
||||
|
|
|
|||
|
|
@ -177,3 +177,35 @@ async def test_abandoned_stream_is_not_cached(local_cache, request_kwargs, monke
|
|||
|
||||
assert len(fake_handler.calls) == 2
|
||||
assert replayed == STREAM_EVENTS
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_stream_replay_logs_once_when_polled_after_exhaustion():
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import (
|
||||
CachedAnthropicMessagesStreamIterator,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.streaming_handler import (
|
||||
PassThroughStreamingHandler,
|
||||
)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
iterator = CachedAnthropicMessagesStreamIterator(
|
||||
events=[event.decode("utf-8") for event in STREAM_EVENTS],
|
||||
litellm_logging_obj=logging_obj,
|
||||
request_body={"model": "claude-sonnet-4-5"},
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
PassThroughStreamingHandler,
|
||||
"_route_streaming_logging_to_handler",
|
||||
new=AsyncMock(),
|
||||
) as mock_route:
|
||||
assert await _collect(iterator) == STREAM_EVENTS
|
||||
for _ in range(2):
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await iterator.__anext__()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
mock_route.assert_called_once()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue