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

feat(passthrough): let a WebSocket relay bill a session from its client text

Browse source 
websocket_passthrough_request takes an optional meter. The relay measures each client text frame once it reaches the upstream, and when the session ends it logs the model, provider and cost the meter prices for the total. A route can then bill an upstream that reports no usage of its own, without the relay knowing what is counted or the client's text being kept

Routes that pass no meter behave exactly as before
This commit is contained in:
shrey kharbanda 2026-10-30 10:49:09 -07:00
parent 84617874b8
commit 68717ecebf
3 changed files with 129 additions and 1 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

@ -148,6 +148,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
EndpointType,
PassthroughStandardLoggingPayload,
WebSocketSessionMeter,
WebSocketSessionUsage,
)
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage
@ -2595,6 +2597,13 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict
return inject_trace_context(headers, parent_span=parent_span)
def _log_session_usage(logging_obj: LiteLLMLoggingObj, usage: WebSocketSessionUsage) -> None:
logging_obj.model = usage.model # rebind-ok: the spend logger reads model and cost off the shared logging object
logging_obj.model_call_details["model"] = usage.model # rebind-ok: same shared logging object
logging_obj.model_call_details["custom_llm_provider"] = usage.custom_llm_provider # rebind-ok: same
logging_obj.model_call_details["response_cost"] = usage.response_cost # rebind-ok: same
async def websocket_passthrough_request(
websocket: WebSocket,
target: str,
@ -2607,6 +2616,7 @@ async def websocket_passthrough_request(
setup_model_rewriter: Callable[[str], str] | None = None,
builtin_provider_route: bool = False,
model: str | None = None,
meter: WebSocketSessionMeter | None = None,
):
"""
WebSocket passthrough request handler.
@ -2622,6 +2632,7 @@ async def websocket_passthrough_request(
setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream
builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider
model: The model the caller named before connecting, logged on a failure before the setup frame arrives
meter: Prices the session from the client text frames delivered upstream, for an upstream that reports no usage
"""
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj
@ -2632,6 +2643,7 @@ async def websocket_passthrough_request(
# Initialize tracking variables
start_time: Final = datetime.now()
websocket_messages: Final[list[dict[str, object]]] = []
metered_total = 0 # rebind-ok: forward_client_to_upstream adds the measure of each client text frame it delivers
litellm_call_id: Final = str(uuid.uuid4())
verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target)
@ -2750,6 +2762,7 @@ async def websocket_passthrough_request(
async def forward_client_to_upstream() -> None:
"""Forward messages from client to upstream WebSocket"""
nonlocal metered_total
try:
while True:
message = await websocket.receive()
@ -2819,6 +2832,8 @@ 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 meter is not None:
metered_total += meter.measure(text_data)
elif bytes_data is not None:
await upstream_ws.send(bytes_data)
except asyncio.CancelledError:
@ -2945,6 +2960,8 @@ async def websocket_passthrough_request(
self.url = target_url
mock_response: Final = MockWebSocketResponse(target)
if meter is not None:
_log_session_usage(logging_obj, meter.price(metered_total))
# Use the same success handler as HTTP passthrough endpoints
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(

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

@ -1,5 +1,6 @@
from dataclasses import dataclass
from enum import Enum
from typing import Final
from typing import Final, Protocol
from pydantic import Field
from typing_extensions import TypedDict
@ -75,6 +76,19 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False):
"""
@dataclass(frozen=True, slots=True)
class WebSocketSessionUsage:
model: str
custom_llm_provider: str
response_cost: float | None
class WebSocketSessionMeter(Protocol):
def measure(self, client_text: str) -> int: ...
def price(self, total: int) -> WebSocketSessionUsage: ...
class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase):
"""
Settings for a specific guardrail on a passthrough endpoint.

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

@ -77,6 +77,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
EndpointType,
PassthroughStandardLoggingPayload,
WebSocketSessionUsage,
)
from tests._master_key import MASTER_KEY
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@ -7736,6 +7737,102 @@ async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text
upstream_ws.close.assert_awaited_once()
@dataclass(frozen=True, slots=True)
class _LengthMeter:
rate: float
def measure(self, client_text: str) -> int:
return len(client_text)
def price(self, total: int) -> WebSocketSessionUsage:
return WebSocketSessionUsage(model="voice-1", custom_llm_provider="tts-provider", response_cost=total * self.rate)
@pytest.mark.asyncio
async def test_websocket_passthrough_logs_the_meters_price_for_the_client_text_it_delivered():
"""Text frames reach the upstream unchanged and are measured, binary frames are not, and the session is logged at
the model, provider and cost the meter prices for the total."""
frames = ("hello", b"\x00\x01binary", "world!")
upstream_ws = RecordingUpstreamWebSocket()
websocket = _client_websocket(
AsyncMock(
side_effect=[
*(
{"type": "websocket.receive", "bytes": frame}
if isinstance(frame, bytes)
else {"type": "websocket.receive", "text": frame}
for frame in frames
),
{"type": "websocket.disconnect"},
]
)
)
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 websocket_passthrough_request(
websocket=websocket,
target="wss://upstream.example/v1/stream",
custom_headers={},
user_api_key_dict=UserAPIKeyAuth(),
forward_headers=False,
endpoint="/provider/v1/stream",
accept_websocket=False,
meter=_LengthMeter(rate=0.5),
)
assert [call.args[0] for call in upstream_ws.send.await_args_list] == list(frames)
logging_obj = success_handler.call_args.kwargs["logging_obj"]
assert logging_obj.model == "voice-1"
assert logging_obj.model_call_details["model"] == "voice-1"
assert logging_obj.model_call_details["custom_llm_provider"] == "tts-provider"
assert logging_obj.model_call_details["response_cost"] == (len("hello") + len("world!")) * 0.5
@pytest.mark.asyncio
async def test_websocket_passthrough_meters_only_client_text_the_upstream_accepted():
upstream_ws = RecordingUpstreamWebSocket()
upstream_ws.send = AsyncMock(side_effect=[None, ConnectionError("upstream went away")])
websocket = _client_websocket(
AsyncMock(
side_effect=[
{"type": "websocket.receive", "text": "hello"},
{"type": "websocket.receive", "text": "never delivered"},
{"type": "websocket.disconnect"},
]
)
)
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 websocket_passthrough_request(
websocket=websocket,
target="wss://upstream.example/v1/stream",
custom_headers={},
user_api_key_dict=UserAPIKeyAuth(),
forward_headers=False,
endpoint="/provider/v1/stream",
accept_websocket=False,
meter=_LengthMeter(rate=0.5),
)
logging_obj = success_handler.call_args.kwargs["logging_obj"]
assert [call.args[0] for call in upstream_ws.send.await_args_list] == ["hello", "never delivered"]
assert logging_obj.model_call_details["response_cost"] == len("hello") * 0.5
@pytest.mark.asyncio
async def test_websocket_passthrough_vertex_live_setup_ack_names_the_model_but_is_not_billed_as_usage():
"""Vertex Live keeps its special first frame: the setup acknowledgement is forwarded verbatim, read for the