mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
fix(speech)!: correctly propagate spend logging and usage for Gemini TTS bridge (#44551)
* fix(speech): correctly propagate spend logging and usage for Gemini TTS bridge * fix(speech): call Gemini once per TTS request and log the bridge's token usage aspeech re-ran speech whenever the first result was not a coroutine, so the speech-to-completion bridge sent every Gemini TTS request upstream twice. It now returns the sync result. The bridge's inner completion runs as an internal sub call, so only the speech call fires success callbacks, and the completion's usage rides on the audio response so the spend log and usage_object carry real token counts. Replaces the separate inner Logging object and the cost_calculator fallback, which never reached the spend row. * fix(speech): flag a cached Gemini TTS reply as a cache hit so it logs at no charge * fix(speech): price a cached Gemini TTS reply's savings from its token usage and keep its cache key * test(integration): cover Gemini TTS speech billing, logging, cache hits and deployment TPM * test(integration): start the speech TPM cell in a minute's first 10 s so the two requests never cross a minute --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
6f66b7a201
commit
f4347d4f50
8 changed files with 497 additions and 16 deletions
|
|
@ -1103,6 +1103,8 @@ def get_usage_object(
|
|||
|
||||
if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
|
||||
return None if completion_response.usage is None else _decisions_usage(completion_response.usage)
|
||||
if isinstance(completion_response, HttpxBinaryResponseContent):
|
||||
return completion_response.usage
|
||||
if usage_obj is None:
|
||||
return None
|
||||
if isinstance(usage_obj, Usage):
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ from typing import TYPE_CHECKING, Final
|
|||
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from litellm._internal_context import internal_sub_call
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import LiteLLMLoggingObj
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
|
|
@ -111,9 +113,10 @@ class SpeechToCompletionBridgeHandler:
|
|||
voice=voice,
|
||||
)
|
||||
|
||||
result: Final = completion(
|
||||
**request_data,
|
||||
)
|
||||
with internal_sub_call():
|
||||
result: Final = completion(
|
||||
**request_data,
|
||||
)
|
||||
|
||||
requested_response_format: Final = optional_params.get("response_format")
|
||||
if isinstance(result, ModelResponse):
|
||||
|
|
|
|||
|
|
@ -5,11 +5,12 @@ from typing import TYPE_CHECKING, Final, cast
|
|||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS
|
||||
from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_param
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.llms.openai import ChatCompletionUserMessage, HttpxBinaryResponseContent
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
|
||||
def _completion_response_cost(model_response: "ModelResponse") -> float | None:
|
||||
|
|
@ -20,6 +21,13 @@ def _completion_response_cost(model_response: "ModelResponse") -> float | None:
|
|||
return response_cost if isinstance(response_cost, float) else None
|
||||
|
||||
|
||||
def _completion_usage(model_response: "ModelResponse") -> "Usage | None":
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
usage: Final = getattr(model_response, "usage", None)
|
||||
return usage if isinstance(usage, Usage) else None
|
||||
|
||||
|
||||
GEMINI_TTS_CHAT_AUDIO_FORMAT: Final = "pcm16"
|
||||
GEMINI_TTS_RAW_RESPONSE_FORMAT: Final = "pcm"
|
||||
GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: Final = frozenset({"wav", GEMINI_TTS_RAW_RESPONSE_FORMAT})
|
||||
|
|
@ -179,4 +187,9 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
)
|
||||
binary_response: Final = HttpxBinaryResponseContent(response)
|
||||
binary_response.set_response_cost(_completion_response_cost(model_response))
|
||||
binary_response.usage = _completion_usage(model_response)
|
||||
cached_hidden_params: Final = get_hidden_params(model_response) if served_from_cache(model_response) else None
|
||||
if cached_hidden_params is not None:
|
||||
set_hidden_param(binary_response, "cache_hit", True)
|
||||
set_hidden_param(binary_response, "cache_key", cached_hidden_params.get("cache_key"))
|
||||
return binary_response
|
||||
|
|
|
|||
|
|
@ -8500,14 +8500,8 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent:
|
|||
api_base=api_base_hint,
|
||||
)
|
||||
|
||||
# Await normally
|
||||
init_response: Final = await loop.run_in_executor(None, func_with_context)
|
||||
if asyncio.iscoroutine(init_response):
|
||||
response = await init_response
|
||||
else:
|
||||
# Call the synchronous function using run_in_executor
|
||||
response = await loop.run_in_executor(None, func_with_context)
|
||||
return response
|
||||
return await init_response if asyncio.iscoroutine(init_response) else init_response
|
||||
except Exception as e:
|
||||
custom_llm_provider = custom_llm_provider or "openai"
|
||||
raise exception_type(
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import builtins
|
|||
from collections.abc import Iterable, Mapping
|
||||
from enum import Enum
|
||||
from os import PathLike
|
||||
from typing import IO, Any, Final, Generic, Literal, Optional, TypeAlias, Union
|
||||
from typing import IO, TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, Union
|
||||
|
||||
import httpx
|
||||
from openai import Omit
|
||||
|
|
@ -96,6 +96,9 @@ from litellm.types.responses.main import (
|
|||
|
||||
from .base import CachedTokensDetails
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
FileContent = IO[bytes] | bytes | PathLike
|
||||
|
||||
FileTypes = (
|
||||
|
|
@ -119,6 +122,7 @@ class BinaryResponseSummary(TypedDict):
|
|||
object: ReadOnly[Literal["binary"]]
|
||||
content_type: ReadOnly[str | None]
|
||||
num_bytes: ReadOnly[int]
|
||||
usage: NotRequired[ReadOnly[Mapping[str, builtins.object]]]
|
||||
|
||||
|
||||
_ResponseT = TypeVar("_ResponseT", bound=httpx.Response | SDKResponse, default=httpx.Response)
|
||||
|
|
@ -126,6 +130,7 @@ _ResponseT = TypeVar("_ResponseT", bound=httpx.Response | SDKResponse, default=h
|
|||
|
||||
class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT]):
|
||||
_hidden_params: dict
|
||||
usage: "Usage | None"
|
||||
response: _ResponseT # pyright: ignore[reportIncompatibleVariableOverride] # SDK accepts both backends at runtime
|
||||
|
||||
@property
|
||||
|
|
@ -139,6 +144,7 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT
|
|||
def __init__(self, response: _ResponseT) -> None:
|
||||
super().__init__(response) # pyright: ignore[reportArgumentType] # SDK accepts both backends at runtime
|
||||
self._hidden_params = {}
|
||||
self.usage = None
|
||||
|
||||
@property
|
||||
def content_type(self) -> str | None:
|
||||
|
|
@ -146,11 +152,12 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT
|
|||
return headers.get("content-type")
|
||||
|
||||
def logging_summary(self) -> BinaryResponseSummary:
|
||||
return {
|
||||
summary: Final[BinaryResponseSummary] = {
|
||||
"object": "binary",
|
||||
"content_type": self.content_type,
|
||||
"num_bytes": self._num_bytes(),
|
||||
}
|
||||
return summary if self.usage is None else {**summary, "usage": self.usage.model_dump()}
|
||||
|
||||
def _num_bytes(self) -> int:
|
||||
try:
|
||||
|
|
|
|||
69
tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py
Normal file
69
tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import eventually
|
||||
from integration._support.wire import wire_server
|
||||
from integration.providers._gemini_tts_voice_support import (
|
||||
AUDIO_TOKENS,
|
||||
BACKEND,
|
||||
GEMINI_API_KEY,
|
||||
PCM,
|
||||
PROMPT_TOKENS,
|
||||
gemini_peer,
|
||||
marker,
|
||||
received_text,
|
||||
wav_payload,
|
||||
)
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
BINARY_RESPONSE: Final = "HttpxBinaryResponseContent"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Logged:
|
||||
call_type: object
|
||||
response_type: str
|
||||
tokens: tuple[object, object]
|
||||
|
||||
|
||||
class _Recorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._lock: Final = threading.Lock()
|
||||
self.events: tuple[_Logged, ...] = ()
|
||||
|
||||
def log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
payload: Final = kwargs.get("standard_logging_object")
|
||||
logged: Final = payload if isinstance(payload, dict) else {}
|
||||
event: Final = _Logged(
|
||||
call_type=kwargs.get("call_type"),
|
||||
response_type=type(response_obj).__name__,
|
||||
tokens=(logged.get("prompt_tokens"), logged.get("completion_tokens")),
|
||||
)
|
||||
with self._lock:
|
||||
self.events = (*self.events, event)
|
||||
|
||||
|
||||
def test_sync_speech_logs_one_success_event_carrying_the_reply_token_usage(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
recorder: Final = _Recorder()
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
with wire_server(gemini_peer) as wire:
|
||||
monkeypatch.setenv("GEMINI_API_BASE", wire.url)
|
||||
text: Final = marker()
|
||||
audio: Final = litellm.speech(model=f"gemini/{BACKEND}", input=text, voice="Kore", api_key=GEMINI_API_KEY)
|
||||
assert wav_payload(audio.content) == PCM
|
||||
assert sum(1 for request in wire.drain() if received_text(request) == text) == 1
|
||||
events: Final = eventually(
|
||||
lambda: recorder.events,
|
||||
lambda found: any(event.response_type == BINARY_RESPONSE for event in found),
|
||||
seconds=30,
|
||||
)
|
||||
assert events == (_Logged("speech", BINARY_RESPONSE, (PROMPT_TOKENS, AUDIO_TOKENS)),)
|
||||
334
tests/integration/spend/test_gemini_tts_speech_billing.py
Normal file
334
tests/integration/spend/test_gemini_tts_speech_billing.py
Normal file
|
|
@ -0,0 +1,334 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import JSON_OBJECT, eventually, gateway_from_environment, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Wire, wire_server
|
||||
from integration.providers._gemini_tts_voice_support import (
|
||||
AUDIO_TOKENS,
|
||||
BACKEND,
|
||||
GEMINI_API_KEY,
|
||||
PCM,
|
||||
PROMPT_TOKENS,
|
||||
REJECTED_PREFIX,
|
||||
VERTEX_LOCATION,
|
||||
VERTEX_PROJECT,
|
||||
gemini_peer,
|
||||
marker,
|
||||
received_text,
|
||||
wav_payload,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
Provider = Literal["gemini", "vertex_ai"]
|
||||
MODELS: Final[Mapping[Provider, str]] = {"gemini": "tts-billing-gemini", "vertex_ai": "tts-billing-vertex"}
|
||||
TPM_MODEL: Final = "tts-billing-tpm"
|
||||
TPM_LIMIT: Final = PROMPT_TOKENS + AUDIO_TOKENS
|
||||
BINARY_RESPONSE: Final = "HttpxBinaryResponseContent"
|
||||
CALLBACK_LOG: Final = "TTS_BILLING_CALLBACK_LOG"
|
||||
CALLBACK_MODULE: Final = """
|
||||
import json
|
||||
import os
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
def _event(kind, kwargs, response_obj):
|
||||
payload = kwargs.get("standard_logging_object") or {}
|
||||
metadata = payload.get("metadata") or {}
|
||||
response = payload.get("response")
|
||||
usage = response.get("usage") if isinstance(response, dict) else None
|
||||
return {
|
||||
"kind": kind,
|
||||
"key": metadata.get("user_api_key_hash"),
|
||||
"call_type": kwargs.get("call_type"),
|
||||
"response_type": type(response_obj).__name__,
|
||||
"response_cost": payload.get("response_cost"),
|
||||
"prompt_tokens": payload.get("prompt_tokens"),
|
||||
"completion_tokens": payload.get("completion_tokens"),
|
||||
"cache_hit": payload.get("cache_hit"),
|
||||
"saved_cache_cost": payload.get("saved_cache_cost"),
|
||||
"cache_key": payload.get("cache_key"),
|
||||
"logged_usage_total": usage.get("total_tokens") if isinstance(usage, dict) else None,
|
||||
}
|
||||
|
||||
|
||||
class SpeechEvents(CustomLogger):
|
||||
def _write(self, line):
|
||||
with open(os.environ["TTS_BILLING_CALLBACK_LOG"], "a") as sink:
|
||||
sink.write(json.dumps(line) + "\\n")
|
||||
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self._write(_event("sync_success", kwargs, response_obj))
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self._write(_event("async_success", kwargs, response_obj))
|
||||
|
||||
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self._write(_event("sync_failure", kwargs, response_obj))
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self._write(_event("async_failure", kwargs, response_obj))
|
||||
|
||||
|
||||
speech_events = SpeechEvents()
|
||||
"""
|
||||
SPEND_ROWS_SQL: Final = (
|
||||
'SELECT call_type, spend, prompt_tokens, completion_tokens, cache_hit FROM "LiteLLM_SpendLogs"'
|
||||
' WHERE api_key = %s ORDER BY "startTime"'
|
||||
)
|
||||
KEY_SPEND_SQL: Final = 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s'
|
||||
|
||||
Event = dict[str, JsonValue]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BillingProxy:
|
||||
proxy: OwnedProxy
|
||||
callback_log: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Sent:
|
||||
status: int
|
||||
body: bytes
|
||||
cache_key: str
|
||||
|
||||
|
||||
def _hashed(key: str) -> str:
|
||||
return sha256(key.encode()).hexdigest()
|
||||
|
||||
|
||||
def _model_list(token_url: str) -> list[JsonValue]:
|
||||
return [
|
||||
{"model_name": MODELS["gemini"], "litellm_params": {"model": f"gemini/{BACKEND}", "api_key": GEMINI_API_KEY}},
|
||||
{
|
||||
"model_name": TPM_MODEL,
|
||||
"litellm_params": {"model": f"gemini/{BACKEND}", "api_key": GEMINI_API_KEY, "tpm": TPM_LIMIT},
|
||||
},
|
||||
{
|
||||
"model_name": MODELS["vertex_ai"],
|
||||
"litellm_params": {
|
||||
"model": f"vertex_ai/{BACKEND}",
|
||||
"vertex_project": VERTEX_PROJECT,
|
||||
"vertex_location": VERTEX_LOCATION,
|
||||
"vertex_credentials": service_account_json(VERTEX_PROJECT, token_url.rstrip("/")),
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _config(directory: Path, token_url: str, *, cached: bool) -> Path:
|
||||
stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
settings: Final = object_value(stock["litellm_settings"])
|
||||
kept: Final = {name: value for name, value in settings.items() if name not in {"cache", "cache_params"}}
|
||||
cache: Final[dict[str, JsonValue]] = (
|
||||
{"cache": True, "cache_params": {**object_value(settings["cache_params"]), "namespace": marker()}}
|
||||
if cached
|
||||
else {"cache": False}
|
||||
)
|
||||
routing: Final[dict[str, JsonValue]] = (
|
||||
{}
|
||||
if cached
|
||||
else {
|
||||
"routing_strategy": "usage-based-routing-v2",
|
||||
"redis_host": os.environ["REDIS_HOST"],
|
||||
"redis_port": int(os.environ["REDIS_PORT"]),
|
||||
}
|
||||
)
|
||||
(directory / "tts_billing_callbacks.py").write_text(CALLBACK_MODULE)
|
||||
path: Final = directory / "tts-billing.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**stock,
|
||||
"model_list": _model_list(token_url),
|
||||
"litellm_settings": {**kept, **cache, "callbacks": "tts_billing_callbacks.speech_events"},
|
||||
"router_settings": {**object_value(stock["router_settings"]), "num_retries": 0, **routing},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def _billing_proxy(wire: Wire, directory: Path, *, cached: bool) -> Iterator[BillingProxy]:
|
||||
callback_log: Final = directory / "callback_events.jsonl"
|
||||
callback_log.touch()
|
||||
with gateway_from_environment() as gateway:
|
||||
config: Final = _config(directory, gateway.upstream_url, cached=cached)
|
||||
overrides: Final = {"GEMINI_API_BASE": wire.url, "VERTEXAI_API_BASE": wire.url, CALLBACK_LOG: str(callback_log)}
|
||||
with owned_proxy_process(gateway, directory, overrides, config=config, workers=2) as proxy:
|
||||
yield BillingProxy(proxy, callback_log)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def tts_wire() -> Iterator[Wire]:
|
||||
with wire_server(gemini_peer) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def uncached(tts_wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[BillingProxy]:
|
||||
yield from _billing_proxy(tts_wire, tmp_path_factory.mktemp("tts-billing-uncached"), cached=False)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cached(tts_wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[BillingProxy]:
|
||||
yield from _billing_proxy(tts_wire, tmp_path_factory.mktemp("tts-billing-cached"), cached=True)
|
||||
|
||||
|
||||
def _speech(billing: BillingProxy, model: str, text: str, key: str, voice: str = "Kore") -> Sent:
|
||||
response: Final = billing.proxy.gateway.request(
|
||||
"POST", "/v1/audio/speech", {"model": model, "input": text, "voice": voice}, key=key
|
||||
)
|
||||
return Sent(response.status_code, response.content, response.headers.get("x-litellm-cache-key", ""))
|
||||
|
||||
|
||||
def _upstream_calls(wire: Wire, text: str) -> int:
|
||||
return sum(1 for request in wire.drain() if received_text(request) == text)
|
||||
|
||||
|
||||
def _spend_rows(key: str, count: int) -> list[dict[str, JsonValue]]:
|
||||
return eventually(lambda: read_rows(SPEND_ROWS_SQL, (_hashed(key),)), lambda rows: len(rows) == count, seconds=70)
|
||||
|
||||
|
||||
def _key_spend(key: str, expected: float) -> JsonValue:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(KEY_SPEND_SQL, (_hashed(key),)),
|
||||
lambda found: len(found) == 1 and found[0]["spend"] == pytest.approx(expected),
|
||||
seconds=70,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
return rows[0]["spend"]
|
||||
|
||||
|
||||
def _chat_spend(billing: BillingProxy, model: str, key: str) -> float:
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": marker()}],
|
||||
"modalities": ["audio"],
|
||||
"audio": {"voice": "Kore", "format": "pcm16"},
|
||||
}
|
||||
response: Final = billing.proxy.gateway.request("POST", "/v1/chat/completions", body, key=key)
|
||||
assert response.status_code == 200, response.text
|
||||
row: Final = _spend_rows(key, 1)[0]
|
||||
assert (row["prompt_tokens"], row["completion_tokens"]) == (PROMPT_TOKENS, AUDIO_TOKENS), row
|
||||
spend: Final = row["spend"]
|
||||
assert isinstance(spend, float) and spend > 0, row
|
||||
return spend
|
||||
|
||||
|
||||
def _events(billing: BillingProxy, key: str) -> tuple[Event, ...]:
|
||||
complete_lines: Final = billing.callback_log.read_text().split("\n")[:-1]
|
||||
events: Final = tuple(JSON_OBJECT.validate_json(line) for line in complete_lines)
|
||||
return tuple(event for event in events if event["key"] == _hashed(key))
|
||||
|
||||
|
||||
def _speech_events(billing: BillingProxy, key: str, count: int) -> tuple[Event, ...]:
|
||||
return eventually(
|
||||
lambda: _events(billing, key),
|
||||
lambda found: sum(1 for event in found if event["response_type"] == BINARY_RESPONSE) >= count,
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def _summary(event: Event) -> tuple[JsonValue, ...]:
|
||||
return (event["call_type"], event["response_type"], event["prompt_tokens"], event["completion_tokens"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"])
|
||||
def test_speech_calls_google_once_and_bills_and_logs_its_token_usage_once(
|
||||
uncached: BillingProxy, tts_wire: Wire, provider: Provider
|
||||
) -> None:
|
||||
with uncached.proxy.gateway.scenario() as scenario:
|
||||
speech_key: Final = scenario.key()
|
||||
text: Final = marker()
|
||||
sent: Final = _speech(uncached, MODELS[provider], text, speech_key)
|
||||
assert sent.status == 200, sent.body[:300]
|
||||
assert wav_payload(sent.body) == PCM
|
||||
assert _upstream_calls(tts_wire, text) == 1
|
||||
chat_spend: Final = _chat_spend(uncached, MODELS[provider], scenario.key())
|
||||
row: Final = _spend_rows(speech_key, 1)[0]
|
||||
assert row["call_type"] == "aspeech", row
|
||||
assert (row["prompt_tokens"], row["completion_tokens"]) == (PROMPT_TOKENS, AUDIO_TOKENS), row
|
||||
assert row["spend"] == pytest.approx(chat_spend), row
|
||||
assert _key_spend(speech_key, chat_spend) == pytest.approx(chat_spend)
|
||||
events: Final = _speech_events(uncached, speech_key, 1)
|
||||
assert [_summary(event) for event in events] == [("aspeech", BINARY_RESPONSE, PROMPT_TOKENS, AUDIO_TOKENS)]
|
||||
assert events[0]["response_cost"] == pytest.approx(chat_spend), events
|
||||
assert events[0]["logged_usage_total"] == PROMPT_TOKENS + AUDIO_TOKENS, events
|
||||
|
||||
|
||||
def test_cached_speech_repeat_is_a_free_cache_hit_after_a_billed_first_call(
|
||||
cached: BillingProxy, tts_wire: Wire
|
||||
) -> None:
|
||||
model: Final = MODELS["gemini"]
|
||||
with cached.proxy.gateway.scenario() as scenario:
|
||||
speech_key: Final = scenario.key()
|
||||
text: Final = marker()
|
||||
first: Final = _speech(cached, model, text, speech_key)
|
||||
repeat: Final = _speech(cached, model, text, speech_key)
|
||||
assert (first.status, repeat.status) == (200, 200), (first.body[:300], repeat.body[:300])
|
||||
assert wav_payload(first.body) == PCM and repeat.body == first.body
|
||||
assert _upstream_calls(tts_wire, text) == 1
|
||||
assert repeat.cache_key, "the cached repeat carried no x-litellm-cache-key header"
|
||||
chat_spend: Final = _chat_spend(cached, model, scenario.key())
|
||||
billed, hit = _spend_rows(speech_key, 2)
|
||||
assert billed["cache_hit"] != "True" and billed["spend"] == pytest.approx(chat_spend), billed
|
||||
assert (billed["prompt_tokens"], billed["completion_tokens"]) == (PROMPT_TOKENS, AUDIO_TOKENS), billed
|
||||
assert (hit["cache_hit"], hit["spend"]) == ("True", 0.0), hit
|
||||
assert _key_spend(speech_key, chat_spend) == pytest.approx(chat_spend)
|
||||
events: Final = _speech_events(cached, speech_key, 2)
|
||||
assert len(events) == 2, events
|
||||
hits: Final = tuple(event for event in events if event["cache_hit"] is True)
|
||||
assert len(hits) == 1, events
|
||||
assert (hits[0]["response_cost"], hits[0]["cache_key"]) == (0.0, repeat.cache_key), hits
|
||||
assert hits[0]["saved_cache_cost"] == pytest.approx(chat_spend), hits
|
||||
paid: Final = tuple(event for event in events if event["cache_hit"] is not True)
|
||||
assert len(paid) == 1 and paid[0]["response_cost"] == pytest.approx(chat_spend), events
|
||||
|
||||
|
||||
def test_speech_rejected_by_google_reaches_the_caller_and_logs_one_failure_after_one_call(
|
||||
uncached: BillingProxy, tts_wire: Wire
|
||||
) -> None:
|
||||
with uncached.proxy.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
text: Final = marker()
|
||||
sent: Final = _speech(uncached, MODELS["gemini"], text, key, voice="Custom-Voice")
|
||||
assert sent.status == 400, sent.body[:300]
|
||||
assert f"{REJECTED_PREFIX}Custom-Voice" in sent.body.decode(), sent.body[:300]
|
||||
assert _upstream_calls(tts_wire, text) == 1
|
||||
events: Final = eventually(lambda: _events(uncached, key), lambda found: len(found) >= 2, seconds=70)
|
||||
logged: Final = sorted((str(event["kind"]), event["call_type"]) for event in events)
|
||||
assert logged == [("async_failure", "aspeech"), ("sync_failure", "aspeech")], events
|
||||
|
||||
|
||||
def test_a_speech_reply_fills_the_deployment_tpm_limit_so_the_next_request_gets_429(
|
||||
uncached: BillingProxy, tts_wire: Wire
|
||||
) -> None:
|
||||
eventually(lambda: datetime.now(UTC), lambda now: now.second < 10, seconds=65)
|
||||
with uncached.proxy.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
first_text: Final = marker()
|
||||
first: Final = _speech(uncached, TPM_MODEL, first_text, key)
|
||||
assert first.status == 200, first.body[:300]
|
||||
_spend_rows(key, 1)
|
||||
second_text: Final = marker()
|
||||
second: Final = _speech(uncached, TPM_MODEL, second_text, key)
|
||||
assert second.status == 429, second.body[:300]
|
||||
assert b"No deployments available" in second.body, second.body[:300]
|
||||
received: Final = tuple(received_text(request) for request in tts_wire.drain())
|
||||
assert (received.count(first_text), received.count(second_text)) == (1, 0)
|
||||
|
|
@ -4425,18 +4425,27 @@ class _RecordedSpeechSuccess:
|
|||
spend_metadata: Mapping[str, object]
|
||||
response_cost: float | None
|
||||
logged_response_cost: float | None
|
||||
logged_tokens: tuple[object, object]
|
||||
logged_cache_hit: object
|
||||
logged_saved_cache_cost: object
|
||||
logged_cache_key: object
|
||||
|
||||
|
||||
def _record_speech_success(payload: dict[str, object]) -> _RecordedSpeechSuccess:
|
||||
call_type: Final = payload.get("call_type")
|
||||
response_cost: Final = payload.get("response_cost")
|
||||
logging_payload: Final = payload.get("standard_logging_object")
|
||||
logged_cost: Final = logging_payload.get("response_cost") if isinstance(logging_payload, dict) else None
|
||||
logged: Final[Mapping[str, object]] = logging_payload if isinstance(logging_payload, dict) else {}
|
||||
logged_cost: Final = logged.get("response_cost")
|
||||
return _RecordedSpeechSuccess(
|
||||
call_type=call_type if isinstance(call_type, str) else None,
|
||||
spend_metadata=get_litellm_metadata_from_kwargs(payload),
|
||||
response_cost=response_cost if isinstance(response_cost, float) else None,
|
||||
logged_response_cost=logged_cost if isinstance(logged_cost, float) else None,
|
||||
logged_tokens=(logged.get("prompt_tokens"), logged.get("completion_tokens")),
|
||||
logged_cache_hit=logged.get("cache_hit"),
|
||||
logged_saved_cache_cost=logged.get("saved_cache_cost"),
|
||||
logged_cache_key=logged.get("cache_key"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -4445,6 +4454,11 @@ class _SuccessEventRecorder(CustomLogger):
|
|||
super().__init__()
|
||||
self.events: list[_RecordedSpeechSuccess] = [] # mutable-ok: test recorder of success-callback events
|
||||
|
||||
def log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
self.events.append(_record_speech_success(kwargs))
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
|
|
@ -4459,6 +4473,16 @@ async def _wait_for_success_event(recorder: _SuccessEventRecorder, call_type: st
|
|||
pytest.fail(f"no {call_type} success event; got {[e.call_type for e in recorder.events]}")
|
||||
|
||||
|
||||
async def _wait_for_success_events(
|
||||
recorder: _SuccessEventRecorder, call_type: str, count: int
|
||||
) -> tuple[_RecordedSpeechSuccess, ...]:
|
||||
for _ in range(100):
|
||||
if len(events := tuple(e for e in recorder.events if e.call_type == call_type)) >= count:
|
||||
return events
|
||||
await asyncio.sleep(0.05)
|
||||
pytest.fail(f"expected {count} {call_type} success events; got {[e.call_type for e in recorder.events]}")
|
||||
|
||||
|
||||
def _gemini_tts_generate_content_response() -> dict[str, object]:
|
||||
return {
|
||||
"candidates": [
|
||||
|
|
@ -4490,7 +4514,7 @@ def _gemini_tts_generate_content_response() -> dict[str, object]:
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking(
|
||||
async def test_aspeech_gemini_bridge_bills_one_call_with_proxy_metadata_and_completion_tokens(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
|
@ -4510,9 +4534,11 @@ async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking(
|
|||
metadata={"user_api_key": "hashed-virtual-key", "user_api_key_user_id": "user-1"},
|
||||
)
|
||||
|
||||
assert mock_route.called
|
||||
assert mock_route.call_count == 1
|
||||
assert mock_route.calls.last.request.headers["x-goog-api-key"] == "fake-gemini-key"
|
||||
speech_event: Final = await _wait_for_success_event(recorder, call_type="aspeech")
|
||||
assert recorder.events == [speech_event]
|
||||
assert speech_event.logged_tokens == (5, 60)
|
||||
assert speech_event.spend_metadata["user_api_key"] == "hashed-virtual-key"
|
||||
assert speech_event.spend_metadata["user_api_key_user_id"] == "user-1"
|
||||
expected_prompt_cost, expected_completion_cost = litellm.cost_per_token(
|
||||
|
|
@ -4525,6 +4551,39 @@ async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking(
|
|||
assert speech_event.logged_response_cost == pytest.approx(expected_cost)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aspeech_gemini_bridge_logs_a_cached_reply_as_a_free_cache_hit(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.delenv("GEMINI_API_KEY", raising=False)
|
||||
monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
|
||||
recorder: Final = _SuccessEventRecorder()
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
|
||||
mock_route: Final = respx_mock.post(
|
||||
url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*"
|
||||
).mock(return_value=httpx.Response(200, json=_gemini_tts_generate_content_response()))
|
||||
|
||||
for _ in range(2):
|
||||
await litellm.aspeech(
|
||||
model="gemini/gemini-2.5-flash-preview-tts",
|
||||
input="cache hit check",
|
||||
voice="Kore",
|
||||
api_key="fake-gemini-key",
|
||||
)
|
||||
|
||||
first, cached = await _wait_for_success_events(recorder, call_type="aspeech", count=2)
|
||||
assert mock_route.call_count == 1
|
||||
assert len(recorder.events) == 2
|
||||
assert first.logged_cache_hit is not True
|
||||
assert first.logged_response_cost is not None and first.logged_response_cost > 0
|
||||
assert cached.logged_cache_hit is True
|
||||
assert cached.logged_response_cost == 0.0
|
||||
assert cached.logged_saved_cache_cost == pytest.approx(first.logged_response_cost)
|
||||
assert isinstance(cached.logged_cache_key, str) and cached.logged_cache_key
|
||||
|
||||
|
||||
def _stream_builder_text_chunk(model: str, content: str, finish_reason: str | None = None) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-cost",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue