diff --git a/Dockerfile b/Dockerfile index 9ad9ab31b65..68d7b14d19f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,8 +1,8 @@ # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/backend/Dockerfile b/backend/Dockerfile index c08014fc0ef..d969be69a20 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index c84003a065f..94e53bafdd9 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,8 +1,8 @@ # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 2729babb6d6..9304dd3784f 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,6 +1,6 @@ # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a diff --git a/gateway/Dockerfile b/gateway/Dockerfile index a2ca3d3f83f..041fd11678b 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 18321c70f6e..482bf88cbc4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2951,7 +2951,12 @@ class Logging(LiteLLMLoggingBaseClass): ) self.model_call_details["end_time"] = end_time self.model_call_details.setdefault("original_response", None) - self.model_call_details["response_cost"] = 0 + # A stream interrupted mid-flight still billed the provider for the + # chunks already delivered; the router stashes that recovered usage as + # ``combined_usage_object`` and pre-computes its cost, so preserve it + # here instead of zeroing the spend on an otherwise-failed request. + if self.model_call_details.get("combined_usage_object") is None: + self.model_call_details["response_cost"] = 0 if hasattr(exception, "headers") and isinstance(exception.headers, dict): self.model_call_details.setdefault("litellm_params", {}) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 8da66d4600d..413ddb71bf8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -6,6 +6,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS +from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, @@ -339,8 +340,7 @@ class StandardBuiltInToolCostTracking: # and _handle_web_search_cost() is never called. if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True return False @@ -352,8 +352,7 @@ class StandardBuiltInToolCostTracking: elif usage is not None: if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True elif ( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 59d0465e6d4..ba8b0764741 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1,7 +1,7 @@ # What is this? ## Helper utilities for cost_per_token() -from typing import Literal, Optional, Tuple, TypedDict, cast +from typing import Any, Literal, Optional, Tuple, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -30,6 +30,26 @@ _IMAGE_RESPONSE_CALL_TYPES = frozenset( ) +def _get_web_search_requests(server_tool_use: Any) -> Optional[int]: + """ + Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value + that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance, + or any other object supporting attribute access. + + Returns ``None`` when the value cannot be resolved — callers can + distinguish "absent" from "zero" using ``is None``. + + See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder`` + historically left this as a plain ``dict``, which broke direct attribute + access in cost calculation. + """ + if server_tool_use is None: + return None + if isinstance(server_tool_use, dict): + return server_tool_use.get("web_search_requests") + return getattr(server_tool_use, "web_search_requests", None) + + def _is_above_128k(tokens: float) -> bool: if tokens > 128000: return True diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c3842..c2a17ae8dcc 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -588,7 +588,18 @@ class ChunkProcessor: hasattr(usage_chunk, "server_tool_use") and usage_chunk.server_tool_use is not None ): - server_tool_use = usage_chunk.server_tool_use + # Coerce dict to ServerToolUse so downstream cost-calc code + # (which accesses .web_search_requests as an attribute) + # doesn't raise AttributeError. Some providers / streaming + # paths leave server_tool_use as a plain dict on the chunk. + if isinstance(usage_chunk.server_tool_use, dict): + server_tool_use = ServerToolUse(**usage_chunk.server_tool_use) + elif isinstance(usage_chunk.server_tool_use, ServerToolUse): + server_tool_use = usage_chunk.server_tool_use + else: + server_tool_use = ServerToolUse.model_validate( + usage_chunk.server_tool_use + ) if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 29c0d0629e8..5ed8cea9122 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1923,11 +1923,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: @@ -2152,11 +2170,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: @@ -2233,6 +2267,7 @@ class CustomStreamWrapper: litellm.request_timeout ) if self.logging_obj is not None: + self._record_partial_usage_for_failure() ## LOGGING threading.Thread( target=self.logging_obj.failure_handler, @@ -2246,6 +2281,7 @@ class CustomStreamWrapper: except Exception as e: traceback_exception = traceback.format_exc() if self.logging_obj is not None: + self._record_partial_usage_for_failure() ## LOGGING threading.Thread( target=self.logging_obj.failure_handler, @@ -2257,6 +2293,33 @@ class CustomStreamWrapper: ) self._handle_stream_fallback_error(e) + def _record_partial_usage_for_failure(self) -> None: + """ + A stream that breaks mid-flight still billed the provider for the chunks + already delivered. Recover that partial usage from the chunks seen so + far and stash it, with its cost, on the logging object so the failure + handler records the real partial spend instead of zero. A request that + later recovers via a router fallback overwrites this with the combined + success log on the same request id, so this never double counts. + """ + if self.logging_obj is None or not self.chunks: + return + try: + partial_response = litellm.stream_chunk_builder(chunks=self.chunks) + usage = cast(Optional[Usage], getattr(partial_response, "usage", None)) + if usage is None: + return + self.logging_obj.model_call_details["combined_usage_object"] = usage + self.logging_obj.model_call_details["response_cost"] = ( + self.logging_obj._response_cost_calculator(result=partial_response) + or 0.0 + ) + except Exception as recover_error: + verbose_logger.debug( + "could not recover partial usage for interrupted stream: %s", + recover_error, + ) + def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn": """ Common error handling for both __next__ and __anext__. diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index 3882d8f978c..6a031498dae 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Optional, Tuple from litellm.litellm_core_utils.llm_cost_calc.utils import ( _get_token_base_cost, + _get_web_search_requests, _parse_prompt_tokens_details, calculate_cache_writing_cost, generic_cost_per_token, @@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search( if model_info is None: return 0.0 - if ( - usage is None - or usage.server_tool_use is None - or usage.server_tool_use.web_search_requests is None - ): + if usage is None: + return 0.0 + web_search_requests = _get_web_search_requests( + getattr(usage, "server_tool_use", None) + ) + if web_search_requests is None: return 0.0 ## Get the cost per web search request @@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search( return 0.0 ## Calculate the total cost - total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests + total_cost = cost_per_web_search_request * web_search_requests return total_cost diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b1a109169c2..c07e510452d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3448,9 +3448,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}", ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3688f25ac44..0a85b181fba 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -40,7 +40,7 @@ class _ProxyDBLogger(CustomLogger): kwargs, response_obj, start_time, end_time ) - async def async_post_call_failure_hook( + async def async_post_call_failure_hook( # noqa: PLR0915 self, request_data: dict, original_exception: Exception, @@ -162,9 +162,20 @@ class _ProxyDBLogger(CustomLogger): if obj_start is not None: actual_start_time = obj_start + # A stream that broke mid-flight still billed the provider for the + # chunks already delivered. ``post_call_failure_hook`` lifts that + # recovered cost onto request_data (the usage rides along in + # ``combined_usage_object`` for the token columns), so attribute the + # real partial spend to this failure row instead of zero. + recovered_response_cost = 0.0 + if isinstance(request_data.get("combined_usage_object"), litellm.Usage): + recovered_response_cost = max( + float(request_data.get("response_cost") or 0.0), 0.0 + ) + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, - response_cost=0.0, + response_cost=recovered_response_cost, user_id=user_api_key_dict.user_id, end_user_id=user_api_key_dict.end_user_id, team_id=user_api_key_dict.team_id, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a912a88a993..49aa8c0072c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -6,18 +6,29 @@ 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 ( + get_content_from_model_response, +) 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 @@ -136,6 +147,84 @@ class AnthropicPassthroughLoggingHandler: return model return None + @staticmethod + def _stream_was_interrupted( + all_chunks: Sequence[Union[str, bytes]], + ) -> bool: + """ + Anthropic ends a stream with ``content_block_stop`` -> ``message_delta`` + -> ``message_stop``; a client disconnect leaves the last event mid + ``content_block_delta``. Scan from the tail and decide on the first + terminal-region event, so the common completed case is O(1) rather than + re-deserializing every line of the stream. + """ + for raw in reversed(all_chunks): + text = raw.decode("utf-8") if isinstance(raw, bytes) else raw + for line in reversed(text.splitlines()): + if not line.startswith("data:"): + continue + try: + data = json.loads(line[len("data:") :].strip()) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(data, dict): + continue + etype = data.get("type") + if etype == "message_delta": + return False + if etype in ( + "content_block_delta", + "content_block_stop", + "message_start", + ): + return True + return True + + @staticmethod + def _recover_interrupted_stream_output_tokens( + response: Union[ModelResponse, TextCompletionResponse], + all_chunks: Sequence[Union[str, bytes]], + model: str, + ) -> None: + """ + An Anthropic stream interrupted before its terminal ``message_delta`` + (client disconnect) carries only the ``message_start`` ``output_tokens`` + placeholder (typically 1-3), so completion tokens and spend are + undercounted ~20x. Re-tokenize the buffered output text to recover a + realistic ``output_tokens`` for usage/cost. Completed streams are + untouched because their terminal ``message_delta`` short-circuits here. + """ + if not isinstance(response, ModelResponse): + return + if not AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks): + return + usage = getattr(response, "usage", None) + if usage is None: + return + output_text = get_content_from_model_response(response) + if not output_text: + return + try: + recovered_output_tokens = litellm.token_counter( + model=model, text=output_text, count_response_tokens=True + ) + except Exception: + verbose_proxy_logger.warning( + "Could not re-tokenize interrupted stream output; " + "keeping placeholder completion token count." + ) + return + if recovered_output_tokens <= (usage.completion_tokens or 0): + return + usage.completion_tokens = recovered_output_tokens + usage.total_tokens = (usage.prompt_tokens or 0) + recovered_output_tokens + # Anthropic costing reads completion_tokens_details.text_tokens, so the + # stale message_start placeholder there must be corrected too or spend + # stays undercounted even after completion_tokens is fixed. + details = getattr(usage, "completion_tokens_details", None) + if details is not None and getattr(details, "text_tokens", None) is not None: + details.text_tokens = recovered_output_tokens + @staticmethod def _create_anthropic_response_logging_payload( litellm_model_response: Union[ModelResponse, TextCompletionResponse], @@ -191,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") ) @@ -262,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..." @@ -277,6 +398,11 @@ class AnthropicPassthroughLoggingHandler: "result": None, "kwargs": {}, } + AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( + response=complete_streaming_response, + all_chunks=all_chunks, + model=model, + ) kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( litellm_model_response=complete_streaming_response, model=model, @@ -550,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( # noqa: PLR0915 + 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( # noqa: PLR0915 httpx_response: httpx.Response, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py index b9df8ecede3..a7ec2f0d368 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py @@ -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") ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 33a6b719280..7a725472dd7 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -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()] diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index e2881faca0d..26d8d748d2e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -265,6 +265,13 @@ def get_logging_payload( # noqa: PLR0915 elif isinstance(_usage, dict): usage = _usage + # A request that failed mid-stream has no usable response_obj usage, but the + # streaming handler may have recovered the usage from the chunks already + # delivered. Honor that override so the partial usage lands in spend tracking. + _combined_usage = kwargs.get("combined_usage_object") + if not usage and isinstance(_combined_usage, litellm.Usage): + usage = _combined_usage.model_dump() + id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs) standard_logging_payload = cast( Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index effd62da717..f0da87e1271 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2006,12 +2006,21 @@ class ProxyLogging: # compute preprocessing latency after the logging object is popped. _logging_obj = request_data.get("litellm_logging_obj") if _logging_obj is not None: - _first_handoff = getattr(_logging_obj, "model_call_details", {}).get( - "first_api_call_start_time" - ) + _model_call_details = getattr(_logging_obj, "model_call_details", {}) + _first_handoff = _model_call_details.get("first_api_call_start_time") if _first_handoff is not None: request_data["first_api_call_start_time"] = _first_handoff + # A stream that broke mid-flight still billed the provider for the + # chunks already delivered; the streaming handler stashes that + # recovered usage and cost here. Lift them onto request_data so the + # failure-path spend callbacks (which run after the logging object + # is popped) record the real partial spend instead of zero. + _recovered_usage = _model_call_details.get("combined_usage_object") + if _recovered_usage is not None: + request_data["combined_usage_object"] = _recovered_usage + request_data["response_cost"] = _model_call_details.get("response_cost") + # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c3bf729238f..3f0fe99d836 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1561,7 +1561,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, ): @@ -1662,6 +1662,9 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details=_prompt_tokens_details or None, ) + if isinstance(server_tool_use, dict): + server_tool_use = ServerToolUse(**server_tool_use) + if server_tool_use is not None: self.server_tool_use = server_tool_use else: # maintain openai compatibility in usage object if possible diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 2160514251a..6e79922a97a 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py new file mode 100644 index 00000000000..4eee6b59d34 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py @@ -0,0 +1,88 @@ +""" +Tests that the cost-tracking call sites tolerate ``server_tool_use`` being +either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, + _get_web_search_requests, +) +from litellm.types.utils import ModelResponse, ServerToolUse, Usage + + +class _UsageWithDictServerToolUse: + """ + Tiny stand-in that mimics the broken streaming-rebuild shape: + ``server_tool_use`` is a plain dict. + """ + + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + self.prompt_tokens_details = None + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 5}) == 5 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + stu = ServerToolUse(web_search_requests=7) + assert _get_web_search_requests(stu) == 7 + + +def test_get_web_search_requests_handles_pydantic_with_none_value(): + stu = ServerToolUse() + assert _get_web_search_requests(stu) is None + + +def test_response_object_includes_web_search_call_with_dict_server_tool_use(): + """ + The exact bug: ``usage.server_tool_use`` is a dict and the check in + ``response_object_includes_web_search_call`` used to crash with + ``AttributeError``. + """ + response = ModelResponse() + usage = _UsageWithDictServerToolUse({"web_search_requests": 2}) + + # Must not raise — and must correctly detect the web search call. + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_pydantic_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(ServerToolUse(web_search_requests=2)) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_none_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(None) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is False diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index b64cb7c6905..f5677a3fcf2 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3078,3 +3078,46 @@ class TestFirstApiCallStartTimeSetOnce: assert obj.model_call_details["api_call_start_time"] > first assert obj.model_call_details["first_api_call_start_time"] == first assert user_meta == {} + + +def test_failure_handler_records_recovered_partial_spend(logging_obj): + """A stream interrupted mid-flight still billed the provider for the chunks + already delivered. When the router stashes that recovered usage as + ``combined_usage_object`` and pre-computes ``response_cost``, the failure + handler must preserve them so the failure row carries the real partial + spend instead of zero. + """ + from litellm.types.utils import Usage + + logging_obj.model_call_details["combined_usage_object"] = Usage( + prompt_tokens=17, completion_tokens=9, total_tokens=26 + ) + logging_obj.model_call_details["response_cost"] = 0.00012 + + logging_obj._failure_handler_helper_fn( + exception=Exception("Connection lost"), + traceback_exception="Traceback ...", + ) + + payload = logging_obj.model_call_details["standard_logging_object"] + assert payload["status"] == "failure" + assert payload["response_cost"] == 0.00012 + assert payload["prompt_tokens"] == 17 + assert payload["completion_tokens"] == 9 + assert payload["total_tokens"] == 26 + + +def test_failure_handler_zeroes_spend_without_recovered_usage(logging_obj): + """A failure with no recovered partial usage keeps the existing behavior of + recording zero spend, so the partial-spend preservation does not leak into + ordinary failures. + """ + logging_obj._failure_handler_helper_fn( + exception=Exception("boom"), + traceback_exception="Traceback ...", + ) + + payload = logging_obj.model_call_details["standard_logging_object"] + assert payload["status"] == "failure" + assert payload["response_cost"] == 0 + assert payload["total_tokens"] == 0 diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py new file mode 100644 index 00000000000..4e28d5ba7d2 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py @@ -0,0 +1,130 @@ +""" +Regression tests for https://github.com/BerriAI/litellm/issues/26153 + +``stream_chunk_builder`` used to leave ``usage.server_tool_use`` as a plain +``dict`` when reconstructing a streaming response. Downstream cost-calculation +code (``StandardBuiltInToolCostTracking.response_object_includes_web_search_call`` +and ``get_cost_for_anthropic_web_search``) accesses +``usage.server_tool_use.web_search_requests`` as an attribute, which raised +``AttributeError: 'dict' object has no attribute 'web_search_requests'``. + +These tests reconstruct streaming chunks for an Anthropic-style web_search +response and assert: + +1. ``stream_chunk_builder`` returns ``ServerToolUse`` (not ``dict``) for + ``usage.server_tool_use``. +2. ``completion_cost`` runs end-to-end on the rebuilt response without + raising ``AttributeError``. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm import completion_cost, stream_chunk_builder +from litellm.types.utils import ( + Delta, + ModelResponseStream, + ServerToolUse, + StreamingChoices, + Usage, +) + + +def _make_text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(role="assistant", content=text), + ) + ], + ) + + +def _make_finish_chunk_with_usage_dict_server_tool_use() -> ModelResponseStream: + """Final chunk where server_tool_use is a *dict* — reproduces the bug shape.""" + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=Usage( + prompt_tokens=42, + completion_tokens=11, + total_tokens=53, + # NOTE: passed as a dict on purpose — this is the shape that + # historically slipped through stream_chunk_builder unchanged. + server_tool_use={"web_search_requests": 3}, + ), + ) + + +def test_stream_chunk_builder_coerces_server_tool_use_to_pydantic(): + """ + Regression: stream_chunk_builder must produce ServerToolUse, not dict. + """ + chunks = [ + _make_text_chunk("Otters "), + _make_text_chunk("are great."), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + + assert rebuilt is not None + assert rebuilt.usage is not None # type: ignore[attr-defined] + server_tool_use = rebuilt.usage.server_tool_use # type: ignore[attr-defined] + + assert ( + server_tool_use is not None + ), "server_tool_use should be carried through from the final chunk" + assert isinstance(server_tool_use, ServerToolUse), ( + f"expected ServerToolUse, got {type(server_tool_use).__name__}: " + f"{server_tool_use!r}" + ) + # Attribute access must not raise (this is exactly what was broken). + assert server_tool_use.web_search_requests == 3 + + +def test_completion_cost_does_not_raise_on_streaming_web_search_response(): + """ + Regression: completion_cost(...) must not raise AttributeError when the + response was reconstructed by stream_chunk_builder from a streaming + Anthropic web_search call. + """ + chunks = [ + _make_text_chunk("hello"), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + assert rebuilt is not None + + # The exact dollar amount depends on the model-pricing table; what matters + # for this regression is that it does NOT raise AttributeError on + # `dict has no attribute 'web_search_requests'`. + try: + cost = completion_cost(completion_response=rebuilt) + except AttributeError as e: # pragma: no cover - regression guard + pytest.fail( + "completion_cost raised AttributeError after stream_chunk_builder " + f"(issue #26153 regression): {e}" + ) + + assert isinstance(cost, (int, float)) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd9..35aca525f6c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -520,7 +520,10 @@ 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 + # server_tool_use must be a ServerToolUse pydantic so downstream cost-calc + # (which uses attribute access) works. See issue #26153. + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 2 def test_sort_chunks_handles_dict_hidden_params_created_at(): diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 63e2cb7f35c..6988a2f50ad 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2118,3 +2118,207 @@ def test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum(): f"Expected 'tool_calls' but got {final.choices[0].finish_reason!r}. " "STOP enum was not normalised through map_finish_reason()." ) + + +def test_record_partial_usage_for_failure_stashes_usage_and_cost(): + """A stream that breaks mid-flight must surface the usage assembled from the + chunks already delivered, plus its cost, on the logging object so the + failure handler records the real partial spend instead of zero. + """ + logging_obj = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="partial-usage-1", + function_id="1245", + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-partial-1", + created=1742056047, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="The Roman Empire began when", role="assistant" + ), + ) + ], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.prompt_tokens == 30 + assert stashed.completion_tokens == 1 + assert stashed.total_tokens == 31 + assert isinstance(logging_obj.model_call_details["response_cost"], float) + + +def test_record_partial_usage_for_failure_noop_without_chunks(): + """With no chunks delivered there is nothing billed to recover, so the + failure stash must stay absent and not force a zero-usage row. + """ + logging_obj = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="partial-usage-2", + function_id="1245", + ) + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.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 diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py new file mode 100644 index 00000000000..70fef0162e6 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py @@ -0,0 +1,94 @@ +""" +Tests that ``get_cost_for_anthropic_web_search`` tolerates ``server_tool_use`` +being either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.anthropic.cost_calculation import ( + _get_web_search_requests, + get_cost_for_anthropic_web_search, +) +from litellm.types.utils import ModelInfo, ServerToolUse + + +class _UsageWithServerToolUse: + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + + +def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo: + info: ModelInfo = { # type: ignore[typeddict-item] + "search_context_cost_per_query": { + "search_context_size_low": cost_per_query, + "search_context_size_medium": cost_per_query, + "search_context_size_high": cost_per_query, + } + } + return info + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 4}) == 4 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2 + + +def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use(): + """ + Regression: ``server_tool_use`` was a dict from ``stream_chunk_builder`` and + direct attribute access on it raised ``AttributeError``. + """ + usage = _UsageWithServerToolUse({"web_search_requests": 3}) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use(): + usage = _UsageWithServerToolUse(ServerToolUse(web_search_requests=3)) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_none_server_tool_use(): + usage = _UsageWithServerToolUse(None) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == 0.0 + + +def test_get_cost_for_anthropic_web_search_with_no_usage(): + info = _make_model_info(cost_per_query=0.01) + cost = get_cost_for_anthropic_web_search(model_info=info, usage=None) + assert cost == 0.0 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 116ba83f42e..d151e267b47 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -3370,3 +3370,54 @@ async def test_resolve_end_user_reraises_budget_exceeded( prisma_client=MagicMock(), user_api_key_cache=cache, ) + + +@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, + ) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 771e10a54a0..0cbf308076c 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1067,3 +1067,40 @@ async def test_failure_hook_drops_error_information_traceback_when_env_set( assert "traceback" not in error_information assert error_information["error_class"] == "RuntimeError" assert error_information["error_message"] == "boom-with-traceback" + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_records_recovered_partial_spend(): + """A stream that broke mid-flight still billed the provider. The failure + hook lifts the recovered cost onto request_data as ``response_cost``; this + hook must pass it through to update_database so the failure row records the + real partial spend instead of the hardcoded zero. + """ + from litellm.types.utils import Usage + + logger = _ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key", user_id="u", team_id="t") + + request_data = { + "model": "anthropic/claude-haiku-4-5", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "proxy_server_request": {"request_id": "rid"}, + "response_cost": 3.5e-05, + "combined_usage_object": Usage( + prompt_tokens=30, completion_tokens=1, total_tokens=31 + ), + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("MidStreamFallbackError: read timeout"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 1083c3e9252..f20fb4c40ec 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -1376,6 +1376,147 @@ class TestPureTextFastPathParity: ) +class TestInterruptedStreamOutputTokenRecovery: + """ + When an Anthropic pass-through stream is interrupted (client disconnect) + before the terminal ``message_delta``, the only usage signal is the + ``message_start`` ``output_tokens`` placeholder (typically 1-3), so + completion tokens and spend are undercounted ~20x. The handler must + re-tokenize the buffered ``content_block_delta`` text to recover a + realistic ``output_tokens``; completed streams must stay untouched. + """ + + @staticmethod + def _sse(event, data): + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + _MODEL = "claude-3-5-haiku-20241022" + _OUTPUT_TEXT = ( + "The history of computing spans centuries, beginning with mechanical " + "calculators and the abacus, advancing through Charles Babbage's " + "analytical engine, Ada Lovelace's first algorithm, Alan Turing's " + "theoretical machine, and the electronic computers of the twentieth " + "century that gave rise to the modern information age." + ) + + def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2): + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + words = self._OUTPUT_TEXT.split(" ") + frames = [ + self._sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_interrupted", + "type": "message", + "role": "assistant", + "model": self._MODEL, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": { + "input_tokens": 29, + "output_tokens": placeholder_output_tokens, + }, + }, + }, + ), + self._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + ] + for i, word in enumerate(words): + text = word if i == 0 else " " + word + frames.append( + self._sse( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": text}, + }, + ) + ) + # Client disconnects here: no content_block_stop / message_delta / + # message_stop are ever received. + return list(PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames)) + + def _completed_chunks(self, *, final_output_tokens: int = 80): + chunks = self._interrupted_chunks() + chunks.append( + "data: " + + json.dumps( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": final_output_tokens}, + } + ) + ) + chunks.append('data: {"type": "message_stop"}') + return chunks + + def _run(self, all_chunks): + logging_obj = MagicMock() + logging_obj.model_call_details = {"model": self._MODEL, "stream": True} + logging_obj.litellm_call_id = "test-call-id" + logging_obj.litellm_params = {} + logging_obj.get_router_model_id.return_value = None + + return AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": self._MODEL, "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=all_chunks, + end_time=datetime.now(), + ) + + def test_interrupted_stream_retokenizes_buffered_output(self): + import litellm + + placeholder = 2 + result = self._run( + self._interrupted_chunks(placeholder_output_tokens=placeholder) + ) + usage = result["result"].usage + + expected = litellm.token_counter( + model=self._MODEL, + text=self._OUTPUT_TEXT, + count_response_tokens=True, + ) + + assert expected > placeholder * 5 + assert usage.completion_tokens == expected + assert usage.completion_tokens > placeholder + assert usage.total_tokens == usage.prompt_tokens + expected + # Anthropic spend is priced off completion_tokens_details.text_tokens; if the + # placeholder leaks through here, cost stays undercounted even though + # completion_tokens looks right. + assert usage.completion_tokens_details.text_tokens == expected + + def test_completed_stream_keeps_message_delta_tokens(self): + final = 80 + result = self._run(self._completed_chunks(final_output_tokens=final)) + usage = result["result"].usage + + # Terminal message_delta present: recovery must not fire; the authoritative + # provider count is preserved verbatim. + assert usage.completion_tokens == final + + class TestStreamFalseDeduplication: """ Regression tests for the duplicate-callback bug where a streaming pass-through @@ -1706,3 +1847,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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index f73aee77cc1..38990644154 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -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) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5ca058fc8d9..5272b105eb5 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2009,3 +2009,50 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form( assert sanitized is not None assert "leaked-via-pydantic-msg" not in sanitized["error_message"] assert REDACTED_BY_LITELM_STRING in sanitized["error_message"] + + +def test_get_logging_payload_uses_recovered_combined_usage_on_failure(): + """A request that fails mid-stream has no usable response_obj usage, but the + streaming handler recovers the usage from the chunks already delivered and + the failure hook surfaces it as ``combined_usage_object``. The spend-log + payload must record those token counts instead of zero. + """ + from litellm.types.utils import Usage + + kwargs = { + "model": "anthropic/claude-haiku-4-5", + "call_type": "acompletion", + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + "combined_usage_object": Usage( + prompt_tokens=30, completion_tokens=1, total_tokens=31 + ), + } + response_obj = Exception("MidStreamFallbackError: read timeout") + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert payload["prompt_tokens"] == 30 + assert payload["completion_tokens"] == 1 + assert payload["total_tokens"] == 31 + + +def test_get_logging_payload_failure_without_recovered_usage_is_zero(): + """A failure with no recovered usage keeps zero token counts, so the + combined-usage override never invents tokens for ordinary failures. + """ + kwargs = { + "model": "anthropic/claude-haiku-4-5", + "call_type": "acompletion", + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = Exception("BadRequestError") + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert payload["total_tokens"] == 0 diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 539a32db57d..b2a81a0cfaf 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -435,3 +435,52 @@ class TestPostCallFailureHookProxyExceptionLogging: ) is False ) +class TestPostCallFailureHookLiftsRecoveredPartialSpend: + """A stream that broke mid-flight still billed the provider for the chunks + already delivered. The streaming handler stashes that recovered usage and + cost on the logging object; post_call_failure_hook must lift them onto + request_data before the logging object is popped, so the failure-path spend + callbacks (which run after the pop) record the real partial spend. + """ + + async def _run(self, request_data): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import UserAPIKeyAuth + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("boom"), + user_api_key_dict=UserAPIKeyAuth(), + ) + + @pytest.mark.asyncio + async def test_lifts_recovered_usage_and_cost(self): + from litellm.types.utils import Usage + + recovered_usage = Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31) + logging_obj = MagicMock() + logging_obj.model_call_details = { + "combined_usage_object": recovered_usage, + "response_cost": 3.5e-05, + } + request_data = {"litellm_logging_obj": logging_obj, "metadata": {}} + await self._run(request_data) + + assert request_data["combined_usage_object"] is recovered_usage + assert request_data["response_cost"] == 3.5e-05 + assert "litellm_logging_obj" not in request_data + + @pytest.mark.asyncio + async def test_no_recovered_usage_is_noop(self): + logging_obj = MagicMock() + logging_obj.model_call_details = {} + request_data = {"litellm_logging_obj": logging_obj, "metadata": {}} + await self._run(request_data) + assert "combined_usage_object" not in request_data + assert "response_cost" not in request_data + + diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5e636b86ed6..7a28eacca8d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3340,6 +3340,151 @@ def test_combine_fallback_usage(): assert chunk.usage.total_tokens == 15 +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure(): + """A mid-stream failure with no successful fallback raises and is logged as + a failure, so the router must never dispatch it as a success. Partial-spend + recovery for the failure row happens in the streaming handler, not here, so + this guards only against reintroducing a success log for a failed stream. + """ + from litellm.exceptions import MidStreamFallbackError + from litellm.types.utils import Delta, StreamingChoices, Usage + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, + }, + ], + set_verbose=True, + ) + + error = MidStreamFallbackError( + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="The Roman Empire began when", + ) + + def _make_interrupted_model_response(): + partial_chunk = litellm.ModelResponseStream( + id="chatcmpl-partial-1", + created=1742056047, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="The Roman Empire began when", role="assistant" + ), + ) + ], + usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), + ) + + class _RaisingStream: + def __init__(self): + self.index = 0 + self.chunks = [partial_chunk] + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index == 0: + self.index += 1 + return partial_chunk + raise error + + stream = _RaisingStream() + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.model_call_details = {} + setattr(stream, "model", "gpt-4") + setattr(stream, "custom_llm_provider", "openai") + setattr(stream, "logging_obj", logging_obj) + return stream, logging_obj + + messages = [{"role": "user", "content": "Hello"}] + initial_kwargs = {"model": "gpt-4", "stream": True} + + # Terminal path: no successful fallback -> the error propagates and the + # router never dispatches a success for the failed stream. + model_response, logging_obj = _make_interrupted_model_response() + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=error), + ): + result = await router._acompletion_streaming_iterator( + model_response=model_response, + messages=messages, + initial_kwargs=dict(initial_kwargs), + ) + collected = [] + with pytest.raises(MidStreamFallbackError): + async for chunk in result: + collected.append(chunk) + + assert len(collected) == 1 + logging_obj.dispatch_success_handlers.assert_not_called() + + # Fallback success: the fallback stream owns success accounting via + # _combine_fallback_usage, so this iterator must not dispatch its own. + model_response, logging_obj = _make_interrupted_model_response() + + class _FallbackStream: + def __init__(self, items): + self.items = items + self.index = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index >= len(self.items): + raise StopAsyncIteration + item = self.items[self.index] + self.index += 1 + return item + + fallback_stream = _FallbackStream( + [ + litellm.ModelResponseStream( + id="chatcmpl-fallback-1", + model="gpt-3.5-turbo", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=" continued", role="assistant"), + ) + ], + ) + ] + ) + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ): + result = await router._acompletion_streaming_iterator( + model_response=model_response, + messages=messages, + initial_kwargs=dict(initial_kwargs), + ) + collected = [] + async for chunk in result: + collected.append(chunk) + + assert len(collected) == 2 + logging_obj.dispatch_success_handlers.assert_not_called() + + @pytest.mark.asyncio async def test_team_scoped_model_fallback(): """