diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4fcd3bde843..8c6e2fead2d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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): diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index 6c33621ec89..00d112773f1 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -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): diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 2ed140c0208..9c1e5ae7a29 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -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 diff --git a/litellm/main.py b/litellm/main.py index 38591441e01..9430af34401 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 8cc9e912e80..8dbc0d97d54 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -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: diff --git a/tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py b/tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py new file mode 100644 index 00000000000..3df10da8fda --- /dev/null +++ b/tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py @@ -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)),) diff --git a/tests/integration/spend/test_gemini_tts_speech_billing.py b/tests/integration/spend/test_gemini_tts_speech_billing.py new file mode 100644 index 00000000000..f1bf3de91b3 --- /dev/null +++ b/tests/integration/spend/test_gemini_tts_speech_billing.py @@ -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) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 63f440d4aba..aa15893e395 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -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",