mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(elevenlabs): bill a WebSocket session by the characters ElevenLabs spoke, not the text sent
ElevenLabs bills the text it synthesized, which its alignment data lists, so text a session sent but never heard is no longer billed Audio that comes back with no alignment at all still bills every character the client sent
This commit is contained in:
mateo-berri
parent
9799a1f16a
commit
eaf0940c59
8 changed files with 212 additions and 97 deletions
|
|
@ -6,7 +6,7 @@ from typing import Final, Literal, TypeAlias
|
|||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import ELEVENLABS_DEFAULT_API_BASE, ELEVENLABS_TTS_WEBSOCKET_DEFAULT_MODEL
|
||||
|
|
@ -26,6 +26,7 @@ ELEVENLABS_TTS_WEBSOCKET_ROUTE_SUFFIXES: Final = ("/stream-input", "/multi-strea
|
|||
ELEVENLABS_TTS_CALLER_CREDENTIAL_PARAMS: Final = frozenset({"authorization", "single_use_token"})
|
||||
ELEVENLABS_TTS_CALLER_CREDENTIAL_FIELDS: Final = frozenset({"xi_api_key", "xi-api-key", "authorization"})
|
||||
ELEVENLABS_TTS_AUDIO_FIELDS: Final = frozenset({"audio", "alignment", "normalizedAlignment"})
|
||||
ELEVENLABS_TTS_ALIGNED_CHARACTERS_KEY: Final = "aligned_characters"
|
||||
_MODEL_ID_PARAM: Final = "model_id"
|
||||
_VOICE_ID: Final = re.compile(r"[A-Za-z0-9_-]+")
|
||||
_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"})
|
||||
|
|
@ -122,9 +123,38 @@ def elevenlabs_tts_client_text(text: str) -> ClientTextDecision:
|
|||
return ForwardedClientText(text=forwarded, recorded={"text": frame_text} if frame_text is not None else None)
|
||||
|
||||
|
||||
class _Alignment(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
chars: tuple[str, ...]
|
||||
|
||||
|
||||
_ALIGNMENT: Final = TypeAdapter(_Alignment)
|
||||
|
||||
|
||||
def _aligned_characters(alignment: object) -> int:
|
||||
try:
|
||||
return len(_ALIGNMENT.validate_python(alignment).chars)
|
||||
except ValidationError:
|
||||
return 0
|
||||
|
||||
|
||||
def elevenlabs_tts_recorded_upstream_frame(frame: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return {key: value for key, value in frame.items() if key not in ELEVENLABS_TTS_AUDIO_FIELDS}
|
||||
without_audio: Final = {key: value for key, value in frame.items() if key not in ELEVENLABS_TTS_AUDIO_FIELDS}
|
||||
if not frame.get("audio"):
|
||||
return without_audio
|
||||
return {**without_audio, ELEVENLABS_TTS_ALIGNED_CHARACTERS_KEY: _aligned_characters(frame.get("alignment"))}
|
||||
|
||||
|
||||
def elevenlabs_tts_characters(client_frames: Sequence[Mapping[str, object]]) -> int:
|
||||
def elevenlabs_tts_characters(
|
||||
client_frames: Sequence[Mapping[str, object]], upstream_frames: Sequence[Mapping[str, object]]
|
||||
) -> int:
|
||||
"""ElevenLabs bills the characters it synthesized, which its alignment lists, so text a session sent but never
|
||||
heard (interrupted, unflushed, timed out) is free; audio that came back with no alignment at all bills the text
|
||||
the client sent instead."""
|
||||
audio_frames: Final = tuple(frame for frame in upstream_frames if ELEVENLABS_TTS_ALIGNED_CHARACTERS_KEY in frame)
|
||||
aligned: Final = sum(
|
||||
count for frame in audio_frames if isinstance(count := frame[ELEVENLABS_TTS_ALIGNED_CHARACTERS_KEY], int)
|
||||
)
|
||||
if aligned or not audio_frames:
|
||||
return aligned
|
||||
return sum(len(text) for frame in client_frames if isinstance(text := frame.get("text"), str))
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
|
@ -47,7 +47,7 @@ class ElevenLabsTextToSpeechPassthroughLoggingHandler:
|
|||
def elevenlabs_tts_passthrough_handler(
|
||||
self,
|
||||
request_body: Mapping[str, object],
|
||||
upstream_frame_count: int,
|
||||
upstream_frames: Sequence[Mapping[str, object]],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
upstream_url: str,
|
||||
kwargs: Mapping[str, object] = MappingProxyType({}),
|
||||
|
|
@ -56,7 +56,7 @@ class ElevenLabsTextToSpeechPassthroughLoggingHandler:
|
|||
request_body.get(WEBSOCKET_PASSTHROUGH_CLIENT_FRAMES_KEY, ())
|
||||
)
|
||||
model: Final = elevenlabs_tts_model(upstream_url)
|
||||
characters: Final = elevenlabs_tts_characters(client_frames)
|
||||
characters: Final = elevenlabs_tts_characters(client_frames, upstream_frames)
|
||||
response_cost: Final = _speech_cost(upstream_url, characters)
|
||||
|
||||
provider: Final = litellm.LlmProviders.ELEVENLABS.value
|
||||
|
|
@ -76,7 +76,7 @@ class ElevenLabsTextToSpeechPassthroughLoggingHandler:
|
|||
"model": model,
|
||||
"characters": characters,
|
||||
"client_text_frames": len(client_frames),
|
||||
"upstream_frames": upstream_frame_count,
|
||||
"upstream_frames": len(upstream_frames),
|
||||
}
|
||||
),
|
||||
"kwargs": {
|
||||
|
|
|
|||
|
|
@ -100,7 +100,7 @@ def _audio_websocket_handler_result(
|
|||
case "elevenlabs_tts":
|
||||
return ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=request_body,
|
||||
upstream_frame_count=len(upstream_frames),
|
||||
upstream_frames=tuple(frame for frame in upstream_frames if isinstance(frame, dict)),
|
||||
logging_obj=logging_obj,
|
||||
upstream_url=str(httpx_response.request.url),
|
||||
kwargs=kwargs,
|
||||
|
|
|
|||
|
|
@ -197,7 +197,7 @@ def test_a_session_reaches_elevenlabs_with_the_proxy_key_and_is_billed_per_chara
|
|||
{"text": ""},
|
||||
)
|
||||
rows: Final = _spend_rows(key)
|
||||
characters: Final = len(" ") + sum(len(text) for text in SENTENCES)
|
||||
characters: Final = sum(len(text) for text in SENTENCES)
|
||||
assert float(rows[0]["spend"]) == pytest.approx(characters * _per_character(FLASH))
|
||||
assert rows[0]["model"] == FLASH
|
||||
assert rows[0]["custom_llm_provider"] == "elevenlabs"
|
||||
|
|
@ -223,5 +223,5 @@ def test_an_unpriced_model_is_refused_before_elevenlabs_is_contacted_and_the_def
|
|||
)
|
||||
assert elevenlabs.upgrades.empty()
|
||||
rows: Final = _spend_rows(key)
|
||||
assert float(rows[0]["spend"]) == pytest.approx((1 + len(SENTENCES[0])) * _per_character(DEFAULT_MODEL))
|
||||
assert float(rows[0]["spend"]) == pytest.approx(len(SENTENCES[0]) * _per_character(DEFAULT_MODEL))
|
||||
assert rows[0]["model"] == DEFAULT_MODEL
|
||||
|
|
|
|||
|
|
@ -230,16 +230,28 @@ def test_client_text_the_billing_parser_cannot_read_is_refused_instead_of_forwar
|
|||
assert decision.code == 1007
|
||||
|
||||
|
||||
def test_recorded_upstream_frame_drops_audio_and_alignment():
|
||||
frame = {
|
||||
def _audio_frame(spoken: str | None) -> dict[str, object]:
|
||||
alignment = None if spoken is None else {"chars": list(spoken), "charStartTimesMs": [0], "charDurationsMs": [50]}
|
||||
return {
|
||||
"audio": "SUQzBAAAAAAAI1RTU0UAAAAPAAADTGF2ZjYxLjcuMTAw",
|
||||
"alignment": {"chars": ["H"], "charStartTimesMs": [0], "charDurationsMs": [50]},
|
||||
"normalizedAlignment": {"chars": ["H"], "charStartTimesMs": [0], "charDurationsMs": [50]},
|
||||
"alignment": alignment,
|
||||
"normalizedAlignment": alignment,
|
||||
"contextId": "ctx",
|
||||
"isFinal": None,
|
||||
}
|
||||
|
||||
assert elevenlabs_tts_recorded_upstream_frame(frame) == {"contextId": "ctx", "isFinal": None}
|
||||
|
||||
def test_recorded_upstream_frame_drops_audio_and_alignment_but_keeps_the_aligned_character_count():
|
||||
assert elevenlabs_tts_recorded_upstream_frame(_audio_frame("Hello")) == {
|
||||
"contextId": "ctx",
|
||||
"isFinal": None,
|
||||
"aligned_characters": 5,
|
||||
}
|
||||
assert elevenlabs_tts_recorded_upstream_frame(_audio_frame(None)) == {
|
||||
"contextId": "ctx",
|
||||
"isFinal": None,
|
||||
"aligned_characters": 0,
|
||||
}
|
||||
assert elevenlabs_tts_recorded_upstream_frame({"isFinal": True}) == {"isFinal": True}
|
||||
assert elevenlabs_tts_recorded_upstream_frame({"message": "Invalid model", "error": "invalid_request"}) == {
|
||||
"message": "Invalid model",
|
||||
|
|
@ -247,15 +259,40 @@ def test_recorded_upstream_frame_drops_audio_and_alignment():
|
|||
}
|
||||
|
||||
|
||||
def test_characters_count_every_text_character_including_whitespace():
|
||||
frames = (
|
||||
{"text": " "},
|
||||
{"text": "Hello world. "},
|
||||
{"text": "¿Qué tal? "},
|
||||
{"close_context": True},
|
||||
{"text": None},
|
||||
{"text": ""},
|
||||
def _session(client_texts: tuple[str, ...], *upstream: dict[str, object]) -> int:
|
||||
return elevenlabs_tts_characters(
|
||||
tuple({"text": text} for text in client_texts),
|
||||
tuple(elevenlabs_tts_recorded_upstream_frame(frame) for frame in upstream),
|
||||
)
|
||||
|
||||
assert elevenlabs_tts_characters(frames) == 1 + len("Hello world. ") + len("¿Qué tal? ")
|
||||
assert elevenlabs_tts_characters(()) == 0
|
||||
|
||||
def test_characters_are_the_ones_elevenlabs_spoke_not_the_ones_sent():
|
||||
sent = (" ", "Hello world. ", "¿Qué tal? ", "This part was cut off before it was spoken. ")
|
||||
|
||||
spoken = (_audio_frame("Hello world."), _audio_frame(None), _audio_frame("¿Qué tal?"), {"isFinal": True})
|
||||
|
||||
assert _session(sent, *spoken) == len("Hello world.") + len("¿Qué tal?")
|
||||
|
||||
|
||||
def test_a_session_that_got_no_audio_back_bills_nothing():
|
||||
assert _session((" ", "Text that never reached the generation threshold "), {"isFinal": True}) == 0
|
||||
assert (
|
||||
_session(
|
||||
(" ", "Text the caller dropped before any audio "),
|
||||
)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
def test_audio_without_any_alignment_bills_every_character_sent():
|
||||
sent = (" ", "Hello world. ", "¿Qué tal? ", "")
|
||||
|
||||
assert _session(sent, _audio_frame(None), _audio_frame(None), {"isFinal": True}) == (
|
||||
1 + len("Hello world. ") + len("¿Qué tal? ")
|
||||
)
|
||||
assert (
|
||||
elevenlabs_tts_characters(
|
||||
({"close_context": True}, {"text": None}), (elevenlabs_tts_recorded_upstream_frame(_audio_frame(None)),)
|
||||
)
|
||||
== 0
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import pytest
|
|||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.elevenlabs.common_utils import elevenlabs_tts_recorded_upstream_frame
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.elevenlabs_tts_passthrough_logging_handler import (
|
||||
ElevenLabsTextToSpeechPassthroughLoggingHandler,
|
||||
|
|
@ -19,12 +20,24 @@ pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map")
|
|||
|
||||
FLASH_URL: Final = "wss://api.elevenlabs.io/v1/text-to-speech/voice1/stream-input?model_id=eleven_flash_v2_5"
|
||||
SESSION_TEXT: Final = (" ", "Hello world. ", "¿Qué tal? ", "")
|
||||
SPOKEN_TEXT: Final = ("Hello world.", "¿Qué tal?")
|
||||
|
||||
|
||||
def _client_frames(*texts: str) -> dict[str, object]:
|
||||
return {"client_frames": [{"text": text} for text in texts]}
|
||||
|
||||
|
||||
def _spoken(*texts: str) -> tuple[dict[str, object], ...]:
|
||||
return tuple(
|
||||
dict(
|
||||
elevenlabs_tts_recorded_upstream_frame(
|
||||
{"audio": "AAAA", "alignment": {"chars": list(text)}, "isFinal": None}
|
||||
)
|
||||
)
|
||||
for text in texts
|
||||
)
|
||||
|
||||
|
||||
def _registry_cost(model_id: str, characters: int) -> float:
|
||||
per_character: Final = litellm.model_cost[f"elevenlabs/{model_id}"]["input_cost_per_character"]
|
||||
assert per_character > 0
|
||||
|
|
@ -62,13 +75,14 @@ def test_is_elevenlabs_tts_websocket_route(url_route, expected):
|
|||
assert ElevenLabsTextToSpeechPassthroughLoggingHandler.is_elevenlabs_tts_websocket_route(url_route) is expected
|
||||
|
||||
|
||||
def test_handler_bills_every_sent_character_at_the_registry_rate_and_names_the_model():
|
||||
def test_handler_bills_the_characters_elevenlabs_spoke_at_the_registry_rate_and_names_the_model():
|
||||
logging_obj = _logging_obj()
|
||||
characters = sum(len(text) for text in SESSION_TEXT)
|
||||
characters = sum(len(text) for text in SPOKEN_TEXT)
|
||||
upstream_frames = (*_spoken(*SPOKEN_TEXT), {"isFinal": True})
|
||||
|
||||
handler_result = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames(*SESSION_TEXT),
|
||||
upstream_frame_count=4,
|
||||
upstream_frames=upstream_frames,
|
||||
logging_obj=logging_obj,
|
||||
upstream_url=FLASH_URL,
|
||||
kwargs={"litellm_params": {"metadata": {}}},
|
||||
|
|
@ -88,7 +102,7 @@ def test_handler_bills_every_sent_character_at_the_registry_rate_and_names_the_m
|
|||
"model": "eleven_flash_v2_5",
|
||||
"characters": characters,
|
||||
"client_text_frames": len(SESSION_TEXT),
|
||||
"upstream_frames": 4,
|
||||
"upstream_frames": len(upstream_frames),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -96,13 +110,13 @@ def test_handler_bills_every_sent_character_at_the_registry_rate_and_names_the_m
|
|||
def test_handler_bills_the_model_the_session_ran_on():
|
||||
flash = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("x" * 1000),
|
||||
upstream_frame_count=1,
|
||||
upstream_frames=_spoken("x" * 1000),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=FLASH_URL,
|
||||
)
|
||||
multilingual = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("x" * 1000),
|
||||
upstream_frame_count=1,
|
||||
upstream_frames=_spoken("x" * 1000),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=FLASH_URL.replace("eleven_flash_v2_5", "eleven_multilingual_v2"),
|
||||
)
|
||||
|
|
@ -112,13 +126,16 @@ def test_handler_bills_the_model_the_session_ran_on():
|
|||
assert multilingual["kwargs"]["model"] == "eleven_multilingual_v2"
|
||||
|
||||
|
||||
def test_handler_charges_in_proportion_to_the_text_sent():
|
||||
def test_handler_charges_in_proportion_to_the_text_spoken():
|
||||
short = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("abcde"), upstream_frame_count=1, logging_obj=_logging_obj(), upstream_url=FLASH_URL
|
||||
request_body=_client_frames("abcde"),
|
||||
upstream_frames=_spoken("abcde"),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=FLASH_URL,
|
||||
)
|
||||
long = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("abcde", "fghij", "klmno"),
|
||||
upstream_frame_count=3,
|
||||
upstream_frames=_spoken("abcde", "fghij", "klmno"),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=FLASH_URL,
|
||||
)
|
||||
|
|
@ -132,19 +149,22 @@ def test_handler_records_no_cost_rather_than_a_stand_in_rate_for_an_unpriced_mod
|
|||
|
||||
handler_result = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("Hello "),
|
||||
upstream_frame_count=1,
|
||||
upstream_frames=_spoken("Hello"),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=unpriced_url,
|
||||
)
|
||||
|
||||
assert handler_result["kwargs"]["response_cost"] is None
|
||||
assert handler_result["kwargs"]["model"] == "eleven_not_in_registry"
|
||||
assert handler_result["result"]["response"]["characters"] == len("Hello ")
|
||||
assert handler_result["result"]["response"]["characters"] == len("Hello")
|
||||
|
||||
|
||||
def test_handler_with_no_client_text_bills_nothing():
|
||||
def test_handler_with_no_audio_back_bills_nothing():
|
||||
handler_result = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body={}, upstream_frame_count=0, logging_obj=_logging_obj(), upstream_url=FLASH_URL
|
||||
request_body=_client_frames(*SESSION_TEXT),
|
||||
upstream_frames=({"isFinal": True},),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=FLASH_URL,
|
||||
)
|
||||
|
||||
assert handler_result["kwargs"]["response_cost"] == 0
|
||||
|
|
@ -168,7 +188,7 @@ async def test_success_handler_dispatches_elevenlabs_tts_and_logs_character_spen
|
|||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
logging_obj = _logging_obj("call-el-e2e")
|
||||
request_body = _client_frames(*SESSION_TEXT)
|
||||
upstream_frames = [{"isFinal": None}, {"isFinal": None}, {"isFinal": True}]
|
||||
upstream_frames = [*_spoken(*SPOKEN_TEXT), {"isFinal": True}]
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", team_id="team-tts", user_id="user-1")
|
||||
metadata = {
|
||||
"user_api_key": user_api_key_dict.api_key,
|
||||
|
|
@ -209,14 +229,14 @@ async def test_success_handler_dispatches_elevenlabs_tts_and_logs_character_spen
|
|||
assert payload["model"] == "eleven_flash_v2_5"
|
||||
assert payload["custom_llm_provider"] == "elevenlabs"
|
||||
assert payload["response_cost"] == pytest.approx(
|
||||
_registry_cost("eleven_flash_v2_5", sum(len(text) for text in SESSION_TEXT))
|
||||
_registry_cost("eleven_flash_v2_5", sum(len(text) for text in SPOKEN_TEXT))
|
||||
)
|
||||
assert payload["metadata"]["user_api_key_team_id"] == "team-tts"
|
||||
assert payload["id"] == "call-el-e2e"
|
||||
assert payload["response"] == {
|
||||
"response": {
|
||||
"model": "eleven_flash_v2_5",
|
||||
"characters": sum(len(text) for text in SESSION_TEXT),
|
||||
"characters": sum(len(text) for text in SPOKEN_TEXT),
|
||||
"client_text_frames": len(SESSION_TEXT),
|
||||
"upstream_frames": len(upstream_frames),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -144,7 +144,9 @@ async def test_stream_input_relays_to_elevenlabs_with_only_the_proxy_key(monkeyp
|
|||
assert json.loads(decision.text) == {"text": "Hi ", "flush": True}
|
||||
assert decision.recorded == {"text": "Hi "}
|
||||
assert isinstance(call.frame_policy.client_text('{"text": "Hi ", "x": "\\ud800"}'), RefusedClientText)
|
||||
assert call.frame_policy.recorded_upstream_frame({"audio": "AAAA", "isFinal": None}) == {"isFinal": None}
|
||||
assert call.frame_policy.recorded_upstream_frame(
|
||||
{"audio": "AAAA", "alignment": {"chars": ["H", "i"]}, "isFinal": None}
|
||||
) == {"isFinal": None, "aligned_characters": 2}
|
||||
assert websocket.accepts == [None]
|
||||
assert websocket.closed is None
|
||||
|
||||
|
|
|
|||
|
|
@ -144,9 +144,7 @@ def test_is_multipart():
|
|||
),
|
||||
],
|
||||
)
|
||||
def test_get_endpoint_type_classifies_anthropic_messages_routes(
|
||||
url: str, expected_endpoint_type: EndpointType
|
||||
) -> None:
|
||||
def test_get_endpoint_type_classifies_anthropic_messages_routes(url: str, expected_endpoint_type: EndpointType) -> None:
|
||||
assert HttpPassThroughEndpointHelpers.get_endpoint_type(url) == expected_endpoint_type
|
||||
|
||||
|
||||
|
|
@ -3576,9 +3574,7 @@ async def test_pass_through_request_preserves_target_query_without_client_query(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_pass_through_without_configured_headers_preserves_no_authorization(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
async def test_config_pass_through_without_configured_headers_preserves_no_authorization(tmp_path, monkeypatch):
|
||||
proxy = await _boot_db_backed_proxy(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
|
|
@ -3645,9 +3641,7 @@ async def test_pass_through_rpm_limit_is_sequential_and_scoped_per_key(tmp_path,
|
|||
proxy = await _boot_db_backed_proxy(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
config_pass_through_endpoints=[
|
||||
{"path": "/rpm-rerank", "target": "http://upstream.test/rerank", "auth": True}
|
||||
],
|
||||
config_pass_through_endpoints=[{"path": "/rpm-rerank", "target": "http://upstream.test/rerank", "auth": True}],
|
||||
db_pass_through_endpoints=[],
|
||||
master_key=MASTER_KEY,
|
||||
)
|
||||
|
|
@ -4824,6 +4818,7 @@ from litellm.exceptions import (
|
|||
GuardrailRaisedException,
|
||||
)
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
||||
_PT_MODULE = "litellm.proxy.pass_through_endpoints.pass_through_endpoints"
|
||||
|
||||
|
||||
|
|
@ -7755,7 +7750,7 @@ async def test_websocket_passthrough_elevenlabs_policy_keeps_caller_keys_upstrea
|
|||
expected_request_body = {"client_frames": [{"text": " "}, {"text": "Hola, ¿qué tal? "}, {"text": ""}]}
|
||||
assert success_call["request_body"] == expected_request_body
|
||||
assert success_call["passthrough_logging_payload"]["request_body"] == expected_request_body
|
||||
assert success_call["response_body"] == [{"isFinal": None}, {"isFinal": True}]
|
||||
assert success_call["response_body"] == [{"isFinal": None, "aligned_characters": 1}, {"isFinal": True}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -7795,7 +7790,7 @@ async def test_websocket_passthrough_bills_the_session_when_the_client_drops_whi
|
|||
assert client_send_text.await_count == 2
|
||||
success_call = success_handler.call_args.kwargs
|
||||
assert success_call["request_body"] == {"client_frames": [{"text": " "}, {"text": "Hello there "}]}
|
||||
assert success_call["response_body"] == [{"isFinal": None}]
|
||||
assert success_call["response_body"] == [{"isFinal": None, "aligned_characters": 1}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8214,8 +8209,10 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata
|
|||
}
|
||||
)
|
||||
body: Final = {
|
||||
"model": "upstream-only-model", metadata_slot: {
|
||||
"model_group": "managed-model", "customer_label": "retained",
|
||||
"model": "upstream-only-model",
|
||||
metadata_slot: {
|
||||
"model_group": "managed-model",
|
||||
"customer_label": "retained",
|
||||
"user_api_key_team_model_max_budget": budget,
|
||||
},
|
||||
}
|
||||
|
|
@ -8223,48 +8220,84 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata
|
|||
assert await limiter.is_team_within_model_budget("shared-team", budget, None, "managed-model")
|
||||
start: Final = datetime.now()
|
||||
logging_obj: Final = LiteLLMLoggingObj(
|
||||
model="upstream-only-model", messages=[], stream=False, call_type="pass_through_endpoint",
|
||||
start_time=start, litellm_call_id="custom-budget", function_id="custom-budget", kwargs={},
|
||||
model="upstream-only-model",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="pass_through_endpoint",
|
||||
start_time=start,
|
||||
litellm_call_id="custom-budget",
|
||||
function_id="custom-budget",
|
||||
kwargs={},
|
||||
dynamic_async_success_callbacks=[limiter],
|
||||
)
|
||||
payload: Final = {
|
||||
"url": "https://upstream.test/echo", "request_body": body, "request_method": "POST", "cost_per_request": 0.25,
|
||||
"url": "https://upstream.test/echo",
|
||||
"request_body": body,
|
||||
"request_method": "POST",
|
||||
"cost_per_request": 0.25,
|
||||
}
|
||||
kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
|
||||
request=request, user_api_key_dict=auth, passthrough_logging_payload=payload, logging_obj=logging_obj,
|
||||
_parsed_body=body, litellm_call_id="custom-budget",
|
||||
request=request,
|
||||
user_api_key_dict=auth,
|
||||
passthrough_logging_payload=payload,
|
||||
logging_obj=logging_obj,
|
||||
_parsed_body=body,
|
||||
litellm_call_id="custom-budget",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="upstream-only-model", user="unknown", optional_params={},
|
||||
litellm_params=kwargs["litellm_params"], call_type="pass_through_endpoint",
|
||||
model="upstream-only-model",
|
||||
user="unknown",
|
||||
optional_params={},
|
||||
litellm_params=kwargs["litellm_params"],
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
response: Final = httpx.Response(
|
||||
200, request=httpx.Request("POST", "https://upstream.test/echo"), json={"ok": True},
|
||||
200,
|
||||
request=httpx.Request("POST", "https://upstream.test/echo"),
|
||||
json={"ok": True},
|
||||
)
|
||||
await PassThroughEndpointLogging().pass_through_async_success_handler(
|
||||
httpx_response=response, response_body={"ok": True}, request_body=body, logging_obj=logging_obj,
|
||||
url_route="https://upstream.test/echo", result=response.text, start_time=start, end_time=datetime.now(),
|
||||
cache_hit=False, **kwargs,
|
||||
httpx_response=response,
|
||||
response_body={"ok": True},
|
||||
request_body=body,
|
||||
logging_obj=logging_obj,
|
||||
url_route="https://upstream.test/echo",
|
||||
result=response.text,
|
||||
start_time=start,
|
||||
end_time=datetime.now(),
|
||||
cache_hit=False,
|
||||
**kwargs,
|
||||
)
|
||||
assert logging_obj.model_call_details["response_cost"] == 0.25
|
||||
assert await limiter.is_team_within_model_budget("shared-team", budget, None, "managed-model")
|
||||
metadata: Final = kwargs["litellm_params"]["metadata"]
|
||||
assert (metadata["model_group"], metadata["customer_label"]) == ("managed-model", "retained")
|
||||
assert metadata.keys().isdisjoint({
|
||||
"user_api_key_model_max_budget", "user_api_key_team_model_max_budget",
|
||||
"user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget",
|
||||
})
|
||||
assert metadata.keys().isdisjoint(
|
||||
{
|
||||
"user_api_key_model_max_budget",
|
||||
"user_api_key_team_model_max_budget",
|
||||
"user_api_key_user_model_max_budget",
|
||||
"user_api_key_end_user_model_max_budget",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata_slot", ["metadata", "litellm_metadata"])
|
||||
def test_builtin_passthrough_pins_model_group_to_the_resolved_model(metadata_slot: str) -> None:
|
||||
request: Final = Request({
|
||||
"type": "http", "method": "POST", "path": "/gemini/v1beta/models/gemini-2.5-flash:generateContent",
|
||||
"headers": [], "query_string": b"",
|
||||
})
|
||||
request: Final = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/gemini/v1beta/models/gemini-2.5-flash:generateContent",
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
}
|
||||
)
|
||||
kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
|
||||
request=request, user_api_key_dict=UserAPIKeyAuth(token="hash", user_id="u-1"),
|
||||
passthrough_logging_payload=MagicMock(), logging_obj=MagicMock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(token="hash", user_id="u-1"),
|
||||
passthrough_logging_payload=MagicMock(),
|
||||
logging_obj=MagicMock(),
|
||||
_parsed_body={"contents": [], metadata_slot: {"model_group": "unbounded-client-choice"}},
|
||||
)
|
||||
assert kwargs["litellm_params"]["metadata"]["model_group"] == "gemini-2.5-flash"
|
||||
|
|
@ -8630,7 +8663,9 @@ def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_bo
|
|||
the call to (LIT-1761: passthrough successes carried model_id="")."""
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.url = httpx.URL("http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent")
|
||||
mock_request.url = httpx.URL(
|
||||
"http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
|
||||
)
|
||||
mock_request.headers = Headers({})
|
||||
mock_request.scope = {}
|
||||
mock_request.state = SimpleNamespace(
|
||||
|
|
@ -9661,12 +9696,14 @@ async def _drain_logging_worker():
|
|||
await GLOBAL_LOGGING_WORKER.stop()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def _vcr_outcome_gate(request, vcr):
|
||||
install_live_call_probe(request, vcr)
|
||||
yield
|
||||
record_vcr_outcome(request, vcr)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate")
|
||||
def test_update_pass_through_route_updates_registry():
|
||||
"""
|
||||
|
|
@ -9756,9 +9793,7 @@ def mock_request():
|
|||
return self._dict.values()
|
||||
|
||||
class MockRequest:
|
||||
def __init__(
|
||||
self, headers=None, method="POST", request_body: Optional[dict] = None
|
||||
):
|
||||
def __init__(self, headers=None, method="POST", request_body: Optional[dict] = None):
|
||||
self.headers = headers or {}
|
||||
self.query_params = QueryParams()
|
||||
self.method = method
|
||||
|
|
@ -9866,9 +9901,9 @@ def test_pass_through_routes_support_all_methods():
|
|||
path = route.path
|
||||
methods = set(route.methods)
|
||||
allowed = PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES.get(path, expected_methods)
|
||||
assert (
|
||||
methods == allowed
|
||||
), f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}"
|
||||
assert methods == allowed, (
|
||||
f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}"
|
||||
)
|
||||
|
||||
check_router_methods(llm_router)
|
||||
|
||||
|
|
@ -9898,19 +9933,11 @@ def test_is_bedrock_agent_runtime_route():
|
|||
|
||||
# Test agent runtime endpoints (should return True)
|
||||
assert _is_bedrock_agent_runtime_route("/knowledgebases/kb-123/retrieve") is True
|
||||
assert (
|
||||
_is_bedrock_agent_runtime_route("/agents/knowledgebases/kb-123/retrieve")
|
||||
is True
|
||||
)
|
||||
assert _is_bedrock_agent_runtime_route("/agents/knowledgebases/kb-123/retrieve") is True
|
||||
|
||||
# Test regular bedrock runtime endpoints (should return False)
|
||||
assert (
|
||||
_is_bedrock_agent_runtime_route("/guardrail/test-id/version/1/apply") is False
|
||||
)
|
||||
assert (
|
||||
_is_bedrock_agent_runtime_route("/model/cohere.command-r-v1:0/converse")
|
||||
is False
|
||||
)
|
||||
assert _is_bedrock_agent_runtime_route("/guardrail/test-id/version/1/apply") is False
|
||||
assert _is_bedrock_agent_runtime_route("/model/cohere.command-r-v1:0/converse") is False
|
||||
assert _is_bedrock_agent_runtime_route("/some/random/endpoint") is False
|
||||
|
||||
|
||||
|
|
@ -9977,9 +10004,7 @@ def test_custom_pricing_used_in_cost_calculation():
|
|||
|
||||
# Verify custom cost is different from standard cost (unless prices happen to match)
|
||||
# This confirms custom pricing is actually being applied
|
||||
assert (
|
||||
custom_cost != standard_cost
|
||||
), "Custom pricing should produce different cost than standard pricing"
|
||||
assert custom_cost != standard_cost, "Custom pricing should produce different cost than standard pricing"
|
||||
|
||||
# Test 3: Custom pricing with cache_read_input_token_cost and input_cost_per_token_batches
|
||||
# This specifically tests the parameters that were causing the original issue
|
||||
|
|
@ -10519,6 +10544,7 @@ def test_init_kwargs_no_authenticated_identity_field_is_client_settable(mock_req
|
|||
survived: Final = {key: metadata.get(key) for key in expected if metadata.get(key) != expected[key]}
|
||||
assert survived == {}, f"client-supplied values survived for: {sorted(survived)}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_direct_passthrough_keeps_mapped_baggage_tags_outside_upstream_body():
|
||||
from litellm.proxy.common_utils.http_parsing_utils import read_request_body
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue