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

fix(passthrough,streaming): recover cost on interrupted and agentic Anthropic streams (#31035)

Browse source 
Streaming and pass-through requests could be logged with $0 cost or dropped from
SpendLogs entirely while the upstream provider still billed every token. This
closes the leak paths not already covered by #30160, #30787 and #30788.

- Catch a stream_chunk_builder raise in the core CustomStreamWrapper (sync and
  async). Large agentic tool-use / thinking streams can make assembly re-raise
  as APIError from inside the except-StopIteration handler, where the sibling
  except does not catch it, so it escaped __next__/__anext__ and dropped the
  request; recover best-effort usage from the raw chunks instead
- Add a usage-only fallback for Anthropic streaming pass-through: when
  stream_chunk_builder returns None or raises, rebuild usage from the
  message_start / message_delta SSE events via AnthropicConfig.calculate_usage so
  cache, web-search and geo tokens are priced instead of left at $0
- Decode buffered pass-through bytes with errors="replace" so a stream cut
  mid-multibyte-sequence still logs the usage events already received
- Record response_cost into model_call_details on the pass-through success path
  (it is read from there, not from kwargs), matching the gemini/cohere/openai
  handlers
- Name the key (alias + masked key) in the virtual-key BudgetExceededError so
  operators don't have to reverse-map spend back to a key
This commit is contained in:
Yassin Kortam 2026-06-22 18:51:13 -07:00 • committed by GitHub
parent 1cdb6cd3ac
commit b24b964e04
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 748 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
litellm/litellm_core_utils/streaming_handler.py
View file

@ -2005,11 +2005,29 @@ class CustomStreamWrapper:
except StopIteration:
if self.sent_last_chunk is True:
complete_streaming_response = litellm.stream_chunk_builder(
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,
)
try:
complete_streaming_response = litellm.stream_chunk_builder(
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,
)
except Exception as e:
# stream_chunk_builder can re-raise (as APIError) on large agentic
# streams. The raise originates inside this except-StopIteration block,
# so the sibling `except Exception` below does not catch it; it would
# escape __next__ and drop the request from SpendLogs. Recover
# best-effort usage from the raw chunks so cost is still tracked
verbose_logger.warning(
"stream_chunk_builder raised at end-of-stream (%s); logging "
"best-effort usage from chunks.",
str(e),
)
try:
complete_streaming_response = self.model_response_creator(
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
)
except Exception:
complete_streaming_response = None
response = self.model_response_creator()
if complete_streaming_response is not None:
@ -2234,11 +2252,27 @@ class CustomStreamWrapper:
except (StopAsyncIteration, StopIteration):
if self.sent_last_chunk is True:
# log the final chunk with accurate streaming values
complete_streaming_response = litellm.stream_chunk_builder(
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,
)
try:
complete_streaming_response = litellm.stream_chunk_builder(
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,
)
except Exception as e:
# see sync __next__: a raise from stream_chunk_builder inside this
# except handler escapes __anext__ and drops the request from SpendLogs.
# Recover best-effort usage from the raw chunks so cost is still tracked
verbose_logger.warning(
"stream_chunk_builder raised at end-of-stream (%s); logging "
"best-effort usage from chunks.",
str(e),
)
try:
complete_streaming_response = self.model_response_creator(
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
)
except Exception:
complete_streaming_response = None
response = self.model_response_creator()
if complete_streaming_response is not None:

9
litellm/proxy/auth/auth_checks.py
View file

@ -3663,9 +3663,18 @@ async def _virtual_key_max_budget_check(
# so a NaN max_budget would silently disable enforcement. Treat a
# non-finite max_budget as "no configured limit" rather than as a bypass.
if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget:
# name the key in the error so operators don't have to reverse-map
# spend back to a key; key_name is the masked form (last 4 chars)
key_label = valid_token.key_alias or "key"
key_descriptor = (
f"{key_label} ({valid_token.key_name})"
if valid_token.key_name
else key_label
)
raise litellm.BudgetExceededError(
current_cost=spend,
max_budget=valid_token.max_budget,
message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}",
)

189
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -6,6 +6,7 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -15,12 +16,19 @@ from litellm.llms.anthropic import get_anthropic_config
from litellm.llms.anthropic.chat.handler import (
ModelResponseIterator as AnthropicModelResponseIterator,
)
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
from litellm.types.utils import LiteLLMBatch, ModelResponse, TextCompletionResponse
from litellm.types.utils import (
Choices,
LiteLLMBatch,
Message,
ModelResponse,
TextCompletionResponse,
)
if TYPE_CHECKING:
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
@ -272,6 +280,9 @@ class AnthropicPassthroughLoggingHandler:
kwargs["response_cost"] = response_cost
kwargs["model"] = model
# the pass-through success path reads spend from
# model_call_details["response_cost"], not from kwargs
logging_obj.model_call_details["response_cost"] = response_cost
passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( # type: ignore
kwargs.get("passthrough_logging_payload")
)
@ -343,13 +354,42 @@ class AnthropicPassthroughLoggingHandler:
if chunk_model:
model = chunk_model
complete_streaming_response = (
AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
try:
complete_streaming_response = (
AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
)
)
)
except Exception as e:
# stream_chunk_builder re-raises assembly failures (as litellm.APIError)
# on large agentic tool-use / thinking streams; treat that the same as a
# None result so the usage-only fallback below still recovers cost
verbose_proxy_logger.warning(
"Anthropic passthrough: stream assembly raised (model=%s): %s; falling "
"back to usage-only cost from raw SSE events.",
model,
e,
)
complete_streaming_response = None
if complete_streaming_response is None:
# stream_chunk_builder cannot always reassemble large agentic streams, but
# Anthropic still emits token usage in the message_start / message_delta SSE
# events regardless of content shape; recover usage-only so cost is tracked.
# Guard it too: a raise here would defeat the point and drop the request
try:
complete_streaming_response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=all_chunks,
model=model,
)
except Exception as e:
verbose_proxy_logger.warning(
"Anthropic passthrough: usage-only fallback failed (model=%s): %s",
model,
e,
)
complete_streaming_response = None
if complete_streaming_response is None:
verbose_proxy_logger.error(
"Unable to build complete streaming response for Anthropic passthrough endpoint, not logging..."
@ -636,6 +676,141 @@ class AnthropicPassthroughLoggingHandler:
)
return complete_streaming_response
@staticmethod
def _extract_sse_data(event_str: str) -> Optional[dict]:
"""Parse the JSON object from the ``data:`` line of an Anthropic SSE event."""
for line in event_str.splitlines():
stripped = line.strip()
if stripped.startswith("data:"):
payload = stripped[len("data:") :].strip()
if not payload or payload == "[DONE]":
return None
try:
return cast(dict, json.loads(payload))
except (ValueError, TypeError):
return None
return None
@staticmethod
def _build_usage_only_response_from_chunks(
all_chunks: Sequence[Union[str, bytes]],
model: str,
) -> Optional[ModelResponse]:
"""
Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for
cost tracking when stream_chunk_builder cannot reassemble the stream.
Anthropic emits usage in ``message_start`` (uncached input + cache tokens, and an
initial output_tokens) and the final ``message_delta`` (cumulative output_tokens)
regardless of the content/tool shape, so cost is recoverable even when full
content assembly fails. Returns ``None`` if no usage event is found.
"""
input_tokens = 0
cache_read = 0
cache_creation = 0
cache_creation_5m: Optional[int] = None
cache_creation_1h: Optional[int] = None
output_tokens = 0
web_search_requests: Optional[int] = None
tool_search_requests: Optional[int] = None
inference_geo: Optional[str] = None
stop_reason: Optional[str] = None
found_usage = False
resolved_model = model
for _chunk_str in all_chunks:
for (
event_str
) in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(
_chunk_str
):
data = AnthropicPassthroughLoggingHandler._extract_sse_data(event_str)
if not data:
continue
event_type = data.get("type")
if event_type == "message_start":
message = data.get("message") or {}
if not resolved_model or resolved_model == "unknown":
resolved_model = message.get("model") or resolved_model
usage = message.get("usage") or {}
input_tokens = usage.get("input_tokens") or input_tokens
cache_read = usage.get("cache_read_input_tokens") or cache_read
cache_creation = (
usage.get("cache_creation_input_tokens") or cache_creation
)
_cc = usage.get("cache_creation")
if isinstance(_cc, dict):
cache_creation_5m = _cc.get("ephemeral_5m_input_tokens")
cache_creation_1h = _cc.get("ephemeral_1h_input_tokens")
if usage.get("inference_geo") is not None:
inference_geo = usage.get("inference_geo")
if usage.get("output_tokens") is not None:
output_tokens = usage.get("output_tokens")
found_usage = True
elif event_type == "message_delta":
_delta_stop = (data.get("delta") or {}).get("stop_reason")
if _delta_stop:
stop_reason = _delta_stop
usage = data.get("usage") or {}
if usage.get("output_tokens") is not None:
output_tokens = usage.get("output_tokens")
_stu = usage.get("server_tool_use")
if isinstance(_stu, dict):
if _stu.get("web_search_requests") is not None:
web_search_requests = _stu.get("web_search_requests")
if _stu.get("tool_search_requests") is not None:
tool_search_requests = _stu.get("tool_search_requests")
if usage.get("cache_read_input_tokens") is not None:
cache_read = usage.get("cache_read_input_tokens")
if usage.get("inference_geo") is not None:
inference_geo = usage.get("inference_geo")
found_usage = True
if not found_usage:
return None
# If only the 5m/1h split was provided, derive the cache_creation total from it.
if not cache_creation and (cache_creation_5m or cache_creation_1h):
cache_creation = (cache_creation_5m or 0) + (cache_creation_1h or 0)
# build usage via the same AnthropicConfig.calculate_usage path the success
# cases use, so prompt_tokens are cache-inclusive and cache / server_tool_use /
# inference_geo tokens are priced instead of left at $0
usage_object: dict = {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
}
if cache_read:
usage_object["cache_read_input_tokens"] = cache_read
if cache_creation:
usage_object["cache_creation_input_tokens"] = cache_creation
if cache_creation_5m is not None or cache_creation_1h is not None:
usage_object["cache_creation"] = {
"ephemeral_5m_input_tokens": cache_creation_5m or 0,
"ephemeral_1h_input_tokens": cache_creation_1h or 0,
}
if web_search_requests is not None or tool_search_requests is not None:
_server_tool_use: dict = {}
if web_search_requests is not None:
_server_tool_use["web_search_requests"] = web_search_requests
if tool_search_requests is not None:
_server_tool_use["tool_search_requests"] = tool_search_requests
usage_object["server_tool_use"] = _server_tool_use
if inference_geo is not None:
usage_object["inference_geo"] = inference_geo
usage_obj = AnthropicConfig().calculate_usage(
usage_object=usage_object, reasoning_content=None
)
return ModelResponse(
model=resolved_model,
choices=[
Choices(
finish_reason=(
map_finish_reason(stop_reason) if stop_reason else "stop"
),
index=0,
message=Message(role="assistant", content=""),
)
],
usage=usage_obj,
)
@staticmethod
def batch_creation_handler(
httpx_response: httpx.Response,

3
litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py
View file

@ -116,6 +116,9 @@ class BasePassthroughLoggingHandler(ABC):
kwargs["response_cost"] = response_cost
kwargs["model"] = model
# the pass-through success path reads spend from
# model_call_details["response_cost"], not from kwargs
logging_obj.model_call_details["response_cost"] = response_cost
passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( # type: ignore
kwargs.get("passthrough_logging_payload")
)

6
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -285,8 +285,10 @@ class PassThroughStreamingHandler:
Returns:
List of string lines, with each line being a complete data: {} chunk
"""
# Combine all bytes and decode to string
combined_str = b"".join(raw_bytes).decode("utf-8")
# errors="replace" so a stream cut mid-multibyte-sequence (client disconnect)
# still decodes and logs the usage events already received, instead of raising
# and dropping the whole request from SpendLogs
combined_str = b"".join(raw_bytes).decode("utf-8", errors="replace")
# Split by newlines and filter out empty lines
lines = [line.strip() for line in combined_str.split("\n") if line.strip()]

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

@ -2646,7 +2646,9 @@ def test_chunk_creator_tool_calls_not_dropped_on_finish(
tool_calls=[
ChatCompletionDeltaToolCall(
id="call_abc",
function=Function(name="get_weather", arguments='{"city":"NYC"}'),
function=Function(
name="get_weather", arguments='{"city":"NYC"}'
),
type="function",
index=0,
)
@ -2741,3 +2743,131 @@ def test_record_partial_usage_for_failure_noop_without_chunks():
wrapper._record_partial_usage_for_failure()
assert "combined_usage_object" not in logging_obj.model_call_details
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_stream_chunk_builder_raise_at_end_of_stream_still_recovers_usage(
sync_mode,
):
"""stream_chunk_builder re-raises (as APIError) on large agentic tool-use
streams. That raise originates inside the except-StopIteration handler, so
before the fix it escaped __next__/__anext__ and the request was dropped from
SpendLogs while the provider billed the tokens. The wrapper must catch it and
recover usage from the raw chunks so cost is still tracked."""
final_usage_block = Usage(
completion_tokens=392, prompt_tokens=1799, total_tokens=2191
)
final_chunk = ModelResponseStream(
id="chatcmpl-raise-test",
created=1742056047,
model=None,
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(content="", role="assistant"),
)
],
usage=final_usage_block,
)
test_chunks = bedrock_chunks + [final_chunk]
logging_obj = Logging(
model="bedrock/claude-haiku-4-5-20251001-v1:0",
messages=[{"role": "user", "content": "Hey"}],
stream=True,
call_type="completion",
start_time=time.time(),
litellm_call_id="raise-test",
function_id="1245",
)
response = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(model_responses=test_chunks),
model="bedrock/claude-haiku-4-5-20251001-v1:0",
custom_llm_provider="bedrock",
logging_obj=logging_obj,
stream_options={"include_usage": True},
)
seen_usage = []
with patch.object(
litellm,
"stream_chunk_builder",
side_effect=Exception("simulated assembly failure"),
):
# before the fix this raised and dropped the request; it must not raise now
if sync_mode:
for chunk in response:
if getattr(chunk, "usage", None) is not None:
seen_usage.append(chunk.usage)
else:
async for chunk in response:
if getattr(chunk, "usage", None) is not None:
seen_usage.append(chunk.usage)
assert any(
u.total_tokens == final_usage_block.total_tokens for u in seen_usage
), "usage recovered from raw chunks was not emitted after stream_chunk_builder raised"
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_stream_chunk_builder_raise_and_usage_recovery_failure_does_not_crash(
sync_mode,
):
"""If end-of-stream assembly raises AND best-effort usage recovery from the raw
chunks also fails, the stream must still complete cleanly rather than propagate
the exception to the consumer."""
from litellm.litellm_core_utils import streaming_handler as sh_module
final_chunk = ModelResponseStream(
id="chatcmpl-raise-recover-fail",
created=1742056047,
model=None,
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(content="", role="assistant"),
)
],
usage=Usage(completion_tokens=1, prompt_tokens=1, total_tokens=2),
)
response = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(
model_responses=bedrock_chunks + [final_chunk]
),
model="bedrock/claude-haiku-4-5-20251001-v1:0",
custom_llm_provider="bedrock",
logging_obj=Logging(
model="bedrock/claude-haiku-4-5-20251001-v1:0",
messages=[{"role": "user", "content": "Hey"}],
stream=True,
call_type="completion",
start_time=time.time(),
litellm_call_id="raise-recover-fail",
function_id="1245",
),
stream_options={"include_usage": True},
)
with (
patch.object(
litellm, "stream_chunk_builder", side_effect=Exception("assembly failed")
),
patch.object(
sh_module, "calculate_total_usage", side_effect=Exception("recovery failed")
),
):
# must not raise even though both assembly and recovery fail
if sync_mode:
chunks = [c for c in response]
else:
chunks = [c async for c in response]
assert len(chunks) > 0

51
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -3792,3 +3792,54 @@ async def test_inference_route_still_enforces_team_budget():
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
request=MagicMock(),
)
@pytest.mark.asyncio
async def test_virtual_key_max_budget_error_names_the_key():
"""BudgetExceededError for a virtual key must name the key (alias + masked key)
so operators don't have to reverse-map a spend figure back to a key."""
valid_token = UserAPIKeyAuth(
token="hashed-token",
key_alias="payments-prod",
key_name="sk-...um_g",
max_budget=10.0,
spend=0.0,
)
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=25.0),
):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
message = str(exc_info.value)
assert "payments-prod" in message
assert "sk-...um_g" in message
@pytest.mark.asyncio
async def test_virtual_key_max_budget_not_exceeded_does_not_raise():
"""Spend below the configured budget must not raise."""
valid_token = UserAPIKeyAuth(
token="hashed-token",
key_alias="payments-prod",
max_budget=10.0,
spend=0.0,
)
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=1.0),
):
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)

307
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
View file

@ -1885,3 +1885,310 @@ class TestNonStreamingResponseRedaction:
leaked = logging_obj.model_call_details.get("complete_streaming_response")
assert leaked is None
assert redacted.choices[0].message.content == "redacted-by-litellm"
def _sse_bytes(data: dict) -> bytes:
return f"event: {data['type']}\ndata: {json.dumps(data)}\n\n".encode()
class TestAnthropicUsageOnlyFallback:
"""When stream_chunk_builder cannot reassemble a large/agentic stream (returns
None or raises), Anthropic still emits token usage in the message_start /
message_delta SSE events. The handler must recover usage-only so the request is
priced instead of being dropped from SpendLogs while Anthropic billed the tokens."""
_CHUNKS = [
_sse_bytes(
{
"type": "message_start",
"message": {
"model": "claude-3-5-haiku-20241022",
"usage": {
"input_tokens": 100,
"cache_read_input_tokens": 40,
"cache_creation_input_tokens": 20,
"output_tokens": 1,
},
},
}
),
_sse_bytes(
{
"type": "message_delta",
"usage": {
"output_tokens": 55,
"server_tool_use": {"web_search_requests": 2},
},
}
),
]
def test_build_usage_only_recovers_cache_inclusive_usage(self):
response = (
AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=self._CHUNKS, model="claude-3-5-haiku-20241022"
)
)
assert response is not None
usage = response.usage
# prompt_tokens must be cache-inclusive (input + cache_read + cache_creation)
assert usage.prompt_tokens == 160
assert usage.completion_tokens == 55
assert usage._cache_read_input_tokens == 40
assert usage._cache_creation_input_tokens == 20
assert usage.prompt_tokens_details.cached_tokens == 40
assert usage.server_tool_use.web_search_requests == 2
def test_build_usage_only_returns_none_without_usage_events(self):
chunks = [_sse_bytes({"type": "content_block_delta", "delta": {"text": "hi"}})]
assert (
AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=chunks, model="claude-3-5-haiku-20241022"
)
is None
)
def test_build_usage_only_recovers_cache_split_server_tools_and_model(self):
# the model is "unknown" up-front and only the 5m/1h cache split is sent
# (no flat cache_creation_input_tokens); web/tool-search and geo arrive in
# message_delta. All must be recovered and priced, not left at $0.
chunks = [
"event: ping\ndata: [DONE]\n\n", # ignored sentinel between real events
_sse_bytes(
{
"type": "message_start",
"message": {
"model": "claude-opus-4-6",
"usage": {
"input_tokens": 80,
"output_tokens": 1,
"cache_creation": {
"ephemeral_5m_input_tokens": 12,
"ephemeral_1h_input_tokens": 8,
},
"inference_geo": "us",
},
},
}
),
_sse_bytes(
{
"type": "message_delta",
"delta": {"stop_reason": "tool_use"},
"usage": {
"output_tokens": 40,
"cache_read_input_tokens": 5,
"inference_geo": "us",
"server_tool_use": {
"web_search_requests": 1,
"tool_search_requests": 3,
},
},
}
),
]
response = (
AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=chunks, model="unknown"
)
)
assert response is not None
assert response.model == "claude-opus-4-6"
# the real stop_reason is surfaced, not a hardcoded "stop"
assert response.choices[0].finish_reason == "tool_calls"
usage = response.usage
# 80 input + 20 cache_creation (derived from 12+8) + 5 cache_read
assert usage.prompt_tokens == 105
assert usage.completion_tokens == 40
assert usage._cache_creation_input_tokens == 20
assert usage._cache_read_input_tokens == 5
assert usage.server_tool_use.web_search_requests == 1
assert usage.server_tool_use.tool_search_requests == 3
@pytest.mark.parametrize(
"event_str,expected",
[
("data: [DONE]", None),
("data: ", None),
("data: {not-json", None),
("event: ping", None),
('data: {"a": 1}', {"a": 1}),
],
)
def test_extract_sse_data_handles_malformed_and_sentinel_lines(
self, event_str, expected
):
assert (
AnthropicPassthroughLoggingHandler._extract_sse_data(event_str) == expected
)
def _real_logging_obj(self):
from litellm.litellm_core_utils.litellm_logging import Logging as RealLoggingObj
logging_obj = RealLoggingObj(
model="claude-3-5-haiku-20241022",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="pass_through_endpoint",
start_time=datetime.now(),
litellm_call_id="test-call-id",
function_id="1",
)
logging_obj.model_call_details["litellm_params"] = {}
logging_obj.litellm_params = {}
return logging_obj
@patch("litellm.completion_cost")
@patch.object(
AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response"
)
def test_handler_falls_back_when_assembly_returns_none(
self, mock_assemble, mock_cost
):
mock_assemble.return_value = None
mock_cost.return_value = 0.0021
logging_obj = self._real_logging_obj()
result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/anthropic/v1/messages",
request_body={"model": "claude-3-5-haiku-20241022", "stream": True},
endpoint_type="messages",
start_time=datetime.now(),
all_chunks=list(self._CHUNKS),
end_time=datetime.now(),
)
assert result["result"] is not None
assert result["result"].usage.completion_tokens == 55
assert result["kwargs"]["response_cost"] == 0.0021
@patch("litellm.completion_cost")
@patch.object(
AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response"
)
def test_handler_falls_back_when_assembly_raises(self, mock_assemble, mock_cost):
import litellm
mock_assemble.side_effect = litellm.APIError(
status_code=500,
message="boom",
llm_provider="anthropic",
model="claude-3-5-haiku-20241022",
)
mock_cost.return_value = 0.0021
logging_obj = self._real_logging_obj()
result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/anthropic/v1/messages",
request_body={"model": "claude-3-5-haiku-20241022", "stream": True},
endpoint_type="messages",
start_time=datetime.now(),
all_chunks=list(self._CHUNKS),
end_time=datetime.now(),
)
# a raise from stream_chunk_builder must be treated like a None result,
# not propagate out and drop the request from SpendLogs
assert result["result"] is not None
assert result["result"].usage.completion_tokens == 55
assert result["kwargs"]["response_cost"] == 0.0021
@patch.object(
AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response"
)
def test_handler_returns_none_when_no_usage_recoverable(self, mock_assemble):
# assembly fails AND the chunks carry no usage event, so there is nothing
# to price; the handler must return None rather than fabricate a response
mock_assemble.return_value = None
logging_obj = self._real_logging_obj()
chunks = [_sse_bytes({"type": "content_block_delta", "delta": {"text": "hi"}})]
result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/anthropic/v1/messages",
request_body={"model": "claude-3-5-haiku-20241022", "stream": True},
endpoint_type="messages",
start_time=datetime.now(),
all_chunks=chunks,
end_time=datetime.now(),
)
assert result["result"] is None
assert result["kwargs"] == {}
@patch.object(
AnthropicPassthroughLoggingHandler, "_build_usage_only_response_from_chunks"
)
@patch.object(
AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response"
)
def test_handler_does_not_crash_when_usage_only_fallback_raises(
self, mock_assemble, mock_fallback
):
# if the usage-only fallback itself raises, it must be treated as None and
# drop gracefully, not propagate out and crash the success handler
mock_assemble.return_value = None
mock_fallback.side_effect = Exception("fallback boom")
logging_obj = self._real_logging_obj()
result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/anthropic/v1/messages",
request_body={"model": "claude-3-5-haiku-20241022", "stream": True},
endpoint_type="messages",
start_time=datetime.now(),
all_chunks=list(self._CHUNKS),
end_time=datetime.now(),
)
assert result["result"] is None
assert result["kwargs"] == {}
class TestAnthropicResponseCostRecordedOnModelCallDetails:
"""The pass-through success path reads spend from
model_call_details["response_cost"], not from kwargs, so the streaming payload
builder must record it there or streaming pass-through logs $0."""
def test_create_payload_records_response_cost_on_model_call_details(self):
from litellm.types.utils import Choices, Message, ModelResponse
logging_obj = MagicMock()
logging_obj.model_call_details = {}
logging_obj.get_router_model_id.return_value = None
logging_obj.litellm_params = {}
logging_obj.litellm_call_id = "test-call-id"
response = ModelResponse(
id="test-id",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(content="hello", role="assistant"),
)
],
created=1234567890,
model="claude-3-7-sonnet-20250219",
usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
)
kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
litellm_model_response=response,
model="claude-3-7-sonnet-20250219",
kwargs={},
start_time=datetime.now(),
end_time=datetime.now(),
logging_obj=logging_obj,
)
assert (
logging_obj.model_call_details["response_cost"] == kwargs["response_cost"]
)
assert logging_obj.model_call_details["response_cost"] > 0

17
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
View file

@ -118,3 +118,20 @@ async def test_chunk_processor_does_not_schedule_logging_when_no_chunks():
assert received == []
mock_route.assert_not_called()
def test_convert_raw_bytes_survives_truncated_multibyte_sequence():
"""A stream cut mid-multibyte-sequence (client disconnect) must still decode
via errors="replace" so the usage events already received are logged, instead
of raising UnicodeDecodeError and dropping the whole request from SpendLogs."""
# the 3-byte "☃" (E2 98 83) is cut after 2 bytes, leaving an invalid sequence
# that strict utf-8 decode would raise on, discarding the message_delta line too
truncated_codepoint = "☃".encode("utf-8")[:2]
raw_bytes = [
b'data: {"text": "' + truncated_codepoint,
b'\ndata: {"type": "message_delta"}\n',
]
lines = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes)
assert any('"type": "message_delta"' in line for line in lines)