scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

test(streaming): assert logged cost from success callback for usage-only chunk

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
milan 2026-07-28 23:12:06 +00:00
parent b2d2b29e2f
commit ef26590f72
1 changed files with 35 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -1508,15 +1508,15 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin
@pytest.mark.asyncio
async def test_openrouter_streaming_usage_only_chunk_without_stream_options(
logging_obj: Logging,
):
async def test_openrouter_streaming_usage_only_chunk_without_stream_options():
"""
Regression: OpenRouter's post-finish chunk has `choices: []`. When the caller did not
pass stream_options.include_usage it was dropped before cost tracking, so the
provider-reported cost never reached the assembled response.
"""
from litellm.cost_calculator import get_response_cost_from_hidden_params
import time
from litellm.integrations.custom_logger import CustomLogger
from litellm.utils import ModelResponseListIterator
chunk1 = ModelResponseStream(
@ -1549,30 +1549,48 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options(
),
)
class MockCallback(CustomLogger):
pass
mock_callback = MockCallback()
litellm.success_callback = [mock_callback]
litellm._async_success_callback = [mock_callback]
stream_logging_obj = Logging(
model="openrouter/claude",
messages=[{"role": "user", "content": "Hey"}],
stream=True,
call_type="completion",
start_time=time.time(),
litellm_call_id="12345",
function_id="1245",
)
stream_logging_obj.update_environment_variables(
model="openrouter/claude",
optional_params={},
litellm_params={},
custom_llm_provider="openrouter",
)
response = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(
model_responses=[chunk1, chunk2, usage_only_chunk]
),
model="openrouter/claude",
custom_llm_provider="openrouter",
logging_obj=logging_obj,
logging_obj=stream_logging_obj,
)
collected_chunks = [chunk async for chunk in response]
with patch.object(mock_callback, "async_log_success_event") as mock_success_event:
collected_chunks = [chunk async for chunk in response]
await asyncio.sleep(1)
assert all(getattr(chunk, "usage", None) is None for chunk in collected_chunks)
complete_response = litellm.stream_chunk_builder(
chunks=response.chunks,
messages=[{"role": "user", "content": "Hey"}],
)
assert complete_response is not None
assert complete_response.usage.cost == 0.00025
CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response)
assert (
get_response_cost_from_hidden_params(complete_response._hidden_params) == 0.00025
)
mock_success_event.assert_called_once()
logged_kwargs = mock_success_event.call_args.kwargs["kwargs"]
assert logged_kwargs["response_cost"] == 0.00025
assert logged_kwargs["standard_logging_object"]["response_cost"] == 0.00025
def test_openrouter_streaming_cost_propagates_to_hidden_params():