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

Merge pull request #41171 from BerriAI/litellm_converted_stream_spend_tracking

Browse source 
fix(logging): track spend for streams a deployment hook converted to non-streaming
This commit is contained in:
Mateo Wang 2026-09-15 15:34:13 -07:00 • committed by GitHub
commit c3222ec110
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 342 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/caching/caching_handler.py
View file

@ -35,6 +35,7 @@ from litellm.litellm_core_utils.logging_utils import (
_assemble_complete_response_from_streaming_chunks,
)
from litellm.types.caching import CachedEmbedding
from litellm.types.integrations.custom_logger import converted_stream_requested
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.rerank import RerankResponse
from litellm.types.utils import (
@ -107,17 +108,31 @@ def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
return "choices" in cached_result
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, object]) -> bool:
def _stream_replay_requested(kwargs: Mapping[str, object]) -> bool:
if kwargs.get("stream", False) is True:
return True
return converted_stream_requested(kwargs) and not kwargs.get("_agentic_loop_depth")
def _should_defer_streaming_cache_hit_callbacks(*, cached_result: object) -> bool:
"""
When stream=True, do not run success callbacks at cache-hit time.
When the cache hit is replayed as a stream, do not run success callbacks at cache-hit time.
Cached chat/text completion replay uses CustomStreamWrapper; cached Responses
replay uses CachedResponsesAPIStreamingIterator; cached Anthropic Messages
replay uses CachedAnthropicMessagesStreamIterator. All invoke logging success
handlers when the stream finishes; firing them here too would double-count
spend and callback records.
spend and callback records. A plain (non-stream) replay logs here, since nothing
else will.
"""
return kwargs.get("stream", False) is True
from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import (
CachedAnthropicMessagesStreamIterator,
)
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
return isinstance(
cached_result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator, CachedAnthropicMessagesStreamIterator)
)
def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") -> Mapping[str, object]:
@ -267,7 +282,7 @@ class LLMCachingHandler:
custom_llm_provider=kwargs.get("custom_llm_provider", None),
args=args,
)
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
# LOG SUCCESS
self._async_log_cache_hit_on_callbacks(
logging_obj=logging_obj,
@ -383,7 +398,7 @@ class LLMCachingHandler:
is_async=False,
)
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=cached_result,
start_time=start_time,
@ -823,7 +838,7 @@ class LLMCachingHandler:
if (call_type == CallTypes.acompletion.value or call_type == CallTypes.completion.value) and isinstance(
cached_result, dict
):
if kwargs.get("stream", False) is True:
if _stream_replay_requested(kwargs):
cached_result = self._convert_cached_stream_response(
cached_result=cached_result,
call_type=call_type,
@ -838,7 +853,7 @@ class LLMCachingHandler:
if (
call_type == CallTypes.atext_completion.value or call_type == CallTypes.text_completion.value
) and isinstance(cached_result, dict):
if kwargs.get("stream", False) is True:
if _stream_replay_requested(kwargs):
cached_result = self._convert_cached_stream_response(
cached_result=cached_result,
call_type=call_type,
@ -893,7 +908,7 @@ class LLMCachingHandler:
elif (call_type == "aresponses" or call_type == "responses") and isinstance(cached_result, dict):
use_chat_completion_cache: Final = _is_chat_completion_cached_dict(cached_result)
if use_chat_completion_cache:
if kwargs.get("stream", False) is True:
if _stream_replay_requested(kwargs):
bridge_call_type: Final = (
CallTypes.acompletion.value if call_type == "aresponses" else CallTypes.completion.value
)
@ -921,7 +936,7 @@ class LLMCachingHandler:
):
response_obj._hidden_params["cache_hit"] = True
if kwargs.get("stream", False) is True:
if _stream_replay_requested(kwargs):
cached_result = CachedResponsesAPIStreamingIterator(
response=response_obj,
logging_obj=logging_obj,

5
litellm/responses/streaming_iterator.py
View file

@ -33,6 +33,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
from litellm.types.integrations.custom_logger import converted_stream_requested
from litellm.types.llms.openai import (
PART_UNION_TYPES,
ResponseAPIUsage,
@ -626,7 +627,9 @@ class BaseResponsesAPIStreamingIterator:
return
request_kwargs = getattr(caching_handler, "request_kwargs", None)
if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True:
if not _is_json_object(request_kwargs):
return
if request_kwargs.get("stream") is not True and not converted_stream_requested(request_kwargs):
return
request_kwargs = request_kwargs.copy()
preset_cache_key = getattr(caching_handler, "preset_cache_key", None)

17
litellm/utils.py
View file

@ -846,6 +846,13 @@ def _is_streaming_response_for_correlation(result: object) -> bool:
return isinstance(result, CustomStreamWrapper)
def _is_converted_stream_result(result: object) -> bool:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator))
# Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
def function_setup(
original_function: str,
@ -1889,6 +1896,9 @@ def client(original_function):
_caching_handler_response.cached_result is not None
and _caching_handler_response.final_embedding_cached_response is None
):
if _is_converted_stream_result(_caching_handler_response.cached_result):
logging_obj.stream = True
logging_obj.model_call_details["stream"] = True
return _caching_handler_response.cached_result
elif _caching_handler_response.embedding_all_elements_cache_hit is True:
@ -1946,10 +1956,9 @@ def client(original_function):
raise
end_time = datetime.datetime.now()
if _is_streaming_request(
kwargs=kwargs,
call_type=call_type,
):
if _is_streaming_request(kwargs=kwargs, call_type=call_type) or _is_converted_stream_result(result):
logging_obj.stream = True
logging_obj.model_call_details["stream"] = True
if "complete_response" in kwargs and kwargs["complete_response"] is True:
chunks: Final = []
for idx, chunk in enumerate(result):

24
tests/local_testing/test_caching_handler.py
View file

@ -927,24 +927,14 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request():
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={"stream": True},
)
is True
)
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={"stream": False},
)
is False
)
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={},
)
is False
logging_obj = MagicMock()
logging_obj.model_call_details = {}
stream_replay = CustomStreamWrapper(
completion_stream=iter(()), model="gpt-4o", logging_obj=logging_obj
)
assert _should_defer_streaming_cache_hit_callbacks(cached_result=stream_replay) is True
assert _should_defer_streaming_cache_hit_callbacks(cached_result=ModelResponse()) is False
assert _should_defer_streaming_cache_hit_callbacks(cached_result={"id": "msg_1"}) is False
@pytest.mark.asyncio

87
tests/test_litellm/caching/test_caching_handler.py
View file

@ -693,3 +693,90 @@ async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monke
assert handler.preset_cache_key is not None
assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key
assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key
@pytest.mark.asyncio
async def test_converted_stream_cache_hit_replayed_as_plain_object_logs_at_hit_time(monkeypatch):
import litellm
from litellm.caching.caching import Cache
from litellm.types.utils import CallTypes
async def aanthropic_messages(**kwargs):
return None
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
kwargs = {
"model": "claude-sonnet-5",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 16,
"caching": True,
"stream": False,
"_websearch_interception_converted_stream": True,
}
cached_message = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "hi"}],
}
await litellm.cache.async_add_cache(cached_message, **kwargs)
handler = LLMCachingHandler(original_function=aanthropic_messages, request_kwargs=kwargs, start_time=datetime.now())
logging_obj = _build_logging_obj(CallTypes.aanthropic_messages.value, stream=False)
logging_obj.async_success_handler = AsyncMock()
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
hit = await handler._async_get_cache(
model="claude-sonnet-5",
original_function=aanthropic_messages,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.aanthropic_messages.value,
kwargs=kwargs,
args=(),
)
assert hit is not None and hit.cached_result == cached_message
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
@pytest.mark.asyncio
async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_replays_as_plain_object(monkeypatch):
import litellm
from litellm.caching.caching import Cache
from litellm.types.utils import CallTypes
async def acompletion(**kwargs):
return None
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
kwargs = {
"model": "gpt-5.6",
"messages": [{"role": "user", "content": "run the code"}],
"caching": True,
"stream": False,
"_code_interpreter_interception_converted_stream": True,
"_agentic_loop_depth": 1,
}
await litellm.cache.async_add_cache(
litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]), **kwargs
)
handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now())
logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False)
logging_obj.async_success_handler = AsyncMock()
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
hit = await handler._async_get_cache(
model="gpt-5.6",
original_function=acompletion,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.acompletion.value,
kwargs=kwargs,
args=(),
)
assert hit is not None and isinstance(hit.cached_result, litellm.ModelResponse)
assert hit.cached_result.choices[0].message.content == "done"
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True

206
tests/test_litellm/test_utils.py
View file

@ -20,6 +20,8 @@ from jsonschema import validate
import litellm
from litellm._internal_context import is_internal_call
from litellm.caching.caching import Cache
from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT
from litellm._logging import (
CorrelationContextFilter,
@ -32,6 +34,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor
from litellm.proxy.utils import is_valid_api_key
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY
from litellm.types.utils import (
CallTypes,
Delta,
@ -44,6 +47,7 @@ from litellm.types.utils import (
from litellm.types.utils import all_litellm_params, bedrock_batch_litellm_params
from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams
from litellm.utils import (
CustomStreamWrapper,
ProviderConfigManager,
TextCompletionStreamWrapper,
_check_provider_match,
@ -4953,6 +4957,208 @@ async def test_wrapper_async_restores_originating_task_context_after_success(mon
session_id_var.set("")
class _ConvertStreamDeploymentHook(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict[str, object] | None:
if not kwargs.get("stream"):
return None
return {**kwargs, "stream": False, HEADROOM_CONVERTED_STREAM_KEY: True}
class _SuccessKwargsCapture(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.success_kwargs: list[dict[str, object]] = []
self.stream_event_responses: list[object] = []
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.success_kwargs.append(kwargs)
async def async_log_stream_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.stream_event_responses.append(response_obj)
def _install_converted_stream_callbacks(monkeypatch: pytest.MonkeyPatch) -> _SuccessKwargsCapture:
capture: Final = _SuccessKwargsCapture()
monkeypatch.setattr(litellm, "callbacks", [_ConvertStreamDeploymentHook(), capture])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
return capture
async def _wait_for_success_kwargs(capture: _SuccessKwargsCapture, count: int = 1) -> dict[str, object]:
for _ in range(50):
if len(capture.success_kwargs) >= count and not _PENDING_CACHE_WRITES:
break
await asyncio.sleep(0.05)
await asyncio.sleep(0.2)
assert len(capture.success_kwargs) == count
return capture.success_kwargs[-1]
def _assert_cache_hit_logged_as_stream(capture: _SuccessKwargsCapture, success_kwargs: dict[str, object]) -> None:
standard_logging_object: Final = success_kwargs["standard_logging_object"]
assert isinstance(standard_logging_object, dict)
assert standard_logging_object["cache_hit"] is True
assert standard_logging_object["stream"] is True
assert success_kwargs["stream"] is True
assert capture.stream_event_responses == []
@pytest.mark.asyncio
async def test_wrapper_async_logs_converted_chat_stream_with_standard_logging_object(
monkeypatch: pytest.MonkeyPatch,
) -> None:
capture: Final = _install_converted_stream_callbacks(monkeypatch)
response: Final = await litellm.acompletion(
model="gpt-5.6",
messages=[{"role": "user", "content": "hi"}],
stream=True,
mock_response="converted stream body",
num_retries=0,
)
assert isinstance(response, CustomStreamWrapper)
chunks: Final = [chunk async for chunk in response]
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "converted stream body"
success_kwargs: Final = await _wait_for_success_kwargs(capture)
standard_logging_object: Final = success_kwargs["standard_logging_object"]
assert isinstance(standard_logging_object, dict)
assert standard_logging_object["response_cost"] > 0
assert standard_logging_object["stream"] is True
assert success_kwargs["stream"] is True
@pytest.mark.asyncio
@respx.mock
async def test_wrapper_async_logs_converted_responses_stream_with_standard_logging_object(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
capture: Final = _install_converted_stream_callbacks(monkeypatch)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
respx.post("https://api.openai.com/v1/responses").respond(
json={
"id": "resp_converted",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-5.6",
"output": [
{
"type": "message",
"id": "msg_converted",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "converted stream body", "annotations": []}],
}
],
"usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7},
}
)
response: Final = await litellm.aresponses(
model="openai/gpt-5.6", input="hi", stream=True, api_key="sk-test", num_retries=0
)
assert isinstance(response, BaseResponsesAPIStreamingIterator)
events: Final = [event async for event in response]
assert events[-1].type == "response.completed"
success_kwargs: Final = await _wait_for_success_kwargs(capture)
standard_logging_object: Final = success_kwargs["standard_logging_object"]
assert isinstance(standard_logging_object, dict)
assert standard_logging_object["response_cost"] > 0
assert standard_logging_object["stream"] is True
assert success_kwargs["stream"] is True
@pytest.mark.asyncio
async def test_wrapper_async_replays_cached_converted_chat_stream_as_stream(
monkeypatch: pytest.MonkeyPatch,
) -> None:
capture: Final = _install_converted_stream_callbacks(monkeypatch)
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
request: Final = {
"model": "gpt-5.6",
"messages": [{"role": "user", "content": "replay me from cache"}],
"stream": True,
"mock_response": "converted stream body",
"num_retries": 0,
}
first: Final = await litellm.acompletion(**request)
first_chunks: Final = [chunk async for chunk in first]
assert "".join(chunk.choices[0].delta.content or "" for chunk in first_chunks) == "converted stream body"
await _wait_for_success_kwargs(capture)
replay: Final = await litellm.acompletion(**request)
assert isinstance(replay, CustomStreamWrapper)
replay_chunks: Final = [chunk async for chunk in replay]
assert "".join(chunk.choices[0].delta.content or "" for chunk in replay_chunks) == "converted stream body"
_assert_cache_hit_logged_as_stream(capture, await _wait_for_success_kwargs(capture, count=2))
@pytest.mark.asyncio
@respx.mock
async def test_wrapper_async_replays_cached_converted_responses_stream_as_stream(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
capture: Final = _install_converted_stream_callbacks(monkeypatch)
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
route: Final = respx.post("https://api.openai.com/v1/responses").respond(
json={
"id": "resp_cached_converted",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-5.6",
"output": [
{
"type": "message",
"id": "msg_cached_converted",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "converted stream body", "annotations": []}],
}
],
"usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7},
}
)
request: Final = {
"model": "openai/gpt-5.6",
"input": "replay me from cache",
"stream": True,
"api_key": "sk-test",
"num_retries": 0,
}
first: Final = await litellm.aresponses(**request)
assert [event async for event in first][-1].type == "response.completed"
await _wait_for_success_kwargs(capture)
replay: Final = await litellm.aresponses(**request)
assert isinstance(replay, BaseResponsesAPIStreamingIterator)
assert [event async for event in replay][-1].type == "response.completed"
assert route.call_count == 1
_assert_cache_hit_logged_as_stream(capture, await _wait_for_success_kwargs(capture, count=2))
def test_function_setup_failure_after_logging_construction_restores_context(monkeypatch):
"""If function_setup() constructs Logging() (which already mutated
trace_id_var/session_id_var in __init__) but then raises before returning,