From 16500bdf077744765b7f0acf9f4db54c8491b577 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 21:41:33 +0000 Subject: [PATCH] fix(proxy): cap Transcribe pricing media downloads by size and concurrency Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 + .../transcribe_passthrough_logging_handler.py | 51 ++++++++++++++----- ..._transcribe_passthrough_logging_handler.py | 37 ++++++++++++++ 3 files changed, 77 insertions(+), 13 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 962611e82de..0d681f4788e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1575,6 +1575,8 @@ BASE_MCP_ROUTE: Final = "/mcp" TRANSCRIBE_JOB_POLLING_INTERVAL_SECONDS: Final = 10.0 TRANSCRIBE_JOB_MAX_POLLING_ATTEMPTS: Final = 720 # 2 hours TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS: Final = 28800 # Amazon Transcribe quota: maximum audio file length +TRANSCRIBE_MAX_MEDIA_BYTES: Final = 2 * 1024**3 # Amazon Transcribe quota: maximum audio file size +TRANSCRIBE_MEDIA_DOWNLOAD_CONCURRENCY: Final = 1 TRANSCRIBE_MEDIA_FETCH_ATTEMPTS: Final = 3 TRANSCRIBE_MEASURABLE_MEDIA_FORMATS: Final = frozenset({"flac", "mp3", "ogg", "wav"}) # what libsndfile can read diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py index 35ccd2320a2..2745876229b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py @@ -7,7 +7,7 @@ from datetime import datetime from functools import lru_cache, partial from pathlib import Path from types import MappingProxyType -from typing import Final, Protocol, TypeAlias +from typing import IO, Final, Protocol, TypeAlias from urllib.parse import quote import httpx @@ -19,8 +19,10 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( TRANSCRIBE_JOB_MAX_POLLING_ATTEMPTS, TRANSCRIBE_JOB_POLLING_INTERVAL_SECONDS, + TRANSCRIBE_MAX_MEDIA_BYTES, TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS, TRANSCRIBE_MEASURABLE_MEDIA_FORMATS, + TRANSCRIBE_MEDIA_DOWNLOAD_CONCURRENCY, TRANSCRIBE_MEDIA_FETCH_ATTEMPTS, ) from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration @@ -298,7 +300,17 @@ def s3_media_url(media_uri: str, aws_region_name: str) -> str | None: return f"https://{bucket}.s3.{aws_region_name}.{dns_suffix}/{quote(key)}" -def transcribe_media_duration_probe(aws_region_name: str) -> MediaDurationProbe: +async def write_media_within_limit(response: httpx.Response, media_file: IO[bytes], max_bytes: int) -> bool: + if int(response.headers.get("content-length", "0")) > max_bytes: + return False + async for chunk in response.aiter_bytes(): + _ = media_file.write(chunk) + if media_file.tell() > max_bytes: + return False + return True + + +def transcribe_media_duration_probe(aws_region_name: str, download_slots: asyncio.Semaphore) -> MediaDurationProbe: from botocore.auth import S3SigV4Auth from botocore.awsrequest import AWSRequest @@ -316,24 +328,30 @@ def transcribe_media_duration_probe(aws_region_name: str) -> MediaDurationProbe: return None headers: Final = await run_aws_signing(sign_s3_get, url) client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.PassThroughEndpoint).client - with tempfile.NamedTemporaryFile() as media_file: - async with client.stream("GET", url, headers=headers) as response: - _ = response.raise_for_status() - async for chunk in response.aiter_bytes(): - _ = media_file.write(chunk) - media_file.flush() - return await asyncio.to_thread(calculate_request_duration, Path(media_file.name)) + async with download_slots: + with tempfile.NamedTemporaryFile() as media_file: + async with client.stream("GET", url, headers=headers) as response: + _ = response.raise_for_status() + if not await write_media_within_limit(response, media_file, TRANSCRIBE_MAX_MEDIA_BYTES): + verbose_proxy_logger.warning( + "Transcribe media %s exceeds the size cap, charging maximum", media_uri + ) + return None + media_file.flush() + return await asyncio.to_thread(calculate_request_duration, Path(media_file.name)) return media_seconds -async def price_transcription_job_live(job_name: str, aws_region_name: str, cost_per_second: float) -> float: +async def price_transcription_job_live( + job_name: str, aws_region_name: str, cost_per_second: float, download_slots: asyncio.Semaphore +) -> float: try: return await price_transcription_job( job_name, cost_per_second, get_job=transcribe_job_lookup(aws_region_name), - media_seconds=transcribe_media_duration_probe(aws_region_name), + media_seconds=transcribe_media_duration_probe(aws_region_name, download_slots), ) except Exception as e: # noqa: BLE001 # an unreadable job must still be charged, so fail closed at the maximum verbose_proxy_logger.exception("Pricing Transcribe job %s failed, charging maximum: %s", job_name, e) @@ -341,8 +359,15 @@ async def price_transcription_job_live(job_name: str, aws_region_name: str, cost class TranscribePassthroughLoggingHandler: - def __init__(self, job_pricer: JobPricer = price_transcription_job_live) -> None: - self._job_pricer: Final = job_pricer + def __init__(self, job_pricer: JobPricer | None = None) -> None: + self._job_pricer: Final = ( + job_pricer + if job_pricer is not None + else partial( + price_transcription_job_live, + download_slots=asyncio.Semaphore(TRANSCRIBE_MEDIA_DOWNLOAD_CONCURRENCY), + ) + ) self._pricing_tasks: Final[set[asyncio.Task[None]]] = set() # mutable-ok: asyncio holds tasks weakly @staticmethod diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py index f3fdf50fbe8..80f118f76dd 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py @@ -1,4 +1,5 @@ import asyncio +import io from datetime import datetime from unittest.mock import MagicMock @@ -15,6 +16,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.transcribe_passt transcribe_cost_per_second, transcribe_supported_operations, transcribe_unpriceable_request_reason, + write_media_within_limit, ) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, @@ -213,6 +215,41 @@ class TestS3MediaUrl: ) +class _ChunkedStream(httpx.AsyncByteStream): + def __init__(self, *chunks: bytes) -> None: + self._chunks = chunks + + async def __aiter__(self): + for chunk in self._chunks: + yield chunk + + +def _media_response(*chunks: bytes, content_length: int | None) -> httpx.Response: + headers = {"content-length": str(content_length)} if content_length is not None else {} + return httpx.Response(200, headers=headers, stream=_ChunkedStream(*chunks)) + + +class TestWriteMediaWithinLimit: + @pytest.mark.asyncio + async def test_media_within_the_cap_is_written_whole(self): + media_file = io.BytesIO() + assert await write_media_within_limit(_media_response(b"abc", b"def", content_length=6), media_file, 6) is True + assert media_file.getvalue() == b"abcdef" + + @pytest.mark.asyncio + async def test_advertised_size_over_the_cap_is_refused_before_downloading(self): + media_file = io.BytesIO() + assert await write_media_within_limit(_media_response(b"abcdef", content_length=7), media_file, 6) is False + assert media_file.getvalue() == b"" + + @pytest.mark.asyncio + async def test_stream_growing_past_the_cap_is_cut_off(self): + media_file = io.BytesIO() + response = _media_response(b"abc", b"def", b"ghi", content_length=None) + assert await write_media_within_limit(response, media_file, 5) is False + assert media_file.getvalue() == b"abcdef" + + class TestPriceTranscriptionJob: @pytest.mark.asyncio async def test_polls_until_completed_then_charges_the_media_length_rounded_up(self):