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

@@ -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
 

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

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):