Merge pull request #16739 from BerriAI/litellm_/audio/speech

[Perf] Fix `/audio/speech` performance by using `shared_sessions`
This commit is contained in:
Alexsander Hamir 2025-11-30 17:27:44 -08:00 committed by GitHub
commit b949ec90db
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 48 additions and 28 deletions

View file

@ -255,6 +255,7 @@ TOGETHER_AI_EMBEDDING_350_M = int(os.getenv("TOGETHER_AI_EMBEDDING_350_M", 350))
QDRANT_SCALAR_QUANTILE = float(os.getenv("QDRANT_SCALAR_QUANTILE", 0.99))
QDRANT_VECTOR_SIZE = int(os.getenv("QDRANT_VECTOR_SIZE", 1536))
CACHED_STREAMING_CHUNK_DELAY = float(os.getenv("CACHED_STREAMING_CHUNK_DELAY", 0.02))
AUDIO_SPEECH_CHUNK_SIZE = 8192 # chunk_size for audio speech streaming. Balance between latency and memory usage
MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int(
os.getenv("MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB", 512)
)

View file

@ -1414,6 +1414,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
timeout: Union[float, httpx.Timeout],
aspeech: Optional[bool] = None,
client=None,
shared_session: Optional["ClientSession"] = None,
) -> HttpxBinaryResponseContent:
if aspeech is not None and aspeech is True:
return self.async_audio_speech(
@ -1428,6 +1429,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
max_retries=max_retries,
timeout=timeout,
client=client,
shared_session=shared_session,
) # type: ignore
openai_client = self._get_openai_client(
@ -1437,6 +1439,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
timeout=timeout,
max_retries=max_retries,
client=client,
shared_session=shared_session,
)
response = cast(OpenAI, openai_client).audio.speech.create(
@ -1460,6 +1463,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
max_retries: int,
timeout: Union[float, httpx.Timeout],
client=None,
shared_session: Optional["ClientSession"] = None,
) -> HttpxBinaryResponseContent:
openai_client = cast(
AsyncOpenAI,
@ -1470,6 +1474,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
timeout=timeout,
max_retries=max_retries,
client=client,
shared_session=shared_session,
),
)

View file

@ -5868,6 +5868,7 @@ def speech( # noqa: PLR0915
proxy_server_request = kwargs.get("proxy_server_request", None)
extra_headers = kwargs.get("extra_headers", None)
model_info = kwargs.get("model_info", None)
shared_session = kwargs.get("shared_session", None)
model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base
) # type: ignore
@ -5981,6 +5982,7 @@ def speech( # noqa: PLR0915
timeout=timeout,
client=client, # pass AsyncOpenAI, OpenAI client
aspeech=aspeech,
shared_session=shared_session,
)
elif custom_llm_provider == "azure":
# Check if this is Azure Speech Service (Cognitive Services TTS)

View file

@ -14,6 +14,7 @@ from datetime import datetime, timedelta
from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
Dict,
List,
Literal,
@ -32,6 +33,7 @@ from litellm.constants import (
AIOHTTP_CONNECTOR_LIMIT,
AIOHTTP_KEEPALIVE_TIMEOUT,
AIOHTTP_TTL_DNS_CACHE,
AUDIO_SPEECH_CHUNK_SIZE,
BASE_MCP_ROUTE,
DEFAULT_MAX_RECURSE_DEPTH,
DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
@ -4087,7 +4089,8 @@ async def async_data_generator(
):
verbose_proxy_logger.debug("inside generator")
try:
str_so_far = ""
# Use a list to accumulate response segments to avoid O(n^2) string concatenation
str_so_far_parts: list[str] = []
error_message: Optional[str] = None
async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
@ -4103,12 +4106,12 @@ async def async_data_generator(
user_api_key_dict=user_api_key_dict,
response=chunk,
data=request_data,
str_so_far=str_so_far,
str_so_far="".join(str_so_far_parts),
)
if isinstance(chunk, (ModelResponse, ModelResponseStream)):
response_str = litellm.get_response_string(response_obj=chunk)
str_so_far += response_str
str_so_far_parts.append(response_str)
if isinstance(chunk, BaseModel):
chunk = chunk.model_dump_json(exclude_none=True, exclude_unset=True)
@ -5302,6 +5305,18 @@ async def moderations(
)
async def _audio_speech_chunk_generator(
_response: HttpxBinaryResponseContent,
) -> AsyncGenerator[bytes, None]:
# chunk_size has a big impact on latency, it can't be too small or too large
# too small: latency is high
# too large: latency is low, but memory usage is high
# 8192 is a good compromise
_generator = await _response.aiter_bytes(chunk_size=AUDIO_SPEECH_CHUNK_SIZE)
async for chunk in _generator:
yield chunk
@router.post(
"/v1/audio/speech",
dependencies=[Depends(user_api_key_auth)],
@ -5347,7 +5362,7 @@ async def audio_speech(
### CALL HOOKS ### - modify incoming data / reject request before calling the model
data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=data, call_type="image_generation"
user_api_key_dict=user_api_key_dict, data=data, call_type="aspeech"
)
## ROUTE TO CORRECT ENDPOINT ##
@ -5374,12 +5389,6 @@ async def audio_speech(
response_cost = hidden_params.get("response_cost", None) or ""
litellm_call_id = hidden_params.get("litellm_call_id", None) or ""
# Printing each chunk size
async def generate(_response: HttpxBinaryResponseContent):
_generator = await _response.aiter_bytes(chunk_size=1024)
async for chunk in _generator:
yield chunk
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
model_id=model_id,
@ -5394,21 +5403,20 @@ async def audio_speech(
hidden_params=hidden_params,
)
select_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
request_data=data,
)
# Determine media type based on model type
media_type = "audio/mpeg" # Default for OpenAI TTS
request_model = data.get("model", "")
if "gemini" in request_model.lower() and (
"tts" in request_model.lower() or "preview-tts" in request_model.lower()
):
media_type = "audio/wav" # Gemini TTS returns WAV format after conversion
if request_model:
request_model_lower = request_model.lower()
if "gemini" in request_model_lower and (
"tts" in request_model_lower or "preview-tts" in request_model_lower
):
media_type = "audio/wav" # Gemini TTS returns WAV format after conversion
return StreamingResponse(
generate(response), media_type=media_type, headers=custom_headers # type: ignore
_audio_speech_chunk_generator(response), # type: ignore[arg-type]
media_type=media_type,
headers=custom_headers, # type: ignore
)
except Exception as e:

View file

@ -4441,17 +4441,20 @@ def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream])
responses_api_response = getattr(response_obj, "response", None)
if responses_api_response and hasattr(responses_api_response, "output"):
output_list = responses_api_response.output
response_str = ""
# Use list accumulation to avoid O(n^2) string concatenation:
# repeatedly doing `response_str += part` copies the full string each time
# because Python strings are immutable, so total work grows with n^2.
response_output_parts: List[str] = []
for output_item in output_list:
# Handle output items with content array
if hasattr(output_item, "content"):
for content_part in output_item.content:
if hasattr(content_part, "text"):
response_str += content_part.text
response_output_parts.append(content_part.text)
# Handle output items with direct text field
elif hasattr(output_item, "text"):
response_str += output_item.text
return response_str
response_output_parts.append(output_item.text)
return "".join(response_output_parts)
# Handle Responses API text delta events
if hasattr(response_obj, "type") and hasattr(response_obj, "delta"):
@ -4465,16 +4468,17 @@ def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream])
response_obj.choices
)
response_str = ""
# Use list accumulation to avoid O(n^2) string concatenation across choices
response_parts: List[str] = []
for choice in _choices:
if isinstance(choice, Choices):
if choice.message.content is not None:
response_str += choice.message.content
response_parts.append(str(choice.message.content))
elif isinstance(choice, StreamingChoices):
if choice.delta.content is not None:
response_str += choice.delta.content
response_parts.append(str(choice.delta.content))
return response_str
return "".join(response_parts)
def get_api_key(llm_provider: str, dynamic_api_key: Optional[str]):