mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(proxy): bill partial Responses API streams when the client disconnects (#47611)
Backport of #47611 to stable/1.103.x
Cherry-picked from f0ae445cc4 (main)
This commit is contained in:
devin-ai-integration[bot]
parent
d7e344440c
commit
bb2ec9a422
10 changed files with 537 additions and 22 deletions
|
|
@ -4,7 +4,7 @@ import asyncio
|
|||
import json
|
||||
import traceback
|
||||
from collections import deque
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping
|
||||
from collections.abc import AsyncGenerator, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm import verbose_logger
|
||||
|
|
@ -23,8 +23,11 @@ from .transformation import (
|
|||
LiteLLMAnthropicToResponsesAPIAdapter,
|
||||
)
|
||||
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
|
||||
from litellm.types.llms.openai import ResponseIncompleteEvent
|
||||
|
||||
|
||||
class AnthropicResponsesStreamWrapper:
|
||||
|
|
@ -391,8 +394,11 @@ class AnthropicResponsesStreamWrapper:
|
|||
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]:
|
||||
"""Yield SSE-encoded bytes for each Anthropic event chunk."""
|
||||
def async_anthropic_sse_wrapper(self) -> "AnthropicResponsesSSEByteStream":
|
||||
"""SSE-encoded bytes for each Anthropic event chunk."""
|
||||
return AnthropicResponsesSSEByteStream(self)
|
||||
|
||||
async def _sse_frames(self) -> AsyncGenerator[bytes, None]:
|
||||
async for chunk in self:
|
||||
if isinstance(chunk, dict):
|
||||
event_type: str = str(chunk.get("type", "message"))
|
||||
|
|
@ -400,3 +406,30 @@ class AnthropicResponsesStreamWrapper:
|
|||
yield payload.encode()
|
||||
else:
|
||||
yield chunk
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return partial_responses_incomplete_event_on_disconnect(self.responses_stream)
|
||||
|
||||
|
||||
class AnthropicResponsesSSEByteStream:
|
||||
"""
|
||||
The byte stream the proxy relays for a default-route /v1/messages request.
|
||||
Keeps the Responses API iterator reachable so a client disconnect can bill
|
||||
the partial output the way native /v1/responses streams do.
|
||||
"""
|
||||
|
||||
def __init__(self, wrapper: AnthropicResponsesStreamWrapper) -> None:
|
||||
self._wrapper: Final = wrapper
|
||||
self._frames: Final = wrapper._sse_frames() # pyright: ignore[reportPrivateUsage] # the wrapper's own byte stream
|
||||
|
||||
def __aiter__(self) -> "AnthropicResponsesSSEByteStream":
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
return await self._frames.__anext__()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._frames.aclose()
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return self._wrapper.partial_response_on_disconnect()
|
||||
|
|
|
|||
|
|
@ -95,6 +95,7 @@ from litellm.proxy.dd_span_tagger import DDSpanTagger
|
|||
from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
|
|
@ -396,6 +397,13 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
|
|||
# A natural end-of-stream success event already fired and released the
|
||||
# slot; do not bill again, and let the caller skip the slot release.
|
||||
return True
|
||||
partial_responses_event: Final = partial_responses_incomplete_event_on_disconnect(response)
|
||||
if partial_responses_event is not None:
|
||||
verbose_proxy_logger.debug(
|
||||
"Billing partial streamed Responses API spend after client disconnect, litellm_call_id=%s",
|
||||
logging_obj.litellm_call_id,
|
||||
)
|
||||
return await _dispatch_disconnect_billing_event(logging_obj, partial_responses_event)
|
||||
chunks: Final[object] = getattr(response, "chunks", None)
|
||||
if not isinstance(chunks, list) or not chunks:
|
||||
return False
|
||||
|
|
@ -425,6 +433,10 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
|
|||
partial_usage: Final = getattr(partial_response, "usage", None)
|
||||
if isinstance(partial_usage, Usage):
|
||||
backfill_missing_cache_usage_fields(partial_usage)
|
||||
return await _dispatch_disconnect_billing_event(logging_obj, partial_response)
|
||||
|
||||
|
||||
async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool:
|
||||
try:
|
||||
await logging_obj.dispatch_success_handlers(
|
||||
partial_response,
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from litellm.responses.litellm_completion_transformation.custom_tools import (
|
|||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import (
|
||||
PART_UNION_TYPES,
|
||||
|
|
@ -33,6 +33,7 @@ from litellm.types.llms.openai import (
|
|||
ReasoningSummaryTextDeltaEvent,
|
||||
ReasoningSummaryTextDoneEvent,
|
||||
ResponseCompletedEvent,
|
||||
ResponseIncompleteEvent,
|
||||
ResponseCreatedEvent,
|
||||
ResponseInProgressEvent,
|
||||
ResponseInputParam,
|
||||
|
|
@ -1258,6 +1259,43 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
)
|
||||
return reasoning_aligned
|
||||
|
||||
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
|
||||
"""
|
||||
The response.incomplete event to bill when the client went away mid-stream.
|
||||
Built from the chat completion chunks relayed so far, so the usage is what
|
||||
stream_chunk_builder counts for them.
|
||||
"""
|
||||
if self.completed_response is not None or not self.collected_chat_completion_chunks:
|
||||
return None
|
||||
built: Final = stream_chunk_builder(
|
||||
chunks=self.collected_chat_completion_chunks,
|
||||
messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
input=self.request_input,
|
||||
responses_api_request=dict(self.responses_api_request),
|
||||
),
|
||||
logging_obj=self.litellm_logging_obj,
|
||||
)
|
||||
if not isinstance(built, ModelResponse):
|
||||
return None
|
||||
responses_api_response: Final = (
|
||||
LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
request_input=self.request_input,
|
||||
chat_completion_response=built,
|
||||
responses_api_request=self.responses_api_request,
|
||||
)
|
||||
)
|
||||
partial: Final = self._with_encoded_response_id(
|
||||
responses_api_response.model_copy(
|
||||
update={
|
||||
"id": self._cached_response_id or responses_api_response.id,
|
||||
"status": "incomplete",
|
||||
"output": list(self._output_with_streamed_item_ids(responses_api_response)),
|
||||
}
|
||||
)
|
||||
)
|
||||
stamp_responses_usage_cost(partial, self.litellm_logging_obj)
|
||||
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
|
||||
|
||||
def _emit_response_completed_event(self, litellm_model_response: ModelResponse) -> ResponseCompletedEvent | None:
|
||||
if litellm_model_response:
|
||||
# Transform the response
|
||||
|
|
|
|||
|
|
@ -41,10 +41,12 @@ from litellm.types.integrations.custom_logger import converted_stream_requested
|
|||
from litellm.types.llms.openai import (
|
||||
PART_UNION_TYPES,
|
||||
ResponseAPIUsage,
|
||||
ResponseIncompleteEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
ResponsesAPIStreamingResponse,
|
||||
)
|
||||
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
|
||||
from litellm.types.utils import CallTypes
|
||||
from litellm.utils import async_post_call_success_deployment_hook
|
||||
|
||||
|
|
@ -289,6 +291,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._yielded_first_chunk = False
|
||||
self._generated_content = ""
|
||||
self._generated_tool_arguments = ""
|
||||
self._created_response: ResponsesAPIResponse | None = None
|
||||
self._completed_response_cached = False
|
||||
self._completed_response_logged = False
|
||||
self._completed_response_cache_hit: bool | None = None
|
||||
|
|
@ -450,6 +453,10 @@ class BaseResponsesAPIStreamingIterator:
|
|||
# Store the completed response (also for incomplete/failed so logging still fires)
|
||||
_chunk_type: Final = getattr(openai_responses_api_chunk, "type", None)
|
||||
openai_types = _get_openai_response_types()
|
||||
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED:
|
||||
_created: Final[object] = getattr(openai_responses_api_chunk, "response", None)
|
||||
if isinstance(_created, ResponsesAPIResponse):
|
||||
self._created_response = _created
|
||||
if openai_responses_api_chunk and _chunk_type in (
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
|
|
@ -481,7 +488,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
else openai_responses_api_chunk.model_copy(update={"response": _billed_response})
|
||||
)
|
||||
self.completed_response = _terminal_chunk
|
||||
_stamp_responses_usage_cost(_billed_response, self.logging_obj)
|
||||
stamp_responses_usage_cost(_billed_response, self.logging_obj)
|
||||
|
||||
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
|
||||
self._handle_logging_failed_response()
|
||||
|
|
@ -875,6 +882,34 @@ class BaseResponsesAPIStreamingIterator:
|
|||
pass
|
||||
|
||||
|
||||
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
|
||||
"""
|
||||
The response.incomplete event to bill when the client went away mid-stream,
|
||||
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.
|
||||
"""
|
||||
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:
|
||||
return None
|
||||
usage: Final = _estimate_usage_safely(
|
||||
self.model or "", self.request_data.get("input"), self.request_data, generated
|
||||
)
|
||||
if usage is None:
|
||||
return None
|
||||
partial: Final = created.model_copy(
|
||||
update={
|
||||
"status": "incomplete",
|
||||
"usage": usage,
|
||||
"output": _partial_output_items(self._generated_content),
|
||||
}
|
||||
)
|
||||
stamp_responses_usage_cost(partial, self.logging_obj)
|
||||
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
|
||||
|
||||
async def call_post_streaming_hooks_for_testing(
|
||||
iterator: object, chunk: ResponsesAPIStreamingResponse
|
||||
) -> ResponsesAPIStreamingResponse:
|
||||
|
|
@ -1347,6 +1382,37 @@ def _add_text_like_part_events(
|
|||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _BillsPartialStreamOnDisconnect(Protocol):
|
||||
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: ...
|
||||
|
||||
|
||||
def partial_responses_incomplete_event_on_disconnect(stream: object) -> ResponseIncompleteEvent | None:
|
||||
"""Ask the stream behind a proxy or router wrapper for the partial response to bill, if it can build one."""
|
||||
if not isinstance(stream, _BillsPartialStreamOnDisconnect):
|
||||
return None
|
||||
try:
|
||||
partial: Final = stream.partial_response_on_disconnect()
|
||||
except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown
|
||||
verbose_logger.debug("Failed to build partial Responses API response for disconnect billing: %s", e)
|
||||
return None
|
||||
return partial if isinstance(partial, ResponseIncompleteEvent) else None
|
||||
|
||||
|
||||
def _partial_output_items(generated_text: str) -> Sequence[GenericResponseOutputItem]:
|
||||
if not generated_text:
|
||||
return []
|
||||
return [
|
||||
GenericResponseOutputItem(
|
||||
type="message",
|
||||
id=f"msg_{uuid.uuid4()}",
|
||||
status="incomplete",
|
||||
role="assistant",
|
||||
content=[OutputText(type="output_text", text=generated_text, annotations=[])],
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _billed_terminal_response(
|
||||
response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None
|
||||
) -> ResponsesAPIResponse | None:
|
||||
|
|
@ -1453,7 +1519,7 @@ def _estimate_usage_safely(
|
|||
return None
|
||||
|
||||
|
||||
def _stamp_responses_usage_cost(
|
||||
def stamp_responses_usage_cost(
|
||||
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
|
||||
) -> None:
|
||||
if response_obj is None or logging_obj is None:
|
||||
|
|
@ -1479,7 +1545,7 @@ def build_synthetic_response_events(
|
|||
chunk_size: int,
|
||||
) -> list[ResponsesAPIStreamingResponse]:
|
||||
openai_types: Final = _get_openai_response_types()
|
||||
_stamp_responses_usage_cost(transformed, logging_obj)
|
||||
stamp_responses_usage_cost(transformed, logging_obj)
|
||||
|
||||
events: Final[list[ResponsesAPIStreamingResponse]] = [
|
||||
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed),
|
||||
|
|
|
|||
|
|
@ -132,6 +132,7 @@ from litellm.router_strategy.tag_based_routing import (
|
|||
is_valid_deployment_tag,
|
||||
)
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
HiddenParamsAsyncIteratorWrapper,
|
||||
_HiddenParamsHost,
|
||||
add_fallback_headers_to_response,
|
||||
add_retry_headers_to_response,
|
||||
|
|
@ -335,6 +336,7 @@ if TYPE_CHECKING:
|
|||
from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseAPIUsage,
|
||||
ResponseIncompleteEvent,
|
||||
ResponseInputParam,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
|
|
@ -585,6 +587,11 @@ def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, bu
|
|||
return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS
|
||||
|
||||
|
||||
def _stream_behind_hidden_params_wrapper(stream: object) -> object:
|
||||
"""set_response_headers wraps a bare byte stream before the fallback wrapper sees it; bill the stream inside."""
|
||||
return stream._inner if isinstance(stream, HiddenParamsAsyncIteratorWrapper) else stream # pyright: ignore[reportPrivateUsage] # the wrapper exposes nothing else
|
||||
|
||||
|
||||
class FallbackAwareAnthropicMessagesStream:
|
||||
"""
|
||||
Bare async generators can't carry the `_hidden_params` attribute the
|
||||
|
|
@ -611,6 +618,13 @@ class FallbackAwareAnthropicMessagesStream:
|
|||
self._source_iterator = fallback_response
|
||||
self.fallback_headers_adopted = True
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
|
||||
return partial_responses_incomplete_event_on_disconnect(
|
||||
_stream_behind_hidden_params_wrapper(self._source_iterator)
|
||||
)
|
||||
|
||||
def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream":
|
||||
return self
|
||||
|
||||
|
|
@ -3203,6 +3217,7 @@ class Router:
|
|||
from datetime import datetime
|
||||
|
||||
self._async_generator = async_generator
|
||||
self._billing_source: object = source_iterator
|
||||
# Mirror every attribute BaseResponsesAPIStreamingIterator.__init__
|
||||
# would have set. The wrapper bypasses super().__init__ (it has no
|
||||
# httpx.Response of its own and no provider config to drive), so
|
||||
|
|
@ -3252,9 +3267,28 @@ class Router:
|
|||
def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]:
|
||||
prepared: Final = Router._prepare_fallback_hidden_params(fallback_response)
|
||||
self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} # mutable-ok: stream metadata
|
||||
self._billing_source = fallback_response
|
||||
self.fallback_headers_adopted = True
|
||||
return prepared
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
"""
|
||||
The event to bill on a client disconnect. After a mid-stream hop the primary
|
||||
stream's relayed tokens are folded into the fallback's event, the way
|
||||
_combine_responses_fallback_usage does for a natural end of stream.
|
||||
"""
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
|
||||
current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source)
|
||||
if self._billing_source is source_iterator:
|
||||
return current
|
||||
primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator)
|
||||
if current is None or primary is None:
|
||||
return current or primary
|
||||
if primary.response.usage is not None:
|
||||
Router._combine_responses_fallback_usage(current, primary.response.usage)
|
||||
return current
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
|
|
|
|||
|
|
@ -7416,6 +7416,118 @@ def _apply_stream_usage_tracking(
|
|||
)
|
||||
|
||||
|
||||
async def _start_partial_responses_stream(self, recorder: _RecordingSuccessLogger, *, relay_text: bool = True):
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
|
||||
mocked: Final = await litellm.aresponses(
|
||||
model="gpt-4o-mini",
|
||||
input="tell me a story",
|
||||
mock_response="The codename is AZURE-FALCON-42 and the story is long.",
|
||||
stream=True,
|
||||
api_key="test-key",
|
||||
)
|
||||
logging_obj: Final = mocked.logging_obj
|
||||
logging_obj.dynamic_async_success_callbacks = [recorder]
|
||||
stream: Final = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="gpt-4o-mini",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "gpt-4o-mini", "input": "tell me a story"},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
created: Final = {
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {
|
||||
"id": "resp_partial",
|
||||
"object": "response",
|
||||
"created_at": 1700000000.0,
|
||||
"model": "gpt-4o-mini",
|
||||
"status": "in_progress",
|
||||
"output": [],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"service_tier": "priority",
|
||||
},
|
||||
}
|
||||
delta: Final = {
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_partial",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "The codename is AZURE-FALCON-42",
|
||||
}
|
||||
for event in (created, delta) if relay_text else (created,):
|
||||
stream._process_chunk(json.dumps(event))
|
||||
return stream
|
||||
|
||||
async def _bill_and_wait(self, recorder: _RecordingSuccessLogger, streamed: object, logging_obj: object) -> bool:
|
||||
billed: Final = await _bill_partial_streamed_spend_on_disconnect({"litellm_logging_obj": logging_obj}, streamed)
|
||||
for _ in range(50):
|
||||
if recorder.success_events:
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
await asyncio.sleep(0.5)
|
||||
return billed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_bills_partial_native_responses_stream_with_counted_usage(self):
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
recorder: Final = _RecordingSuccessLogger()
|
||||
stream: Final = await self._start_partial_responses_stream(recorder)
|
||||
billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj)
|
||||
|
||||
assert billed is True
|
||||
assert len(recorder.success_events) == 1
|
||||
partial_response: Final = recorder.success_events[0]["response_obj"]
|
||||
assert isinstance(partial_response, ResponsesAPIResponse)
|
||||
assert partial_response.status == "incomplete"
|
||||
assert getattr(partial_response, "service_tier") == "priority"
|
||||
assert partial_response.usage.output_tokens == litellm.token_counter(
|
||||
model="gpt-4o-mini", text="The codename is AZURE-FALCON-42"
|
||||
)
|
||||
assert partial_response.usage.input_tokens > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_bills_partial_default_route_messages_stream_through_the_responses_adapter(self):
|
||||
from litellm.llms.anthropic.pass_through.responses_adapters.streaming_iterator import (
|
||||
AnthropicResponsesStreamWrapper,
|
||||
)
|
||||
from litellm.router import FallbackAwareAnthropicMessagesStream
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
async def _sse_frames() -> AsyncGenerator[bytes, None]:
|
||||
yield b"event: message_start\n\n"
|
||||
|
||||
recorder: Final = _RecordingSuccessLogger()
|
||||
stream: Final = await self._start_partial_responses_stream(recorder)
|
||||
wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=stream, model="gpt-4o-mini")
|
||||
streamed: Final = FallbackAwareAnthropicMessagesStream(_sse_frames(), wrapper.async_anthropic_sse_wrapper())
|
||||
billed: Final = await self._bill_and_wait(recorder, streamed, stream.logging_obj)
|
||||
|
||||
assert billed is True
|
||||
assert len(recorder.success_events) == 1
|
||||
partial_response: Final = recorder.success_events[0]["response_obj"]
|
||||
assert isinstance(partial_response, ResponsesAPIResponse)
|
||||
assert partial_response.usage.output_tokens == litellm.token_counter(
|
||||
model="gpt-4o-mini", text="The codename is AZURE-FALCON-42"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_before_any_responses_output_bills_nothing(self):
|
||||
recorder: Final = _RecordingSuccessLogger()
|
||||
stream: Final = await self._start_partial_responses_stream(recorder, relay_text=False)
|
||||
billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj)
|
||||
|
||||
assert billed is False
|
||||
assert recorder.success_events == []
|
||||
|
||||
|
||||
class TestApplyStreamUsageTracking:
|
||||
def test_default_injects_usage_and_marks_strip_for_chat_completions(self):
|
||||
data = {"stream": True, "model": "gpt-5.4-nano"}
|
||||
|
|
|
|||
|
|
@ -5183,3 +5183,40 @@ def test_convert_response_function_tool_call_keeps_the_index_it_is_given(index:
|
|||
)
|
||||
|
||||
assert tool_call["index"] == index
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chunks():
|
||||
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
|
||||
LiteLLMCompletionStreamingIterator,
|
||||
)
|
||||
|
||||
chat_stream: Final = await litellm.acompletion(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "tell me a story"}],
|
||||
mock_response="The codename is AZURE-FALCON-42 and the story is long.",
|
||||
stream=True,
|
||||
api_key="test-key",
|
||||
)
|
||||
iterator: Final = LiteLLMCompletionStreamingIterator(
|
||||
model="gpt-4o-mini",
|
||||
litellm_custom_stream_wrapper=chat_stream,
|
||||
request_input="tell me a story",
|
||||
responses_api_request={},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
assert iterator.partial_response_on_disconnect() is None, "nothing relayed yet, nothing to bill"
|
||||
events: Final = [await iterator.__anext__() for _ in range(6)]
|
||||
relayed_text: Final = "".join(
|
||||
getattr(event, "delta", "") for event in events if event.type == "response.output_text.delta"
|
||||
)
|
||||
assert relayed_text, "the bridge should have relayed some text by now"
|
||||
|
||||
incomplete_event: Final = iterator.partial_response_on_disconnect()
|
||||
|
||||
assert incomplete_event is not None
|
||||
assert incomplete_event.type == "response.incomplete"
|
||||
assert incomplete_event.response.status == "incomplete"
|
||||
assert incomplete_event.response.usage is not None
|
||||
assert incomplete_event.response.usage.output_tokens == litellm.token_counter(model="gpt-4o-mini", text=relayed_text)
|
||||
assert incomplete_event.response.usage.input_tokens > 0
|
||||
|
|
|
|||
|
|
@ -7,7 +7,10 @@ crashed with exit 139 whenever any CustomLogger was registered).
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -17,6 +20,7 @@ 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.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
|
|
@ -156,3 +160,99 @@ async def test_sync_callbacks_run_only_after_async_handler_completes(recording_e
|
|||
submit_times = recording_executor.submit_times_for(logging_obj)
|
||||
assert len(submit_times) == 1
|
||||
assert submit_times[0] >= recorder.async_hook_finished
|
||||
|
||||
|
||||
_CREATED_RESPONSE_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = {
|
||||
"id": "resp_partial",
|
||||
"object": "response",
|
||||
"created_at": 1700000000.0,
|
||||
"model": "gpt-5.4-nano",
|
||||
"status": "in_progress",
|
||||
"output": [],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"service_tier": "priority",
|
||||
}
|
||||
_CREATED_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = {
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": _CREATED_RESPONSE_WITH_PRIORITY_TIER,
|
||||
}
|
||||
|
||||
|
||||
def _text_delta(text: str, sequence_number: int) -> Mapping[str, object]:
|
||||
return {
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": sequence_number,
|
||||
"item_id": "msg_partial",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": text,
|
||||
}
|
||||
|
||||
|
||||
def _completed_after(text: str) -> Mapping[str, object]:
|
||||
return {
|
||||
"type": "response.completed",
|
||||
"sequence_number": 9,
|
||||
"response": {
|
||||
**_CREATED_RESPONSE_WITH_PRIORITY_TIER,
|
||||
"status": "completed",
|
||||
"usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _live_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator:
|
||||
return ResponsesAPIStreamingIterator(
|
||||
response=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"},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
|
||||
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)):
|
||||
iterator._process_chunk(json.dumps(event))
|
||||
|
||||
incomplete_event: Final = iterator.partial_response_on_disconnect()
|
||||
|
||||
assert incomplete_event is not None
|
||||
assert incomplete_event.type == "response.incomplete"
|
||||
partial: Final = incomplete_event.response
|
||||
assert partial.status == "incomplete"
|
||||
assert getattr(partial, "service_tier") == "priority"
|
||||
assert partial.output[0].content[0].text == "Once upon a time a fox slept."
|
||||
assert partial.usage.output_tokens == litellm.token_counter(
|
||||
model="gpt-5.4-nano", text="Once upon a time a fox slept."
|
||||
)
|
||||
assert partial.usage.input_tokens == litellm.token_counter(
|
||||
model="gpt-5.4-nano", messages=[{"role": "user", "content": "tell me a story"}]
|
||||
)
|
||||
assert partial.usage.total_tokens == partial.usage.input_tokens + partial.usage.output_tokens
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"events",
|
||||
[
|
||||
(),
|
||||
(_CREATED_WITH_PRIORITY_TIER,),
|
||||
(_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1), _completed_after("Once upon a time")),
|
||||
],
|
||||
ids=["nothing_relayed", "created_only", "already_completed"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_response_on_disconnect_is_none_when_there_is_nothing_to_bill_or_the_stream_billed_itself(
|
||||
events: tuple[Mapping[str, object], ...],
|
||||
):
|
||||
iterator: Final = _live_iterator(_make_logging_obj())
|
||||
for event in events:
|
||||
iterator._process_chunk(json.dumps(event))
|
||||
|
||||
assert iterator.partial_response_on_disconnect() is None
|
||||
await asyncio.sleep(0)
|
||||
|
|
|
|||
|
|
@ -436,27 +436,27 @@ def _responses_api_response_with_usage() -> ResponsesAPIResponse:
|
|||
)
|
||||
|
||||
|
||||
def test_stamp_responses_usage_cost_stamps_computed_cost():
|
||||
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
|
||||
def teststamp_responses_usage_cost_stamps_computed_cost():
|
||||
from litellm.responses.streaming_iterator import stamp_responses_usage_cost
|
||||
|
||||
response = _responses_api_response_with_usage()
|
||||
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
||||
logging_obj._response_cost_calculator.return_value = 0.000704
|
||||
|
||||
_stamp_responses_usage_cost(response, logging_obj)
|
||||
stamp_responses_usage_cost(response, logging_obj)
|
||||
|
||||
assert getattr(response.usage, "cost", None) == pytest.approx(0.000704)
|
||||
logging_obj._response_cost_calculator.assert_called_once_with(result=response)
|
||||
|
||||
|
||||
def test_stamp_responses_usage_cost_keeps_provider_reported_cost():
|
||||
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
|
||||
def teststamp_responses_usage_cost_keeps_provider_reported_cost():
|
||||
from litellm.responses.streaming_iterator import stamp_responses_usage_cost
|
||||
|
||||
response = _responses_api_response_with_usage()
|
||||
setattr(response.usage, "cost", 0.5)
|
||||
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
||||
|
||||
_stamp_responses_usage_cost(response, logging_obj)
|
||||
stamp_responses_usage_cost(response, logging_obj)
|
||||
|
||||
assert getattr(response.usage, "cost", None) == pytest.approx(0.5)
|
||||
logging_obj._response_cost_calculator.assert_not_called()
|
||||
|
|
@ -475,8 +475,8 @@ def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse:
|
|||
)
|
||||
|
||||
|
||||
def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
|
||||
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
|
||||
def teststamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
|
||||
from litellm.responses.streaming_iterator import stamp_responses_usage_cost
|
||||
response = _unvalidated_response_with_dict_usage(
|
||||
{
|
||||
"input_tokens": 29,
|
||||
|
|
@ -488,7 +488,7 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
|
|||
)
|
||||
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
||||
|
||||
_stamp_responses_usage_cost(response, logging_obj)
|
||||
stamp_responses_usage_cost(response, logging_obj)
|
||||
|
||||
assert isinstance(response.usage, ResponseAPIUsage)
|
||||
assert response.usage.cost == pytest.approx(3e-05)
|
||||
|
|
@ -496,27 +496,27 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
|
|||
logging_obj._response_cost_calculator.assert_not_called()
|
||||
|
||||
|
||||
def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost():
|
||||
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
|
||||
def teststamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost():
|
||||
from litellm.responses.streaming_iterator import stamp_responses_usage_cost
|
||||
response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149})
|
||||
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
||||
logging_obj._response_cost_calculator.return_value = 0.000704
|
||||
|
||||
_stamp_responses_usage_cost(response, logging_obj)
|
||||
stamp_responses_usage_cost(response, logging_obj)
|
||||
|
||||
assert isinstance(response.usage, ResponseAPIUsage)
|
||||
assert response.usage.cost == pytest.approx(0.000704)
|
||||
logging_obj._response_cost_calculator.assert_called_once_with(result=response)
|
||||
|
||||
|
||||
def test_stamp_responses_usage_cost_survives_calculator_failure():
|
||||
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
|
||||
def teststamp_responses_usage_cost_survives_calculator_failure():
|
||||
from litellm.responses.streaming_iterator import stamp_responses_usage_cost
|
||||
|
||||
response = _responses_api_response_with_usage()
|
||||
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
||||
logging_obj._response_cost_calculator.side_effect = RuntimeError("cost map unavailable")
|
||||
|
||||
_stamp_responses_usage_cost(response, logging_obj)
|
||||
stamp_responses_usage_cost(response, logging_obj)
|
||||
|
||||
assert getattr(response.usage, "cost", None) is None
|
||||
|
||||
|
|
|
|||
|
|
@ -12755,6 +12755,24 @@ def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream:
|
|||
return FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object())
|
||||
|
||||
|
||||
def test_anthropic_messages_fallback_wrapper_bills_the_stream_behind_the_hidden_params_wrapper():
|
||||
"""set_response_headers wraps the bare byte stream, the disconnect bill must still reach the stream inside it."""
|
||||
from litellm.router_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper
|
||||
from litellm.types.llms.openai import ResponseIncompleteEvent
|
||||
|
||||
event: Final = MagicMock(spec=ResponseIncompleteEvent)
|
||||
|
||||
class _BillableByteStream:
|
||||
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent:
|
||||
return event
|
||||
|
||||
wrapped: Final = FallbackAwareAnthropicMessagesStream(
|
||||
_anthropic_messages_empty_generator(), HiddenParamsAsyncIteratorWrapper(_BillableByteStream())
|
||||
)
|
||||
|
||||
assert wrapped.partial_response_on_disconnect() is event
|
||||
|
||||
|
||||
def _anthropic_messages_make_router() -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
|
|
@ -16919,3 +16937,68 @@ class TestMemberAutoRouterInference:
|
|||
monkeypatch.setitem(sys.modules, "fastapi", None)
|
||||
monkeypatch.delitem(sys.modules, "litellm.proxy.auth.auto_router_checks", raising=False)
|
||||
assert (await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}})).model == "restricted-model"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens():
|
||||
"""After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone."""
|
||||
import json
|
||||
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
|
||||
def _relayed_iterator(text: str) -> ResponsesAPIStreamingIterator:
|
||||
logging_obj: Final = MagicMock(spec=LiteLLMLogging)
|
||||
logging_obj.model_call_details = {"litellm_params": {}}
|
||||
logging_obj.completion_start_time = None
|
||||
iterator: Final = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="gpt-5.4",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "gpt-5.4", "input": "Hi"},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
created: Final = {
|
||||
"id": "resp_1",
|
||||
"object": "response",
|
||||
"created_at": 1700000000.0,
|
||||
"model": "gpt-5.4",
|
||||
"status": "in_progress",
|
||||
"output": [],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
}
|
||||
iterator._process_chunk(json.dumps({"type": "response.created", "sequence_number": 0, "response": created}))
|
||||
iterator._process_chunk(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": text,
|
||||
}
|
||||
)
|
||||
)
|
||||
return iterator
|
||||
|
||||
router: Final = Router(
|
||||
model_list=[{"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k"}}]
|
||||
)
|
||||
primary: Final = _relayed_iterator("Once upon a time there was a lighthouse keeper")
|
||||
fallback: Final = _relayed_iterator("who lit the lamp")
|
||||
wrapped: Final = await router._aresponses_streaming_iterator(
|
||||
response=primary,
|
||||
initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": AsyncMock()},
|
||||
)
|
||||
primary_tokens: Final = primary.partial_response_on_disconnect().response.usage.output_tokens
|
||||
fallback_tokens: Final = fallback.partial_response_on_disconnect().response.usage.output_tokens
|
||||
|
||||
assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens
|
||||
|
||||
wrapped.adopt_fallback_headers(fallback)
|
||||
|
||||
assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + fallback_tokens
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue