mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(responses)!: bill partial usage on mid-stream failures and report bridged /v1/messages failures to the proxy (#47537)
A Responses API stream that fails partway now bills the usage it already consumed, ends with an error instead of a clean stop, and honors stream_timeout between chunks. Bridged /v1/messages calls report those failures to the proxy's failure hooks, so spend, budgets and logs see them
This commit is contained in:
Mateo Wang
GitHub
parent
a604c89059
commit
dfd9ad4103
12 changed files with 1155 additions and 47 deletions
|
|
@ -842,6 +842,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self._async_success_enqueued: bool = False
|
||||
self._sync_success_scheduled: bool = False
|
||||
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
|
||||
self._stream_failure_claimed: bool = False
|
||||
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
|
||||
|
||||
def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging":
|
||||
|
|
@ -2561,6 +2562,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
},
|
||||
)
|
||||
|
||||
def claim_stream_failure_report(self) -> bool:
|
||||
if self._stream_failure_claimed:
|
||||
return False
|
||||
self._stream_failure_claimed = True
|
||||
return True
|
||||
|
||||
def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None:
|
||||
"""Bill a fully streamed response on the failure log when a post-call hook rejects it."""
|
||||
usage: Final = getattr(assembled, "usage", None)
|
||||
|
|
|
|||
|
|
@ -2,15 +2,16 @@
|
|||
## Translates OpenAI call to Anthropic `/v1/messages` format
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from collections.abc import AsyncGenerator, Iterator, Mapping
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import anyio
|
||||
from pydantic import ConfigDict, field_validator
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm._logging import redact_internal_details_from_client_message
|
||||
from litellm._uuid import uuid
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.exceptions import APIError, MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
encrypted_reasoning_signature,
|
||||
)
|
||||
|
|
@ -22,6 +23,7 @@ from litellm.llms.anthropic.messages.native.utils import (
|
|||
from litellm.llms.anthropic.messages.safeguards import StreamedSafeguardResults
|
||||
from litellm.llms.anthropic.messages.utils import anthropic_sse_frame
|
||||
from litellm.responses.streaming_iterator import (
|
||||
map_stream_error_to_exception,
|
||||
partial_responses_incomplete_event_on_disconnect,
|
||||
stream_error_status_and_message,
|
||||
)
|
||||
|
|
@ -130,10 +132,12 @@ class AnthropicResponsesStreamWrapper:
|
|||
) -> None:
|
||||
self.responses_stream = responses_stream
|
||||
self.model = model
|
||||
self.litellm_logging_obj = litellm_logging_obj
|
||||
self.safeguards_evaluator: Final = safeguards_evaluator
|
||||
self._message_id: str = f"msg_{uuid.uuid4()}"
|
||||
if litellm_logging_obj is not None:
|
||||
litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id)
|
||||
self._unreported_failure: Exception | None = None
|
||||
self._current_block_index: int = -1
|
||||
# Map item_id -> content_block_index so we can stop the right block later
|
||||
self._item_id_to_block_index: dict[str, int] = {}
|
||||
|
|
@ -377,16 +381,24 @@ class AnthropicResponsesStreamWrapper:
|
|||
|
||||
if event_type == "response.failed":
|
||||
failed: Final = _FailedResponseEvent.model_validate(event)
|
||||
status_code, message = stream_error_status_and_message(
|
||||
failed.response.error if failed.response is not None else None
|
||||
)
|
||||
upstream_error: Final = failed.response.error if failed.response is not None else None
|
||||
status_code, message = stream_error_status_and_message(upstream_error)
|
||||
verbose_logger.error(
|
||||
"AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed (%s): %s",
|
||||
self.model,
|
||||
status_code,
|
||||
message,
|
||||
)
|
||||
self._fail_stream(status_code, message)
|
||||
self._fail_stream(
|
||||
status_code,
|
||||
message,
|
||||
failure=map_stream_error_to_exception(
|
||||
upstream_error,
|
||||
self.model,
|
||||
getattr(self.responses_stream, "custom_llm_provider", None) or "",
|
||||
logging_obj=self.litellm_logging_obj,
|
||||
),
|
||||
)
|
||||
return
|
||||
|
||||
# ---- response completed -> message_delta + message_stop ----
|
||||
|
|
@ -439,10 +451,36 @@ class AnthropicResponsesStreamWrapper:
|
|||
self._sent_message_stop = True
|
||||
return
|
||||
|
||||
def _fail_stream(self, status_code: int, message: str) -> None:
|
||||
def _fail_stream(self, status_code: int, message: str, failure: Exception | None = None) -> None:
|
||||
self._stream_failed = True
|
||||
self._unreported_failure = failure or APIError(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
llm_provider=getattr(self.responses_stream, "custom_llm_provider", None),
|
||||
model=self.model,
|
||||
)
|
||||
self._chunk_queue.append(_anthropic_error_chunk(status_code, message))
|
||||
|
||||
async def _report_unreported_failure(self) -> None:
|
||||
failure: Final = self._unreported_failure
|
||||
if failure is None:
|
||||
return
|
||||
self._unreported_failure = None
|
||||
on_failure: Final[Callable[[Exception], Awaitable[None]] | None] = getattr(
|
||||
self.litellm_logging_obj, "_on_detached_stream_failure", None
|
||||
)
|
||||
if on_failure is None:
|
||||
return
|
||||
try:
|
||||
with anyio.CancelScope(shield=True):
|
||||
await on_failure(failure)
|
||||
except Exception as hook_failure: # noqa: BLE001 # a failing proxy hook must not hide the error event from the client
|
||||
verbose_logger.warning(
|
||||
"AnthropicResponsesStreamWrapper: stream failure hook raised: %s(%s)",
|
||||
type(hook_failure).__name__,
|
||||
hook_failure,
|
||||
)
|
||||
|
||||
def __aiter__(self) -> "AnthropicResponsesStreamWrapper":
|
||||
return self
|
||||
|
||||
|
|
@ -450,6 +488,7 @@ class AnthropicResponsesStreamWrapper:
|
|||
if self._chunk_queue:
|
||||
return self._chunk_queue.popleft()
|
||||
if self._stream_failed:
|
||||
await self._report_unreported_failure()
|
||||
raise StopAsyncIteration
|
||||
|
||||
if not self._sent_message_start:
|
||||
|
|
@ -478,7 +517,7 @@ class AnthropicResponsesStreamWrapper:
|
|||
verbose_logger.exception(
|
||||
"AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed", self.model
|
||||
)
|
||||
self._fail_stream(*_failure_status_and_message(e))
|
||||
self._fail_stream(*_failure_status_and_message(e), failure=_original_failure(e))
|
||||
|
||||
if not self._chunk_queue and not self._sent_message_stop and not self._stream_failed:
|
||||
verbose_logger.error(
|
||||
|
|
@ -498,8 +537,14 @@ class AnthropicResponsesStreamWrapper:
|
|||
|
||||
async def _sse_frames(self) -> AsyncGenerator[bytes, None]:
|
||||
safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator)
|
||||
async for chunk in self:
|
||||
yield anthropic_sse_frame(await safeguard_results.observe(chunk))
|
||||
try:
|
||||
async for chunk in self:
|
||||
yield anthropic_sse_frame(await safeguard_results.observe(chunk))
|
||||
finally:
|
||||
from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt
|
||||
|
||||
if not closing_superseded_attempt():
|
||||
await self._report_unreported_failure()
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return partial_responses_incomplete_event_on_disconnect(self.responses_stream)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import json
|
|||
import logging
|
||||
import math
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
from itertools import chain
|
||||
|
|
@ -136,7 +137,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
|
|||
_check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
check_and_merge_model_level_guardrails,
|
||||
)
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
|
||||
|
|
@ -2006,6 +2007,49 @@ def _content_text(fields: Mapping[str, object]) -> str:
|
|||
return str(fields.get("content", ""))
|
||||
|
||||
|
||||
async def _no_stream_failure_to_release() -> None:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _ClientStreamFailureReport:
|
||||
logging_obj: LiteLLMLoggingObj
|
||||
report: Callable[[Exception], Awaitable[None]]
|
||||
held: Exception | None = None
|
||||
|
||||
@classmethod
|
||||
def hold_while_streaming(cls, logging_obj: LiteLLMLoggingObj | None) -> Callable[[], Awaitable[None]]:
|
||||
if logging_obj is None or logging_obj.on_detached_stream_failure is None:
|
||||
return _no_stream_failure_to_release
|
||||
return cls(logging_obj, logging_obj.on_detached_stream_failure)._start_holding()
|
||||
|
||||
def _start_holding(self) -> Callable[[], Awaitable[None]]:
|
||||
self.logging_obj.on_detached_stream_failure = self.hold
|
||||
return self.release
|
||||
|
||||
async def hold(self, failure: Exception) -> None:
|
||||
if self.held is None:
|
||||
self.held = failure
|
||||
|
||||
async def release(self) -> None:
|
||||
self.logging_obj.on_detached_stream_failure = self.report
|
||||
if self.held is None:
|
||||
return
|
||||
await asyncio.shield(
|
||||
spawn_logging_task(
|
||||
asyncio.get_running_loop(), self._report_held(self.held), task_name="stream failure report"
|
||||
)
|
||||
)
|
||||
|
||||
async def _report_held(self, failure: Exception) -> None:
|
||||
try:
|
||||
await self.report(failure)
|
||||
except Exception as hook_failure: # noqa: BLE001 # a failing failure hook must not skip the stream cleanup
|
||||
verbose_proxy_logger.warning(
|
||||
"Stream failure hook raised: %s(%s)", type(hook_failure).__name__, sanitize_for_log(hook_failure)
|
||||
)
|
||||
|
||||
|
||||
class ProxyBaseLLMRequestProcessing:
|
||||
def __init__(self, data: dict, served_model: str | None = None):
|
||||
self.data = data
|
||||
|
|
@ -3498,6 +3542,8 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data: Final = self.data
|
||||
|
||||
async def _on_detached_stream_failure(exc: Exception) -> None:
|
||||
if not logging_obj.claim_stream_failure_report():
|
||||
return
|
||||
try:
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -4448,11 +4494,15 @@ class ProxyBaseLLMRequestProcessing:
|
|||
client_disconnected = False
|
||||
delivered_chunk = False
|
||||
recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the yielded bytes
|
||||
stream_logging_obj: Final[object] = request_data.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # legacy request data is narrowed to Logging below
|
||||
guarded_stream: Final[AsyncGenerator[object, None]] = proxy_logging_obj.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
)
|
||||
release_stream_failure: Final = _ClientStreamFailureReport.hold_while_streaming(
|
||||
stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None
|
||||
)
|
||||
try:
|
||||
str_so_far = ""
|
||||
async for chunk in guarded_stream:
|
||||
|
|
@ -4523,10 +4573,18 @@ class ProxyBaseLLMRequestProcessing:
|
|||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.async_data_generator(): Exception occured - %s", e
|
||||
)
|
||||
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data=request_data,
|
||||
already_reported: Final = (
|
||||
isinstance(stream_logging_obj, LiteLLMLoggingObj)
|
||||
and not stream_logging_obj.claim_stream_failure_report()
|
||||
)
|
||||
transformed_exception: Final = (
|
||||
None
|
||||
if already_reported
|
||||
else await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data=request_data,
|
||||
)
|
||||
)
|
||||
if transformed_exception is not None:
|
||||
e = transformed_exception
|
||||
|
|
@ -4549,15 +4607,18 @@ class ProxyBaseLLMRequestProcessing:
|
|||
seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail)
|
||||
yield seal + error_frame if seal else error_frame
|
||||
finally:
|
||||
await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
|
||||
request=request,
|
||||
request_data=request_data,
|
||||
response=response,
|
||||
stream_completed=stream_completed,
|
||||
client_disconnected=client_disconnected,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
try:
|
||||
await release_stream_failure()
|
||||
finally:
|
||||
await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
|
||||
request=request,
|
||||
request_data=request_data,
|
||||
response=response,
|
||||
stream_completed=stream_completed,
|
||||
client_disconnected=client_disconnected,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def async_sse_data_generator(
|
||||
|
|
|
|||
|
|
@ -41,7 +41,11 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError
|
|||
from litellm.litellm_core_utils.asyncify import run_async_function
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base
|
||||
from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation
|
||||
from litellm.litellm_core_utils.internal_call_metadata import (
|
||||
is_background_response,
|
||||
requests_background,
|
||||
without_budget_reservation,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
|
||||
|
|
@ -65,7 +69,7 @@ from litellm.types.llms.openai import (
|
|||
ResponsesAPIStreamingResponse,
|
||||
)
|
||||
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
|
||||
from litellm.types.utils import CallTypes
|
||||
from litellm.types.utils import CallTypes, ModelResponse
|
||||
from litellm.utils import async_post_call_success_deployment_hook
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -288,7 +292,7 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]:
|
|||
return _status_code_for_error_fields(error_type, error_code), message
|
||||
|
||||
|
||||
def _map_stream_error_to_exception(
|
||||
def map_stream_error_to_exception(
|
||||
error_obj: object,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
|
|
@ -578,6 +582,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
self._generated_content + self._generated_tool_arguments,
|
||||
use_default_image_token_count=False,
|
||||
)
|
||||
if _estimate_wanted
|
||||
else None
|
||||
|
|
@ -731,7 +736,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._earlier_rounds_spend = earlier
|
||||
|
||||
def _map_error_event_exception(self, error_obj: object) -> Exception:
|
||||
return _map_stream_error_to_exception(
|
||||
return map_stream_error_to_exception(
|
||||
error_obj,
|
||||
self.model or "",
|
||||
self.custom_llm_provider or "",
|
||||
|
|
@ -981,6 +986,39 @@ class BaseResponsesAPIStreamingIterator:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
def _record_partial_usage_for_failure(self) -> None:
|
||||
generated_text: Final = self._generated_content + self._generated_tool_arguments
|
||||
if (
|
||||
not generated_text
|
||||
or requests_background(self.request_data)
|
||||
or is_background_response(self._created_response)
|
||||
or self._completed_response_logged
|
||||
or self.logging_obj.model_call_details.get("combined_usage_object") is not None
|
||||
):
|
||||
return
|
||||
estimate: Final = _estimate_usage_safely(
|
||||
self.model or "",
|
||||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
generated_text,
|
||||
use_default_image_token_count=True,
|
||||
)
|
||||
if estimate is None:
|
||||
return
|
||||
usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the success path above uses
|
||||
estimate
|
||||
)
|
||||
served_tier: Final[object] = getattr(self._created_response, "service_tier", None)
|
||||
assembled: Final = ModelResponse(
|
||||
model=self.model, usage=usage, service_tier=served_tier if isinstance(served_tier, str) else None
|
||||
)
|
||||
if self._earlier_rounds_spend is None:
|
||||
self.logging_obj.record_assembled_response_for_failure(assembled)
|
||||
return
|
||||
self.logging_obj.record_failure_after_earlier_rounds(
|
||||
self._earlier_rounds_spend, usage, self.logging_obj.response_cost_calculator(result=assembled) or 0.0
|
||||
)
|
||||
|
||||
def _handle_failure(self, exception: Exception):
|
||||
"""
|
||||
Trigger failure handlers before bubbling the exception.
|
||||
|
|
@ -990,6 +1028,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
if self._failure_handled:
|
||||
return
|
||||
self._failure_handled = True
|
||||
self._record_partial_usage_for_failure()
|
||||
|
||||
traceback_exception: Final = traceback.format_exc()
|
||||
end_time: Final = datetime.now()
|
||||
|
|
@ -1060,15 +1099,20 @@ class BaseResponsesAPIStreamingIterator:
|
|||
the same shape the success handler assembles for a natural end of stream.
|
||||
Its response is the response.created payload (it carries the served
|
||||
service_tier) plus the text and tool arguments relayed so far, with counted
|
||||
usage since the provider never sent any. None when nothing was relayed yet
|
||||
or the stream already reached its terminal event and billed itself.
|
||||
usage since the provider never sent any. None when nothing was relayed yet,
|
||||
the stream already reached its terminal event and billed itself, or its failure
|
||||
log already billed the relayed output.
|
||||
"""
|
||||
created: Final = self._created_response
|
||||
generated: Final = self._generated_content + self._generated_tool_arguments
|
||||
if self.completed_response is not None or created is None or not generated:
|
||||
if created is None or not generated or self._billed_without_disconnect():
|
||||
return None
|
||||
usage: Final = _estimate_usage_safely(
|
||||
self.model or "", self.request_data.get("input"), self.request_data, generated
|
||||
self.model or "",
|
||||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
generated,
|
||||
use_default_image_token_count=False,
|
||||
)
|
||||
if usage is None:
|
||||
return None
|
||||
|
|
@ -1082,6 +1126,9 @@ class BaseResponsesAPIStreamingIterator:
|
|||
stamp_responses_usage_cost(partial, self.logging_obj)
|
||||
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
|
||||
|
||||
def _billed_without_disconnect(self) -> bool:
|
||||
return self.completed_response is not None or self._failure_handled
|
||||
|
||||
def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None:
|
||||
self._yielded_first_chunk = True
|
||||
if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES:
|
||||
|
|
@ -1218,6 +1265,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
raise StopAsyncIteration from e
|
||||
self._raise_for_transport_error(e)
|
||||
except httpx.TimeoutException as e:
|
||||
if self.completed_response is not None:
|
||||
self.finished = True
|
||||
raise StopAsyncIteration from e
|
||||
self._raise_for_transport_timeout(e)
|
||||
except httpx.HTTPError as e:
|
||||
# Handle HTTP errors
|
||||
|
|
@ -1303,6 +1353,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
raise StopIteration from e
|
||||
self._raise_for_transport_error(e)
|
||||
except httpx.TimeoutException as e:
|
||||
if self.completed_response is not None:
|
||||
self.finished = True
|
||||
raise StopIteration from e
|
||||
self._raise_for_transport_timeout(e)
|
||||
except httpx.HTTPError as e:
|
||||
# Handle HTTP errors
|
||||
|
|
@ -1727,6 +1780,8 @@ def _estimate_usage_from_text(
|
|||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
*,
|
||||
use_default_image_token_count: bool = False,
|
||||
) -> ResponseAPIUsage:
|
||||
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union
|
||||
|
|
@ -1734,7 +1789,7 @@ def _estimate_usage_from_text(
|
|||
only_since_compaction=True,
|
||||
)
|
||||
input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, messages=messages
|
||||
model=model, messages=messages, use_default_image_token_count=use_default_image_token_count
|
||||
)
|
||||
output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, text=generated_text, count_response_tokens=True
|
||||
|
|
@ -1751,6 +1806,8 @@ def _estimate_usage_safely(
|
|||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
*,
|
||||
use_default_image_token_count: bool,
|
||||
) -> ResponseAPIUsage | None:
|
||||
try:
|
||||
return _estimate_usage_from_text(
|
||||
|
|
@ -1758,6 +1815,7 @@ def _estimate_usage_safely(
|
|||
request_input=request_input,
|
||||
responses_api_request=responses_api_request,
|
||||
generated_text=generated_text,
|
||||
use_default_image_token_count=use_default_image_token_count,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e)
|
||||
|
|
@ -2276,7 +2334,7 @@ class _ResponsesWsTurn:
|
|||
)
|
||||
if failed_event is None:
|
||||
return None
|
||||
return _map_stream_error_to_exception(
|
||||
return map_stream_error_to_exception(
|
||||
_ws_event_error(failed_event), model, custom_llm_provider, logging_obj=connection_logging_obj
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -6456,6 +6456,9 @@ class Router:
|
|||
hop_kwargs: Final = mid_stream_fallback_hop_kwargs(
|
||||
model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs
|
||||
)
|
||||
shared_logging: Final = kwargs.get("litellm_logging_obj")
|
||||
if isinstance(shared_logging, LiteLLMLogging):
|
||||
shared_logging.carry_failed_attempt_usage()
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks_helper(
|
||||
model=model, original_generic_function=original_generic_function, **kwargs
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1516,9 +1516,9 @@ def test_bedrock_mantle_openai_envelope_context_overflow_maps_to_context_window_
|
|||
|
||||
|
||||
def test_bedrock_mantle_streamed_context_overflow_event_maps_to_context_window_exceeded():
|
||||
from litellm.responses.streaming_iterator import _map_stream_error_to_exception
|
||||
from litellm.responses.streaming_iterator import map_stream_error_to_exception
|
||||
|
||||
mapped_exception = _map_stream_error_to_exception(
|
||||
mapped_exception = map_stream_error_to_exception(
|
||||
{
|
||||
"code": "context_length_exceeded",
|
||||
"message": "Your input exceeds the context window of this model. Please adjust your input and try again.",
|
||||
|
|
|
|||
|
|
@ -4,20 +4,25 @@ Tests for AnthropicResponsesStreamWrapper
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from types import SimpleNamespace
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../..")))
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
encrypted_reasoning_signature,
|
||||
)
|
||||
from litellm.llms.anthropic.messages.native.streaming_iterator import aclose_superseded_attempt
|
||||
from litellm.llms.anthropic.messages.native.utils import INCOMPLETE_STREAM_ERROR_MESSAGE
|
||||
from litellm.llms.anthropic.messages.responses_adapters.streaming_iterator import (
|
||||
AnthropicResponsesStreamWrapper,
|
||||
|
|
@ -622,3 +627,194 @@ class TestUpstreamFailureEndsStreamWithErrorEvent:
|
|||
chunks = _collect(iter(()))
|
||||
assert [chunk["type"] for chunk in chunks] == ["message_start", "error"]
|
||||
assert chunks[1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}
|
||||
|
||||
|
||||
class _FailureHookRecorder:
|
||||
def __init__(self) -> None:
|
||||
self.received: tuple[Exception, ...] = ()
|
||||
|
||||
async def __call__(self, failure: Exception) -> None:
|
||||
self.received = (*self.received, failure)
|
||||
|
||||
|
||||
def _logging_obj_with_failure_hook(on_failure: Callable[[Exception], Awaitable[None]]) -> LitellmLogging:
|
||||
logging_obj = LitellmLogging(
|
||||
model="m",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id="lit-6871-test",
|
||||
function_id="lit-6871-test",
|
||||
)
|
||||
logging_obj._on_detached_stream_failure = on_failure
|
||||
return logging_obj
|
||||
|
||||
|
||||
def _collect_with_logging(
|
||||
stream: AsyncIterator[dict[str, object]], logging_obj: LitellmLogging
|
||||
) -> list[dict[str, object]]:
|
||||
async def _run() -> list[dict[str, object]]:
|
||||
wrapper = AnthropicResponsesStreamWrapper(responses_stream=stream, model="m", litellm_logging_obj=logging_obj)
|
||||
return [chunk async for chunk in wrapper]
|
||||
|
||||
return asyncio.run(_run())
|
||||
|
||||
|
||||
class _DeliveryCountingConsumer:
|
||||
def __init__(self) -> None:
|
||||
self.delivered: tuple[dict[str, object], ...] = ()
|
||||
self.delivered_when_hook_fired: tuple[int, ...] = ()
|
||||
|
||||
async def hook(self, failure: Exception) -> None:
|
||||
self.delivered_when_hook_fired = (*self.delivered_when_hook_fired, len(self.delivered))
|
||||
|
||||
async def drain(self, wrapper: AnthropicResponsesStreamWrapper) -> None:
|
||||
async for chunk in wrapper:
|
||||
self.delivered = (*self.delivered, chunk)
|
||||
|
||||
|
||||
def _failing_after_one_delta() -> AsyncIterator[dict[str, object]]:
|
||||
async def _gen() -> AsyncIterator[dict[str, object]]:
|
||||
yield {"type": "response.created"}
|
||||
yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}
|
||||
yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"}
|
||||
raise litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m")
|
||||
|
||||
return _gen()
|
||||
|
||||
|
||||
class TestUpstreamFailureReachesTheProxyFailureHook:
|
||||
def test_raised_upstream_failure_reaches_the_hook_unwrapped_once_the_error_event_was_consumed(self):
|
||||
rate_limit = litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m")
|
||||
wrapped = MidStreamFallbackError(
|
||||
message=str(rate_limit),
|
||||
model="m",
|
||||
llm_provider="openai",
|
||||
original_exception=rate_limit,
|
||||
is_pre_first_chunk=False,
|
||||
)
|
||||
recorder = _FailureHookRecorder()
|
||||
|
||||
async def _gen() -> AsyncIterator[dict[str, object]]:
|
||||
yield {"type": "response.created"}
|
||||
yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}
|
||||
yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"}
|
||||
raise wrapped
|
||||
|
||||
chunks = _collect_with_logging(_gen(), _logging_obj_with_failure_hook(recorder))
|
||||
assert [chunk["type"] for chunk in chunks] == [
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"error",
|
||||
]
|
||||
assert recorder.received == (rate_limit,)
|
||||
|
||||
def test_response_failed_event_reaches_the_hook_as_an_exception_carrying_the_upstream_status(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
failed = SimpleNamespace(
|
||||
status="failed",
|
||||
output=[],
|
||||
usage=None,
|
||||
error={"code": "rate_limit_exceeded", "message": "Rate limit reached for gpt-5.5, try again in 20s."},
|
||||
)
|
||||
|
||||
async def _gen() -> AsyncIterator[dict[str, object]]:
|
||||
yield {"type": "response.created"}
|
||||
yield {"type": "response.failed", "response": failed}
|
||||
|
||||
chunks = _collect_with_logging(_gen(), _logging_obj_with_failure_hook(recorder))
|
||||
assert chunks[-1]["type"] == "error"
|
||||
(failure,) = recorder.received
|
||||
assert isinstance(failure, litellm.RateLimitError)
|
||||
assert failure.status_code == 429
|
||||
assert "Rate limit reached for gpt-5.5, try again in 20s." in failure.message
|
||||
|
||||
def test_a_raising_hook_never_hides_the_error_event_from_the_client(self):
|
||||
async def _exploding_hook(failure: Exception) -> None:
|
||||
raise RuntimeError("spend writer is down")
|
||||
|
||||
logging_obj = _logging_obj_with_failure_hook(_exploding_hook)
|
||||
|
||||
async def _gen() -> AsyncIterator[dict[str, object]]:
|
||||
yield {"type": "response.created"}
|
||||
raise ConnectionResetError("Response payload is not completed")
|
||||
|
||||
chunks = _collect_with_logging(_gen(), logging_obj)
|
||||
assert [chunk["type"] for chunk in chunks] == ["message_start", "error"]
|
||||
assert chunks[-1]["error"] == {"type": "api_error", "message": "Response payload is not completed"}
|
||||
|
||||
def test_the_hook_fires_only_after_the_consumer_took_the_error_event(self):
|
||||
consumer = _DeliveryCountingConsumer()
|
||||
|
||||
async def _run() -> None:
|
||||
wrapper = AnthropicResponsesStreamWrapper(
|
||||
responses_stream=_failing_after_one_delta(),
|
||||
model="m",
|
||||
litellm_logging_obj=_logging_obj_with_failure_hook(consumer.hook),
|
||||
)
|
||||
await consumer.drain(wrapper)
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await wrapper.__anext__()
|
||||
|
||||
asyncio.run(_run())
|
||||
assert [chunk["type"] for chunk in consumer.delivered] == [
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"error",
|
||||
]
|
||||
assert consumer.delivered_when_hook_fired == (len(consumer.delivered),)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("close", "expected_reports"),
|
||||
[
|
||||
pytest.param(lambda sse: sse.aclose(), 1, id="client_closes"),
|
||||
pytest.param(aclose_superseded_attempt, 0, id="fallback_takes_over"),
|
||||
],
|
||||
)
|
||||
def test_closing_the_stream_on_the_error_event_reports_unless_a_fallback_takes_over(self, close, expected_reports):
|
||||
recorder = _FailureHookRecorder()
|
||||
|
||||
async def _close_on_the_error_event() -> None:
|
||||
wrapper = AnthropicResponsesStreamWrapper(
|
||||
responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(recorder)
|
||||
)
|
||||
sse = wrapper.async_anthropic_sse_wrapper()
|
||||
async for frame in sse:
|
||||
if frame.startswith(b"event: error\n"):
|
||||
await close(sse)
|
||||
return
|
||||
pytest.fail("the stream ended without an error event")
|
||||
|
||||
asyncio.run(_close_on_the_error_event())
|
||||
assert len(recorder.received) == expected_reports
|
||||
|
||||
def test_a_client_disconnect_during_the_failure_hook_lets_the_hook_finish(self):
|
||||
hook_started = anyio.Event()
|
||||
hook_released = anyio.Event()
|
||||
finished: list[Exception] = []
|
||||
|
||||
async def _held_hook(failure: Exception) -> None:
|
||||
hook_started.set()
|
||||
await hook_released.wait()
|
||||
finished.append(failure)
|
||||
|
||||
async def _disconnect_mid_hook() -> None:
|
||||
wrapper = AnthropicResponsesStreamWrapper(
|
||||
responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(_held_hook)
|
||||
)
|
||||
|
||||
async def _consume() -> None:
|
||||
async for _ in wrapper.async_anthropic_sse_wrapper():
|
||||
pass
|
||||
|
||||
async with anyio.create_task_group() as task_group:
|
||||
task_group.start_soon(_consume)
|
||||
await hook_started.wait()
|
||||
task_group.cancel_scope.cancel()
|
||||
hook_released.set()
|
||||
|
||||
asyncio.run(_disconnect_mid_hook())
|
||||
assert len(finished) == 1
|
||||
|
|
|
|||
|
|
@ -10826,12 +10826,19 @@ class TestResolveLitellmCallId:
|
|||
class _FailureHookRecorder:
|
||||
"""Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it."""
|
||||
|
||||
def __init__(self, raises: Optional[Exception] = None):
|
||||
def __init__(self, raises: Optional[Exception] = None, gate: Optional[asyncio.Event] = None):
|
||||
self.calls = []
|
||||
self.entered = asyncio.Event()
|
||||
self.reported = asyncio.Event()
|
||||
self._raises = raises
|
||||
self._gate = gate
|
||||
|
||||
async def post_call_failure_hook(self, **kwargs):
|
||||
self.entered.set()
|
||||
if self._gate is not None:
|
||||
await self._gate.wait()
|
||||
self.calls.append(kwargs)
|
||||
self.reported.set()
|
||||
if self._raises is not None:
|
||||
raise self._raises
|
||||
|
||||
|
|
@ -10932,6 +10939,171 @@ class TestDetachedStreamFailureHook:
|
|||
|
||||
assert [call["original_exception"] for call in recorder.calls] == [failure]
|
||||
|
||||
@staticmethod
|
||||
def _armed_stream(recorder: _FailureHookRecorder, iterator_hook_for, response: Optional[MagicMock] = None):
|
||||
from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth
|
||||
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-6.1-sol",
|
||||
messages=[{"role": "user", "content": "Write a story."}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=None,
|
||||
litellm_call_id="call-lit6871",
|
||||
function_id="function-lit6871",
|
||||
)
|
||||
user_api_key_dict = RealUserAPIKeyAuth(api_key="sk-test")
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj})
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
|
||||
proxy_logging_obj.post_call_failure_hook = recorder.post_call_failure_hook
|
||||
proxy_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook_for(logging_obj)
|
||||
processing_obj._arm_detached_stream_failure_hook(
|
||||
logging_obj=logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
stream = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=MagicMock() if response is None else response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=processing_obj.data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
serialize_chunk=lambda chunk: "data: chunk\n\n",
|
||||
serialize_error=lambda error: "data: error\n\n",
|
||||
)
|
||||
return logging_obj, stream
|
||||
|
||||
@staticmethod
|
||||
def _hook_reporting(upstream_failure: Exception, after_report=None):
|
||||
def iterator_hook_for(logging_obj):
|
||||
async def iterator_hook(**_kwargs):
|
||||
yield b"event: message_start\n\n"
|
||||
await logging_obj.on_detached_stream_failure(upstream_failure)
|
||||
yield b"event: error\n\n"
|
||||
yield b"event: message_stop\n\n"
|
||||
if after_report is not None:
|
||||
raise after_report
|
||||
|
||||
return iterator_hook
|
||||
|
||||
return iterator_hook_for
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_iterator_hook_failure_after_a_bridged_failure_is_the_one_the_failure_hook_gets(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
hook_failure = RuntimeError("iterator hook rejected the stream")
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure))
|
||||
|
||||
frames = [frame async for frame in stream]
|
||||
|
||||
assert frames == ["data: chunk\n\n"] * 3 + ["data: error\n\n"]
|
||||
assert [call["original_exception"] for call in recorder.calls] == [hook_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_hook_replacing_the_iterator_hook_failure_aborts_the_stream(self):
|
||||
replacement = HTTPException(status_code=418, detail="rejected by a custom failure hook")
|
||||
recorder = _FailureHookRecorder(raises=replacement)
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
hook_failure = RuntimeError("iterator hook rejected the stream")
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure))
|
||||
|
||||
frames = [await stream.__anext__() for _ in range(3)]
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await stream.__anext__()
|
||||
|
||||
assert raised.value is replacement
|
||||
assert frames == ["data: chunk\n\n"] * 3
|
||||
assert [call["original_exception"] for call in recorder.calls] == [hook_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_after_the_stream_ends(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure))
|
||||
|
||||
frames = []
|
||||
async for frame in stream:
|
||||
frames.append(frame)
|
||||
assert recorder.calls == []
|
||||
|
||||
assert frames == ["data: chunk\n\n"] * 3
|
||||
assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_report_and_stream_cleanup_survive_the_client_task_being_cancelled_mid_report(self):
|
||||
gate = asyncio.Event()
|
||||
recorder = _FailureHookRecorder(gate=gate)
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
response = MagicMock()
|
||||
response.aclose = AsyncMock()
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure), response=response)
|
||||
|
||||
async def client_reads_the_whole_stream():
|
||||
return [frame async for frame in stream]
|
||||
|
||||
client = asyncio.create_task(client_reads_the_whole_stream())
|
||||
await asyncio.wait_for(recorder.entered.wait(), timeout=5)
|
||||
client.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await client
|
||||
|
||||
response.aclose.assert_awaited_once()
|
||||
gate.set()
|
||||
await asyncio.wait_for(recorder.reported.wait(), timeout=5)
|
||||
assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_when_the_client_leaves_early(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure))
|
||||
|
||||
await stream.__anext__()
|
||||
await stream.__anext__()
|
||||
await stream.aclose()
|
||||
|
||||
assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_reported_after_the_stream_ends_still_reaches_the_failure_hook(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed after the client stream closed")
|
||||
|
||||
async def iterator_hook(**_kwargs):
|
||||
yield b"event: message_start\n\n"
|
||||
|
||||
logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook)
|
||||
|
||||
frames = [frame async for frame in stream]
|
||||
await logging_obj.on_detached_stream_failure(upstream_failure)
|
||||
|
||||
assert frames == ["data: chunk\n\n"]
|
||||
assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_reported_by_an_iterator_hook_is_not_reported_again_when_the_stream_closes(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
hook_failure = RuntimeError("iterator hook rejected the error frame")
|
||||
|
||||
async def iterator_hook(**_kwargs):
|
||||
yield b"event: message_start\n\n"
|
||||
raise hook_failure
|
||||
|
||||
logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook)
|
||||
|
||||
frames = [frame async for frame in stream]
|
||||
await logging_obj.on_detached_stream_failure(upstream_failure)
|
||||
|
||||
assert frames == ["data: chunk\n\n", "data: error\n\n"]
|
||||
assert [call["original_exception"] for call in recorder.calls] == [hook_failure]
|
||||
|
||||
|
||||
class TestPostCallMaskedOutputReachesDeferredLogging:
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -2668,6 +2668,65 @@ async def test_stream_auto_execute_bills_the_paid_first_round_on_its_success_row
|
|||
assert _request_failure_row(logging_obj, recorder.logged[0]).get("response_cost", 0.0) == 0.0
|
||||
|
||||
|
||||
class _CutAfterStreamedText(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
opened: Final = {**json.loads(json.dumps(_ROUND_TWO_BODY, default=dict)), "status": "in_progress"}
|
||||
yield _sse(
|
||||
(
|
||||
{"type": "response.created", "response": opened, "sequence_number": 0},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_round_two",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "PING " * 40,
|
||||
"sequence_number": 1,
|
||||
},
|
||||
)
|
||||
)
|
||||
raise httpx.ReadError("connection reset mid-stream")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_bills_a_follow_up_cut_after_streamed_text_with_no_second_fixed_margin(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}})
|
||||
_serve_responses(
|
||||
respx_mock,
|
||||
monkeypatch,
|
||||
(_streamed(_ROUND_ONE_BODY), httpx.Response(200, stream=_CutAfterStreamedText(), headers=SSE_HEADERS)),
|
||||
)
|
||||
recorder: Final = _RoundRecorder()
|
||||
logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-cut-after-text")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await _stream_auto_execute(logging_obj)
|
||||
ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
|
||||
await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0)
|
||||
|
||||
_assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY)
|
||||
assert len(recorder.failed) == 1
|
||||
failed: Final = recorder.failed[0]
|
||||
assert failed["completion_tokens"] > 0
|
||||
priced_with_its_own_margin: Final = litellm.completion_cost(
|
||||
completion_response=ModelResponse(
|
||||
model="gpt-5.6",
|
||||
usage=Usage(
|
||||
prompt_tokens=failed["prompt_tokens"],
|
||||
completion_tokens=failed["completion_tokens"],
|
||||
total_tokens=failed["prompt_tokens"] + failed["completion_tokens"],
|
||||
),
|
||||
),
|
||||
model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
assert failed["response_cost"] == pytest.approx(priced_with_its_own_margin - _FIXED_MARGIN)
|
||||
_assert_breakdown_agrees_with_the_billed_cost(failed)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
|
|
|
|||
|
|
@ -9,20 +9,23 @@ crashed with exit 139 whenever any CustomLogger was registered).
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module
|
||||
from litellm.responses import streaming_iterator as responses_streaming_iterator_module
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.responses import streaming_iterator as responses_streaming_iterator_module
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
class RecordingCustomLogger(CustomLogger):
|
||||
|
|
@ -169,6 +172,240 @@ async def test_sync_callbacks_run_only_after_async_handler_completes(recording_e
|
|||
assert submit_times[0] >= recorder.async_hook_finished
|
||||
|
||||
|
||||
class FailureRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.failure_payloads: tuple[StandardLoggingPayload, ...] = ()
|
||||
self.success_payloads: tuple[StandardLoggingPayload, ...] = ()
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.failure_payloads = (
|
||||
*self.failure_payloads,
|
||||
cast(StandardLoggingPayload, kwargs["standard_logging_object"]),
|
||||
)
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.success_payloads = (
|
||||
*self.success_payloads,
|
||||
cast(StandardLoggingPayload, kwargs["standard_logging_object"]),
|
||||
)
|
||||
|
||||
|
||||
class _ChunkRecorder:
|
||||
def __init__(self) -> None:
|
||||
self.chunks: tuple[ResponsesAPIStreamingResponse, ...] = ()
|
||||
|
||||
async def drain(self, iterator: ResponsesAPIStreamingIterator) -> None:
|
||||
async for chunk in iterator:
|
||||
self.chunks = (*self.chunks, chunk)
|
||||
|
||||
|
||||
_SEA_DELTAS = ("The sea ", "is wide ", "and deep.")
|
||||
_SEA_PROMPT = "Write a 300 word story about the sea."
|
||||
_UNREACHABLE_IMAGE_INPUT: list[dict[str, object]] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "Describe this picture."},
|
||||
{"type": "input_image", "image_url": "http://images.invalid/photo.png", "detail": "high"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
_TOOL_ARGUMENT_DELTAS = ('{"city": ', '"Lisbon", ', '"unit": "celsius"}')
|
||||
|
||||
|
||||
class _UpstreamThatTimesOutAfterThreeDeltas:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
def __init__(self, event_type: str = "response.output_text.delta", deltas: tuple[str, ...] = _SEA_DELTAS) -> None:
|
||||
self.event_type = event_type
|
||||
self.deltas = deltas
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
for sequence_number, delta in enumerate(self.deltas):
|
||||
event = {
|
||||
"type": self.event_type,
|
||||
"item_id": "item_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": delta,
|
||||
"sequence_number": sequence_number,
|
||||
}
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
raise httpx.ReadTimeout("Timeout on reading data from socket")
|
||||
|
||||
|
||||
class _UpstreamThatFailsWithUsageAfterThreeDeltas:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
for sequence_number, delta in enumerate(_SEA_DELTAS):
|
||||
event = {
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": delta,
|
||||
"sequence_number": sequence_number,
|
||||
}
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
failed = {
|
||||
"type": "response.failed",
|
||||
"sequence_number": len(_SEA_DELTAS),
|
||||
"response": {
|
||||
"id": "resp_1",
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"status": "failed",
|
||||
"model": "gpt-5.4-nano",
|
||||
"output": [],
|
||||
"error": {"code": "server_error", "message": "The model had an internal error"},
|
||||
"usage": {"input_tokens": 400, "output_tokens": 9, "total_tokens": 409},
|
||||
},
|
||||
}
|
||||
yield f"data: {json.dumps(failed)}\n\n".encode()
|
||||
|
||||
|
||||
async def _failure_payload_after_mid_stream_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
request_input: str | list[dict[str, object]],
|
||||
upstream: _UpstreamThatTimesOutAfterThreeDeltas | None = None,
|
||||
background_fields: Mapping[str, object] | None = None,
|
||||
) -> tuple[FailureRecorder, StandardLoggingPayload]:
|
||||
upstream_response: Final = upstream or _UpstreamThatTimesOutAfterThreeDeltas()
|
||||
recorder = FailureRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "success_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [recorder])
|
||||
logging_obj = _make_logging_obj()
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
optional_params={"stream": True},
|
||||
litellm_params={"custom_llm_provider": "openai", "aresponses": True},
|
||||
)
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=cast(httpx.Response, upstream_response),
|
||||
model="gpt-5.4-nano",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai",
|
||||
request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, **(background_fields or {})},
|
||||
)
|
||||
received = _ChunkRecorder()
|
||||
|
||||
with pytest.raises(litellm.Timeout, match="Timeout on reading data from socket"):
|
||||
await received.drain(iterator)
|
||||
assert [getattr(chunk, "delta", None) for chunk in received.chunks] == list(upstream_response.deltas)
|
||||
|
||||
await iterator._await_pending_logging()
|
||||
(payload,) = recorder.failure_payloads
|
||||
return recorder, payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_read_timeout_logs_failure_with_partial_usage_and_no_success(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
recorder, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["prompt_tokens"] == litellm.token_counter(
|
||||
model="gpt-5.4-nano", messages=[{"role": "user", "content": _SEA_PROMPT}]
|
||||
)
|
||||
assert payload["completion_tokens"] == litellm.token_counter(
|
||||
model="gpt-5.4-nano", text="".join(_SEA_DELTAS), count_response_tokens=True
|
||||
)
|
||||
assert payload["response_cost"] > 0
|
||||
assert recorder.success_payloads == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_read_timeout_bills_the_tool_call_arguments_already_streamed(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
upstream: Final = _UpstreamThatTimesOutAfterThreeDeltas(
|
||||
event_type="response.function_call_arguments.delta", deltas=_TOOL_ARGUMENT_DELTAS
|
||||
)
|
||||
_, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT, upstream)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["completion_tokens"] == litellm.token_counter(
|
||||
model="gpt-5.4-nano", text="".join(_TOOL_ARGUMENT_DELTAS), count_response_tokens=True
|
||||
)
|
||||
assert payload["response_cost"] > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_of_a_background_response_bills_nothing_since_retrieving_it_bills_the_whole_run(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_, payload = await _failure_payload_after_mid_stream_timeout(
|
||||
monkeypatch, _SEA_PROMPT, background_fields={"background": True}
|
||||
)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["completion_tokens"] == 0
|
||||
assert payload["response_cost"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_usage_estimate_never_fetches_image_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _UNREACHABLE_IMAGE_INPUT)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["prompt_tokens"] == litellm.token_counter(
|
||||
model="gpt-5.4-nano",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this picture."},
|
||||
{"type": "image_url", "image_url": {"url": "http://images.invalid/photo.png", "detail": "high"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
use_default_image_token_count=True,
|
||||
)
|
||||
assert payload["completion_tokens"] > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_keeps_the_usage_the_provider_reported(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
recorder = FailureRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
logging_obj = _make_logging_obj()
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
optional_params={"stream": True},
|
||||
litellm_params={"custom_llm_provider": "openai", "aresponses": True},
|
||||
)
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=cast(httpx.Response, _UpstreamThatFailsWithUsageAfterThreeDeltas()),
|
||||
model="gpt-5.4-nano",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai",
|
||||
request_data={"model": "gpt-5.4-nano", "input": _SEA_PROMPT, "stream": True},
|
||||
)
|
||||
|
||||
with pytest.raises(MidStreamFallbackError, match="The model had an internal error"):
|
||||
await _ChunkRecorder().drain(iterator)
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
(payload,) = recorder.failure_payloads
|
||||
assert (payload["prompt_tokens"], payload["completion_tokens"]) == (400, 9)
|
||||
|
||||
|
||||
_CREATED_RESPONSE_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = {
|
||||
"id": "resp_partial",
|
||||
"object": "response",
|
||||
|
|
@ -211,17 +448,135 @@ def _completed_after(text: str) -> Mapping[str, object]:
|
|||
}
|
||||
|
||||
|
||||
def _live_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator:
|
||||
def _live_iterator(
|
||||
logging_obj: LitellmLogging, upstream: object = None, background: bool = False
|
||||
) -> ResponsesAPIStreamingIterator:
|
||||
return ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
response=cast(httpx.Response, upstream or httpx.Response(200)),
|
||||
model="gpt-5.4-nano",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "gpt-5.4-nano", "input": "tell me a story"},
|
||||
request_data={"model": "gpt-5.4-nano", "input": "tell me a story", "background": background},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
|
||||
class _UpstreamThatTimesOutAfterRelayingText:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)):
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
raise httpx.ReadTimeout("Timeout on reading data from socket")
|
||||
|
||||
|
||||
class _UpstreamThatCompletesThenTimesOut:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
events: Final = (
|
||||
_CREATED_WITH_PRIORITY_TIER,
|
||||
*(_text_delta(delta, sequence_number) for sequence_number, delta in enumerate(_SEA_DELTAS, start=1)),
|
||||
_completed_after("".join(_SEA_DELTAS)),
|
||||
)
|
||||
for event in events:
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
raise httpx.ReadTimeout("Timeout on reading data from socket")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_after_completed_response_ends_the_stream_and_bills_it_once(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
recorder: Final = FailureRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "success_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [recorder])
|
||||
logging_obj: Final = _make_logging_obj()
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
optional_params={"stream": True},
|
||||
litellm_params={"custom_llm_provider": "openai", "aresponses": True},
|
||||
)
|
||||
iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatCompletesThenTimesOut())
|
||||
|
||||
chunks: Final = _ChunkRecorder()
|
||||
await chunks.drain(iterator)
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
assert chunks.chunks[-1].type == "response.completed"
|
||||
assert len(recorder.success_payloads) == 1
|
||||
success_payload: Final = recorder.success_payloads[0]
|
||||
response_payload: Final = success_payload["response"]
|
||||
assert isinstance(response_payload, dict)
|
||||
usage: Final = cast(dict[str, object], response_payload["usage"])
|
||||
assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == (5, 7, 12)
|
||||
assert recorder.failure_payloads == ()
|
||||
assert logging_obj.model_call_details.get("combined_usage_object") is None
|
||||
|
||||
|
||||
class _UpstreamThatTimesOutAfterCreated:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
def __init__(self, **created_fields: object) -> None:
|
||||
self.created_fields = created_fields
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
created: Final = {
|
||||
**_CREATED_WITH_PRIORITY_TIER,
|
||||
"response": {**_CREATED_RESPONSE_WITH_PRIORITY_TIER, **self.created_fields},
|
||||
}
|
||||
for event in (created, _text_delta("Once upon a time", 1)):
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
raise httpx.ReadTimeout("Timeout on reading data from socket")
|
||||
|
||||
|
||||
async def _mid_stream_failure_payload(
|
||||
monkeypatch: pytest.MonkeyPatch, requested_tier: str | None, **created_fields: object
|
||||
) -> StandardLoggingPayload:
|
||||
recorder: Final = FailureRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
logging_obj: Final = _make_logging_obj()
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
optional_params={"stream": True}
|
||||
if requested_tier is None
|
||||
else {"stream": True, "service_tier": requested_tier},
|
||||
litellm_params={"custom_llm_provider": "openai", "aresponses": True},
|
||||
)
|
||||
iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatTimesOutAfterCreated(**created_fields))
|
||||
|
||||
with pytest.raises(litellm.Timeout):
|
||||
await _ChunkRecorder().drain(iterator)
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
(payload,) = recorder.failure_payloads
|
||||
return payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_bills_the_relayed_text_at_the_tier_the_provider_served(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
flex_requested: Final = await _mid_stream_failure_payload(monkeypatch, "flex", service_tier="default")
|
||||
nothing_requested: Final = await _mid_stream_failure_payload(monkeypatch, None, service_tier="default")
|
||||
|
||||
assert nothing_requested["response_cost"] > 0
|
||||
assert flex_requested["response_cost"] == nothing_requested["response_cost"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_bills_nothing_when_the_provider_runs_the_response_in_background(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
payload: Final = await _mid_stream_failure_payload(monkeypatch, None, background=True)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert (payload["completion_tokens"], payload["response_cost"]) == (0, 0)
|
||||
|
||||
|
||||
def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier():
|
||||
iterator: Final = _live_iterator(_make_logging_obj())
|
||||
for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)):
|
||||
|
|
@ -263,3 +618,25 @@ async def test_partial_response_on_disconnect_is_none_when_there_is_nothing_to_b
|
|||
|
||||
assert iterator.partial_response_on_disconnect() is None
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_response_on_disconnect_is_none_once_the_failure_log_billed_the_relayed_text():
|
||||
iterator: Final = _live_iterator(_make_logging_obj(), upstream=_UpstreamThatTimesOutAfterRelayingText())
|
||||
|
||||
with pytest.raises(litellm.Timeout):
|
||||
await _ChunkRecorder().drain(iterator)
|
||||
|
||||
assert iterator.partial_response_on_disconnect() is None
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
|
||||
def test_partial_response_on_disconnect_reports_the_relayed_text_of_a_background_response():
|
||||
iterator: Final = _live_iterator(_make_logging_obj(), background=True)
|
||||
for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)):
|
||||
iterator._process_chunk(json.dumps(event))
|
||||
|
||||
incomplete_event: Final = iterator.partial_response_on_disconnect()
|
||||
|
||||
assert incomplete_event is not None
|
||||
assert incomplete_event.response.output[0].content[0].text == "Once upon a time"
|
||||
|
|
|
|||
|
|
@ -210,9 +210,11 @@ _TRAILING_ERRORS = [
|
|||
httpx.RemoteProtocolError("peer closed connection without sending complete message body"),
|
||||
]
|
||||
|
||||
_TRAILING_ERRORS_AFTER_COMPLETED = [*_TRAILING_ERRORS, httpx.ReadTimeout("Timeout on reading data from socket")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type)
|
||||
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS_AFTER_COMPLETED, ids=type)
|
||||
async def test_transport_error_after_completed_event_ends_stream_cleanly(trailing_error):
|
||||
"""A sloppy connection close after `response.completed` must not turn a
|
||||
complete stream into an error (regression guard for the transport no longer
|
||||
|
|
@ -243,7 +245,7 @@ async def test_transport_error_before_completed_event_raises():
|
|||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type)
|
||||
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS_AFTER_COMPLETED, ids=type)
|
||||
def test_sync_transport_error_after_completed_event_ends_stream_cleanly(trailing_error):
|
||||
iterator = _make_sync_iterator(
|
||||
sse_events=_COMPLETE_STREAM_EVENTS,
|
||||
|
|
@ -1205,7 +1207,9 @@ async def test_completed_event_survives_a_failing_usage_estimate():
|
|||
falls back to usage None."""
|
||||
malformed_input: Final = [{"type": "message", "role": "user", "content": 42}]
|
||||
with pytest.raises(ValueError, match="Invalid content type"):
|
||||
_estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world")
|
||||
_estimate_usage_from_text(
|
||||
"gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world"
|
||||
)
|
||||
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
|
|
|
|||
|
|
@ -3774,6 +3774,132 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t
|
|||
)
|
||||
|
||||
|
||||
def _responses_sse(*events: Mapping[str, object]) -> str:
|
||||
return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events)
|
||||
|
||||
|
||||
def _responses_envelope(response_id: str, model: str, status: str, **fields: object) -> dict[str, object]:
|
||||
return {
|
||||
"id": response_id,
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"status": status,
|
||||
"model": model,
|
||||
"output": [],
|
||||
**fields,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fallback_usage(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
recorder: Final = _SpendRecorder()
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
respx_mock.post("https://primary.example/v1/responses").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
text=_responses_sse(
|
||||
{"type": "response.created", "response": _responses_envelope("resp_primary", "gpt-4o", "in_progress")},
|
||||
*(
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_primary",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": word,
|
||||
}
|
||||
for word in ("Once ", "upon ", "a ", "time ", "there")
|
||||
),
|
||||
{
|
||||
"type": "response.failed",
|
||||
"response": _responses_envelope(
|
||||
"resp_primary",
|
||||
"gpt-4o",
|
||||
"failed",
|
||||
error={"code": "server_error", "message": "The model had an internal error"},
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
respx_mock.post("https://backup.example/v1/responses").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
text=_responses_sse(
|
||||
{"type": "response.created", "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress")},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": _responses_envelope(
|
||||
"resp_backup",
|
||||
"gpt-4o-mini",
|
||||
"completed",
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_backup",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "Once upon a time", "annotations": []}],
|
||||
}
|
||||
],
|
||||
usage={"input_tokens": 1000, "output_tokens": 700, "total_tokens": 1700},
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "primary",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "k1", "api_base": "https://primary.example/v1"},
|
||||
},
|
||||
{
|
||||
"model_name": "backup",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "k2",
|
||||
"api_base": "https://backup.example/v1",
|
||||
},
|
||||
},
|
||||
],
|
||||
fallbacks=[{"primary": ["backup"]}],
|
||||
num_retries=0,
|
||||
)
|
||||
shared_logging, _ = litellm.utils.function_setup(
|
||||
original_function="aresponses",
|
||||
rules_obj=litellm.utils.Rules(),
|
||||
start_time=datetime.now(),
|
||||
model="primary",
|
||||
input="Tell me a story",
|
||||
stream=True,
|
||||
litellm_call_id="shared-call",
|
||||
)
|
||||
|
||||
response: Final = await router.aresponses(
|
||||
model="primary",
|
||||
input="Tell me a story",
|
||||
stream=True,
|
||||
litellm_call_id="shared-call",
|
||||
litellm_logging_obj=shared_logging,
|
||||
)
|
||||
_ = [event async for event in response]
|
||||
await asyncio.wait_for(recorder.succeeded.wait(), timeout=10)
|
||||
|
||||
success: Final = recorder.successes[-1]
|
||||
assert (success["prompt_tokens"], success["completion_tokens"]) == (1000, 700)
|
||||
spend_row: Final = recorder.spend_rows[-1]
|
||||
assert (spend_row["prompt_tokens"], spend_row["completion_tokens"], spend_row["spend"]) == (
|
||||
1000,
|
||||
700,
|
||||
success["response_cost"],
|
||||
)
|
||||
|
||||
|
||||
def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited():
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue