mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(proxy): bill each websocket lane cut off at session end on its own row
When a client leaves while turns on several stream_id lanes are still running, each unfinished turn now gets its own spend row under the groups and model its own frame asked for, instead of one merged row. Also pins the PR's new single-assigned test variables as Final and swaps datetime.now() in the logging-handler and success-handler fixtures for fixed timestamps.
This commit is contained in:
mateo-berri
parent
af0ad899de
commit
1d70354c06
6 changed files with 206 additions and 167 deletions
|
|
@ -3313,7 +3313,10 @@ async def websocket_passthrough_request(
|
|||
lanes_in_flight: Final = (
|
||||
() if upstream_frame_kind is None else _lanes_in_flight(upstream_frame_kind, websocket_messages)
|
||||
)
|
||||
log_turn(next(iter(lanes_in_flight), None), _every_frame)
|
||||
*earlier_lanes, last_lane = lanes_in_flight or (None,)
|
||||
for lane in earlier_lanes:
|
||||
log_turn(lane, _frames_of_lane(lane))
|
||||
log_turn(last_lane, _every_frame)
|
||||
|
||||
# Call the proxy logging success hook
|
||||
if proxy_logging_obj:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import json
|
||||
from typing import Any, cast
|
||||
from typing import Any, Final, cast
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -42,7 +42,7 @@ def test_realtime_request_body_caches_each_model_separately():
|
|||
|
||||
|
||||
def test_realtime_request_body_keeps_a_model_name_with_quotes_inside_its_own_field():
|
||||
body = realtime_request_body('gpt-4o", "user": "someone-else')
|
||||
body: Final = realtime_request_body('gpt-4o", "user": "someone-else')
|
||||
|
||||
assert json.loads(body) == {"model": 'gpt-4o", "user": "someone-else'}
|
||||
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket
|
|||
REALTIME_MODEL: Final = "gpt-realtime-2.1"
|
||||
RESPONSES_MODEL: Final = "gpt-5.4-mini"
|
||||
REASONING_MODEL: Final = "gpt-5.4"
|
||||
SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0)
|
||||
|
||||
|
||||
def _logging_obj() -> MagicMock:
|
||||
|
|
@ -86,13 +87,13 @@ def _handle(
|
|||
websocket_messages=messages,
|
||||
logging_obj=logging_obj if logging_obj is not None else _logging_obj(),
|
||||
url_route=url_route,
|
||||
start_time=datetime.now(),
|
||||
start_time=SESSION_START,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio_out: int) -> float:
|
||||
prices = litellm.model_cost[model]
|
||||
prices: Final = litellm.model_cost[model]
|
||||
return (
|
||||
text_in * prices["input_cost_per_token"]
|
||||
+ audio_in * prices["input_cost_per_audio_token"]
|
||||
|
|
@ -102,16 +103,16 @@ def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio
|
|||
|
||||
|
||||
def test_realtime_session_is_priced_on_the_model_the_session_reported():
|
||||
logging_obj = _logging_obj()
|
||||
messages = [
|
||||
logging_obj: Final = _logging_obj()
|
||||
messages: Final = [
|
||||
_session_created(REALTIME_MODEL),
|
||||
{"type": "response.created", "response": {"id": "resp_fake"}},
|
||||
_realtime_response_done(text_in=33, audio_in=120, text_out=16, audio_out=240),
|
||||
]
|
||||
|
||||
handled = _handle(messages, logging_obj=logging_obj)
|
||||
handled: Final = _handle(messages, logging_obj=logging_obj)
|
||||
|
||||
expected_cost = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240)
|
||||
expected_cost: Final = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240)
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["kwargs"]["model"] == REALTIME_MODEL
|
||||
assert handled["kwargs"]["custom_llm_provider"] == "openai"
|
||||
|
|
@ -122,8 +123,8 @@ def test_realtime_session_is_priced_on_the_model_the_session_reported():
|
|||
|
||||
|
||||
def test_realtime_cost_covers_every_turn_of_the_session():
|
||||
one_turn = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)])
|
||||
two_turns = _handle(
|
||||
one_turn: Final = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)])
|
||||
two_turns: Final = _handle(
|
||||
[
|
||||
_session_created(REALTIME_MODEL),
|
||||
_realtime_response_done(10, 0, 5, 0),
|
||||
|
|
@ -136,13 +137,13 @@ def test_realtime_cost_covers_every_turn_of_the_session():
|
|||
|
||||
|
||||
def test_responses_turn_is_priced_on_the_model_the_response_reported():
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
prices: Final = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["kwargs"]["model"] == RESPONSES_MODEL
|
||||
assert handled["result"].usage.prompt_tokens == 13
|
||||
|
|
@ -150,7 +151,7 @@ def test_responses_turn_is_priced_on_the_model_the_response_reported():
|
|||
|
||||
|
||||
def test_responses_cost_covers_every_turn_of_the_connection():
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5),
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7),
|
||||
|
|
@ -158,8 +159,8 @@ def test_responses_cost_covers_every_turn_of_the_connection():
|
|||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"]
|
||||
prices: Final = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost: Final = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"]
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["result"].usage.total_tokens == 45
|
||||
|
||||
|
|
@ -170,9 +171,9 @@ def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_
|
|||
the provider's own frames may set the price of a session, so a model the
|
||||
caller put in the request cannot stand in for one the provider never named.
|
||||
"""
|
||||
logging_obj = _logging_obj()
|
||||
logging_obj: Final = _logging_obj()
|
||||
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0)],
|
||||
kwargs=MappingProxyType({"model": REALTIME_MODEL}),
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -189,7 +190,7 @@ def test_each_responses_turn_is_priced_on_the_model_that_turn_ran():
|
|||
starts on a big model and finishes on a small one has to pay the big
|
||||
model's price for the turn that used it.
|
||||
"""
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[
|
||||
_responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5),
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7),
|
||||
|
|
@ -197,9 +198,9 @@ def test_each_responses_turn_is_priced_on_the_model_that_turn_ran():
|
|||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
first = litellm.model_cost[REASONING_MODEL]
|
||||
second = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = (
|
||||
first: Final = litellm.model_cost[REASONING_MODEL]
|
||||
second: Final = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost: Final = (
|
||||
13 * first["input_cost_per_token"]
|
||||
+ 5 * first["output_cost_per_token"]
|
||||
+ 20 * second["input_cost_per_token"]
|
||||
|
|
@ -210,7 +211,7 @@ def test_each_responses_turn_is_priced_on_the_model_that_turn_ran():
|
|||
|
||||
|
||||
def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model():
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[
|
||||
_session_created(RESPONSES_MODEL),
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}),
|
||||
|
|
@ -218,14 +219,14 @@ def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model():
|
|||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
prices: Final = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
|
||||
|
||||
def test_a_session_nobody_named_a_model_for_is_left_unpriced():
|
||||
logging_obj = _logging_obj()
|
||||
handled = _handle(
|
||||
logging_obj: Final = _logging_obj()
|
||||
handled: Final = _handle(
|
||||
[{"type": "response.done", "response": {"usage": {"input_tokens": 10, "output_tokens": 5}}}],
|
||||
kwargs=MappingProxyType({"litellm_call_id": "call-fake"}),
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -237,9 +238,9 @@ def test_a_session_nobody_named_a_model_for_is_left_unpriced():
|
|||
|
||||
|
||||
def test_a_session_that_reported_no_usage_still_records_its_model():
|
||||
logging_obj = _logging_obj()
|
||||
logging_obj: Final = _logging_obj()
|
||||
|
||||
handled = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj)
|
||||
handled: Final = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj)
|
||||
|
||||
assert handled["result"] is None
|
||||
assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"}
|
||||
|
|
@ -248,7 +249,7 @@ def test_a_session_that_reported_no_usage_still_records_its_model():
|
|||
|
||||
|
||||
def test_a_responses_turn_without_readable_usage_is_not_billed():
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[{"type": "response.completed", "response": {"model": RESPONSES_MODEL, "status": "completed"}}],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
|
@ -264,7 +265,7 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row():
|
|||
price it already worked out or the cheaper model pays for both turns. The
|
||||
row it lands in names the model the connection opened on.
|
||||
"""
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5),
|
||||
_responses_completed(REASONING_MODEL, input_tokens=20, output_tokens=7),
|
||||
|
|
@ -272,7 +273,7 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row():
|
|||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
logged_cost = litellm.response_cost_calculator(
|
||||
logged_cost: Final = litellm.response_cost_calculator(
|
||||
response_object=handled["result"],
|
||||
model=handled["kwargs"]["model"],
|
||||
custom_llm_provider="openai",
|
||||
|
|
@ -280,8 +281,8 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row():
|
|||
optional_params={},
|
||||
)
|
||||
|
||||
session_prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
priced_on_the_session_model = (
|
||||
session_prices: Final = litellm.model_cost[RESPONSES_MODEL]
|
||||
priced_on_the_session_model: Final = (
|
||||
33 * session_prices["input_cost_per_token"] + 12 * session_prices["output_cost_per_token"]
|
||||
)
|
||||
assert logged_cost == pytest.approx(handled["kwargs"]["response_cost"])
|
||||
|
|
@ -295,9 +296,9 @@ def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran():
|
|||
bill, but response.created already named the model, so the row says which
|
||||
model answered instead of falling back to unknown.
|
||||
"""
|
||||
logging_obj = _logging_obj()
|
||||
logging_obj: Final = _logging_obj()
|
||||
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[{"type": "response.created", "response": {"id": "resp_fake", "model": REASONING_MODEL}}],
|
||||
logging_obj=logging_obj,
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
|
|
@ -314,19 +315,19 @@ def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(even
|
|||
A turn cut short by max_output_tokens or a provider failure still reports
|
||||
the tokens it burned, and OpenAI bills them, so the row has to as well.
|
||||
"""
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, event=event)],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
prices: Final = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["result"].usage.total_tokens == 18
|
||||
|
||||
|
||||
def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed():
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[{"type": "response.failed", "response": {"model": RESPONSES_MODEL, "status": "failed", "usage": None}}],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
|
@ -353,21 +354,21 @@ def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_ha
|
|||
billed at that model's price instead of $0, while a turn whose own model
|
||||
has a price keeps it.
|
||||
"""
|
||||
handled = _handle(
|
||||
handled: Final = _handle(
|
||||
[_responses_completed(turn_model, input_tokens=13, output_tokens=5)],
|
||||
kwargs=MappingProxyType({"litellm_params": {"metadata": {"model_group": requested_model}}}),
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[priced_as]
|
||||
expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
prices: Final = litellm.model_cost[priced_as]
|
||||
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
assert expected_cost > 0
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["kwargs"]["model"] == turn_model
|
||||
|
||||
|
||||
def _http_passthrough_cost(response: Mapping[str, object]) -> float:
|
||||
http_response = httpx.Response(
|
||||
http_response: Final = httpx.Response(
|
||||
200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses")
|
||||
)
|
||||
_, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost(
|
||||
|
|
@ -390,10 +391,10 @@ def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_r
|
|||
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)
|
||||
turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served)
|
||||
plain_turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)
|
||||
|
||||
handled = _handle([turn], url_route="/openai_passthrough/v1/responses")
|
||||
handled: Final = _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"])
|
||||
|
|
|
|||
|
|
@ -7697,7 +7697,7 @@ class RecordingFrameModelGate:
|
|||
|
||||
async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict:
|
||||
self.checked.append(model)
|
||||
refusal = self.refusals.get(model)
|
||||
refusal: Final = self.refusals.get(model)
|
||||
if refusal is not None:
|
||||
return WebsocketFrameModelRefused(message=refusal, error_type=self.error_type)
|
||||
return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ()))
|
||||
|
|
@ -7714,8 +7714,8 @@ async def _relay_client_frames(
|
|||
gate: WebsocketFrameModelGate | None,
|
||||
user_api_key_dict: UserAPIKeyAuth | None = None,
|
||||
) -> SimpleNamespace:
|
||||
upstream_ws = RecordingUpstreamWebSocket()
|
||||
websocket = _client_websocket(
|
||||
upstream_ws: Final = RecordingUpstreamWebSocket()
|
||||
websocket: Final = _client_websocket(
|
||||
AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}])
|
||||
)
|
||||
|
||||
|
|
@ -7758,9 +7758,9 @@ async def _relay_client_frames(
|
|||
],
|
||||
)
|
||||
async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(frame, expected_model):
|
||||
gate = RecordingFrameModelGate()
|
||||
gate: Final = RecordingFrameModelGate()
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
relayed: Final = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
|
||||
assert gate.checked == [expected_model]
|
||||
assert [call.args[0] for call in relayed.upstream.send.await_args_list] == [json.dumps(frame)]
|
||||
|
|
@ -7769,10 +7769,10 @@ async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(fr
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cannot_call():
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
|
||||
frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}})
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
|
||||
frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}})
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
relayed: Final = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert json.loads(relayed.client.send_text.await_args.args[0]) == {
|
||||
|
|
@ -7790,10 +7790,10 @@ async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the
|
|||
reads as a truncated number. The first sentence names the reason whole, and
|
||||
the error event before the close still carries every word with its type.
|
||||
"""
|
||||
refusal = "Budget has been exceeded! Key=" + "f" * 64 + " Spend=1.23456789, Budget=1.0"
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": refusal}), error_type="budget_exceeded")
|
||||
refusal: Final = "Budget has been exceeded! Key=" + "f" * 64 + " Spend=1.23456789, Budget=1.0"
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": refusal}), error_type="budget_exceeded")
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps({"type": "response.create", "model": "gpt-5.4"})], gate)
|
||||
relayed: Final = await _relay_client_frames([json.dumps({"type": "response.create", "model": "gpt-5.4"})], gate)
|
||||
|
||||
assert relayed.client.close.await_args_list == [call(code=1008, reason="Budget has been exceeded!")]
|
||||
assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == {
|
||||
|
|
@ -7804,15 +7804,15 @@ async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied():
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL}))
|
||||
frame = json.dumps(
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL}))
|
||||
frame: Final = json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {"model": "gpt-realtime-2.1", "input_audio_transcription": {"model": "gpt-4o-transcribe"}},
|
||||
}
|
||||
)
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
relayed: Final = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
|
|
@ -7821,13 +7821,13 @@ async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_mode
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_stops_relaying_after_a_refused_frame():
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
|
||||
frames = [
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
|
||||
frames: Final = [
|
||||
json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}),
|
||||
json.dumps({"type": "response.create"}),
|
||||
]
|
||||
|
||||
relayed = await _relay_client_frames(frames, gate)
|
||||
relayed: Final = await _relay_client_frames(frames, gate)
|
||||
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert gate.checked == ["gpt-realtime-2.1"]
|
||||
|
|
@ -7845,9 +7845,9 @@ async def test_websocket_passthrough_stops_relaying_after_a_refused_frame():
|
|||
],
|
||||
)
|
||||
async def test_websocket_passthrough_forwards_frames_that_name_no_model(frame):
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"}))
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"}))
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
relayed: Final = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == []
|
||||
assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame]
|
||||
|
|
@ -7862,13 +7862,13 @@ async def test_websocket_passthrough_never_bills_the_model_a_client_frame_named(
|
|||
session onto that model's price, so only the provider's own frames decide
|
||||
what a session is billed as.
|
||||
"""
|
||||
gate = RecordingFrameModelGate()
|
||||
frames = [
|
||||
gate: Final = RecordingFrameModelGate()
|
||||
frames: Final = [
|
||||
json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}),
|
||||
json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}),
|
||||
]
|
||||
|
||||
relayed = await _relay_client_frames(frames, gate)
|
||||
relayed: Final = await _relay_client_frames(frames, gate)
|
||||
|
||||
assert gate.checked == ["gpt-realtime-2.1", "gpt-realtime-2.1-mini"]
|
||||
assert "model" not in relayed.success_handler.call_args.kwargs
|
||||
|
|
@ -7888,9 +7888,9 @@ async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(fram
|
|||
"\\u006dodel" asks for a model just as plainly as one carrying "model" and
|
||||
has to face the same check.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL}))
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL}))
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
relayed: Final = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == [expected_model]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
|
|
@ -7899,10 +7899,10 @@ async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(fram
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for():
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
|
||||
frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode()
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
|
||||
frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode()
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
relayed: Final = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == ["gpt-realtime-2.1"]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
|
|
@ -7924,9 +7924,9 @@ async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_t
|
|||
a model list could reach any model through one. The gateway cannot see that
|
||||
model, so a key that is not allowed every model cannot run a stored prompt.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL}))
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL}))
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
relayed: Final = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
|
||||
error_message: Final = json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"]
|
||||
close_reason: Final = relayed.client.close.await_args.kwargs["reason"]
|
||||
|
|
@ -7940,9 +7940,9 @@ async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_t
|
|||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES)
|
||||
async def test_websocket_passthrough_forwards_a_stored_prompt_for_a_key_allowed_every_model(frame):
|
||||
gate = RecordingFrameModelGate()
|
||||
gate: Final = RecordingFrameModelGate()
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
relayed: Final = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
|
||||
assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [json.dumps(frame)]
|
||||
assert relayed.client.close.await_args_list == [call()]
|
||||
|
|
@ -7950,12 +7950,12 @@ async def test_websocket_passthrough_forwards_a_stored_prompt_for_a_key_allowed_
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_prompt_is_refused_for_spend():
|
||||
refusal = "Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0"
|
||||
gate = RecordingFrameModelGate(
|
||||
refusal: Final = "Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0"
|
||||
gate: Final = RecordingFrameModelGate(
|
||||
MappingProxyType({STORED_PROMPT_PROBE_MODEL: refusal}), error_type="budget_exceeded"
|
||||
)
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(STORED_PROMPT_FRAMES[0])], gate)
|
||||
relayed: Final = await _relay_client_frames([json.dumps(STORED_PROMPT_FRAMES[0])], gate)
|
||||
|
||||
assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == {
|
||||
"type": "budget_exceeded",
|
||||
|
|
@ -7966,14 +7966,14 @@ async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_promp
|
|||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_checks_the_model_a_hosted_tool_names():
|
||||
"""A hosted tool such as image generation runs its own model, so that model gets the same check as the turn's"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-image-2": KEY_MODEL_REFUSAL}))
|
||||
frame = {
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-image-2": KEY_MODEL_REFUSAL}))
|
||||
frame: Final = {
|
||||
"type": "response.create",
|
||||
"model": "gpt-5.4-mini",
|
||||
"tools": [{"type": "web_search"}, {"type": "image_generation", "model": "gpt-image-2"}],
|
||||
}
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
relayed: Final = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
|
||||
assert gate.checked == ["gpt-5.4-mini", "gpt-image-2"]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
|
|
@ -7989,9 +7989,9 @@ async def test_websocket_passthrough_checks_the_model_a_hosted_tool_names():
|
|||
],
|
||||
)
|
||||
async def test_websocket_passthrough_refuses_an_image_tool_without_a_model_for_a_key_limited_to_some_models(frame):
|
||||
gate = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL}))
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL}))
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
relayed: Final = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
|
||||
assert gate.checked[-1] == STORED_PROMPT_PROBE_MODEL
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
|
|
@ -8001,10 +8001,10 @@ async def test_websocket_passthrough_refuses_an_image_tool_without_a_model_for_a
|
|||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_checks_a_frame_only_the_standard_json_parser_reads():
|
||||
"""A lone surrogate escape is valid JSON syntax the fast parser refuses, so the frame is still checked"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
|
||||
frame = '{"type": "response.create", "model": "gpt-5.4", "metadata": {"note": "\\ud800"}}'
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
|
||||
frame: Final = '{"type": "response.create", "model": "gpt-5.4", "metadata": {"note": "\\ud800"}}'
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
relayed: Final = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == ["gpt-5.4"]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
|
|
@ -8012,10 +8012,10 @@ async def test_websocket_passthrough_checks_a_frame_only_the_standard_json_parse
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_forwards_binary_audio_untouched():
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"}))
|
||||
audio = bytes(range(256)) * 4
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"}))
|
||||
audio: Final = bytes(range(256)) * 4
|
||||
|
||||
relayed = await _relay_client_frames([audio], gate)
|
||||
relayed: Final = await _relay_client_frames([audio], gate)
|
||||
|
||||
assert gate.checked == []
|
||||
assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [audio]
|
||||
|
|
@ -8023,9 +8023,9 @@ async def test_websocket_passthrough_forwards_binary_audio_untouched():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_without_a_frame_gate_forwards_every_frame():
|
||||
frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}})
|
||||
frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}})
|
||||
|
||||
relayed = await _relay_client_frames([frame], None)
|
||||
relayed: Final = await _relay_client_frames([frame], None)
|
||||
|
||||
assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame]
|
||||
assert "model" not in relayed.success_handler.call_args.kwargs
|
||||
|
|
@ -8051,7 +8051,7 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out
|
|||
already used still have to reach spend, including the last frame the client
|
||||
never collected.
|
||||
"""
|
||||
upstream_ws = FakeUpstreamWebSocket(
|
||||
upstream_ws: Final = FakeUpstreamWebSocket(
|
||||
REALTIME_SESSION_CREATED,
|
||||
REALTIME_RESPONSE_DONE,
|
||||
json.dumps({"type": "response.output_text.delta", "delta": "too late"}),
|
||||
|
|
@ -8112,7 +8112,7 @@ class ScriptedUpstreamWebSocket:
|
|||
from websockets.exceptions import ConnectionClosedOK
|
||||
from websockets.frames import Close
|
||||
|
||||
frame = await self._frames.get()
|
||||
frame: Final = await self._frames.get()
|
||||
if frame is None:
|
||||
raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
|
||||
return frame
|
||||
|
|
@ -8126,14 +8126,14 @@ def _turn_taking_client(
|
|||
from starlette.websockets import WebSocketState
|
||||
|
||||
delivered: list[str] = []
|
||||
script = iter(frames)
|
||||
receives = iter(range(len(frames) + 1) if answers_due is None else answers_due)
|
||||
script: Final = iter(frames)
|
||||
receives: Final = iter(range(len(frames) + 1) if answers_due is None else answers_due)
|
||||
|
||||
async def receive():
|
||||
answers_due = next(receives)
|
||||
answers_due: Final = next(receives)
|
||||
while sum(answered_by in text for text in delivered) < answers_due:
|
||||
await asyncio.sleep(0)
|
||||
frame = next(script, None)
|
||||
frame: Final = next(script, None)
|
||||
return {"type": "websocket.disconnect"} if frame is None else _received_frame(frame)
|
||||
|
||||
websocket = _client_websocket(receive)
|
||||
|
|
@ -8208,21 +8208,21 @@ async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_th
|
|||
and be charged to the group that let its model through, so a group budget
|
||||
sees its share during the session and never pays for the other group's turn.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(
|
||||
gate: Final = RecordingFrameModelGate(
|
||||
groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}),
|
||||
)
|
||||
client = _turn_taking_client(
|
||||
client: Final = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}),
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4"}),
|
||||
],
|
||||
answered_by='"response.completed"',
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini"), _responses_turn("gpt-5.4")])
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini"), _responses_turn("gpt-5.4")])
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, gate)
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, gate)
|
||||
|
||||
rows = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [
|
||||
["gpt-5.4-mini", "gpt-5.4-mini"],
|
||||
["gpt-5.4", "gpt-5.4"],
|
||||
|
|
@ -8246,10 +8246,10 @@ async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_a
|
|||
and under the model that frame asked for, even when the provider answers
|
||||
under a dated snapshot name a per-model budget was never set on.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(
|
||||
gate: Final = RecordingFrameModelGate(
|
||||
groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}),
|
||||
)
|
||||
client = _turn_taking_client(
|
||||
client: Final = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}),
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4"}),
|
||||
|
|
@ -8257,13 +8257,13 @@ async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_a
|
|||
answered_by='"response.completed"',
|
||||
answers_due=(0, 0, 2),
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(
|
||||
replies=[_responses_turn("gpt-5.4-mini-2099-01-01"), _responses_turn("gpt-5.4-2099-01-01")]
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, gate)
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, gate)
|
||||
|
||||
rows = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [
|
||||
["gpt-5.4-mini-2099-01-01", "gpt-5.4-mini-2099-01-01"],
|
||||
["gpt-5.4-2099-01-01", "gpt-5.4-2099-01-01"],
|
||||
|
|
@ -8291,10 +8291,10 @@ async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_
|
|||
open. Those frames never run, so each later turn still has to reach spend
|
||||
under the groups and model its own frame asked for.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(
|
||||
gate: Final = RecordingFrameModelGate(
|
||||
groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}),
|
||||
)
|
||||
client = _turn_taking_client(
|
||||
client: Final = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "input": "hi"}),
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4", "previous_response_id": "resp_gone"}),
|
||||
|
|
@ -8304,7 +8304,7 @@ async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_
|
|||
answered_by='"response.completed"',
|
||||
answers_due=(0, 0, 0, 0, 2),
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(
|
||||
replies=[
|
||||
[_responses_error("missing_required_parameter")],
|
||||
[_responses_error("previous_response_not_found")],
|
||||
|
|
@ -8313,9 +8313,9 @@ async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_
|
|||
]
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, gate)
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, gate)
|
||||
|
||||
rows = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"]
|
||||
assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [
|
||||
["tier-a"],
|
||||
|
|
@ -8331,10 +8331,10 @@ async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_fram
|
|||
Each turn still has to reach spend on its own row, under the groups and
|
||||
model its own frame asked for.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(
|
||||
gate: Final = RecordingFrameModelGate(
|
||||
groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}),
|
||||
)
|
||||
client = _turn_taking_client(
|
||||
client: Final = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "stream_id": "slow", "model": "gpt-5.4"}),
|
||||
json.dumps({"type": "response.create", "stream_id": "fast", "model": "gpt-5.4-mini"}),
|
||||
|
|
@ -8342,7 +8342,7 @@ async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_fram
|
|||
answered_by='"response.completed"',
|
||||
answers_due=(0, 0, 2),
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(
|
||||
replies=[
|
||||
[_on_lane(_responses_turn_started("gpt-5.4"), "slow")],
|
||||
[
|
||||
|
|
@ -8352,9 +8352,9 @@ async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_fram
|
|||
]
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, gate)
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, gate)
|
||||
|
||||
rows = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
assert [[(frame["stream_id"], frame["response"]["model"]) for frame in row["response_body"]] for row in rows] == [
|
||||
[("fast", "gpt-5.4-mini"), ("fast", "gpt-5.4-mini")],
|
||||
[("slow", "gpt-5.4"), ("slow", "gpt-5.4")],
|
||||
|
|
@ -8373,8 +8373,8 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k
|
|||
charged to the groups and the model the key opened it with, whatever model
|
||||
a later session.update or response.create names.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(groups=MappingProxyType({"gpt-realtime-2.1-mini": ("tier-b",)}))
|
||||
client = _turn_taking_client(
|
||||
gate: Final = RecordingFrameModelGate(groups=MappingProxyType({"gpt-realtime-2.1-mini": ("tier-b",)}))
|
||||
client: Final = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1-mini"}}),
|
||||
json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}),
|
||||
|
|
@ -8382,7 +8382,7 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k
|
|||
answered_by='"type"',
|
||||
answers_due=(1, 2, 4),
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(
|
||||
REALTIME_SESSION_CREATED,
|
||||
replies=[
|
||||
[json.dumps({"type": "session.updated"})],
|
||||
|
|
@ -8390,7 +8390,7 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k
|
|||
],
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(
|
||||
relayed: Final = await _relay_turns(
|
||||
client,
|
||||
upstream_ws,
|
||||
gate,
|
||||
|
|
@ -8400,7 +8400,7 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k
|
|||
)
|
||||
|
||||
assert gate.checked == ["gpt-realtime-2.1-mini", "gpt-realtime-2.1-mini"]
|
||||
rows = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [["tier-a"]]
|
||||
assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-realtime-2.1"]
|
||||
|
||||
|
|
@ -8408,10 +8408,10 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k
|
|||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_logs_no_spend_row_for_a_session_refused_before_any_answer():
|
||||
"""A session refused at its first frame ran nothing, so its only record is the refusal, not an extra $0 row"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
|
||||
client = _turn_taking_client([json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"error"')
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
|
||||
client: Final = _turn_taking_client([json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"error"')
|
||||
|
||||
relayed = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate)
|
||||
relayed: Final = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate)
|
||||
|
||||
assert relayed.success_handler.call_args_list == []
|
||||
|
||||
|
|
@ -8422,18 +8422,18 @@ async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_ou
|
|||
A client can leave while a later turn is still running. That turn never
|
||||
ends, but it was started and may be billed, so it still gets its own row.
|
||||
"""
|
||||
client = _turn_taking_client(
|
||||
client: Final = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}),
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4"}),
|
||||
],
|
||||
answered_by='"response.created"',
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(
|
||||
replies=[_responses_turn("gpt-5.4-mini"), [_responses_turn_started("gpt-5.4")]]
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate())
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate())
|
||||
|
||||
assert [
|
||||
[json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]]
|
||||
|
|
@ -8441,6 +8441,41 @@ async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_ou
|
|||
] == [["response.created", "response.completed"], ["response.created"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_charges_each_lane_a_client_walked_out_of_to_the_frame_that_started_it():
|
||||
"""
|
||||
A client can leave while turns on several stream_id lanes are still
|
||||
running. Each unfinished turn still gets its own row, under the groups and
|
||||
model its own frame asked for.
|
||||
"""
|
||||
gate: Final = RecordingFrameModelGate(
|
||||
groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}),
|
||||
)
|
||||
client: Final = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "stream_id": "slow", "model": "gpt-5.4"}),
|
||||
json.dumps({"type": "response.create", "stream_id": "fast", "model": "gpt-5.4-mini"}),
|
||||
],
|
||||
answered_by='"response.created"',
|
||||
)
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(
|
||||
replies=[
|
||||
[_on_lane(_responses_turn_started("gpt-5.4"), "slow")],
|
||||
[_on_lane(_responses_turn_started("gpt-5.4-mini"), "fast")],
|
||||
]
|
||||
)
|
||||
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, gate)
|
||||
|
||||
rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
assert [[frame["stream_id"] for frame in row["response_body"]] for row in rows] == [["slow"], ["fast"]]
|
||||
assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4", "gpt-5.4-mini"]
|
||||
assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [
|
||||
["tier-b"],
|
||||
["tier-a"],
|
||||
]
|
||||
|
||||
|
||||
class UpstreamFinishingTheTurnAfterTheClientLeft:
|
||||
"""Starts a turn for each client frame and sends the rest of it only after the client left, and stops once closed"""
|
||||
|
||||
|
|
@ -8474,7 +8509,7 @@ class UpstreamFinishingTheTurnAfterTheClientLeft:
|
|||
from websockets.exceptions import ConnectionClosedOK
|
||||
from websockets.frames import Close
|
||||
|
||||
frame = await self._frames.get()
|
||||
frame: Final = await self._frames.get()
|
||||
if frame is None:
|
||||
raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
|
||||
return frame
|
||||
|
|
@ -8492,16 +8527,16 @@ async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_
|
|||
"""
|
||||
from starlette.websockets import WebSocketDisconnect, WebSocketState
|
||||
|
||||
client_left = asyncio.Event()
|
||||
send_failed = asyncio.Event()
|
||||
client_left: Final = asyncio.Event()
|
||||
send_failed: Final = 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
|
||||
receive: Final = client.receive
|
||||
send_text: Final = client.send_text
|
||||
|
||||
async def receive_until_leaving() -> Mapping[str, object]:
|
||||
message = await receive()
|
||||
message: Final = await receive()
|
||||
if message["type"] == "websocket.disconnect":
|
||||
client_left.set()
|
||||
if not hang_up_read_before_the_next_send:
|
||||
|
|
@ -8517,13 +8552,13 @@ async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_
|
|||
|
||||
client.receive = receive_until_leaving
|
||||
client.send_text = AsyncMock(side_effect=send_text_until_leaving)
|
||||
upstream_ws = UpstreamFinishingTheTurnAfterTheClientLeft(
|
||||
upstream_ws: Final = 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)
|
||||
relayed: Final = 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"]]
|
||||
|
|
@ -8539,13 +8574,13 @@ async def test_websocket_passthrough_logs_no_extra_row_for_what_the_provider_sen
|
|||
turn, such as rate_limits.updated. The session ends with nothing running,
|
||||
so its turn is its only row and the request count stays one per turn.
|
||||
"""
|
||||
client = _turn_taking_client([json.dumps({"type": "response.create"})], answered_by='"rate_limits.updated"')
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
client: Final = _turn_taking_client([json.dumps({"type": "response.create"})], answered_by='"rate_limits.updated"')
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(
|
||||
REALTIME_SESSION_CREATED,
|
||||
replies=[[REALTIME_RESPONSE_DONE, json.dumps({"type": "rate_limits.updated", "rate_limits": []})]],
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime")
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime")
|
||||
|
||||
assert [
|
||||
[frame["type"] for frame in logged.kwargs["response_body"]]
|
||||
|
|
@ -8560,15 +8595,15 @@ async def test_websocket_passthrough_carries_the_session_the_provider_opened_int
|
|||
session, and never again. Each later turn logs on its own, so it has to
|
||||
carry that opening frame or it would be priced as a model nobody named.
|
||||
"""
|
||||
client = _turn_taking_client(
|
||||
client: Final = _turn_taking_client(
|
||||
[json.dumps({"type": "response.create"}), json.dumps({"type": "response.create"})],
|
||||
answered_by='"response.done"',
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(
|
||||
REALTIME_SESSION_CREATED, replies=[[REALTIME_RESPONSE_DONE], [REALTIME_RESPONSE_DONE]]
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime")
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime")
|
||||
|
||||
assert [
|
||||
[frame["type"] for frame in logged.kwargs["response_body"]]
|
||||
|
|
@ -8584,8 +8619,8 @@ async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_lat
|
|||
but it was paid for, so it has to reach spend instead of turning the whole
|
||||
session into a failure that bills nothing.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
|
||||
client = _turn_taking_client(
|
||||
gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
|
||||
client: Final = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}),
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4"}),
|
||||
|
|
@ -8593,9 +8628,9 @@ async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_lat
|
|||
answered_by='"response.created"',
|
||||
)
|
||||
created, completed = _responses_turn("gpt-5.4-mini")
|
||||
upstream_ws = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed])
|
||||
upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed])
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, gate)
|
||||
relayed: Final = await _relay_turns(client, upstream_ws, gate)
|
||||
|
||||
assert [json.loads(text.args[0])["type"] for text in client.send_text.await_args_list] == [
|
||||
"response.created",
|
||||
|
|
@ -9536,7 +9571,7 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
|
|||
|
||||
|
||||
async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object]:
|
||||
served = await _relay_client_frames(
|
||||
served: Final = await _relay_client_frames(
|
||||
[json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})],
|
||||
gate,
|
||||
UserAPIKeyAuth(api_key="hashed-key"),
|
||||
|
|
@ -9548,7 +9583,7 @@ async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object
|
|||
async def test_websocket_session_is_charged_to_the_group_its_frames_were_authorized_through():
|
||||
"""A websocket names its model in a frame, long after the relay built its logging metadata, so the
|
||||
groups the gate matched have to be picked up at the end or the group pays for nothing it allowed."""
|
||||
metadata = await _billed_metadata(RecordingFrameModelGate(groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",)})))
|
||||
metadata: Final = await _billed_metadata(RecordingFrameModelGate(groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",)})))
|
||||
|
||||
assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"]
|
||||
assert metadata["user_api_key"] == "hashed-key"
|
||||
|
|
@ -9556,7 +9591,7 @@ async def test_websocket_session_is_charged_to_the_group_its_frames_were_authori
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_session_that_matched_no_group_is_charged_to_none():
|
||||
metadata = await _billed_metadata(RecordingFrameModelGate())
|
||||
metadata: Final = await _billed_metadata(RecordingFrameModelGate())
|
||||
|
||||
assert not metadata[MODEL_ACCESS_GROUP_METADATA_KEY]
|
||||
|
||||
|
|
|
|||
|
|
@ -677,8 +677,8 @@ def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str,
|
|||
logging_obj=logging_obj,
|
||||
url_route=url_route,
|
||||
result="websocket_connection_successful",
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
start_time=_START,
|
||||
end_time=_END,
|
||||
cache_hit=False,
|
||||
litellm_call_id="call-fake",
|
||||
)
|
||||
|
|
@ -686,11 +686,11 @@ def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str,
|
|||
|
||||
@pytest.mark.parametrize("url_route", ["/openai/v1/realtime", "/openai_passthrough/v1/realtime"])
|
||||
def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route):
|
||||
normalized = _normalize(url_route, REALTIME_FRAMES)
|
||||
normalized: Final = _normalize(url_route, REALTIME_FRAMES)
|
||||
|
||||
assert normalized["kwargs"]["model"] == "gpt-realtime-2.1"
|
||||
assert normalized["kwargs"]["custom_llm_provider"] == "openai"
|
||||
prices = litellm.model_cost["gpt-realtime-2.1"]
|
||||
prices: Final = litellm.model_cost["gpt-realtime-2.1"]
|
||||
assert normalized["kwargs"]["response_cost"] == pytest.approx(
|
||||
33 * prices["input_cost_per_token"] + 16 * prices["output_cost_per_token"]
|
||||
)
|
||||
|
|
@ -698,7 +698,7 @@ def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route):
|
|||
|
||||
|
||||
def test_openai_http_passthrough_bodies_are_left_to_the_http_handlers():
|
||||
normalized = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}})
|
||||
normalized: Final = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}})
|
||||
|
||||
assert normalized["standard_logging_response_object"] is None
|
||||
assert normalized["kwargs"] == {"litellm_call_id": "call-fake"}
|
||||
|
|
|
|||
|
|
@ -137,7 +137,7 @@ async def _serve(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: Mapping[str, object],
|
||||
) -> _Served:
|
||||
served = _Served(relay=_FakeRelay(), frame_model_gate=_FakeFrameModelGate())
|
||||
served: Final = _Served(relay=_FakeRelay(), frame_model_gate=_FakeFrameModelGate())
|
||||
await openai_websocket_proxy_route(
|
||||
websocket=websocket,
|
||||
endpoint=endpoint,
|
||||
|
|
@ -153,7 +153,7 @@ async def _serve(
|
|||
@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"])
|
||||
async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch):
|
||||
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
|
||||
websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="sk-provider"):
|
||||
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
|
||||
|
|
@ -178,10 +178,10 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m
|
|||
async def test_openai_websocket_reports_its_route_without_the_proxy_root_path(monkeypatch):
|
||||
"""Route allowlists and spend rows key on the route the proxy serves, which a root path prefix is not part of"""
|
||||
monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm")
|
||||
websocket = _FakeWebSocket("/litellm/openai/v1/responses", "")
|
||||
websocket: Final = _FakeWebSocket("/litellm/openai/v1/responses", "")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="sk-provider"):
|
||||
served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED)
|
||||
served: Final = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED)
|
||||
|
||||
assert [call.endpoint for call in served.relay.calls] == ["/openai/v1/responses"]
|
||||
|
||||
|
|
@ -263,7 +263,7 @@ async def test_openai_websocket_closes_cleanly_when_workload_identity_exchange_f
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing():
|
||||
websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1")
|
||||
websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value=None):
|
||||
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
|
||||
|
|
@ -279,7 +279,7 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing
|
|||
@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"])
|
||||
@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS)
|
||||
async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings):
|
||||
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
|
||||
websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
|
||||
|
||||
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings)
|
||||
|
||||
|
|
@ -320,11 +320,11 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol():
|
|||
[["gpt-5.4-mini"], ["gpt-realtime-2.1", "gpt-5.4"], []],
|
||||
)
|
||||
async def test_openai_websocket_connects_model_restricted_keys(key_models):
|
||||
websocket = _FakeWebSocket("/openai/v1/responses", "")
|
||||
user_api_key_dict = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake")
|
||||
websocket: Final = _FakeWebSocket("/openai/v1/responses", "")
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="sk-provider"):
|
||||
served = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED)
|
||||
served: Final = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED)
|
||||
|
||||
assert len(served.relay.calls) == 1
|
||||
assert served.relay.calls[0].client_frame_model_gate is served.frame_model_gate
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue