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

Fix Anthropic streaming web search cost usage

Browse source 
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
oss-agent-shin 2026-05-06 17:47:07 +00:00
parent b83d11351f
commit 4853cfd25d
No known key found for this signature in database
3 changed files with 105 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -464,12 +464,15 @@ class ChunkProcessor:
id=id,
)
def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict:
def _usage_chunk_calculation_helper(
self, usage_chunk: Union[Usage, Dict[str, Any]]
) -> dict:
prompt_tokens = 0
completion_tokens = 0
## anthropic prompt caching information ##
cache_creation_input_tokens: Optional[int] = None
cache_read_input_tokens: Optional[int] = None
server_tool_use: Optional[ServerToolUse] = None
completion_tokens_details: Optional[CompletionTokensDetails] = None
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
@ -499,12 +502,19 @@ class ChunkProcessor:
usage_chunk.prompt_tokens_details, PromptTokensDetailsWrapper
):
prompt_tokens_details = usage_chunk.prompt_tokens_details
if "server_tool_use" in usage_chunk:
raw_server_tool_use = usage_chunk.get("server_tool_use")
if isinstance(raw_server_tool_use, dict):
server_tool_use = ServerToolUse(**raw_server_tool_use)
elif isinstance(raw_server_tool_use, ServerToolUse):
server_tool_use = raw_server_tool_use
return {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"cache_creation_input_tokens": cache_creation_input_tokens,
"cache_read_input_tokens": cache_read_input_tokens,
"server_tool_use": server_tool_use,
"completion_tokens_details": completion_tokens_details,
"prompt_tokens_details": prompt_tokens_details,
}
@ -584,11 +594,8 @@ class ChunkProcessor:
completion_tokens_details = usage_chunk_dict[
"completion_tokens_details"
]
if (
hasattr(usage_chunk, "server_tool_use")
and usage_chunk.server_tool_use is not None
):
server_tool_use = usage_chunk.server_tool_use
if usage_chunk_dict["server_tool_use"] is not None:
server_tool_use = usage_chunk_dict["server_tool_use"]
if (
usage_chunk_dict["prompt_tokens_details"] is not None
and getattr(

7
litellm/types/utils.py
View file

@ -1553,7 +1553,7 @@ class Usage(SafeAttributeModel, CompletionUsage):
completion_tokens_details: Optional[
Union[CompletionTokensDetailsWrapper, dict]
] = None,
server_tool_use: Optional[ServerToolUse] = None,
server_tool_use: Optional[Union[ServerToolUse, dict]] = None,
cost: Optional[float] = None,
**params,
):
@ -1655,7 +1655,10 @@ class Usage(SafeAttributeModel, CompletionUsage):
)
if server_tool_use is not None:
self.server_tool_use = server_tool_use
if isinstance(server_tool_use, dict):
self.server_tool_use = ServerToolUse(**server_tool_use)
else:
self.server_tool_use = server_tool_use
else: # maintain openai compatibility in usage object if possible
del self.server_tool_use

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

@ -1,14 +1,11 @@
import json
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm import stream_chunk_builder
from litellm import completion_cost, stream_chunk_builder
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
@ -520,7 +517,92 @@ 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.server_tool_use is not None
assert usage.server_tool_use.web_search_requests == 2
def test_stream_chunk_builder_anthropic_web_search_completion_cost():
chunks = [
ModelResponseStream(
id="chatcmpl-mocked-usage-1",
created=1745513206,
model="claude-sonnet-4-6",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
provider_specific_fields=None,
content="hello",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields=None,
stream_options={"include_usage": True},
usage=Usage(
completion_tokens=0,
prompt_tokens=50,
total_tokens=50,
completion_tokens_details=None,
server_tool_use=ServerToolUse(web_search_requests=2),
prompt_tokens_details=None,
),
),
ModelResponseStream(
id="chatcmpl-mocked-usage-1",
created=1745513207,
model="claude-sonnet-4-6",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(
provider_specific_fields=None,
content=None,
role=None,
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields=None,
stream_options={"include_usage": True},
usage=Usage(
completion_tokens=27,
prompt_tokens=0,
total_tokens=27,
completion_tokens_details=None,
prompt_tokens_details=None,
),
),
]
response = stream_chunk_builder(chunks)
assert response is not None
assert response.usage is not None
assert response.usage.server_tool_use is not None
assert response.usage.server_tool_use.web_search_requests == 2
assert completion_cost(completion_response=response) > 0
dict_response = stream_chunk_builder([chunk.model_dump() for chunk in chunks])
assert dict_response is not None
assert dict_response.usage is not None
assert dict_response.usage.server_tool_use is not None
assert dict_response.usage.server_tool_use.web_search_requests == 2
assert completion_cost(completion_response=dict_response) > 0
def test_sort_chunks_handles_dict_hidden_params_created_at():