scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

fix(proxy): price websocket Responses turns by served tier and hosted tools, and bill turns a client walked out of

Browse source 
A websocket Responses turn is now priced through the same Responses cost path the HTTP passthrough uses, so a priority-tier answer and a hosted web search cost what they cost over HTTP. When the client leaves mid-answer, the gateway keeps reading the provider's socket until that turn ends, up to 10 minutes, so the turn the provider still bills reaches spend
This commit is contained in:
mateo-berri 2026-10-29 10:56:52 -07:00
parent 5185a536db
commit ad4a71c2b9
4 changed files with 182 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py
View file

@ -13,7 +13,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage
from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponsesAPIResponse
from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind
from litellm.types.utils import LlmProviders, ModelResponse, Usage
@ -43,18 +43,11 @@ class _ModelNamingEvent(BaseModel):
response: _NamedModel | None = None
class _EndedResponse(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
model: str | None = None
usage: ResponseAPIUsage | None = None
class _ResponsesTurnEndEvent(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
type: Literal["response.completed", "response.incomplete", "response.failed"]
response: _EndedResponse
response: ResponsesAPIResponse
_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent)
@ -70,6 +63,7 @@ class _Billing:
@dataclass(frozen=True, slots=True)
class _CompletedTurn:
model: str
response: ResponsesAPIResponse
usage: Usage
@ -136,6 +130,7 @@ def _responses_turn(event: Mapping[str, object], session_model: str) -> _Complet
return None
return _CompletedTurn(
model=ended.response.model or session_model,
response=ended.response,
usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter
ended.response.usage
),
@ -149,23 +144,24 @@ def _responses_billing(messages: WebsocketMessages, model: str, requested_model:
usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns])
return _Billing(
usage=usage,
cost=sum(_responses_cost(turn.usage, (turn.model, requested_model)) for turn in turns),
cost=sum(_responses_cost(turn.response, (turn.model, requested_model)) for turn in turns),
)
def _model_cost(usage: Usage, model: str) -> float | None:
def _model_cost(response: ResponsesAPIResponse, model: str) -> float | None:
try:
return litellm.completion_cost(
completion_response=ModelResponse(model=model, usage=usage),
completion_response=response,
model=model,
custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER,
call_type="responses",
)
except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log
return None
def _responses_cost(usage: Usage, price_as: Sequence[str | None]) -> float:
costs: Final = (_model_cost(usage, model) for model in dict.fromkeys(price_as) if model is not None)
def _responses_cost(response: ResponsesAPIResponse, price_as: Sequence[str | None]) -> float:
costs: Final = (_model_cost(response, model) for model in dict.fromkeys(price_as) if model is not None)
cost: Final = next((cost for cost in costs if cost), None)
if cost is None:
verbose_proxy_logger.warning(

42
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -2743,6 +2743,14 @@ def _session_end_needs_a_row(
return "starts_turn" in kinds or (ledger.logged_turns == 0 and bool(kinds - {"opens_session"}))
def _turn_in_flight(
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, unlogged_frames: Sequence[Mapping[str, object]]
) -> bool:
return upstream_frame_kind is not None and any(
upstream_frame_kind(frame) == "starts_turn" for frame in unlogged_frames
)
async def _relay_client_frame(
websocket: WebSocket,
upstream_ws: ClientConnection,
@ -2765,6 +2773,7 @@ async def _relay_client_frame(
_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"})
WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0
async def websocket_passthrough_request(
@ -2779,6 +2788,7 @@ async def websocket_passthrough_request(
setup_model_rewriter: Callable[[str], str] | None = None,
client_frame_model_gate: WebsocketFrameModelGate | None = None,
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None,
walkout_turn_drain_seconds: float = WALKOUT_TURN_DRAIN_SECONDS,
):
"""
WebSocket passthrough request handler.
@ -2795,6 +2805,8 @@ async def websocket_passthrough_request(
client_frame_model_gate: Optional per-frame check of every model a client frame names
upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones
that open the session, so each turn logs its own spend and keeps the session's provider-named model
walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client
left, so the provider's usage for it still reaches spend
"""
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj
@ -3019,7 +3031,8 @@ async def websocket_passthrough_request(
message = await websocket.receive()
message_type = message.get("type")
if message_type == "websocket.disconnect":
await upstream_ws.close()
if not _turn_in_flight(upstream_frame_kind, websocket_messages):
await upstream_ws.close()
break
text_data: str | None = message.get("text")
@ -3147,11 +3160,17 @@ async def websocket_passthrough_request(
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
record_upstream_frame(upstream_message)
if not _client_socket_is_open(websocket):
if not _turn_in_flight(upstream_frame_kind, websocket_messages):
await upstream_ws.close()
return
if isinstance(upstream_message, bytes):
await websocket.send_bytes(upstream_message)
else:
await websocket.send_text(upstream_message)
try:
if isinstance(upstream_message, bytes):
await websocket.send_bytes(upstream_message)
else:
await websocket.send_text(upstream_message)
except WebSocketDisconnect:
if not _turn_in_flight(upstream_frame_kind, websocket_messages):
raise
async def forward_upstream_to_client() -> Close | None:
try:
@ -3176,12 +3195,15 @@ async def websocket_passthrough_request(
raise
# Create tasks for bidirectional message forwarding
tasks: Final = [
asyncio.create_task(forward_client_to_upstream()),
asyncio.create_task(forward_upstream_to_client()),
]
client_task: Final = asyncio.create_task(forward_client_to_upstream())
upstream_task: Final = asyncio.create_task(forward_upstream_to_client())
tasks: Final = (client_task, upstream_task)
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
if client_task.done() and _turn_in_flight(upstream_frame_kind, websocket_messages):
await asyncio.wait((upstream_task,), timeout=walkout_turn_drain_seconds)
done: Final = tuple(task for task in tasks if task.done())
pending: Final = tuple(task for task in tasks if not task.done())
# Cancel remaining tasks
for task in pending:

52
tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py
View file

@ -4,9 +4,13 @@ from types import MappingProxyType
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
import litellm
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
OpenAIWebsocketPassthroughLoggingHandler,
openai_websocket_frame_kind,
@ -45,14 +49,22 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou
def _responses_completed(
model: str, input_tokens: int, output_tokens: int, event: str = "response.completed"
model: str,
input_tokens: int,
output_tokens: int,
event: str = "response.completed",
served: Mapping[str, object] = MappingProxyType({}),
) -> dict[str, object]:
return {
"type": event,
"response": {
"id": "resp_fake",
"object": "response",
"created_at": 1761700000,
"model": model,
"status": "completed",
"output": [],
**served,
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
@ -201,10 +213,7 @@ def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model():
handled = _handle(
[
_session_created(RESPONSES_MODEL),
{
"type": "response.completed",
"response": {"usage": {"input_tokens": 13, "output_tokens": 5, "total_tokens": 18}},
},
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}),
],
url_route="/openai_passthrough/v1/responses",
)
@ -357,6 +366,39 @@ def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_ha
assert handled["kwargs"]["model"] == turn_model
def _http_passthrough_cost(response: Mapping[str, object]) -> float:
http_response = httpx.Response(
200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses")
)
_, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost(
RESPONSES_MODEL, http_response, _logging_obj(), "openai"
)
return cost
@pytest.mark.parametrize(
"served",
[
{"service_tier": "priority"},
{"output": [{"type": "web_search_call", "id": "ws_fake", "status": "completed", "action": {"type": "search"}}]},
],
ids=["priority_tier", "hosted_web_search"],
)
def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_response(served):
"""
The provider bills a turn by the tier that served it and the hosted tools it
called, not only its tokens, so a websocket turn has to cost exactly what the
same response costs over HTTP, never the cheaper token-only price.
"""
turn = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served)
plain_turn = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)
handled = _handle([turn], url_route="/openai_passthrough/v1/responses")
assert handled["kwargs"]["response_cost"] == pytest.approx(_http_passthrough_cost(turn["response"]))
assert handled["kwargs"]["response_cost"] > _http_passthrough_cost(plain_turn["response"])
@pytest.mark.parametrize(
("frame_type", "kind"),
[

93
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -8103,6 +8103,7 @@ async def _relay_turns(
endpoint: str = "/openai_passthrough/v1/responses",
target: str = "wss://api.openai.com/v1/responses",
opening_access_groups: Sequence[str] = (),
walkout_turn_drain_seconds: float = 0.05,
) -> SimpleNamespace:
with (
_patched_websocket_passthrough_environment(upstream_ws) as proxy_logging,
@ -8121,6 +8122,7 @@ async def _relay_turns(
accept_websocket=False,
client_frame_model_gate=gate,
upstream_frame_kind=openai_websocket_frame_kind,
walkout_turn_drain_seconds=walkout_turn_drain_seconds,
),
timeout=10,
)
@ -8297,6 +8299,97 @@ async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_ou
] == [["response.created", "response.completed"], ["response.created"]]
class UpstreamFinishingTheTurnAfterTheClientLeft:
"""Starts a turn for each client frame and sends the rest of it only after the client left, and stops once closed"""
def __init__(self, client_left: asyncio.Event, started: str, rest_of_the_turn: Sequence[str]):
self._frames: asyncio.Queue[str | None] = asyncio.Queue()
self._client_left = client_left
self._started = started
self._rest_of_the_turn = rest_of_the_turn
self._closed = False
self._finishing: list[asyncio.Task[None]] = []
self.send = AsyncMock(side_effect=self._start_turn)
self.close = AsyncMock(side_effect=self._close)
def _close(self, *args: object, **kwargs: object) -> None:
self._closed = True
self._frames.put_nowait(None)
async def _start_turn(self, frame: str | bytes) -> None:
self._frames.put_nowait(self._started)
self._finishing.append(asyncio.create_task(self._finish_once_the_client_left()))
async def _finish_once_the_client_left(self) -> None:
await self._client_left.wait()
for frame in self._rest_of_the_turn:
await asyncio.sleep(0)
if self._closed:
return
self._frames.put_nowait(frame)
async def recv(self, decode: bool | None = None) -> str:
from websockets.exceptions import ConnectionClosedOK
from websockets.frames import Close
frame = await self._frames.get()
if frame is None:
raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
return frame
@pytest.mark.asyncio
@pytest.mark.parametrize("hang_up_read_before_the_next_send", [True, False], ids=["read_first", "found_by_a_send"])
async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_the_provider_ends_it(
hang_up_read_before_the_next_send: bool,
):
"""
Leaving mid-answer must not make a turn free. The provider keeps running and
bills it, so the gateway keeps reading until the turn ends and logs its usage,
whether it reads the hang-up first or finds it by failing to relay a frame.
"""
from starlette.websockets import WebSocketDisconnect, WebSocketState
client_left = asyncio.Event()
send_failed = asyncio.Event()
client = _turn_taking_client(
[json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"response.created"'
)
receive = client.receive
send_text = client.send_text
async def receive_until_leaving() -> Mapping[str, object]:
message = await receive()
if message["type"] == "websocket.disconnect":
client_left.set()
if not hang_up_read_before_the_next_send:
await send_failed.wait()
client.client_state = WebSocketState.DISCONNECTED
return message
async def send_text_until_leaving(text: str) -> None:
if client_left.is_set():
send_failed.set()
raise WebSocketDisconnect(code=1006)
await send_text(text)
client.receive = receive_until_leaving
client.send_text = AsyncMock(side_effect=send_text_until_leaving)
upstream_ws = UpstreamFinishingTheTurnAfterTheClientLeft(
client_left,
_responses_turn_started("gpt-5.4"),
(json.dumps({"type": "response.output_text.delta", "delta": "Hi"}), _responses_turn_ended("gpt-5.4")),
)
relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=5)
assert [
[json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]]
for logged in relayed.success_handler.call_args_list
] == [["response.created", "response.output_text.delta", "response.completed"]]
upstream_ws.close.assert_awaited()
@pytest.mark.asyncio
async def test_websocket_passthrough_logs_no_extra_row_for_what_the_provider_sends_after_the_last_turn():
"""