mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(responses): write one spend-log row per native WebSocket turn
Each response.create turn on a native Responses WebSocket connection now logs its own call as it finishes, with its own usage and the resp_ id the client saw as request_id. A failed turn is booked alone and no longer flips the other turns.
This commit is contained in:
mateo-berri
parent
627356d3a0
commit
63483e6a88
17 changed files with 649 additions and 329 deletions
|
|
@ -2892,6 +2892,7 @@ _RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "re
|
|||
|
||||
|
||||
class _ResponsesWsEventResponse(LiteLLMBaseModel):
|
||||
id: str | None = None
|
||||
usage: Mapping[str, object] | None = None
|
||||
service_tier: str | None = None
|
||||
|
||||
|
|
@ -2943,6 +2944,26 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
|
|||
)
|
||||
return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(list(collected_usage_objects))
|
||||
|
||||
@staticmethod
|
||||
def response_id_from_responses_ws_results(results: Sequence[Mapping[str, object]]) -> str | None:
|
||||
response_ids: Final = tuple(
|
||||
event.response.id
|
||||
for event in (_ResponsesWsEvent.model_validate(result) for result in results)
|
||||
if event.response is not None and event.response.id
|
||||
)
|
||||
return response_ids[-1] if response_ids else None
|
||||
|
||||
@staticmethod
|
||||
def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject:
|
||||
tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results)
|
||||
raw_events: Final = cast(OpenAIRealtimeStreamList, list(results)) # cast-ok: raw WS frames, unvalidated field
|
||||
return LiteLLMRealtimeStreamLoggingObject(
|
||||
id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results),
|
||||
usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results),
|
||||
results=raw_events,
|
||||
service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None,
|
||||
)
|
||||
|
||||
|
||||
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"
|
||||
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ from types import MappingProxyType, TracebackType
|
|||
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
|
||||
|
||||
from httpx import Response
|
||||
from pydantic import BaseModel, JsonValue
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm import _custom_logger_compatible_callbacks_literal
|
||||
|
|
@ -536,6 +536,10 @@ def _resolve_mantle_region_for_cost(
|
|||
return resolve_mantle_region(litellm_params or MappingProxyType({}))
|
||||
|
||||
|
||||
_PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"})
|
||||
_LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _provider_response_id(source: object) -> str | None:
|
||||
candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None)
|
||||
return candidate if isinstance(candidate, str) and candidate else None
|
||||
|
|
@ -760,6 +764,32 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
|
||||
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
|
||||
|
||||
def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging":
|
||||
connection_params: Final = _LITELLM_PARAMS.validate_python(self.litellm_params)
|
||||
litellm_params: Final[dict[str, object]] = {
|
||||
**connection_params,
|
||||
**{
|
||||
key: {**value}
|
||||
for key, value in connection_params.items()
|
||||
if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict)
|
||||
},
|
||||
"litellm_call_id": litellm_call_id,
|
||||
}
|
||||
new_call: Final = copy.copy(self)
|
||||
new_call.litellm_call_id = litellm_call_id
|
||||
new_call.start_time = start_time
|
||||
new_call.litellm_params = litellm_params
|
||||
new_call.model_call_details = {
|
||||
**self.model_call_details,
|
||||
"litellm_call_id": litellm_call_id,
|
||||
"start_time": start_time,
|
||||
"first_api_call_start_time": None,
|
||||
"litellm_params": litellm_params,
|
||||
}
|
||||
new_call.streaming_chunks = []
|
||||
new_call.sync_streaming_chunks = []
|
||||
return new_call
|
||||
|
||||
def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
|
||||
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
|
||||
self.response_timing_metrics = dict(timing_metrics)
|
||||
|
|
@ -2362,20 +2392,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
|
||||
elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped
|
||||
combined_ws_usage: Final = (
|
||||
ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
|
||||
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
|
||||
)
|
||||
)
|
||||
ws_tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(
|
||||
logging_result = ResponsesWebSocketTokenUsageProcessor.create_logging_object(
|
||||
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
|
||||
)
|
||||
ws_service_tier: Final = next(iter(ws_tier_partition)) if len(ws_tier_partition) == 1 else None
|
||||
logging_result = LiteLLMRealtimeStreamLoggingObject(
|
||||
usage=combined_ws_usage,
|
||||
results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
|
||||
service_tier=ws_service_tier,
|
||||
)
|
||||
|
||||
elif (
|
||||
self.call_type == CallTypes.llm_passthrough_route.value
|
||||
|
|
@ -6628,7 +6647,7 @@ def get_standard_logging_object_payload(
|
|||
else raw_usage_dict
|
||||
)
|
||||
|
||||
id = response_obj.get("id", kwargs.get("litellm_call_id"))
|
||||
id = _provider_response_id(response_obj) or kwargs.get("litellm_call_id")
|
||||
|
||||
_model_id: Final = metadata.get("model_info", {}).get("id", "")
|
||||
_model_group: Final = metadata.get("model_group", "")
|
||||
|
|
|
|||
|
|
@ -155,7 +155,10 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.realtime import RealtimeQueryParams
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.responses.main import DeleteResponseResult
|
||||
from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults
|
||||
from litellm.types.responses.streaming_websocket import (
|
||||
ResponsesWebSocketRequestDefaults,
|
||||
ResponsesWebSocketTurnFailureHook,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
|
|
@ -6745,8 +6748,9 @@ class BaseLLMHTTPHandler:
|
|||
custom_llm_provider: str | None = None,
|
||||
first_message: str | None = None,
|
||||
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
|
||||
on_turn_failure: ResponsesWebSocketTurnFailureHook | None = None,
|
||||
**kwargs: object,
|
||||
) -> Exception | None:
|
||||
) -> None:
|
||||
"""
|
||||
Handles Responses API WebSocket mode.
|
||||
|
||||
|
|
@ -6780,7 +6784,7 @@ class BaseLLMHTTPHandler:
|
|||
**kwargs,
|
||||
)
|
||||
await handler.run()
|
||||
return None
|
||||
return
|
||||
|
||||
import websockets
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
|
|
@ -6901,8 +6905,9 @@ class BaseLLMHTTPHandler:
|
|||
authorized_model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
request_defaults=request_defaults,
|
||||
on_turn_failure=on_turn_failure,
|
||||
)
|
||||
return await streaming.bidirectional_forward()
|
||||
await streaming.bidirectional_forward()
|
||||
|
||||
except websockets.exceptions.InvalidStatus as e:
|
||||
verbose_logger.exception("Error connecting to responses WS backend: %s", e)
|
||||
|
|
@ -6916,7 +6921,6 @@ class BaseLLMHTTPHandler:
|
|||
pass
|
||||
else:
|
||||
raise Exception(f"Unexpected error while closing WebSocket: {close_error}")
|
||||
return None
|
||||
|
||||
def image_edit_handler(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -5377,7 +5377,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
that usage instead of refunding it.
|
||||
"""
|
||||
try:
|
||||
stash: Final = get_request_stash()
|
||||
failed_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(request_data)
|
||||
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(failed_request))
|
||||
if stash is None:
|
||||
return
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
|
|
|
|||
|
|
@ -469,11 +469,7 @@ class _ProxyDBLogger(CustomLogger):
|
|||
await _release_budget_reservation(budget_reservation=budget_reservation)
|
||||
# Non-model call types (health checks, afile_delete) have no model or standard_logging_object.
|
||||
# Use .get() for "stream" to avoid KeyError on health checks.
|
||||
# WS session wrappers (_aresponses_websocket, _arealtime) also reach here with
|
||||
# result=None; their per-turn costs are tracked on the inner aresponses/realtime calls.
|
||||
if sl_object is None and (
|
||||
not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime")
|
||||
):
|
||||
if sl_object is None and (not kwargs.get("model") or kwargs.get("call_type") == "_arealtime"):
|
||||
verbose_proxy_logger.warning(
|
||||
"Cost tracking - skipping, no standard_logging_object for call_type=%s",
|
||||
kwargs.get("call_type", "unknown"),
|
||||
|
|
|
|||
|
|
@ -44,9 +44,12 @@ from litellm.types.llms.openai import (
|
|||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.responses.main import DeleteResponseResult
|
||||
from litellm.types.responses.streaming_websocket import ResponsesWebSocketTurnFailureHook
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
|
@ -1436,6 +1439,40 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec
|
|||
return MappingProxyType({key: value for key, value in hints.items() if value is not None})
|
||||
|
||||
|
||||
_RESPONSES_WS_FIRST_FRAME_KEYS: Final = frozenset({"first_message", "input", "previous_response_id"})
|
||||
_RESPONSES_WS_PER_TURN_DICT_KEYS: Final = frozenset({"metadata", "litellm_metadata"})
|
||||
|
||||
|
||||
def _responses_ws_turn_request_data(
|
||||
connection_request_data: Mapping[str, object], turn_logging_obj: "LiteLLMLoggingObj"
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
**{key: value for key, value in connection_request_data.items() if key not in _RESPONSES_WS_FIRST_FRAME_KEYS},
|
||||
**{
|
||||
key: {**value}
|
||||
for key, value in connection_request_data.items()
|
||||
if key in _RESPONSES_WS_PER_TURN_DICT_KEYS and isinstance(value, dict)
|
||||
},
|
||||
"litellm_call_id": turn_logging_obj.litellm_call_id,
|
||||
"litellm_logging_obj": turn_logging_obj,
|
||||
}
|
||||
|
||||
|
||||
def _responses_ws_turn_failure_hook(
|
||||
proxy_logging_obj: "ProxyLogging",
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
connection_request_data: Mapping[str, object],
|
||||
) -> ResponsesWebSocketTurnFailureHook:
|
||||
async def book_turn_failure(turn_logging_obj: "LiteLLMLoggingObj", exception: Exception) -> None:
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=exception,
|
||||
request_data=_responses_ws_turn_request_data(connection_request_data, turn_logging_obj),
|
||||
)
|
||||
|
||||
return book_turn_failure
|
||||
|
||||
|
||||
def _responses_ws_failure_frame(failure: Exception) -> str:
|
||||
raw_status: Final = getattr(failure, "status_code", None)
|
||||
status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500
|
||||
|
|
@ -1584,7 +1621,10 @@ async def _responses_websocket_session(
|
|||
return
|
||||
|
||||
routed_data: Final = dict(
|
||||
data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message)
|
||||
data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
responses_ws_turn_failure_hook=_responses_ws_turn_failure_hook(proxy_logging_obj, user_api_key_dict, data),
|
||||
**_routing_hints_from_first_ws_frame(first_message),
|
||||
)
|
||||
# Phase 2: route to upstream provider
|
||||
try:
|
||||
|
|
@ -1594,13 +1634,7 @@ async def _responses_websocket_session(
|
|||
llm_router=llm_router,
|
||||
user_model=user_model,
|
||||
)
|
||||
failure: Final = await llm_call
|
||||
if isinstance(failure, Exception):
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=failure,
|
||||
request_data=routed_data,
|
||||
)
|
||||
await llm_call
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Responses WebSocket error")
|
||||
with contextlib.suppress(Exception):
|
||||
|
|
@ -1611,6 +1645,8 @@ async def _responses_websocket_session(
|
|||
request_data=routed_data,
|
||||
)
|
||||
await websocket.close(code=1011, reason="Internal server error")
|
||||
finally:
|
||||
await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict)
|
||||
|
||||
|
||||
@router.websocket("/v1/responses")
|
||||
|
|
|
|||
|
|
@ -59,7 +59,10 @@ from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
|||
from litellm.llms.openai.data_residency import infer_openai_data_residency
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.responses.main import *
|
||||
from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults
|
||||
from litellm.types.responses.streaming_websocket import (
|
||||
ResponsesWebSocketRequestDefaults,
|
||||
ResponsesWebSocketTurnFailureHook,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import all_litellm_params
|
||||
from litellm.utils import (
|
||||
|
|
@ -2422,12 +2425,13 @@ async def _aresponses_websocket(
|
|||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
timeout: float | None = None,
|
||||
responses_ws_turn_failure_hook: ResponsesWebSocketTurnFailureHook | None = None,
|
||||
**kwargs,
|
||||
) -> Exception | None:
|
||||
) -> None:
|
||||
"""
|
||||
Private function to handle the Responses API WebSocket mode.
|
||||
|
||||
For PROXY use only. Returns the provider failure that ended the connection, if any.
|
||||
For PROXY use only. ``responses_ws_turn_failure_hook`` books each failed turn of a native socket.
|
||||
|
||||
Resolves the LLM provider from ``model``, looks up the matching
|
||||
``BaseResponsesAPIConfig``, and hands off to
|
||||
|
|
@ -2500,7 +2504,7 @@ async def _aresponses_websocket(
|
|||
deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS}
|
||||
first_message: Final = kwargs.get("first_message")
|
||||
|
||||
return await base_llm_http_handler.async_responses_websocket(
|
||||
await base_llm_http_handler.async_responses_websocket(
|
||||
model=resolved_model,
|
||||
websocket=websocket,
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -2517,5 +2521,6 @@ async def _aresponses_websocket(
|
|||
litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
|
||||
custom_llm_provider=_custom_llm_provider,
|
||||
request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs),
|
||||
on_turn_failure=responses_ws_turn_failure_hook,
|
||||
**remaining_kwargs,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,8 +7,10 @@ import time
|
|||
import traceback
|
||||
import uuid
|
||||
from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, cast, overload, runtime_checkable
|
||||
|
||||
|
|
@ -56,6 +58,7 @@ if TYPE_CHECKING:
|
|||
ResponsesBackendWebSocket,
|
||||
ResponsesClientWebSocket,
|
||||
ResponsesWebSocketRequestDefaults,
|
||||
ResponsesWebSocketTurnFailureHook,
|
||||
)
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
|
||||
|
|
@ -1831,9 +1834,59 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output
|
|||
|
||||
_RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"})
|
||||
|
||||
_RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES
|
||||
|
||||
_RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
|
||||
|
||||
|
||||
def _json_object_or_none(message: object) -> Mapping[str, object] | None:
|
||||
if _is_json_object(message):
|
||||
return message
|
||||
if not isinstance(message, (str, bytes)):
|
||||
return None
|
||||
try:
|
||||
parsed: Final = _load_json_value(message)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
return parsed if _is_json_object(parsed) else None
|
||||
|
||||
|
||||
def _user_texts_in_input_item(item: object) -> tuple[str, ...]:
|
||||
if not _is_json_object(item) or item.get("type") != "message" or item.get("role") != "user":
|
||||
return ()
|
||||
content: Final = item.get("content", [])
|
||||
if isinstance(content, str):
|
||||
return (content,)
|
||||
return tuple(
|
||||
text
|
||||
for block in _json_array_or_empty(content)
|
||||
if _is_json_object(block)
|
||||
and block.get("type") == "input_text"
|
||||
and isinstance(text := block.get("text"), str)
|
||||
and text
|
||||
)
|
||||
|
||||
|
||||
def _logged_user_messages(client_frame: object) -> tuple[dict[str, object], ...]:
|
||||
frame: Final = _json_object_or_none(client_frame)
|
||||
if frame is None or frame.get("type") != "response.create":
|
||||
return ()
|
||||
input_items: Final = frame.get("input", [])
|
||||
if isinstance(input_items, str):
|
||||
return ({"role": "user", "content": input_items},)
|
||||
texts: Final = chain.from_iterable(_user_texts_in_input_item(item) for item in _json_array_or_empty(input_items))
|
||||
return tuple({"role": "user", "content": text} for text in texts)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ResponsesWsTurn:
|
||||
logging_obj: LiteLLMLoggingObj
|
||||
events: tuple[_MutableJsonObject, ...] = ()
|
||||
|
||||
def with_event(self, event: _MutableJsonObject) -> _ResponsesWsTurn:
|
||||
return _ResponsesWsTurn(logging_obj=self.logging_obj, events=(*self.events, event))
|
||||
|
||||
|
||||
def _ws_event_error(event: Mapping[str, object]) -> object:
|
||||
if event.get("type") == "error":
|
||||
return event.get("error")
|
||||
|
|
@ -1916,6 +1969,7 @@ class ResponsesWebSocketStreaming:
|
|||
authorized_model: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
|
||||
on_turn_failure: ResponsesWebSocketTurnFailureHook | None = None,
|
||||
):
|
||||
self.websocket = websocket
|
||||
self.backend_ws = backend_ws
|
||||
|
|
@ -1925,8 +1979,8 @@ class ResponsesWebSocketStreaming:
|
|||
litellm_metadata: Final = self.request_data.get("litellm_metadata")
|
||||
self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {}
|
||||
self.custom_llm_provider: str | None = custom_llm_provider
|
||||
self.messages: list[_MutableJsonObject] = []
|
||||
self.input_messages: list[dict[str, object]] = []
|
||||
self.on_turn_failure: ResponsesWebSocketTurnFailureHook | None = on_turn_failure
|
||||
self._open_turns: tuple[_ResponsesWsTurn, ...] = ()
|
||||
self.first_message = first_message
|
||||
self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or []
|
||||
self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or []
|
||||
|
|
@ -1936,104 +1990,70 @@ class ResponsesWebSocketStreaming:
|
|||
self.authorized_model: str | None = authorized_model
|
||||
self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults
|
||||
|
||||
def _should_store_event(self, event_obj: _MutableJsonObject) -> bool:
|
||||
return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES
|
||||
def _start_turn(self, client_frame: str | None) -> None:
|
||||
turn_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes
|
||||
turn_logging_obj: Final = self.logging_obj.copy_for_new_call(
|
||||
litellm_call_id=str(uuid.uuid4()), start_time=turn_start
|
||||
)
|
||||
turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame))
|
||||
turn_logging_obj.pre_call(input=client_frame, api_key="")
|
||||
self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj))
|
||||
|
||||
def _store_event(self, event: str | bytes | dict[str, object]) -> None:
|
||||
if isinstance(event, bytes):
|
||||
event = event.decode("utf-8")
|
||||
if isinstance(event, str):
|
||||
try:
|
||||
event_obj = _load_json_object(event)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return
|
||||
else:
|
||||
event_obj = event
|
||||
def _record_client_frame(self, client_frame: str) -> None:
|
||||
frame: Final = _json_object_or_none(client_frame)
|
||||
if frame is not None and frame.get("type") == "response.create":
|
||||
self._start_turn(client_frame)
|
||||
|
||||
if self._should_store_event(event_obj):
|
||||
self.messages.append(event_obj)
|
||||
def _record_backend_event(self, event: str) -> None:
|
||||
event_obj: Final = _json_object_or_none(event)
|
||||
if event_obj is None or event_obj.get("type") not in RESPONSES_WS_LOGGED_EVENT_TYPES:
|
||||
return
|
||||
if not self._open_turns:
|
||||
self._start_turn(None)
|
||||
turn: Final = self._open_turns[0].with_event(event_obj)
|
||||
if event_obj.get("type") in _RESPONSES_WS_TERMINAL_EVENT_TYPES:
|
||||
self._open_turns = self._open_turns[1:]
|
||||
self._log_turn(turn)
|
||||
return
|
||||
self._open_turns = (turn, *self._open_turns[1:])
|
||||
|
||||
def _collect_input_from_client_event(self, message: object) -> None:
|
||||
"""Extract user input content from response.create for logging."""
|
||||
try:
|
||||
if isinstance(message, str):
|
||||
msg_obj = _load_json_object(message)
|
||||
elif _is_json_object(message):
|
||||
msg_obj = message
|
||||
else:
|
||||
return
|
||||
def _log_open_turns(self) -> None:
|
||||
open_turns: Final = self._open_turns
|
||||
self._open_turns = ()
|
||||
for turn in open_turns:
|
||||
if turn.events:
|
||||
self._log_turn(turn)
|
||||
|
||||
if msg_obj.get("type") != "response.create":
|
||||
return
|
||||
def _log_turn(self, turn: _ResponsesWsTurn) -> None:
|
||||
_spawn_logging_task(
|
||||
asyncio.get_running_loop(), self._dispatch_turn(turn), task_name="Responses WS turn logging"
|
||||
)
|
||||
|
||||
input_items: Final = msg_obj.get("input", [])
|
||||
if isinstance(input_items, str):
|
||||
self.input_messages.append({"role": "user", "content": input_items})
|
||||
return
|
||||
|
||||
if _is_json_array(input_items):
|
||||
for item in input_items:
|
||||
if not _is_json_object(item):
|
||||
continue
|
||||
if item.get("type") == "message" and item.get("role") == "user":
|
||||
content = item.get("content", [])
|
||||
if isinstance(content, str):
|
||||
self.input_messages.append({"role": "user", "content": content})
|
||||
elif _is_json_array(content):
|
||||
for c in content:
|
||||
if _is_json_object(c) and c.get("type") == "input_text":
|
||||
text = c.get("text", "")
|
||||
if text:
|
||||
self.input_messages.append({"role": "user", "content": text})
|
||||
except (json.JSONDecodeError, AttributeError, TypeError):
|
||||
pass
|
||||
|
||||
def _store_input(self, message: object) -> None:
|
||||
self._collect_input_from_client_event(message)
|
||||
if self.logging_obj:
|
||||
self.logging_obj.pre_call(input=message, api_key="")
|
||||
|
||||
def _failure_exception(self) -> Exception | None:
|
||||
async def _dispatch_turn(self, turn: _ResponsesWsTurn) -> None:
|
||||
events: Final = list(turn.events)
|
||||
failed_event: Final = next(
|
||||
(event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None
|
||||
(event for event in events if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None
|
||||
)
|
||||
if failed_event is None:
|
||||
return None
|
||||
return _map_stream_error_to_exception(
|
||||
await turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True)
|
||||
return
|
||||
exception: Final = _map_stream_error_to_exception(
|
||||
_ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or ""
|
||||
)
|
||||
|
||||
async def _log_messages(self) -> None:
|
||||
if not self.logging_obj:
|
||||
return
|
||||
if self.input_messages:
|
||||
self.logging_obj.model_call_details["messages"] = self.input_messages
|
||||
if not self.messages:
|
||||
return
|
||||
exception: Final = self._failure_exception()
|
||||
if exception is None:
|
||||
asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True))
|
||||
return
|
||||
self._record_usage_for_failure()
|
||||
traceback_exception: Final = "".join(traceback.format_exception(exception))
|
||||
asyncio.create_task(
|
||||
self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True)
|
||||
self._record_usage_for_failure(turn.logging_obj, events)
|
||||
await turn.logging_obj.dispatch_failure_handlers(
|
||||
exception, "".join(traceback.format_exception(exception)), prefer_async_handlers=True
|
||||
)
|
||||
if self.on_turn_failure is not None:
|
||||
await self.on_turn_failure(turn.logging_obj, exception)
|
||||
|
||||
def _record_usage_for_failure(self) -> None:
|
||||
@staticmethod
|
||||
def _record_usage_for_failure(turn_logging_obj: LiteLLMLoggingObj, events: Sequence[_MutableJsonObject]) -> None:
|
||||
from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor
|
||||
from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject
|
||||
|
||||
usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
|
||||
self.messages
|
||||
)
|
||||
tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(self.messages)
|
||||
service_tier: Final = next(iter(tier_partition)) if len(tier_partition) == 1 else None
|
||||
logging_result: Final = LiteLLMRealtimeStreamLoggingObject(
|
||||
usage=usage, results=self.messages, service_tier=service_tier
|
||||
)
|
||||
response_cost: Final = self.logging_obj.response_cost_calculator(result=logging_result) or 0.0
|
||||
self.logging_obj.record_partial_usage_for_failure(usage, response_cost)
|
||||
logging_result: Final = ResponsesWebSocketTokenUsageProcessor.create_logging_object(events)
|
||||
response_cost: Final = turn_logging_obj.response_cost_calculator(result=logging_result) or 0.0
|
||||
turn_logging_obj.record_partial_usage_for_failure(logging_result.usage, response_cost)
|
||||
|
||||
def _wrap_response_event(self, response_str: str) -> str:
|
||||
try:
|
||||
|
|
@ -2092,7 +2112,7 @@ class ResponsesWebSocketStreaming:
|
|||
|
||||
# Log the output-masked form so PII redacted by apply_to_output
|
||||
# guardrails does not appear in success logs.
|
||||
self._store_event(wrapped_str)
|
||||
self._record_backend_event(wrapped_str)
|
||||
|
||||
await self.websocket.send_text(wrapped_str)
|
||||
|
||||
|
|
@ -2101,7 +2121,7 @@ class ResponsesWebSocketStreaming:
|
|||
except Exception as e:
|
||||
verbose_logger.exception("Error in responses WS backend_to_client: %s", e)
|
||||
finally:
|
||||
await self._log_messages()
|
||||
self._log_open_turns()
|
||||
|
||||
def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool:
|
||||
"""
|
||||
|
|
@ -2446,8 +2466,7 @@ class ResponsesWebSocketStreaming:
|
|||
try:
|
||||
if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message):
|
||||
masked_first: Final = await self._mask_response_create(self.first_message)
|
||||
self._store_input(masked_first)
|
||||
self._store_event(masked_first)
|
||||
self._record_client_frame(masked_first)
|
||||
await self.backend_ws.send(masked_first)
|
||||
|
||||
while True:
|
||||
|
|
@ -2455,14 +2474,13 @@ class ResponsesWebSocketStreaming:
|
|||
if not await self._enforce_or_reject_frame(message):
|
||||
continue
|
||||
masked = await self._mask_response_create(message)
|
||||
self._store_input(masked)
|
||||
self._store_event(masked)
|
||||
self._record_client_frame(masked)
|
||||
await self.backend_ws.send(masked)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
|
||||
|
||||
async def bidirectional_forward(self) -> Exception | None:
|
||||
async def bidirectional_forward(self) -> None:
|
||||
forward_task: Final = asyncio.create_task(self.backend_to_client())
|
||||
try:
|
||||
await self.client_to_backend()
|
||||
|
|
@ -2479,7 +2497,6 @@ class ResponsesWebSocketStreaming:
|
|||
await self.backend_ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
return self._failure_exception()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -8406,8 +8406,7 @@ class Router:
|
|||
Track remaining tpm/rpm quota for model in model_list
|
||||
"""
|
||||
try:
|
||||
# WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
|
||||
if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
|
||||
if kwargs.get("call_type") == "_arealtime":
|
||||
return
|
||||
if is_batch_retrieve_call_type(kwargs.get("call_type")):
|
||||
return
|
||||
|
|
|
|||
|
|
@ -495,8 +495,7 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
"""Original method now uses helper functions"""
|
||||
verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event")
|
||||
# WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
|
||||
if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
|
||||
if kwargs.get("call_type") == "_arealtime":
|
||||
return
|
||||
standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
|
||||
if standard_logging_payload is None:
|
||||
|
|
|
|||
|
|
@ -2,10 +2,13 @@ from __future__ import annotations
|
|||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol
|
||||
from typing import TYPE_CHECKING, Protocol
|
||||
|
||||
from litellm.types.guardrails import PresidioPerRequestConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
|
||||
class ResponsesClientWebSocket(Protocol):
|
||||
"""Client-facing websocket surface used by the Responses API websocket handlers."""
|
||||
|
|
@ -43,6 +46,12 @@ class PresidioGuardrailCallback(Protocol):
|
|||
) -> str: ...
|
||||
|
||||
|
||||
class ResponsesWebSocketTurnFailureHook(Protocol):
|
||||
"""Books one failed ``response.create`` turn of a native Responses API socket from that turn's logging object."""
|
||||
|
||||
async def __call__(self, turn_logging_obj: Logging, exception: Exception) -> None: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResponsesWebSocketRequestDefaults:
|
||||
"""Deployment-level request parameters merged into every ``response.create`` frame relayed over a native websocket."""
|
||||
|
|
|
|||
|
|
@ -4478,6 +4478,7 @@ class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase):
|
|||
results: SkipValidation[OpenAIRealtimeStreamList]
|
||||
usage: Usage
|
||||
service_tier: str | None = None
|
||||
id: str | None = None
|
||||
_hidden_params: dict = {}
|
||||
|
||||
@property
|
||||
|
|
|
|||
|
|
@ -1998,34 +1998,34 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(
|
||||
async def test_aresponses_websocket_writes_no_connection_level_success_log(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.responses.main import base_llm_http_handler
|
||||
|
||||
success_events = []
|
||||
logged_call_ids = []
|
||||
|
||||
class CaptureLogger(CustomLogger):
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
success_events.append(response_obj)
|
||||
logged_call_ids.append(kwargs.get("litellm_call_id"))
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()])
|
||||
monkeypatch.setattr(litellm, "failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "success_callback", [])
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [])
|
||||
failure = litellm.BadRequestError(message="invalid_encrypted_content", model="gpt-4o", llm_provider="openai")
|
||||
with patch.object( # test-quality-ok: the provider socket is the seam; how the wrapper treats the relay's outcome is under test
|
||||
base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=failure)
|
||||
with patch.object( # test-quality-ok: the provider socket is the seam; each turn logs inside it, so the wrapper must not
|
||||
base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=None)
|
||||
):
|
||||
outcome = await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test")
|
||||
await litellm._aresponses_websocket(
|
||||
model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test", litellm_call_id="socket-call"
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
with contextlib.suppress(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
|
||||
|
||||
assert outcome is failure
|
||||
assert success_events == []
|
||||
assert "socket-call" not in logged_call_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8078,6 +8078,74 @@ def test_normalize_logging_result_prices_responses_websocket_at_returned_service
|
|||
assert priority_http_cost > default_http_cost
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("completed_response_id", "expected_request_id"),
|
||||
(("resp_wrapped_turn_2", "resp_wrapped_turn_2"), (None, "turn-2-call")),
|
||||
)
|
||||
def test_responses_websocket_turn_is_logged_under_the_response_id_the_client_saw(
|
||||
completed_response_id: str | None, expected_request_id: str
|
||||
):
|
||||
events = [
|
||||
{"type": "response.created", "response": {"id": completed_response_id}},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": completed_response_id,
|
||||
"usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140},
|
||||
},
|
||||
},
|
||||
]
|
||||
logging_obj = _responses_ws_logging_obj()
|
||||
now = datetime.datetime.now()
|
||||
|
||||
payload = litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload(
|
||||
kwargs={"litellm_call_id": "turn-2-call", "model": "gpt-4o", "messages": []},
|
||||
init_response_obj=logging_obj.normalize_logging_result(result=events),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
logging_obj=logging_obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["id"] == expected_request_id
|
||||
|
||||
|
||||
def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata():
|
||||
reservation = {"reserved_cost": 0.5, "finalized": False}
|
||||
connection = _responses_ws_logging_obj()
|
||||
connection_start = connection.start_time
|
||||
connection.litellm_params = {
|
||||
"metadata": {"user_api_key_alias": "codex", "budget_reservation": reservation},
|
||||
"litellm_metadata": {"model_info": {"id": "dep-1"}},
|
||||
"litellm_call_id": connection.litellm_call_id,
|
||||
}
|
||||
connection.model_call_details = {
|
||||
**connection.model_call_details,
|
||||
"litellm_params": connection.litellm_params,
|
||||
"first_api_call_start_time": connection_start,
|
||||
}
|
||||
turn_start = datetime.datetime(2026, 10, 9, 17, 24, 47)
|
||||
|
||||
turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start)
|
||||
turn.litellm_params["metadata"]["status"] = "failure"
|
||||
turn.litellm_params["litellm_metadata"]["model_group"] = "gpt-4o"
|
||||
turn.model_call_details["standard_logging_object"] = {"id": "resp_1"}
|
||||
|
||||
assert (turn.litellm_call_id, turn.start_time) == ("turn-1-call", turn_start)
|
||||
assert turn.model_call_details["litellm_call_id"] == "turn-1-call"
|
||||
assert turn.model_call_details["start_time"] == turn_start
|
||||
assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1-call"
|
||||
assert turn.model_call_details["first_api_call_start_time"] is None
|
||||
assert turn.litellm_params["metadata"]["budget_reservation"] is reservation
|
||||
assert (connection.litellm_call_id, connection.start_time) == ("responses-ws-usage-test", connection_start)
|
||||
assert connection.litellm_params["litellm_call_id"] == "responses-ws-usage-test"
|
||||
assert "status" not in connection.litellm_params["metadata"]
|
||||
assert "model_group" not in connection.litellm_params["litellm_metadata"]
|
||||
assert "standard_logging_object" not in connection.model_call_details
|
||||
assert connection.model_call_details["first_api_call_start_time"] == connection_start
|
||||
|
||||
|
||||
def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj):
|
||||
"""LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead
|
||||
recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs)."""
|
||||
|
|
|
|||
|
|
@ -4294,6 +4294,40 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3():
|
||||
"""
|
||||
A native Responses WebSocket holds one parallel slot for the whole socket and books each failed turn
|
||||
under that turn's own call id. A failed turn must not free the slot while the socket is still open.
|
||||
"""
|
||||
_api_key = hash_token("sk-ws-turn-failure")
|
||||
local_cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
|
||||
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
|
||||
connection_data: Dict[str, Any] = {"model": "gpt-4o-mini", "litellm_call_id": "socket-call"}
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, cache=local_cache, data=connection_data, call_type=""
|
||||
)
|
||||
|
||||
await handler.async_post_call_failure_hook(
|
||||
request_data={"model": "gpt-4o-mini", "litellm_call_id": "turn-2-call"},
|
||||
original_exception=Exception("previous_response_not_found"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
in_flight_after_turn_failure = handler._gauge_in_flight_from_cache_value(
|
||||
await local_cache.async_get_cache(key=counter_key)
|
||||
)
|
||||
await handler.async_post_call_failure_hook(
|
||||
request_data=connection_data,
|
||||
original_exception=Exception("socket closed"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert in_flight_after_turn_failure == 1
|
||||
assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_event_releases_parallel_slot_v3(monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -542,29 +542,6 @@ class TestResponsesWSStreamingFirstMessage:
|
|||
|
||||
|
||||
class TestWSSessionCostTracking:
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_budget_limiter_skips_aresponses_websocket_call_type(self):
|
||||
"""
|
||||
RouterBudgetLimiting.async_log_success_event must not raise when
|
||||
call_type='_aresponses_websocket', even when standard_logging_object is None.
|
||||
Per-turn costs are tracked by individual aresponses calls inside the session;
|
||||
the outer session wrapper fires with result=None.
|
||||
"""
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
|
||||
limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting)
|
||||
kwargs = {
|
||||
"call_type": "_aresponses_websocket",
|
||||
"standard_logging_object": None,
|
||||
"litellm_params": {"custom_llm_provider": "vertex_ai"},
|
||||
}
|
||||
await limiter.async_log_success_event(
|
||||
kwargs=kwargs,
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_budget_limiter_skips_arealtime_call_type(self):
|
||||
"""Same guard applies to _arealtime WS session wrappers."""
|
||||
|
|
@ -820,8 +797,7 @@ class TestResponsesWSFirstFrameModelAuth:
|
|||
ws.close.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider_rejected", [True, False])
|
||||
async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool):
|
||||
async def test_endpoint_books_each_failed_turn_as_its_own_failed_request(self):
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
responses_websocket_endpoint,
|
||||
)
|
||||
|
|
@ -833,23 +809,36 @@ class TestResponsesWSFirstFrameModelAuth:
|
|||
ws.url = "ws://testserver/v1/responses"
|
||||
ws.accept = AsyncMock()
|
||||
ws.receive_text = AsyncMock(
|
||||
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []})
|
||||
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": "first turn"})
|
||||
)
|
||||
ws.close = AsyncMock()
|
||||
|
||||
connection_metadata = {"user_api_key_alias": "codex", "model_group": "gpt-4o-mini"}
|
||||
processor = MagicMock()
|
||||
processor.common_processing_pre_call_logic = AsyncMock(
|
||||
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
|
||||
return_value=(
|
||||
{
|
||||
"model": "gpt-4o-mini",
|
||||
"metadata": connection_metadata,
|
||||
"litellm_metadata": {},
|
||||
"litellm_call_id": "connection-call",
|
||||
},
|
||||
MagicMock(),
|
||||
)
|
||||
)
|
||||
turn_logging_obj = MagicMock()
|
||||
turn_logging_obj.litellm_call_id = "turn-2-call"
|
||||
failure = litellm.BadRequestError(
|
||||
message="invalid_encrypted_content", model="gpt-4o-mini", llm_provider="openai"
|
||||
message="previous_response_not_found", model="gpt-4o-mini", llm_provider="openai"
|
||||
)
|
||||
|
||||
async def fake_llm_call():
|
||||
return failure if provider_rejected else None
|
||||
async def relay_with_one_failed_turn():
|
||||
routed = mock_route_request.await_args.kwargs["data"]
|
||||
await routed["responses_ws_turn_failure_hook"](turn_logging_obj, failure)
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock()
|
||||
user_api_key_dict = MagicMock()
|
||||
|
||||
with (
|
||||
|
|
@ -857,15 +846,15 @@ class TestResponsesWSFirstFrameModelAuth:
|
|||
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint does with the relay's outcome is under test
|
||||
patch( # test-quality-ok: the pre-call processor needs a live proxy; how the endpoint books a failed turn is under test
|
||||
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
|
||||
return_value=processor,
|
||||
),
|
||||
patch( # test-quality-ok: routing is the seam that hands back the relay's outcome
|
||||
patch( # test-quality-ok: routing is the seam that hands the turn failure hook to the relay
|
||||
"litellm.proxy.route_llm_request.route_request",
|
||||
new_callable=AsyncMock,
|
||||
return_value=fake_llm_call(),
|
||||
),
|
||||
side_effect=lambda **_: relay_with_one_failed_turn(),
|
||||
) as mock_route_request,
|
||||
patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj",
|
||||
proxy_logging_obj,
|
||||
|
|
@ -878,14 +867,19 @@ class TestResponsesWSFirstFrameModelAuth:
|
|||
)
|
||||
|
||||
ws.close.assert_not_awaited()
|
||||
if not provider_rejected:
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_awaited()
|
||||
return
|
||||
proxy_logging_obj.post_call_failure_hook.assert_awaited_once()
|
||||
booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs
|
||||
assert booked["original_exception"] is failure
|
||||
assert booked["user_api_key_dict"] is user_api_key_dict
|
||||
assert booked["request_data"]["model"] == "gpt-4o-mini"
|
||||
turn_request = booked["request_data"]
|
||||
assert turn_request["litellm_call_id"] == "turn-2-call"
|
||||
assert turn_request["litellm_logging_obj"] is turn_logging_obj
|
||||
assert turn_request["model"] == "gpt-4o-mini"
|
||||
assert turn_request["metadata"] == connection_metadata
|
||||
assert turn_request["metadata"] is not connection_metadata
|
||||
assert "input" not in turn_request
|
||||
assert "first_message" not in turn_request
|
||||
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self):
|
||||
|
|
@ -914,6 +908,7 @@ class TestResponsesWSFirstFrameModelAuth:
|
|||
)
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None)
|
||||
user_api_key_dict = MagicMock()
|
||||
|
||||
with (
|
||||
|
|
@ -951,6 +946,7 @@ class TestResponsesWSFirstFrameModelAuth:
|
|||
assert booked["original_exception"] is rejection
|
||||
assert booked["user_api_key_dict"] is user_api_key_dict
|
||||
assert booked["request_data"]["model"] == "gpt-4o-mini"
|
||||
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reruns_model_auth_for_first_frame_model(self):
|
||||
|
|
|
|||
|
|
@ -1898,6 +1898,29 @@ def _make_streaming(**kwargs):
|
|||
return ResponsesWebSocketStreaming(**kwargs)
|
||||
|
||||
|
||||
class _TurnLoggingRecorder:
|
||||
"""Stands in for a socket's connection-level Logging object and keeps every per-turn copy it hands out."""
|
||||
|
||||
def __init__(self, response_cost: float = 0.0) -> None:
|
||||
self.response_cost = response_cost
|
||||
self.turns: list[MagicMock] = []
|
||||
self.connection = MagicMock()
|
||||
self.connection.copy_for_new_call = MagicMock(side_effect=self._copy_for_new_call)
|
||||
|
||||
def _copy_for_new_call(self, litellm_call_id: str, start_time: object) -> MagicMock:
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
turn = MagicMock()
|
||||
turn.litellm_call_id = litellm_call_id
|
||||
turn.start_time = start_time
|
||||
turn.model_call_details = {}
|
||||
turn.dispatch_success_handlers = AsyncMock()
|
||||
turn.dispatch_failure_handlers = AsyncMock()
|
||||
turn.response_cost_calculator = MagicMock(return_value=self.response_cost)
|
||||
self.turns.append(turn)
|
||||
return turn
|
||||
|
||||
|
||||
class TestNativeWebSocketGuardrailMasking:
|
||||
"""Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming."""
|
||||
|
||||
|
|
@ -3051,12 +3074,11 @@ class TestNativeWebSocketEncryptedContentAffinity:
|
|||
Exception("stop"),
|
||||
]
|
||||
)
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.dispatch_success_handlers = AsyncMock()
|
||||
recorder = _TurnLoggingRecorder()
|
||||
handler = _make_streaming(
|
||||
websocket=websocket,
|
||||
backend_ws=backend_ws,
|
||||
logging_obj=logging_obj,
|
||||
logging_obj=recorder.connection,
|
||||
request_data={"litellm_metadata": dict(_AFFINITY_METADATA)},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
|
@ -3075,7 +3097,7 @@ class TestNativeWebSocketEncryptedContentAffinity:
|
|||
)
|
||||
assert completed["response"]["output"][0]["encrypted_content"] == wrapped_content
|
||||
await asyncio.sleep(0)
|
||||
logged = logging_obj.dispatch_success_handlers.await_args[0][0]
|
||||
logged = recorder.turns[0].dispatch_success_handlers.await_args[0][0]
|
||||
assert logged[0]["response"]["id"] == completed["response"]["id"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -3163,14 +3185,11 @@ class TestNativeWebSocketEncryptedContentAffinity:
|
|||
Exception("stop"),
|
||||
]
|
||||
)
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.dispatch_success_handlers = AsyncMock()
|
||||
logging_obj.dispatch_failure_handlers = AsyncMock()
|
||||
logging_obj.response_cost_calculator = MagicMock(return_value=0.0)
|
||||
recorder = _TurnLoggingRecorder()
|
||||
handler = _make_streaming(
|
||||
websocket=websocket,
|
||||
backend_ws=backend_ws,
|
||||
logging_obj=logging_obj,
|
||||
logging_obj=recorder.connection,
|
||||
request_data={},
|
||||
authorized_model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
|
|
@ -3179,163 +3198,257 @@ class TestNativeWebSocketEncryptedContentAffinity:
|
|||
await handler.backend_to_client()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
logging_obj.dispatch_success_handlers.assert_not_awaited()
|
||||
logging_obj.dispatch_failure_handlers.assert_awaited_once()
|
||||
exception = logging_obj.dispatch_failure_handlers.await_args[0][0]
|
||||
(turn,) = recorder.turns
|
||||
turn.dispatch_success_handlers.assert_not_awaited()
|
||||
turn.dispatch_failure_handlers.assert_awaited_once()
|
||||
exception = turn.dispatch_failure_handlers.await_args[0][0]
|
||||
assert exception.status_code == expected_status
|
||||
assert failure_frame.get("error", failure_frame.get("response", {}).get("error"))["message"] in str(exception)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_backend_to_client_bills_completed_turns_before_a_failure(self):
|
||||
|
||||
def _completed_turn(response_id: str, input_tokens: int, output_tokens: int) -> list[dict[str, object]]:
|
||||
return [
|
||||
{"type": "response.created", "response": {"id": response_id, "status": "in_progress"}},
|
||||
{"type": "response.output_text.delta", "delta": "hi"},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": response_id,
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"usage": {
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"total_tokens": input_tokens + output_tokens,
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _response_create(text: str) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"type": "response.create",
|
||||
"model": "gpt-5.4-mini",
|
||||
"input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _ScriptedBackend:
|
||||
"""Answers each forwarded response.create with the next scripted turn of upstream events."""
|
||||
|
||||
def __init__(self, turns: list[list[dict[str, object]]]) -> None:
|
||||
import asyncio
|
||||
|
||||
self.turns = list(turns)
|
||||
self.events: asyncio.Queue[str] = asyncio.Queue()
|
||||
self.sent: list[str] = []
|
||||
|
||||
async def send(self, frame: str) -> None:
|
||||
self.sent.append(frame)
|
||||
if json.loads(frame).get("type") == "response.create" and self.turns:
|
||||
for event in self.turns.pop(0):
|
||||
self.events.put_nowait(json.dumps(event))
|
||||
|
||||
async def recv(self, decode: bool = False) -> str:
|
||||
return await self.events.get()
|
||||
|
||||
async def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
async def _until(condition) -> bool:
|
||||
import asyncio
|
||||
|
||||
for _ in range(200):
|
||||
if condition():
|
||||
return True
|
||||
await asyncio.sleep(0)
|
||||
return False
|
||||
|
||||
|
||||
class _ScriptedClient:
|
||||
"""Sends each frame only after the previous turn's log was dispatched, then hangs up."""
|
||||
|
||||
def __init__(self, frames: list[str], recorder: _TurnLoggingRecorder) -> None:
|
||||
self.frames = list(frames)
|
||||
self.recorder = recorder
|
||||
self.sent_text: list[str] = []
|
||||
self.logged_turns_before_each_frame: list[int] = []
|
||||
self.logged_turns_before_hang_up: int | None = None
|
||||
|
||||
def _logged_turns(self) -> int:
|
||||
return sum(
|
||||
1
|
||||
for turn in self.recorder.turns
|
||||
if turn.dispatch_success_handlers.await_count or turn.dispatch_failure_handlers.await_count
|
||||
)
|
||||
|
||||
async def receive_text(self) -> str:
|
||||
sent_so_far = len(self.logged_turns_before_each_frame)
|
||||
await _until(lambda: self._logged_turns() >= sent_so_far)
|
||||
if self.frames:
|
||||
self.logged_turns_before_each_frame.append(self._logged_turns())
|
||||
return self.frames.pop(0)
|
||||
self.logged_turns_before_hang_up = self._logged_turns()
|
||||
raise Exception("client hung up")
|
||||
|
||||
async def send_text(self, text: str) -> None:
|
||||
self.sent_text.append(text)
|
||||
|
||||
|
||||
class TestNativeWebSocketPerTurnLogging:
|
||||
"""One native Responses socket carries many response.create turns, and each turn is its own logged call."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_turn_is_logged_as_its_own_call_while_the_socket_stays_open(self):
|
||||
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
|
||||
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
recorder = _TurnLoggingRecorder()
|
||||
client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder)
|
||||
backend = _ScriptedBackend(
|
||||
[_completed_turn("resp_1", 10, 1), _completed_turn("resp_2", 20, 2), _completed_turn("resp_3", 30, 3)]
|
||||
)
|
||||
handler = _make_streaming(
|
||||
websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={}
|
||||
)
|
||||
|
||||
await handler.bidirectional_forward()
|
||||
|
||||
assert client.logged_turns_before_each_frame == [0, 1, 2]
|
||||
assert client.logged_turns_before_hang_up == 3
|
||||
assert len({turn.litellm_call_id for turn in recorder.turns}) == 3
|
||||
assert recorder.connection.litellm_call_id not in {turn.litellm_call_id for turn in recorder.turns}
|
||||
logged_events = [turn.dispatch_success_handlers.await_args[0][0] for turn in recorder.turns]
|
||||
sent_frames = [json.loads(text) for text in client.sent_text]
|
||||
ids_the_client_saw = [frame["response"]["id"] for frame in sent_frames if frame["type"] == "response.completed"]
|
||||
assert [events[-1]["response"]["id"] for events in logged_events] == ids_the_client_saw
|
||||
assert [
|
||||
ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id)["response_id"]
|
||||
for response_id in ids_the_client_saw
|
||||
] == ["resp_1", "resp_2", "resp_3"]
|
||||
assert [len(events) for events in logged_events] == [2, 2, 2]
|
||||
assert [events[-1]["response"]["usage"]["input_tokens"] for events in logged_events] == [10, 20, 30]
|
||||
assert [turn.model_call_details["messages"] for turn in recorder.turns] == [
|
||||
[{"role": "user", "content": "one"}],
|
||||
[{"role": "user", "content": "two"}],
|
||||
[{"role": "user", "content": "three"}],
|
||||
]
|
||||
for turn in recorder.turns:
|
||||
turn.dispatch_failure_handlers.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_turn_is_booked_alone_and_leaves_the_other_turns_successful(self):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
|
||||
|
||||
websocket = MagicMock()
|
||||
websocket.send_text = AsyncMock()
|
||||
backend_ws = MagicMock()
|
||||
backend_ws.recv = AsyncMock(
|
||||
side_effect=[
|
||||
json.dumps(
|
||||
recorder = _TurnLoggingRecorder()
|
||||
on_turn_failure = AsyncMock()
|
||||
client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder)
|
||||
backend = _ScriptedBackend(
|
||||
[
|
||||
_completed_turn("resp_1", 10, 1),
|
||||
[
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_1",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
"type": "error",
|
||||
"status": 400,
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"code": "previous_response_not_found",
|
||||
"message": "Previous response with id 'resp_x' not found.",
|
||||
},
|
||||
}
|
||||
),
|
||||
json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": "bad turn"}}),
|
||||
Exception("stop"),
|
||||
],
|
||||
_completed_turn("resp_3", 30, 3),
|
||||
]
|
||||
)
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.dispatch_success_handlers = AsyncMock()
|
||||
logging_obj.dispatch_failure_handlers = AsyncMock()
|
||||
logging_obj.response_cost_calculator = MagicMock(return_value=0.01)
|
||||
handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, logging_obj=logging_obj, request_data={})
|
||||
handler = _make_streaming(
|
||||
websocket=client,
|
||||
backend_ws=backend,
|
||||
logging_obj=recorder.connection,
|
||||
request_data={},
|
||||
authorized_model="gpt-5.4-mini",
|
||||
custom_llm_provider="openai",
|
||||
on_turn_failure=on_turn_failure,
|
||||
)
|
||||
|
||||
await handler.backend_to_client()
|
||||
await asyncio.sleep(0)
|
||||
await handler.bidirectional_forward()
|
||||
|
||||
logging_obj.record_partial_usage_for_failure.assert_called_once()
|
||||
usage, response_cost = logging_obj.record_partial_usage_for_failure.call_args[0]
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
|
||||
first, failed, third = recorder.turns
|
||||
for successful in (first, third):
|
||||
successful.dispatch_success_handlers.assert_awaited_once()
|
||||
successful.dispatch_failure_handlers.assert_not_awaited()
|
||||
failed.dispatch_success_handlers.assert_not_awaited()
|
||||
failed.dispatch_failure_handlers.assert_awaited_once()
|
||||
on_turn_failure.assert_awaited_once()
|
||||
booked_turn, booked_exception = on_turn_failure.await_args[0]
|
||||
assert booked_turn is failed
|
||||
assert booked_exception is failed.dispatch_failure_handlers.await_args[0][0]
|
||||
assert booked_exception.status_code == 400
|
||||
assert "not found" in str(booked_exception)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_turn_records_only_its_own_partial_usage(self):
|
||||
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
|
||||
|
||||
recorder = _TurnLoggingRecorder(response_cost=0.01)
|
||||
client = _ScriptedClient([_response_create(text) for text in ("one", "two")], recorder)
|
||||
backend = _ScriptedBackend(
|
||||
[
|
||||
_completed_turn("resp_1", 10, 5),
|
||||
[
|
||||
{"type": "response.created", "response": {"id": "resp_2", "status": "in_progress"}},
|
||||
{
|
||||
"type": "response.failed",
|
||||
"response": {
|
||||
"id": "resp_2",
|
||||
"status": "failed",
|
||||
"error": {"code": "server_error", "message": "upstream blew up"},
|
||||
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
|
||||
},
|
||||
},
|
||||
],
|
||||
]
|
||||
)
|
||||
handler = _make_streaming(
|
||||
websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={}
|
||||
)
|
||||
|
||||
await handler.bidirectional_forward()
|
||||
|
||||
first, failed = recorder.turns
|
||||
first.record_partial_usage_for_failure.assert_not_called()
|
||||
usage, response_cost = failed.record_partial_usage_for_failure.call_args[0]
|
||||
assert (usage.prompt_tokens, usage.completion_tokens) == (0, 0)
|
||||
assert response_cost == 0.01
|
||||
logging_obj.dispatch_success_handlers.assert_not_awaited()
|
||||
logging_obj.dispatch_failure_handlers.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bidirectional_forward_returns_the_provider_failure(self):
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
async def test_a_turn_still_in_flight_when_the_socket_closes_is_logged_with_what_it_received(self):
|
||||
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
|
||||
|
||||
backend_drained = asyncio.Event()
|
||||
backend_events = [
|
||||
json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}),
|
||||
json.dumps(
|
||||
{
|
||||
"type": "error",
|
||||
"status": 400,
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_encrypted_content",
|
||||
"message": "could not be verified",
|
||||
},
|
||||
}
|
||||
),
|
||||
]
|
||||
|
||||
async def recv(decode=False):
|
||||
if backend_events:
|
||||
return backend_events.pop(0)
|
||||
backend_drained.set()
|
||||
raise Exception("stop")
|
||||
|
||||
async def receive_text():
|
||||
await backend_drained.wait()
|
||||
raise Exception("client gone")
|
||||
|
||||
websocket = MagicMock()
|
||||
websocket.send_text = AsyncMock()
|
||||
websocket.receive_text = receive_text
|
||||
backend_ws = MagicMock()
|
||||
backend_ws.recv = recv
|
||||
backend_ws.send = AsyncMock()
|
||||
backend_ws.close = AsyncMock()
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.dispatch_success_handlers = AsyncMock()
|
||||
logging_obj.dispatch_failure_handlers = AsyncMock()
|
||||
logging_obj.response_cost_calculator = MagicMock(return_value=0.0)
|
||||
recorder = _TurnLoggingRecorder()
|
||||
client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder)
|
||||
backend = _ScriptedBackend(
|
||||
[
|
||||
_completed_turn("resp_1", 10, 1),
|
||||
[{"type": "response.created", "response": {"id": "resp_2", "status": "in_progress"}}],
|
||||
]
|
||||
)
|
||||
handler = _make_streaming(
|
||||
websocket=websocket,
|
||||
backend_ws=backend_ws,
|
||||
logging_obj=logging_obj,
|
||||
request_data={},
|
||||
authorized_model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={}
|
||||
)
|
||||
|
||||
failure = await handler.bidirectional_forward()
|
||||
await handler.bidirectional_forward()
|
||||
await _until(lambda: recorder.turns[1].dispatch_success_handlers.await_count == 1)
|
||||
|
||||
assert isinstance(failure, Exception)
|
||||
assert failure.status_code == 400
|
||||
assert "could not be verified" in str(failure)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bidirectional_forward_returns_none_after_a_completed_turn(self):
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
|
||||
|
||||
backend_drained = asyncio.Event()
|
||||
backend_events = [
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_1",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
},
|
||||
}
|
||||
),
|
||||
]
|
||||
|
||||
async def recv(decode=False):
|
||||
if backend_events:
|
||||
return backend_events.pop(0)
|
||||
backend_drained.set()
|
||||
raise Exception("stop")
|
||||
|
||||
async def receive_text():
|
||||
await backend_drained.wait()
|
||||
raise Exception("client gone")
|
||||
|
||||
websocket = MagicMock()
|
||||
websocket.send_text = AsyncMock()
|
||||
websocket.receive_text = receive_text
|
||||
backend_ws = MagicMock()
|
||||
backend_ws.recv = recv
|
||||
backend_ws.send = AsyncMock()
|
||||
backend_ws.close = AsyncMock()
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.dispatch_success_handlers = AsyncMock()
|
||||
logging_obj.dispatch_failure_handlers = AsyncMock()
|
||||
handler = _make_streaming(
|
||||
websocket=websocket,
|
||||
backend_ws=backend_ws,
|
||||
logging_obj=logging_obj,
|
||||
request_data={},
|
||||
authorized_model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
assert await handler.bidirectional_forward() is None
|
||||
first, in_flight, unanswered = recorder.turns
|
||||
first.dispatch_success_handlers.assert_awaited_once()
|
||||
in_flight.dispatch_success_handlers.assert_awaited_once()
|
||||
assert [event["type"] for event in in_flight.dispatch_success_handlers.await_args[0][0]] == ["response.created"]
|
||||
unanswered.dispatch_success_handlers.assert_not_awaited()
|
||||
unanswered.dispatch_failure_handlers.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -56,7 +56,9 @@ async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object])
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"])
|
||||
@pytest.mark.parametrize(
|
||||
"call_type", ["aresponses", "_aresponses_websocket", "anthropic_messages", "aembedding", "arerank"]
|
||||
)
|
||||
async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type):
|
||||
"""Non-chat surfaces carry the provider only on the standard logging payload."""
|
||||
limiter = RouterBudgetLimiting(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue