fix(speech): keep proxy metadata and completion cost through the TTS completion bridge

This commit is contained in:
mateo-berri 2026-08-26 15:12:41 -07:00
parent f57e4b812c
commit 3418d7baf9
5 changed files with 103 additions and 3 deletions

View file

@ -8,6 +8,14 @@ if TYPE_CHECKING:
from litellm.types.utils import ModelResponse
def _completion_response_cost(model_response: "ModelResponse") -> float | None:
hidden_params: Final = getattr(model_response, "_hidden_params", None)
if not isinstance(hidden_params, dict):
return None
response_cost: Final = hidden_params.get("response_cost")
return response_cost if isinstance(response_cost, float) else None
class SpeechToCompletionBridgeTransformationHandler:
def transform_request(
self,
@ -123,4 +131,6 @@ class SpeechToCompletionBridgeTransformationHandler:
# Create an httpx.Response object
response: Final = httpx.Response(status_code=200, content=binary_data, headers=headers)
return HttpxBinaryResponseContent(response)
binary_response: Final = HttpxBinaryResponseContent(response)
binary_response.set_response_cost(_completion_response_cost(model_response))
return binary_response

View file

@ -1590,7 +1590,7 @@ class Logging(LiteLLMLoggingBaseClass):
if transformed_result is not None:
result = transformed_result
if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"):
if isinstance(result, (BaseModel, HttpxBinaryResponseContent)) and hasattr(result, "_hidden_params"):
hidden_params: Final = getattr(result, "_hidden_params", {})
if (
"response_cost" in hidden_params and hidden_params["response_cost"] is not None

View file

@ -8013,7 +8013,7 @@ def speech(
if max_retries is None:
max_retries = litellm.num_retries or openai.DEFAULT_MAX_RETRIES
litellm_params_dict: Final = get_litellm_params(**kwargs)
litellm_params_dict: Final = get_litellm_params(metadata=metadata, **kwargs)
# Get provider-specific text-to-speech config and map parameters
text_to_speech_provider_config = ProviderConfigManager.get_provider_text_to_speech_config(

View file

@ -109,6 +109,9 @@ EmbeddingInput = str | list[str]
class HttpxBinaryResponseContent(_HttpxBinaryResponseContent):
_hidden_params: dict = {}
def set_response_cost(self, response_cost: float | None) -> None:
self._hidden_params = {"response_cost": response_cost}
class NotGiven:
"""

View file

@ -1,7 +1,10 @@
import asyncio
import base64
import contextlib
import copy
import json
import os
from typing import Any, Final
import httpx
import pytest
@ -14,6 +17,9 @@ from unittest.mock import MagicMock, patch
import litellm
from litellm import main as litellm_main
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
from litellm.types.utils import Usage
async def _async_fake_bedrock_image_details(image_url):
@ -2957,3 +2963,84 @@ async def test_acompletion_resolves_provider_from_api_base():
)
assert response.choices[0].message.content == "resolved"
class _SuccessEventRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.events: list[dict[str, Any]] = [] # mutable-ok: test recorder of success-callback kwargs
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
self.events.append(kwargs)
async def _wait_for_success_event(recorder: _SuccessEventRecorder, call_type: str) -> dict[str, Any]:
for _ in range(100):
if (event := next((e for e in recorder.events if e.get("call_type") == call_type), None)) is not None:
return event
await asyncio.sleep(0.05)
pytest.fail(f"no {call_type} success event; got {[e.get('call_type') for e in recorder.events]}")
def _gemini_tts_generate_content_response() -> dict[str, Any]:
return {
"candidates": [
{
"content": {
"parts": [
{
"inlineData": {
"mimeType": "audio/L16;codec=pcm;rate=24000",
"data": base64.b64encode(b"pcm-audio-bytes").decode(),
}
}
],
"role": "model",
},
"finishReason": "STOP",
"index": 0,
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 60,
"totalTokenCount": 65,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 5}],
"candidatesTokensDetails": [{"modality": "AUDIO", "tokenCount": 60}],
},
"modelVersion": "gemini-2.5-flash-preview-tts",
}
@pytest.mark.asyncio
async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
recorder: Final = _SuccessEventRecorder()
monkeypatch.setattr(litellm, "callbacks", [recorder])
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()))
await litellm.aspeech(
model="gemini/gemini-2.5-flash-preview-tts",
input="spend tracking check",
voice="Kore",
api_key="fake-gemini-key",
metadata={"user_api_key": "hashed-virtual-key", "user_api_key_user_id": "user-1"},
)
assert mock_route.called
speech_event: Final = await _wait_for_success_event(recorder, call_type="aspeech")
spend_metadata: Final = get_litellm_metadata_from_kwargs(speech_event)
assert spend_metadata["user_api_key"] == "hashed-virtual-key"
assert spend_metadata["user_api_key_user_id"] == "user-1"
expected_prompt_cost, expected_completion_cost = litellm.cost_per_token(
model="gemini/gemini-2.5-flash-preview-tts",
usage_object=Usage(prompt_tokens=5, completion_tokens=60, total_tokens=65),
)
expected_cost: Final = expected_prompt_cost + expected_completion_cost
assert expected_cost > 0
assert speech_event["response_cost"] == pytest.approx(expected_cost)
assert speech_event["standard_logging_object"]["response_cost"] == pytest.approx(expected_cost)