mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(passthrough): let a WebSocket route count client text for billing
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
parent
25fd335223
commit
6e6dd3f5bc
3 changed files with 163 additions and 2 deletions
|
|
@ -137,6 +137,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,
|
||||
)
|
||||
|
|
@ -2476,13 +2477,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.
|
||||
|
|
@ -2496,6 +2498,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
|
||||
|
|
@ -2506,6 +2511,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)
|
||||
|
|
@ -2624,6 +2630,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()
|
||||
|
|
@ -2693,8 +2700,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:
|
||||
|
|
@ -2835,7 +2846,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,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -67,6 +67,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,
|
||||
)
|
||||
|
|
@ -7822,6 +7823,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():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue