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 d1e9c521eb into 850fe595ac

Browse source
This commit is contained in:
Dhruv Yadav 2026-04-17 12:26:56 -07:00 • committed by GitHub
commit cb075d2d74
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 303 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -471,6 +471,7 @@ class ChunkProcessor:
cache_read_input_tokens: Optional[int] = None
completion_tokens_details: Optional[CompletionTokensDetails] = None
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
cost: Optional[float] = None
if "prompt_tokens" in usage_chunk:
prompt_tokens = usage_chunk.get("prompt_tokens", 0) or 0
@ -480,6 +481,8 @@ class ChunkProcessor:
cache_creation_input_tokens = usage_chunk.get("cache_creation_input_tokens")
if "cache_read_input_tokens" in usage_chunk:
cache_read_input_tokens = usage_chunk.get("cache_read_input_tokens")
if "cost" in usage_chunk:
cost = usage_chunk.get("cost")
if hasattr(usage_chunk, "completion_tokens_details"):
if isinstance(usage_chunk.completion_tokens_details, dict):
completion_tokens_details = CompletionTokensDetails(
@ -506,6 +509,7 @@ class ChunkProcessor:
"cache_read_input_tokens": cache_read_input_tokens,
"completion_tokens_details": completion_tokens_details,
"prompt_tokens_details": prompt_tokens_details,
"cost": cost,
}
def count_reasoning_tokens(self, response: ModelResponse) -> Optional[int]:
@ -543,9 +547,13 @@ class ChunkProcessor:
web_search_requests: Optional[int] = None
completion_tokens_details: Optional[CompletionTokensDetails] = None
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
cost: Optional[float] = None
for chunk in chunks:
usage_chunk: Optional[Usage] = None
if "usage" in chunk:
if hasattr(chunk, "usage") and chunk.usage is not None:
usage_chunk = chunk.usage
elif "usage" in chunk:
usage_chunk = chunk["usage"]
elif (
isinstance(chunk, ModelResponse)
@ -604,6 +612,9 @@ class ChunkProcessor:
prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"]
if usage_chunk_dict["cost"] is not None:
cost = usage_chunk_dict["cost"]
return UsagePerChunk(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
@ -613,6 +624,7 @@ class ChunkProcessor:
web_search_requests=web_search_requests,
completion_tokens_details=completion_tokens_details,
prompt_tokens_details=prompt_tokens_details,
cost=cost,
)
def calculate_usage(
@ -652,6 +664,7 @@ class ChunkProcessor:
prompt_tokens_details: Optional[
PromptTokensDetailsWrapper
] = calculated_usage_per_chunk["prompt_tokens_details"]
cost: Optional[float] = calculated_usage_per_chunk["cost"]
try:
returned_usage.prompt_tokens = prompt_tokens or token_counter(
@ -720,6 +733,9 @@ class ChunkProcessor:
web_search_requests
)
if cost is not None:
setattr(returned_usage, "cost", cost)
# Return a new usage object with the new values
returned_usage = Usage(**returned_usage.model_dump())

66
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1024,10 +1024,14 @@ class CustomStreamWrapper:
if self.custom_llm_provider == "bedrock" and "trace" in model_response:
return model_response
# Default - return StopIteration
if hasattr(model_response, "usage"):
self.chunks.append(model_response)
raise StopIteration
# Don't raise StopIteration here - some providers (like OpenRouter)
# send usage/cost data in chunks after the finish_reason chunk
if (
hasattr(model_response, "usage")
and model_response.usage is not None
):
return model_response
return
# flush any remaining holding chunk
if len(self.holding_chunk) > 0:
if model_response.choices[0].delta.content is None:
@ -1577,12 +1581,16 @@ class CustomStreamWrapper:
self.tool_call = True
if hasattr(chunk, "usage") and chunk.usage is not None:
model_response.usage = chunk.usage
## RETURN ARG
return self.return_processed_chunk_logic(
result = self.return_processed_chunk_logic(
completion_obj=completion_obj,
model_response=model_response, # type: ignore
response_obj=response_obj,
)
return result
except StopIteration:
raise StopIteration
@ -1820,6 +1828,23 @@ class CustomStreamWrapper:
model_response.choices[0].finish_reason = "tool_calls"
return model_response
@staticmethod
def _propagate_usage_cost_to_hidden_params(
response: "ModelResponse",
) -> None:
"""
If the assembled response carries a provider-reported cost on
usage.cost, copy it into _hidden_params so litellm's cost
calculator uses it instead of a token-based estimate.
"""
_usage = getattr(response, "usage", None)
if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None:
if "additional_headers" not in response._hidden_params:
response._hidden_params["additional_headers"] = {}
response._hidden_params["additional_headers"][
"llm_provider-x-litellm-response-cost"
] = float(_usage.cost)
def __next__(self) -> "ModelResponseStream": # noqa: PLR0915
cache_hit = False
if (
@ -1882,6 +1907,10 @@ class CustomStreamWrapper:
if hasattr(
response, "usage"
): # remove usage from chunk, only send on final chunk
usage_to_preserve = response.usage
if usage_to_preserve:
response._hidden_params["usage"] = usage_to_preserve
# Convert the object to a dictionary
obj_dict = response.model_dump()
@ -1920,6 +1949,10 @@ class CustomStreamWrapper:
response = self.model_response_creator()
if complete_streaming_response is not None:
self._propagate_usage_cost_to_hidden_params(
complete_streaming_response
)
setattr(
response,
"usage",
@ -2146,6 +2179,10 @@ class CustomStreamWrapper:
response = self.model_response_creator()
if complete_streaming_response is not None:
self._propagate_usage_cost_to_hidden_params(
complete_streaming_response
)
setattr(
response,
"usage",
@ -2363,12 +2400,16 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage:
"""Assume most recent usage chunk has total usage uptil then."""
prompt_tokens: int = 0
completion_tokens: int = 0
latest_usage_chunk = None
for chunk in chunks:
if "usage" in chunk and chunk["usage"] is not None:
if "prompt_tokens" in chunk["usage"]:
prompt_tokens = chunk["usage"].get("prompt_tokens", 0) or 0
if "completion_tokens" in chunk["usage"]:
completion_tokens = chunk["usage"].get("completion_tokens", 0) or 0
usage = chunk["usage"]
latest_usage_chunk = usage
if "prompt_tokens" in usage:
prompt_tokens = usage.get("prompt_tokens", 0) or 0
if "completion_tokens" in usage:
completion_tokens = usage.get("completion_tokens", 0) or 0
returned_usage_chunk = Usage(
prompt_tokens=prompt_tokens,
@ -2376,6 +2417,13 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage:
total_tokens=prompt_tokens + completion_tokens,
)
if (
latest_usage_chunk
and hasattr(latest_usage_chunk, "cost")
and latest_usage_chunk.cost is not None
):
returned_usage_chunk.cost = latest_usage_chunk.cost
return returned_usage_chunk

1
litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -14,3 +14,4 @@ class UsagePerChunk(TypedDict):
web_search_requests: Optional[int]
completion_tokens_details: Optional[CompletionTokensDetails]
prompt_tokens_details: Optional[PromptTokensDetailsWrapper]
cost: Optional[float]

10
litellm/types/utils.py
View file

@ -1829,16 +1829,19 @@ class ModelResponseStream(ModelResponseBase):
else:
created = created
usage_to_set = None
if "usage" in kwargs and kwargs["usage"] is not None:
if isinstance(kwargs["usage"], dict):
kwargs["usage"] = Usage(**kwargs["usage"])
usage_to_set = Usage(**kwargs["usage"])
kwargs["usage"] = usage_to_set
elif isinstance(kwargs["usage"], BaseModel):
dump = (
kwargs["usage"].model_dump()
if hasattr(kwargs["usage"], "model_dump")
else kwargs["usage"].dict()
)
kwargs["usage"] = Usage(**dump)
usage_to_set = Usage(**dump)
kwargs["usage"] = usage_to_set
kwargs["id"] = id
kwargs["created"] = created
@ -1847,6 +1850,9 @@ class ModelResponseStream(ModelResponseBase):
super().__init__(**kwargs)
if usage_to_set is not None:
self.usage = usage_to_set
def __contains__(self, key):
# Define custom behavior for the 'in' operator
return hasattr(self, key)

68
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -183,7 +183,11 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks():
make_chunk(role="assistant", content=None),
make_chunk(
thinking_blocks=[
{"type": "thinking", "thinking": "Step 1 analysis...", "signature": None}
{
"type": "thinking",
"thinking": "Step 1 analysis...",
"signature": None,
}
]
),
make_chunk(
@ -201,7 +205,11 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks():
),
make_chunk(
thinking_blocks=[
{"type": "thinking", "thinking": "Step 2 analysis...", "signature": None}
{
"type": "thinking",
"thinking": "Step 2 analysis...",
"signature": None,
}
]
),
make_chunk(
@ -402,7 +410,7 @@ def test_stream_chunk_builder_litellm_usage_chunks():
def test_get_model_from_chunks_azure_model_router():
"""
Test that _get_model_from_chunks finds the actual model from Azure Model Router chunks.
Azure Model Router returns the request model (e.g., 'azure-model-router') in the first chunk,
but subsequent chunks contain the actual model (e.g., 'gpt-4.1-nano-2025-04-14').
This is important for accurate cost calculation.
@ -413,24 +421,24 @@ def test_get_model_from_chunks_azure_model_router():
{"model": "gpt-4.1-nano-2025-04-14", "id": "chatcmpl-123", "choices": []},
{"model": "gpt-4.1-nano-2025-04-14", "id": "chatcmpl-123", "choices": []},
]
result = ChunkProcessor._get_model_from_chunks(
chunks=chunks, first_chunk_model="azure-model-router"
)
# Should return the actual model, not the request model
assert result == "gpt-4.1-nano-2025-04-14"
# Test when all chunks have the same model (non-router case)
chunks_same_model = [
{"model": "gpt-4", "id": "chatcmpl-456", "choices": []},
{"model": "gpt-4", "id": "chatcmpl-456", "choices": []},
]
result_same = ChunkProcessor._get_model_from_chunks(
chunks=chunks_same_model, first_chunk_model="gpt-4"
)
# Should return the first chunk's model when all are the same
assert result_same == "gpt-4"
@ -511,8 +519,8 @@ def test_stream_chunk_builder_anthropic_web_search():
assert usage.prompt_tokens == 50
assert usage.completion_tokens == 27
assert usage.total_tokens == 77
assert usage.server_tool_use['web_search_requests'] == 2
assert usage.total_tokens == 77
assert usage.server_tool_use["web_search_requests"] == 2
def test_sort_chunks_handles_dict_hidden_params_created_at():
@ -602,4 +610,42 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields():
response = stream_chunk_builder(chunks=[chunk_dict])
assert response is not None
assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default"
assert (
response._hidden_params["provider_specific_fields"]["traffic_type"] == "default"
)
def test_cost_field_in_usage_chunks():
chunk1_usage = Usage(completion_tokens=1, prompt_tokens=10, total_tokens=11)
chunk1 = ModelResponseStream(
id="chatcmpl-1",
created=1745513206,
model="openrouter/claude",
choices=[
StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
],
usage=chunk1_usage,
)
chunk2_usage = Usage(
completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
)
chunk2 = ModelResponseStream(
id="chatcmpl-1",
created=1745513207,
model="openrouter/claude",
choices=[
StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
],
usage=chunk2_usage,
)
processor = ChunkProcessor(chunks=[chunk1, chunk2])
usage = processor.calculate_usage(
chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi"
)
assert hasattr(usage, "cost")
assert usage.cost == 0.00025
assert usage.prompt_tokens == 10
assert usage.completion_tokens == 5

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

@ -1171,6 +1171,169 @@ def test_has_any_special_delta_attributes(
assert result is False
def test_calculate_total_usage_with_cost():
from litellm.litellm_core_utils.streaming_handler import calculate_total_usage
chunk1_usage = Usage(completion_tokens=1, prompt_tokens=10, total_tokens=11)
chunk1 = ModelResponseStream(
id="test-1",
created=1745513206,
model="openrouter/test",
choices=[
StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
],
usage=chunk1_usage,
)
chunk2_usage = Usage(
completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
)
chunk2 = ModelResponseStream(
id="test-1",
created=1745513207,
model="openrouter/test",
choices=[
StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
],
usage=chunk2_usage,
)
usage = calculate_total_usage([chunk1, chunk2])
assert hasattr(usage, "cost")
assert usage.cost == 0.00025
assert usage.prompt_tokens == 10
assert usage.completion_tokens == 5
@pytest.mark.asyncio
async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Logging):
from litellm.utils import ModelResponseListIterator
chunk1 = ModelResponseStream(
id="chatcmpl-or",
created=1742056047,
model="openrouter/claude",
choices=[
StreamingChoices(
finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant")
)
],
usage=None,
)
chunk2 = ModelResponseStream(
id="chatcmpl-or",
created=1742056048,
model="openrouter/claude",
choices=[
StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
],
usage=None,
)
chunk3_usage = Usage(
completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
)
chunk3 = ModelResponseStream(
id="chatcmpl-or",
created=1742056049,
model="openrouter/claude",
choices=[
StreamingChoices(finish_reason=None, index=0, delta=Delta(content=""))
],
usage=chunk3_usage,
)
completion_stream = ModelResponseListIterator(
model_responses=[chunk1, chunk2, chunk3]
)
response = CustomStreamWrapper(
completion_stream=completion_stream,
model="openrouter/claude",
custom_llm_provider="openrouter",
logging_obj=logging_obj,
stream_options={"include_usage": True},
)
collected_chunks = []
async for chunk in response:
collected_chunks.append(chunk)
usage_chunks = [c for c in collected_chunks if hasattr(c, "usage") and c.usage]
assert len(usage_chunks) > 0
assert hasattr(usage_chunks[-1].usage, "cost")
assert usage_chunks[-1].usage.cost == 0.00025
def test_openrouter_streaming_cost_propagates_to_hidden_params():
"""
Verify that provider-reported cost from usage.cost flows into
_hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"]
on the complete streaming response, so litellm's cost calculator uses it.
"""
import litellm
chunk1 = ModelResponseStream(
id="chatcmpl-or",
created=1742056047,
model="openrouter/claude",
choices=[
StreamingChoices(
finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant")
)
],
usage=None,
)
chunk2 = ModelResponseStream(
id="chatcmpl-or",
created=1742056048,
model="openrouter/claude",
choices=[
StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
],
usage=None,
)
chunk3 = ModelResponseStream(
id="chatcmpl-or",
created=1742056049,
model="openrouter/claude",
choices=[
StreamingChoices(finish_reason=None, index=0, delta=Delta(content=""))
],
usage=Usage(
completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
),
)
# Build the complete response as stream_chunk_builder does
complete_response = litellm.stream_chunk_builder(
chunks=[chunk1, chunk2, chunk3],
messages=[{"role": "user", "content": "test"}],
)
assert complete_response is not None
assert hasattr(complete_response.usage, "cost")
assert complete_response.usage.cost == 0.00025
# Use the real propagation method from CustomStreamWrapper
CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response)
assert "additional_headers" in complete_response._hidden_params
assert (
complete_response._hidden_params["additional_headers"][
"llm_provider-x-litellm-response-cost"
]
== 0.00025
)
# Verify the cost calculator would pick this up
from litellm.cost_calculator import get_response_cost_from_hidden_params
provider_cost = get_response_cost_from_hidden_params(
complete_response._hidden_params
)
assert provider_cost == 0.00025
def test_handle_special_delta_attributes(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):