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

test(passthrough): bill the speak relay tests through the real success handler

Browse source 
The relay tests reuse the hang-up PR's helper, which runs the real success handler and returns what a spend logger receives, instead of mocking our own handler. tests/unit/AGENTS.md rules out doubles on our own code. The empty request body test is dropped, since the spend a session logs is now the assertion
This commit is contained in:
shrey kharbanda 2026-10-26 19:05:55 -07:00
parent 71d23db938
commit 660c43cb9f
2 changed files with 38 additions and 109 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_speak_passthrough_logging_handler.py
View file

@ -113,7 +113,7 @@ class _CapturingLogger(CustomLogger):
super().__init__()
self.payloads: list[StandardLoggingPayload] = [] # mutable-ok: the logger appends each payload it receives
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
self.payloads.append(kwargs["standard_logging_object"])

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

@ -68,7 +68,6 @@ 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,
)
@ -7653,10 +7652,11 @@ class _RelayedSession:
async def _relay_and_capture_spend(
monkeypatch: pytest.MonkeyPatch,
upstream_ws: FakeUpstreamWebSocket,
upstream_ws: "FakeUpstreamWebSocket | SpeakingUpstreamWebSocket",
websocket: MagicMock,
target: str,
endpoint: str,
count_client_text: Callable[[str | bytes], int] | None = None,
) -> _RelayedSession:
"""Runs the relay with the real success handler and returns what a spend logger receives"""
capturing_logger: Final = _CapturingSpendLogger()
@ -7687,6 +7687,7 @@ async def _relay_and_capture_spend(
forward_headers=False,
endpoint=endpoint,
accept_websocket=False,
count_client_text=count_client_text,
),
timeout=RELAY_HANG_GUARD_SECONDS,
)
@ -7834,7 +7835,6 @@ SPEAK_CLOSE_FRAME: Final = json.dumps({"type": "Close"})
CLIENT_BINARY_FRAME: Final = b"\x00\xffclient-binary"
SPEAK_SESSION_FRAMES: Final = ({"text": SPEAK_FRAME}, {"bytes": CLIENT_BINARY_FRAME}, {"text": FLUSH_FRAME})
SPEAK_SESSION_CLIENT_LENGTH: Final = sum(map(len, (SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME)))
RELAY_HANG_GUARD_SECONDS: Final = 5.0
class SpeakingUpstreamWebSocket:
@ -7891,67 +7891,35 @@ def _speak_client(frames: tuple[dict[str, str | bytes], ...] = SPEAK_SESSION_FRA
return websocket
async def _relay_deepgram_speak(
upstream_ws: SpeakingUpstreamWebSocket,
websocket: MagicMock,
count_client_text: Callable[[str | bytes], int] | None,
) -> AsyncMock:
"""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
def _aura_2_usd_per_character() -> float:
return litellm.model_cost["deepgram/aura-2"]["input_cost_per_character"]
@pytest.mark.asyncio
async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed() -> None:
@pytest.mark.usefixtures("local_model_cost_map")
async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""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"""
the session is billed for the summed count of the text and binary frames"""
upstream_ws: Final = SpeakingUpstreamWebSocket()
websocket: Final = _speak_client()
success_handler: Final = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len)
session: Final = await _relay_and_capture_spend(
monkeypatch, upstream_ws, websocket, DEEPGRAM_SPEAK_TARGET, "/deepgram/v1/speak", 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
}
assert [log["response_cost"] for log in session.spend_logs] == [
pytest.approx(SPEAK_SESSION_CLIENT_LENGTH * _aura_2_usd_per_character())
]
websocket.close.assert_awaited_once_with()
@pytest.mark.asyncio
async def test_websocket_passthrough_without_a_counter_reports_an_empty_request_body() -> None:
"""Listen, realtime and every other caller pass no counter, so their success handler sees what it saw before"""
upstream_ws: Final = SpeakingUpstreamWebSocket()
success_handler: Final = 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.usefixtures("local_model_cost_map")
@pytest.mark.parametrize(
("close_code", "close_reason"),
[
@ -7960,26 +7928,19 @@ async def test_websocket_passthrough_without_a_counter_reports_an_empty_request_
],
)
async def test_websocket_passthrough_counts_client_text_when_the_upstream_ends_with_an_error_close(
close_code: int, close_reason: str
monkeypatch: pytest.MonkeyPatch, close_code: int, close_reason: str
) -> None:
upstream_ws: Final = SpeakingUpstreamWebSocket(close_code=close_code, close_reason=close_reason)
websocket: Final = _speak_client()
success_handler: Final = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len)
session: Final = await _relay_and_capture_spend(
monkeypatch, upstream_ws, websocket, DEEPGRAM_SPEAK_TARGET, "/deepgram/v1/speak", 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
}
class _CapturingSpendLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.payloads: list[StandardLoggingPayload] = [] # mutable-ok: the logger appends each payload it receives
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.payloads.append(kwargs["standard_logging_object"])
assert [log["response_cost"] for log in session.spend_logs] == [
pytest.approx(SPEAK_SESSION_CLIENT_LENGTH * _aura_2_usd_per_character())
]
@pytest.mark.asyncio
@ -7987,52 +7948,20 @@ class _CapturingSpendLogger(CustomLogger):
async def test_deepgram_speak_session_logs_character_spend_once_and_never_the_spoken_text(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Runs the relay with the speak counter and the real success handler, so the spend logger sees what production
writes for one Speak message"""
capturing_logger: Final = _CapturingSpendLogger()
monkeypatch.setattr(litellm, "_async_success_callback", [capturing_logger])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "callbacks", [])
queued: Final[list[Coroutine[object, object, None]]] = [] # mutable-ok: the fake worker keeps each coroutine
upstream_ws: Final = SpeakingUpstreamWebSocket()
session: Final = await _relay_and_capture_spend(
monkeypatch,
SpeakingUpstreamWebSocket(),
_speak_client(),
DEEPGRAM_SPEAK_TARGET,
"/deepgram/v1/speak",
count_client_text=deepgram_speak_characters,
)
with (
patch("litellm.proxy.proxy_server.proxy_logging_obj") as proxy_logging,
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect",
return_value=FakeUpstreamConnect(upstream_ws),
),
patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker,
):
proxy_logging.pre_call_hook = AsyncMock(return_value={})
proxy_logging.post_call_success_hook = AsyncMock()
proxy_logging.post_call_failure_hook = AsyncMock()
worker.ensure_initialized_and_enqueue = MagicMock(
side_effect=lambda async_coroutine: queued.append(async_coroutine)
)
await asyncio.wait_for(
websocket_passthrough_request(
websocket=_speak_client(),
target=DEEPGRAM_SPEAK_TARGET,
custom_headers={"Authorization": "Token dg-provider-key"},
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", team_id="team-tts"),
forward_headers=False,
endpoint="/deepgram/v1/speak",
accept_websocket=False,
count_client_text=deepgram_speak_characters,
),
timeout=RELAY_HANG_GUARD_SECONDS,
)
for coroutine in queued:
await coroutine
assert len(capturing_logger.payloads) == 1
payload: Final = capturing_logger.payloads[0]
assert payload["model"] == "aura-2-thalia-en"
assert payload["custom_llm_provider"] == "deepgram"
# Aura-2 pay as you go rate, $0.030 per 1,000 characters, from https://deepgram.com/pricing, checked 2026-10-26
assert payload["response_cost"] == pytest.approx(len(SPOKEN_TEXT) * 0.030 / 1000)
assert SPOKEN_TEXT not in json.dumps(payload, default=str)
assert [(log["model"], log["custom_llm_provider"]) for log in session.spend_logs] == [
("aura-2-thalia-en", "deepgram")
]
assert session.spend_logs[0]["response_cost"] == pytest.approx(len(SPOKEN_TEXT) * _aura_2_usd_per_character())
assert SPOKEN_TEXT not in json.dumps(session.spend_logs[0], default=str)
@pytest.mark.asyncio