scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

feat(passthrough): let a WebSocket route count client text for billing

Browse source 
websocket_passthrough_request takes an optional count_client_text callback.
It is applied to each client frame, text or binary, after the frame
reaches the upstream, and only the summed integer reaches the success
handler, under WEBSOCKET_CLIENT_TEXT_COUNT_KEY in request_body. The frames
are never kept, so a route can bill on client text it must not store.
Binary frames are counted too, since an upstream may parse them the same
way as text frames

Frames are relayed unchanged in both directions. Callers that pass no
counter keep request_body={}. custom_headers is typed as the read-only
Mapping[str, str] the relay actually uses

Refs LIT-10070
This commit is contained in:
shrey kharbanda 2026-10-26 13:41:38 -07:00
parent 763cb04df2
commit 4ae796053f
3 changed files with 163 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -139,6 +139,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
WEBSOCKET_CLIENT_TEXT_COUNT_KEY,
EndpointType,
PassthroughStandardLoggingPayload,
)
@ -2489,13 +2490,14 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict
async def websocket_passthrough_request(
websocket: WebSocket,
target: str,
custom_headers: dict,
custom_headers: Mapping[str, str],
user_api_key_dict: UserAPIKeyAuth,
forward_headers: bool | None = False,
endpoint: str | None = None,
cost_per_request: float | None = None,
accept_websocket: bool = True,
setup_model_rewriter: Callable[[str], str] | None = None,
count_client_text: Callable[[str | bytes], int] | None = None,
):
"""
WebSocket passthrough request handler.
@ -2509,6 +2511,9 @@ async def websocket_passthrough_request(
endpoint: The endpoint path (for logging purposes)
cost_per_request: Optional field - cost per request to the target endpoint
setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream
count_client_text: Optional read-only count of each client frame, text or binary, after it reaches the
upstream, such as billable characters. Only the sum reaches the success handler, under WEBSOCKET_CLIENT_TEXT_COUNT_KEY
in request_body, so routes whose client text must not be stored can still bill it
"""
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj
@ -2519,6 +2524,7 @@ async def websocket_passthrough_request(
# Initialize tracking variables
start_time: Final = datetime.now()
websocket_messages: Final[list[dict[str, object]]] = []
client_text_count: int = 0 # rebind-ok: forward_client_to_upstream adds each client frame's count
litellm_call_id: Final = str(uuid.uuid4())
verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target)
@ -2637,6 +2643,7 @@ async def websocket_passthrough_request(
async def forward_client_to_upstream() -> None:
"""Forward messages from client to upstream WebSocket"""
nonlocal client_text_count
try:
while True:
message = await websocket.receive()
@ -2706,8 +2713,12 @@ async def websocket_passthrough_request(
# Not a JSON message or doesn't contain setup data
await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter))
if count_client_text is not None:
client_text_count += count_client_text(text_data)
elif bytes_data is not None:
await upstream_ws.send(bytes_data)
if count_client_text is not None:
client_text_count += count_client_text(bytes_data)
except asyncio.CancelledError:
raise
except Exception:
@ -2844,7 +2855,9 @@ async def websocket_passthrough_request(
end_time=end_time,
logging_obj=logging_obj,
cache_hit=False,
request_body={},
request_body={}
if count_client_text is None
else {WEBSOCKET_CLIENT_TEXT_COUNT_KEY: client_text_count},
**success_kwargs,
)
)

2
litellm/types/passthrough_endpoints/pass_through_endpoints.py
View file

@ -27,6 +27,8 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__
SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response"
WEBSOCKET_CLIENT_TEXT_COUNT_KEY: Final = "litellm_websocket_client_text_count"
class BackgroundResponseSettlement(str, Enum):
SETTLES = "settles"

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

@ -66,6 +66,7 @@ from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
WEBSOCKET_CLIENT_TEXT_COUNT_KEY,
EndpointType,
PassthroughStandardLoggingPayload,
)
@ -7808,6 +7809,151 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock:
return increment_spend_counters
DEEPGRAM_SPEAK_TARGET = "wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en"
SPEAK_AUDIO_CHUNK = bytes(range(256)) * 8
SPOKEN_TEXT = "The quick brown fox jumps over the lazy dog."
SPEAK_FRAME = json.dumps({"type": "Speak", "text": SPOKEN_TEXT})
FLUSH_FRAME = json.dumps({"type": "Flush"})
FLUSHED_FRAME = json.dumps({"type": "Flushed", "sequence_id": 0})
SPEAK_CLOSE_FRAME = json.dumps({"type": "Close"})
CLIENT_BINARY_FRAME = b"\x00\xffclient-binary"
SPEAK_SESSION_FRAMES = ({"text": SPEAK_FRAME}, {"bytes": CLIENT_BINARY_FRAME}, {"text": FLUSH_FRAME})
SPEAK_SESSION_CLIENT_LENGTH = sum(map(len, (SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME)))
RELAY_HANG_GUARD_SECONDS = 5.0
class SpeakingUpstreamWebSocket:
"""Answers each Speak with audio as soon as it arrives and ends the session on Close, the way Deepgram speak
does, so a relay that held audio back until the end would never see the Close it waits for"""
def __init__(self, close_code: int = 1000, close_reason: str = ""):
self.received: list[str | bytes] = []
self.close = AsyncMock()
self._close_code = close_code
self._close_reason = close_reason
self._replies: asyncio.Queue[str | bytes | None] = asyncio.Queue()
async def send(self, message: str | bytes) -> None:
self.received.append(message)
match json.loads(message).get("type") if isinstance(message, str) else None:
case "Speak":
self._replies.put_nowait(SPEAK_AUDIO_CHUNK)
case "Flush":
self._replies.put_nowait(FLUSHED_FRAME)
case "Close":
self._replies.put_nowait(None)
async def recv(self, decode: bool | None = None) -> str | bytes:
from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK
from websockets.frames import Close
reply = await self._replies.get()
if reply is not None:
return reply
close = Close(self._close_code, self._close_reason)
closed = ConnectionClosedOK if self._close_code == 1000 else ConnectionClosedError
raise closed(rcvd=close, sent=close, rcvd_then_sent=True)
def _speak_client(frames: tuple[dict[str, str | bytes], ...] = SPEAK_SESSION_FRAMES):
"""A client socket that sends ``frames``, then Close once the first audio has reached it"""
first_audio = asyncio.Event()
async def script():
for frame in frames:
yield {"type": "websocket.receive", **frame}
await first_audio.wait()
yield {"type": "websocket.receive", "text": SPEAK_CLOSE_FRAME}
await asyncio.Event().wait()
steps = script()
async def receive():
return await anext(steps)
websocket = _client_websocket(receive)
websocket.send_bytes = AsyncMock(side_effect=lambda data: first_audio.set())
return websocket
async def _relay_deepgram_speak(upstream_ws, websocket, count_client_text):
"""Runs the generic relay the way the Deepgram speak route does and returns the success handler mock"""
with (
_patched_websocket_passthrough_environment(upstream_ws),
patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam
"litellm.proxy.pass_through_endpoints.pass_through_endpoints."
"pass_through_endpoint_logging.pass_through_async_success_handler",
new=AsyncMock(),
) as success_handler,
):
await asyncio.wait_for(
websocket_passthrough_request(
websocket=websocket,
target=DEEPGRAM_SPEAK_TARGET,
custom_headers={"Authorization": "Token dg-provider-key"},
user_api_key_dict=UserAPIKeyAuth(),
forward_headers=False,
endpoint="/deepgram/v1/speak",
accept_websocket=False,
count_client_text=count_client_text,
),
timeout=RELAY_HANG_GUARD_SECONDS,
)
return success_handler
@pytest.mark.asyncio
async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed():
"""Audio reaches the client while the session is open, every client frame reaches the upstream unchanged, and
the success handler gets only the summed count of the text and binary frames, never the frames"""
upstream_ws = SpeakingUpstreamWebSocket()
websocket = _speak_client()
success_handler = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len)
assert upstream_ws.received == [SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME]
assert [call.args[0] for call in websocket.send_bytes.await_args_list] == [SPEAK_AUDIO_CHUNK]
assert [call.args[0] for call in websocket.send_text.await_args_list] == [FLUSHED_FRAME]
assert success_handler.call_args.kwargs["request_body"] == {
WEBSOCKET_CLIENT_TEXT_COUNT_KEY: SPEAK_SESSION_CLIENT_LENGTH
}
websocket.close.assert_awaited_once_with()
@pytest.mark.asyncio
async def test_websocket_passthrough_without_a_counter_reports_an_empty_request_body():
"""Listen, realtime and every other caller pass no counter, so their success handler sees what it saw before"""
upstream_ws = SpeakingUpstreamWebSocket()
success_handler = await _relay_deepgram_speak(upstream_ws, _speak_client(), count_client_text=None)
assert upstream_ws.received == [SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME]
assert success_handler.call_args.kwargs["request_body"] == {}
assert success_handler.call_args.kwargs["response_body"] == [json.loads(FLUSHED_FRAME)]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("close_code", "close_reason"),
[
pytest.param(1008, "DATA-0000: invalid text", id="policy"),
pytest.param(1011, "NET-0001: internal", id="internal"),
],
)
async def test_websocket_passthrough_counts_client_text_when_the_upstream_ends_with_an_error_close(
close_code: int, close_reason: str
):
upstream_ws = SpeakingUpstreamWebSocket(close_code=close_code, close_reason=close_reason)
websocket = _speak_client()
success_handler = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len)
websocket.close.assert_awaited_once_with(code=close_code, reason=close_reason)
assert success_handler.call_args.kwargs["request_body"] == {
WEBSOCKET_CLIENT_TEXT_COUNT_KEY: SPEAK_SESSION_CLIENT_LENGTH
}
@pytest.mark.asyncio
async def test_passthrough_success_reconciles_budget_reservation():
"""