diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py index dc2d4408c20..c2fa1388557 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py @@ -1,7 +1,7 @@ import re from collections.abc import AsyncIterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, cast import litellm from litellm._logging import verbose_logger @@ -50,6 +50,24 @@ class AnthropicMessagesStreamCacheWriter: def has_buffered_provider_output(self) -> bool: return getattr(self.stream, "has_buffered_provider_output", False) is True + @property + def chunks(self) -> list | None: + return cast( # cast-ok: the billing helper itself treats chunks as an opaque getattr + "list | None", getattr(self.stream, "chunks", None) + ) + + @property + def messages(self) -> list | None: + return cast( # cast-ok: messages is a plain list on the inner stream + "list | None", getattr(self.stream, "messages", None) + ) + + @property + def model(self) -> str | None: + return cast( # cast-ok: model is a str on the inner stream + "str | None", getattr(self.stream, "model", None) + ) + def __aiter__(self) -> "AnthropicMessagesStreamCacheWriter": return self diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py index 22d14614108..fe22890f1b5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py @@ -278,6 +278,40 @@ class _HeldBackStream: raise StopAsyncIteration +class _AttributedStream: + """Stream stub carrying the billing attributes the disconnect helper reads.""" + + def __init__(self, chunks: list) -> None: + self.chunks = [object()] + self.messages = [{"role": "user", "content": "hi"}] + self.model = "gpt-4o-mini" + self._pending = list(chunks) + + def __aiter__(self) -> "_AttributedStream": + return self + + async def __anext__(self) -> bytes: + if not self._pending: + raise StopAsyncIteration + return self._pending.pop(0) + + +@pytest.mark.asyncio +async def test_cache_writer_exposes_inner_stream_billing_attributes(request_kwargs): + caching_handler = LLMCachingHandler( + original_function=handler.anthropic_messages, + request_kwargs=dict(request_kwargs), + start_time=datetime.datetime.now(), + ) + inner = _AttributedStream(STREAM_EVENTS) + writer = AnthropicMessagesStreamCacheWriter(stream=inner, caching_handler=caching_handler) + + assert writer.chunks is inner.chunks + assert writer.messages is inner.messages + assert writer.model == "gpt-4o-mini" + assert await _collect(writer) == STREAM_EVENTS + + def test_cache_writer_forwards_has_buffered_provider_output(request_kwargs): caching_handler = LLMCachingHandler( original_function=handler.anthropic_messages,