mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(proxy): scope Transcribe jobs to the key that started them and charge rewritten media the maximum
Standard jobs are tagged litellm-owner on StartTranscriptionJob so GetTranscriptionJob and DeleteTranscriptionJob only work for the owner or a proxy admin, and account-wide operations need a proxy admin. Media rewritten after job creation is charged the eight hour maximum, and the success handler takes an injected log dispatch instead of tests patching its private method Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
4885594a1e
commit
ce735f586c
6 changed files with 410 additions and 45 deletions
|
|
@ -1578,6 +1578,7 @@ TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS: Final = 28800 # Amazon Transcribe quota:
|
|||
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_MEDIA_LAST_MODIFIED_TOLERANCE_SECONDS: Final = 1.0 # S3 Last-Modified carries whole seconds only
|
||||
TRANSCRIBE_MEASURABLE_MEDIA_FORMATS: Final = frozenset({"flac", "mp3", "ogg", "wav"}) # what libsndfile can read
|
||||
|
||||
BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour
|
||||
|
|
|
|||
|
|
@ -1329,16 +1329,26 @@ async def transcribe_proxy_route(
|
|||
"""
|
||||
Pass-through for the Amazon Transcribe API, e.g. `POST /transcribe/StartTranscriptionJob`.
|
||||
|
||||
The request body is forwarded as-is to the AWS JSON 1.1 API and signed with SigV4
|
||||
using the proxy's AWS credentials. Streaming transcription (`transcribestreaming`)
|
||||
uses a separate HTTP/2 event-stream protocol and is not served by this route.
|
||||
The request body is forwarded to the AWS JSON 1.1 API and signed with SigV4 using the
|
||||
proxy's AWS credentials. Standard jobs are tagged with the calling key's owner so that
|
||||
only that owner (or a proxy admin) can read or delete them; account-wide operations
|
||||
such as ListTranscriptionJobs are limited to proxy admins. Streaming transcription
|
||||
(`transcribestreaming`) uses a separate HTTP/2 event-stream protocol and is not served
|
||||
by this route.
|
||||
|
||||
[Docs](https://docs.litellm.ai/docs/pass_through/transcribe)
|
||||
"""
|
||||
from .llm_provider_handlers.transcribe_passthrough_logging_handler import (
|
||||
TRANSCRIBE_CUSTOM_LLM_PROVIDER,
|
||||
TRANSCRIBE_OWNED_JOB_OPERATIONS,
|
||||
TRANSCRIBE_PRICED_OPERATION,
|
||||
TRANSCRIBE_TARGET_PREFIX,
|
||||
TranscribeRefusal,
|
||||
transcribe_admin_only_refusal,
|
||||
transcribe_cost_per_second,
|
||||
transcribe_job_access_refusal,
|
||||
transcribe_job_lookup,
|
||||
transcribe_owned_start_request,
|
||||
transcribe_supported_operations,
|
||||
transcribe_unpriceable_request_reason,
|
||||
)
|
||||
|
|
@ -1371,6 +1381,23 @@ async def transcribe_proxy_route(
|
|||
unpriceable_reason: Final = transcribe_unpriceable_request_reason(operation, data, transcribe_cost_per_second())
|
||||
if unpriceable_reason is not None:
|
||||
raise HTTPException(status_code=400, detail=unpriceable_reason)
|
||||
admin_only_refusal: Final = transcribe_admin_only_refusal(operation, user_api_key_dict)
|
||||
if admin_only_refusal is not None:
|
||||
raise HTTPException(status_code=admin_only_refusal.status_code, detail=admin_only_refusal.detail)
|
||||
request_body: Final = (
|
||||
transcribe_owned_start_request(data, user_api_key_dict) if operation == TRANSCRIBE_PRICED_OPERATION else data
|
||||
)
|
||||
if isinstance(request_body, TranscribeRefusal):
|
||||
raise HTTPException(status_code=request_body.status_code, detail=request_body.detail)
|
||||
access_refusal: Final = (
|
||||
await transcribe_job_access_refusal(
|
||||
data.get("TranscriptionJobName"), user_api_key_dict, transcribe_job_lookup(aws_region_name)
|
||||
)
|
||||
if operation in TRANSCRIBE_OWNED_JOB_OPERATIONS
|
||||
else None
|
||||
)
|
||||
if access_refusal is not None:
|
||||
raise HTTPException(status_code=access_refusal.status_code, detail=access_refusal.detail)
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post
|
||||
|
||||
|
|
@ -1381,7 +1408,7 @@ async def transcribe_proxy_route(
|
|||
service_name="transcribe",
|
||||
aws_region_name=aws_region_name,
|
||||
url=target_url,
|
||||
body=json.dumps(data),
|
||||
body=json.dumps(request_body),
|
||||
headers=MappingProxyType(
|
||||
{
|
||||
"Content-Type": "application/x-amz-json-1.1",
|
||||
|
|
@ -1396,7 +1423,7 @@ async def transcribe_proxy_route(
|
|||
custom_headers=prepped.headers,
|
||||
custom_llm_provider=TRANSCRIBE_CUSTOM_LLM_PROVIDER,
|
||||
)
|
||||
setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, data)
|
||||
setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, request_body)
|
||||
setattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, prepped.body)
|
||||
return await endpoint_func(request, fastapi_response, user_api_key_dict)
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ import json
|
|||
import math
|
||||
import tempfile
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from email.utils import parsedate_to_datetime
|
||||
from functools import lru_cache, partial
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
|
|
@ -24,6 +26,7 @@ from litellm.constants import (
|
|||
TRANSCRIBE_MEASURABLE_MEDIA_FORMATS,
|
||||
TRANSCRIBE_MEDIA_DOWNLOAD_CONCURRENCY,
|
||||
TRANSCRIBE_MEDIA_FETCH_ATTEMPTS,
|
||||
TRANSCRIBE_MEDIA_LAST_MODIFIED_TOLERANCE_SECONDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
|
|
@ -32,7 +35,16 @@ from litellm.litellm_core_utils.litellm_logging import (
|
|||
get_standard_logging_object_payload,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingResultValues, PassThroughEndpointLoggingTypedDict
|
||||
from litellm.proxy._types import (
|
||||
PassThroughEndpointLoggingResultValues,
|
||||
PassThroughEndpointLoggingTypedDict,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.resource_ownership import (
|
||||
get_primary_resource_owner_scope,
|
||||
is_proxy_admin,
|
||||
user_can_access_resource_owner,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.utils import StandardPassThroughResponseObject
|
||||
|
||||
|
|
@ -45,9 +57,11 @@ TRANSCRIBE_UNPRICED_OPERATIONS: Final = frozenset(
|
|||
)
|
||||
TRANSCRIBE_SURCHARGE_MEMBERS: Final = ("ContentRedaction", "ToxicityDetection")
|
||||
TRANSCRIBE_TERMINAL_JOB_STATUSES: Final = frozenset({"COMPLETED", "FAILED"})
|
||||
TRANSCRIBE_OWNER_TAG: Final = "litellm-owner"
|
||||
TRANSCRIBE_OWNED_JOB_OPERATIONS: Final = frozenset({"GetTranscriptionJob", "DeleteTranscriptionJob"})
|
||||
|
||||
JobLookup: TypeAlias = Callable[[str], Awaitable[Mapping[str, object]]] # mutable-ok: Callable parameter syntax
|
||||
MediaDurationProbe: TypeAlias = Callable[[str], Awaitable[float | None]] # mutable-ok: Callable parameter syntax
|
||||
MediaDurationProbe: TypeAlias = Callable[[str, float], Awaitable[float | None]] # mutable-ok: Callable parameter syntax
|
||||
JobPricer: TypeAlias = Callable[[str, str, float], Awaitable[float]] # mutable-ok: Callable parameter syntax
|
||||
|
||||
|
||||
|
|
@ -60,10 +74,18 @@ class _MediaRef(BaseModel):
|
|||
MediaFileUri: str | None = None
|
||||
|
||||
|
||||
class _JobTag(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
Key: str | None = None
|
||||
Value: str | None = None
|
||||
|
||||
|
||||
class _TranscriptionJob(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
TranscriptionJobStatus: str | None = None
|
||||
CreationTime: float | None = None
|
||||
Media: _MediaRef | None = None
|
||||
Tags: tuple[_JobTag, ...] = ()
|
||||
|
||||
|
||||
class _GetTranscriptionJobResponse(BaseModel):
|
||||
|
|
@ -77,6 +99,13 @@ class _PricedCostMapEntry(BaseModel):
|
|||
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
_JSON_OBJECTS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TranscribeRefusal:
|
||||
status_code: int
|
||||
detail: str
|
||||
|
||||
|
||||
class PassThroughLogDispatch(Protocol):
|
||||
|
|
@ -178,6 +207,60 @@ def requested_media_format(request_body: Mapping[str, object]) -> str | None:
|
|||
return suffix.lower() if dot else None
|
||||
|
||||
|
||||
def transcribe_admin_only_refusal(operation: str, user_api_key_dict: UserAPIKeyAuth) -> TranscribeRefusal | None:
|
||||
if (
|
||||
operation == TRANSCRIBE_PRICED_OPERATION
|
||||
or operation in TRANSCRIBE_OWNED_JOB_OPERATIONS
|
||||
or is_proxy_admin(user_api_key_dict)
|
||||
):
|
||||
return None
|
||||
return TranscribeRefusal(
|
||||
403,
|
||||
f"{operation} reaches every Amazon Transcribe resource in the AWS account, so only a proxy admin may call it;"
|
||||
f" other keys may {TRANSCRIBE_PRICED_OPERATION} and {' or '.join(sorted(TRANSCRIBE_OWNED_JOB_OPERATIONS))}"
|
||||
" for the jobs they started",
|
||||
)
|
||||
|
||||
|
||||
def transcribe_owned_start_request(
|
||||
request_body: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> dict[str, object] | TranscribeRefusal:
|
||||
owner: Final = get_primary_resource_owner_scope(user_api_key_dict)
|
||||
if owner is None:
|
||||
return TranscribeRefusal(400, "The calling key has no identity to record as the owner of the transcription job")
|
||||
try:
|
||||
tags: Final = _JSON_OBJECTS.validate_python(request_body.get("Tags", ()))
|
||||
except ValidationError:
|
||||
return TranscribeRefusal(400, "Tags must be a list of objects with Key and Value members")
|
||||
if any(tag.get("Key") == TRANSCRIBE_OWNER_TAG for tag in tags):
|
||||
return TranscribeRefusal(
|
||||
400, f"The {TRANSCRIBE_OWNER_TAG} tag is assigned by LiteLLM and cannot be supplied by the caller"
|
||||
)
|
||||
owner_tag: Final = _JobTag(Key=TRANSCRIBE_OWNER_TAG, Value=owner).model_dump()
|
||||
return {**request_body, "Tags": (*tags, owner_tag)} # mutable-ok: json.dumps and the body state key take a dict
|
||||
|
||||
|
||||
async def transcribe_job_access_refusal(
|
||||
job_name: object, user_api_key_dict: UserAPIKeyAuth, get_job: JobLookup
|
||||
) -> TranscribeRefusal | None:
|
||||
if is_proxy_admin(user_api_key_dict):
|
||||
return None
|
||||
if not isinstance(job_name, str):
|
||||
return TranscribeRefusal(400, "TranscriptionJobName must be a string")
|
||||
not_found: Final = TranscribeRefusal(
|
||||
404, f"No transcription job named {job_name} was started through this proxy by the calling key"
|
||||
)
|
||||
try:
|
||||
job: Final = _GetTranscriptionJobResponse.model_validate(await get_job(job_name)).TranscriptionJob
|
||||
except Exception as e: # noqa: BLE001 # a job that cannot be read cannot be shown to belong to the caller
|
||||
verbose_proxy_logger.warning("Looking up Transcribe job %s for an ownership check failed: %s", job_name, e)
|
||||
return not_found
|
||||
owner: Final = (
|
||||
next((tag.Value for tag in job.Tags if tag.Key == TRANSCRIBE_OWNER_TAG), None) if job is not None else None
|
||||
)
|
||||
return None if user_can_access_resource_owner(owner, user_api_key_dict) else not_found
|
||||
|
||||
|
||||
def transcription_job_cost(audio_seconds: float, cost_per_second: float) -> float:
|
||||
return math.ceil(audio_seconds) * cost_per_second
|
||||
|
||||
|
|
@ -211,13 +294,14 @@ async def await_transcription_job(
|
|||
|
||||
async def measure_media_seconds(
|
||||
media_uri: str,
|
||||
job_created_at: float,
|
||||
media_seconds: MediaDurationProbe,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
attempts: int = TRANSCRIBE_MEDIA_FETCH_ATTEMPTS,
|
||||
) -> float | None:
|
||||
for attempt in range(1, attempts + 1):
|
||||
try:
|
||||
return await media_seconds(media_uri)
|
||||
return await media_seconds(media_uri, job_created_at)
|
||||
except Exception as e: # noqa: BLE001 # the media is retried, then charged at the maximum if still unreadable
|
||||
verbose_proxy_logger.warning("Measuring Transcribe media %s failed (attempt %d): %s", media_uri, attempt, e)
|
||||
if attempt < attempts:
|
||||
|
|
@ -236,7 +320,9 @@ async def price_transcription_job(
|
|||
"""
|
||||
Amazon Transcribe bills every second of the media file, silence included, and reports no
|
||||
duration itself, so the job is polled to completion and the media it transcribed is measured.
|
||||
Anything that stops the duration from being read is charged as the longest media AWS accepts.
|
||||
The measurement only counts when the object has not been rewritten since the job was created,
|
||||
which is what ties it to the bytes Transcribe read. Anything that stops the duration from
|
||||
being read is charged as the longest media AWS accepts.
|
||||
"""
|
||||
job: Final = await await_transcription_job(job_name, get_job, sleep=sleep, max_attempts=max_attempts)
|
||||
if job is None:
|
||||
|
|
@ -245,9 +331,9 @@ async def price_transcription_job(
|
|||
if job.TranscriptionJobStatus == "FAILED":
|
||||
return 0.0
|
||||
media_uri: Final = job.Media.MediaFileUri if job.Media is not None else None
|
||||
if media_uri is None:
|
||||
if media_uri is None or job.CreationTime is None:
|
||||
return transcribe_max_job_cost(cost_per_second)
|
||||
audio_seconds: Final = await measure_media_seconds(media_uri, media_seconds, sleep=sleep)
|
||||
audio_seconds: Final = await measure_media_seconds(media_uri, job.CreationTime, media_seconds, sleep=sleep)
|
||||
if audio_seconds is None:
|
||||
return transcribe_max_job_cost(cost_per_second)
|
||||
return transcription_job_cost(audio_seconds, cost_per_second)
|
||||
|
|
@ -303,6 +389,14 @@ 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 media_predates_job(headers: Mapping[str, str], job_created_at: float) -> bool:
|
||||
try:
|
||||
modified_at: Final = parsedate_to_datetime(headers["last-modified"]).timestamp()
|
||||
except (KeyError, TypeError, ValueError):
|
||||
return False
|
||||
return modified_at <= job_created_at + TRANSCRIBE_MEDIA_LAST_MODIFIED_TOLERANCE_SECONDS
|
||||
|
||||
|
||||
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
|
||||
|
|
@ -325,7 +419,7 @@ def transcribe_media_duration_probe(aws_region_name: str, download_slots: asynci
|
|||
S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
|
||||
return dict(aws_request.prepare().headers.items()) # mutable-ok: httpx request headers take a dict
|
||||
|
||||
async def media_seconds(media_uri: str) -> float | None:
|
||||
async def media_seconds(media_uri: str, job_created_at: float) -> float | None:
|
||||
url: Final = s3_media_url(media_uri, aws_region_name)
|
||||
if url is None:
|
||||
return None
|
||||
|
|
@ -335,6 +429,11 @@ def transcribe_media_duration_probe(aws_region_name: str, download_slots: asynci
|
|||
with tempfile.NamedTemporaryFile() as media_file:
|
||||
async with client.stream("GET", url, headers=headers) as response:
|
||||
_ = response.raise_for_status()
|
||||
if not media_predates_job(response.headers, job_created_at):
|
||||
verbose_proxy_logger.warning(
|
||||
"Transcribe media %s was rewritten after the job was created, charging maximum", media_uri
|
||||
)
|
||||
return None
|
||||
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
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ from .llm_provider_handlers.gemini_passthrough_logging_handler import (
|
|||
)
|
||||
from .llm_provider_handlers.transcribe_passthrough_logging_handler import (
|
||||
TRANSCRIBE_CUSTOM_LLM_PROVIDER,
|
||||
PassThroughLogDispatch,
|
||||
TranscribePassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.vertex_passthrough_logging_handler import (
|
||||
|
|
@ -52,10 +53,15 @@ def _safe_response_text(httpx_response: httpx.Response) -> str:
|
|||
|
||||
|
||||
class PassThroughEndpointLogging:
|
||||
def __init__(self, transcribe_handler: TranscribePassthroughLoggingHandler | None = None):
|
||||
def __init__(
|
||||
self,
|
||||
transcribe_handler: TranscribePassthroughLoggingHandler | None = None,
|
||||
log_dispatch: PassThroughLogDispatch | None = None,
|
||||
):
|
||||
self.transcribe_passthrough_logging_handler: Final = (
|
||||
transcribe_handler if transcribe_handler is not None else TranscribePassthroughLoggingHandler()
|
||||
)
|
||||
self._log_dispatch: Final = log_dispatch if log_dispatch is not None else self._handle_logging
|
||||
self.TRACKED_VERTEX_METHOD_ROUTES = (
|
||||
"generateContent",
|
||||
"streamGenerateContent",
|
||||
|
|
@ -351,7 +357,7 @@ class PassThroughEndpointLogging:
|
|||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
request_body=request_body,
|
||||
log=self._handle_logging,
|
||||
log=self._log_dispatch,
|
||||
standard_pass_through_logging_payload=passthrough_logging_payload,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
@ -385,7 +391,7 @@ class PassThroughEndpointLogging:
|
|||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
await self._handle_logging(
|
||||
await self._log_dispatch(
|
||||
logging_obj=logging_obj,
|
||||
standard_logging_response_object=standard_logging_response_object,
|
||||
result=result,
|
||||
|
|
|
|||
|
|
@ -7,13 +7,20 @@ import httpx
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.transcribe_passthrough_logging_handler import (
|
||||
TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS,
|
||||
TRANSCRIBE_OWNER_TAG,
|
||||
TranscribePassthroughLoggingHandler,
|
||||
TranscribeRefusal,
|
||||
media_predates_job,
|
||||
price_transcription_job,
|
||||
requested_media_format,
|
||||
s3_media_url,
|
||||
transcribe_admin_only_refusal,
|
||||
transcribe_cost_per_second,
|
||||
transcribe_job_access_refusal,
|
||||
transcribe_owned_start_request,
|
||||
transcribe_supported_operations,
|
||||
transcribe_unpriceable_request_reason,
|
||||
write_media_within_limit,
|
||||
|
|
@ -46,23 +53,27 @@ async def _no_sleep(_: float) -> None:
|
|||
|
||||
|
||||
MEDIA_URI = "s3://b/a.wav"
|
||||
CREATED_AT = 1_789_682_363.696
|
||||
|
||||
|
||||
def _job(status: str, media_uri: str | None = MEDIA_URI) -> dict[str, object]:
|
||||
def _job(
|
||||
status: str, media_uri: str | None = MEDIA_URI, created_at: float | None = CREATED_AT, **members: object
|
||||
) -> dict[str, object]:
|
||||
media = {"Media": {"MediaFileUri": media_uri}} if media_uri else {}
|
||||
return {"TranscriptionJob": {"TranscriptionJobStatus": status, **media}}
|
||||
created = {"CreationTime": created_at} if created_at is not None else {}
|
||||
return {"TranscriptionJob": {"TranscriptionJobStatus": status, **media, **created, **members}}
|
||||
|
||||
|
||||
async def _no_media(uri: str) -> float | None:
|
||||
async def _no_media(uri: str, created_at: float) -> float | None:
|
||||
raise AssertionError("the media must not be measured on this path")
|
||||
|
||||
|
||||
def _media_probe(*durations: float | None | Exception):
|
||||
remaining = list(durations)
|
||||
measured: list[str] = []
|
||||
measured: list[tuple[str, float]] = []
|
||||
|
||||
async def media_seconds(uri: str) -> float | None:
|
||||
measured.append(uri)
|
||||
async def media_seconds(uri: str, created_at: float) -> float | None:
|
||||
measured.append((uri, created_at))
|
||||
outcome = remaining.pop(0) if len(remaining) > 1 else remaining[0]
|
||||
if isinstance(outcome, Exception):
|
||||
raise outcome
|
||||
|
|
@ -266,7 +277,7 @@ class TestPriceTranscriptionJob:
|
|||
|
||||
assert cost == pytest.approx(18 * COST_PER_SECOND)
|
||||
assert seen == ["job-1", "job-1", "job-1"]
|
||||
assert measured == [MEDIA_URI]
|
||||
assert measured == [(MEDIA_URI, CREATED_AT)]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_poll_is_retried_instead_of_ending_pricing(self):
|
||||
|
|
@ -310,7 +321,7 @@ class TestPriceTranscriptionJob:
|
|||
cost = await price_transcription_job("job-1", COST_PER_SECOND, get_job, media_seconds, sleep=_no_sleep)
|
||||
|
||||
assert cost == pytest.approx(TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS * COST_PER_SECOND)
|
||||
assert measured == [MEDIA_URI]
|
||||
assert measured == [(MEDIA_URI, CREATED_AT)]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_media_fetch_is_retried_then_charged_the_maximum(self):
|
||||
|
|
@ -340,6 +351,157 @@ class TestPriceTranscriptionJob:
|
|||
|
||||
assert cost == pytest.approx(TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS * COST_PER_SECOND)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_job_without_creation_time_is_charged_the_maximum_unmeasured(self):
|
||||
get_job, _ = _sequence(_job("COMPLETED", created_at=None))
|
||||
media_seconds, measured = _media_probe(60.0)
|
||||
|
||||
cost = await price_transcription_job("job-1", COST_PER_SECOND, get_job, media_seconds, sleep=_no_sleep)
|
||||
|
||||
assert cost == pytest.approx(TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS * COST_PER_SECOND)
|
||||
assert measured == []
|
||||
|
||||
|
||||
class TestMediaPredatesJob:
|
||||
LAST_MODIFIED = "Thu, 17 Sep 2026 17:45:00 GMT"
|
||||
LAST_MODIFIED_EPOCH = 1_789_667_100.0
|
||||
|
||||
def test_object_written_before_the_job_counts(self):
|
||||
assert media_predates_job(httpx.Headers({"Last-Modified": self.LAST_MODIFIED}), self.LAST_MODIFIED_EPOCH + 30)
|
||||
|
||||
def test_object_written_in_the_same_second_as_the_job_counts(self):
|
||||
assert media_predates_job(httpx.Headers({"Last-Modified": self.LAST_MODIFIED}), self.LAST_MODIFIED_EPOCH - 0.4)
|
||||
|
||||
def test_object_rewritten_after_the_job_does_not_count(self):
|
||||
assert not media_predates_job(
|
||||
httpx.Headers({"Last-Modified": self.LAST_MODIFIED}), self.LAST_MODIFIED_EPOCH - 30
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("headers", [{}, {"Last-Modified": "yesterday"}])
|
||||
def test_unknown_modification_time_does_not_count(self, headers: dict[str, str]):
|
||||
assert not media_predates_job(httpx.Headers(headers), self.LAST_MODIFIED_EPOCH + 30)
|
||||
|
||||
|
||||
VIRTUAL_KEY = UserAPIKeyAuth(api_key="hashed-key-a", user_id="user-a", team_id="team-a")
|
||||
OTHER_VIRTUAL_KEY = UserAPIKeyAuth(api_key="hashed-key-b", user_id="user-b", team_id="team-b")
|
||||
ADMIN_KEY = UserAPIKeyAuth(api_key="hashed-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
|
||||
class TestTranscribeAdminOnlyRefusal:
|
||||
@pytest.mark.parametrize("operation", ["StartTranscriptionJob", "GetTranscriptionJob", "DeleteTranscriptionJob"])
|
||||
def test_job_scoped_operations_are_open_to_virtual_keys(self, operation: str):
|
||||
assert transcribe_admin_only_refusal(operation, VIRTUAL_KEY) is None
|
||||
|
||||
@pytest.mark.parametrize("operation", ["ListTranscriptionJobs", "ListVocabularies", "DeleteVocabulary"])
|
||||
def test_account_wide_operations_are_refused_for_virtual_keys(self, operation: str):
|
||||
refusal = transcribe_admin_only_refusal(operation, VIRTUAL_KEY)
|
||||
|
||||
assert refusal is not None
|
||||
assert refusal.status_code == 403
|
||||
assert operation in refusal.detail
|
||||
|
||||
@pytest.mark.parametrize("operation", ["ListTranscriptionJobs", "DeleteVocabulary"])
|
||||
def test_account_wide_operations_are_open_to_proxy_admins(self, operation: str):
|
||||
assert transcribe_admin_only_refusal(operation, ADMIN_KEY) is None
|
||||
|
||||
|
||||
class TestTranscribeOwnedStartRequest:
|
||||
def test_the_caller_identity_is_appended_to_the_job_tags(self):
|
||||
body = {"TranscriptionJobName": "j", "Tags": [{"Key": "env", "Value": "qa"}]}
|
||||
|
||||
owned = transcribe_owned_start_request(body, VIRTUAL_KEY)
|
||||
|
||||
assert owned == {
|
||||
"TranscriptionJobName": "j",
|
||||
"Tags": ({"Key": "env", "Value": "qa"}, {"Key": TRANSCRIBE_OWNER_TAG, "Value": "user-a"}),
|
||||
}
|
||||
assert body == {"TranscriptionJobName": "j", "Tags": [{"Key": "env", "Value": "qa"}]}
|
||||
|
||||
def test_a_request_without_tags_gets_the_owner_tag(self):
|
||||
owned = transcribe_owned_start_request({"TranscriptionJobName": "j"}, VIRTUAL_KEY)
|
||||
|
||||
assert owned == {"TranscriptionJobName": "j", "Tags": ({"Key": TRANSCRIBE_OWNER_TAG, "Value": "user-a"},)}
|
||||
|
||||
def test_the_caller_cannot_supply_the_owner_tag(self):
|
||||
owned = transcribe_owned_start_request(
|
||||
{"TranscriptionJobName": "j", "Tags": [{"Key": TRANSCRIBE_OWNER_TAG, "Value": "user-b"}]}, VIRTUAL_KEY
|
||||
)
|
||||
|
||||
assert isinstance(owned, TranscribeRefusal)
|
||||
assert owned.status_code == 400
|
||||
|
||||
@pytest.mark.parametrize("tags", ["env=qa", ["env"], {"Key": "env"}])
|
||||
def test_malformed_tags_are_refused(self, tags: object):
|
||||
owned = transcribe_owned_start_request({"TranscriptionJobName": "j", "Tags": tags}, VIRTUAL_KEY)
|
||||
|
||||
assert isinstance(owned, TranscribeRefusal)
|
||||
assert owned.status_code == 400
|
||||
|
||||
def test_a_key_without_any_identity_is_refused(self):
|
||||
owned = transcribe_owned_start_request({"TranscriptionJobName": "j"}, UserAPIKeyAuth())
|
||||
|
||||
assert isinstance(owned, TranscribeRefusal)
|
||||
assert owned.status_code == 400
|
||||
|
||||
|
||||
def _tagged(owner: str | None) -> dict[str, object]:
|
||||
tags = {"Tags": [{"Key": TRANSCRIBE_OWNER_TAG, "Value": owner}]} if owner is not None else {}
|
||||
return _job("COMPLETED", **tags)
|
||||
|
||||
|
||||
class TestTranscribeJobAccessRefusal:
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_key_that_started_the_job_may_read_it(self):
|
||||
get_job, seen = _sequence(_tagged("user-a"))
|
||||
|
||||
assert await transcribe_job_access_refusal("job-1", VIRTUAL_KEY, get_job) is None
|
||||
assert seen == ["job-1"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_job_started_by_another_key_is_reported_missing(self):
|
||||
get_job, _ = _sequence(_tagged("user-b"))
|
||||
|
||||
refusal = await transcribe_job_access_refusal("job-1", VIRTUAL_KEY, get_job)
|
||||
|
||||
assert refusal is not None
|
||||
assert refusal.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_job_started_outside_the_proxy_is_reported_missing(self):
|
||||
get_job, _ = _sequence(_tagged(None))
|
||||
|
||||
refusal = await transcribe_job_access_refusal("job-1", OTHER_VIRTUAL_KEY, get_job)
|
||||
|
||||
assert refusal is not None
|
||||
assert refusal.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_job_that_cannot_be_looked_up_is_reported_missing(self):
|
||||
async def get_job(job_name: str) -> dict[str, object]:
|
||||
raise httpx.HTTPStatusError("boom", request=MagicMock(), response=MagicMock())
|
||||
|
||||
refusal = await transcribe_job_access_refusal("job-1", VIRTUAL_KEY, get_job)
|
||||
|
||||
assert refusal is not None
|
||||
assert refusal.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_non_string_job_name_is_refused_before_any_lookup(self):
|
||||
get_job, seen = _sequence(_tagged("user-a"))
|
||||
|
||||
refusal = await transcribe_job_access_refusal(["job-1"], VIRTUAL_KEY, get_job)
|
||||
|
||||
assert refusal is not None
|
||||
assert refusal.status_code == 400
|
||||
assert seen == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_proxy_admin_reads_any_job_without_a_lookup(self):
|
||||
get_job, seen = _sequence(_tagged("user-b"))
|
||||
|
||||
assert await transcribe_job_access_refusal("job-1", ADMIN_KEY, get_job) is None
|
||||
assert seen == []
|
||||
|
||||
|
||||
class TestTranscribePassthroughHandler:
|
||||
def test_records_model_provider_and_the_given_cost(self):
|
||||
|
|
@ -453,13 +615,14 @@ class TestStartTranscriptionJobIsLoggedAtJobCost:
|
|||
scheduled.append(job_name)
|
||||
return 0.0
|
||||
|
||||
logging = PassThroughEndpointLogging(TranscribePassthroughLoggingHandler(job_pricer=job_pricer))
|
||||
immediate: list[dict[str, object]] = []
|
||||
|
||||
async def handle_logging(**kwargs: object) -> None:
|
||||
async def log_dispatch(**kwargs: object) -> None:
|
||||
immediate.append(kwargs)
|
||||
|
||||
logging._handle_logging = handle_logging # rebind-ok: the shared dispatch is the observable under test
|
||||
logging = PassThroughEndpointLogging(
|
||||
TranscribePassthroughLoggingHandler(job_pricer=job_pricer), log_dispatch=log_dispatch
|
||||
)
|
||||
|
||||
await logging.pass_through_async_success_handler(
|
||||
httpx_response=_make_response("StartTranscriptionJob"),
|
||||
|
|
|
|||
|
|
@ -5287,10 +5287,17 @@ def transcribe_client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]:
|
|||
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual"))
|
||||
monkeypatch.setitem(
|
||||
app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual", user_id="user-a")
|
||||
)
|
||||
yield TestClient(app)
|
||||
|
||||
|
||||
def _owned_job(owner: str | None, status: str = "COMPLETED") -> dict[str, object]:
|
||||
tags = {"Tags": [{"Key": "litellm-owner", "Value": owner}]} if owner is not None else {}
|
||||
return {"TranscriptionJob": {"TranscriptionJobName": "litellm-job-1", "TranscriptionJobStatus": status, **tags}}
|
||||
|
||||
|
||||
class TestTranscribeProxyRoute:
|
||||
START_JOB_BODY: Final = MappingProxyType(
|
||||
{
|
||||
|
|
@ -5299,6 +5306,7 @@ class TestTranscribeProxyRoute:
|
|||
"Media": {"MediaFileUri": "s3://bucket/audio.wav"},
|
||||
}
|
||||
)
|
||||
OWNER_TAG: Final = MappingProxyType({"Key": "litellm-owner", "Value": "user-a"})
|
||||
|
||||
def test_signs_and_forwards_start_transcription_job(self, transcribe_client: TestClient) -> None:
|
||||
upstream_body = {
|
||||
|
|
@ -5317,17 +5325,27 @@ class TestTranscribeProxyRoute:
|
|||
assert targets[0] == "Transcribe.StartTranscriptionJob"
|
||||
assert set(targets[1:]) <= {"Transcribe.GetTranscriptionJob"}
|
||||
sent = route.calls[0].request
|
||||
assert json.loads(sent.content) == dict(self.START_JOB_BODY)
|
||||
assert json.loads(sent.content) == {**dict(self.START_JOB_BODY), "Tags": [dict(self.OWNER_TAG)]}
|
||||
assert sent.headers["content-type"] == "application/x-amz-json-1.1"
|
||||
assert sent.headers["authorization"].startswith("AWS4-HMAC-SHA256 Credential=test-access-key/")
|
||||
assert "/us-west-2/transcribe/aws4_request" in sent.headers["authorization"]
|
||||
assert "x-amz-date" in sent.headers
|
||||
|
||||
def test_the_caller_cannot_forge_the_owner_tag(self, transcribe_client: TestClient) -> None:
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
route = upstream.post(TRANSCRIBE_UPSTREAM)
|
||||
response = transcribe_client.post(
|
||||
"/transcribe/StartTranscriptionJob",
|
||||
json={**dict(self.START_JOB_BODY), "Tags": [{"Key": "litellm-owner", "Value": "user-b"}]},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "litellm-owner" in response.json()["detail"]
|
||||
assert not route.called
|
||||
|
||||
def test_sdk_route_reads_operation_from_x_amz_target_and_resigns(self, transcribe_client: TestClient) -> None:
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
route = upstream.post(TRANSCRIBE_UPSTREAM).mock(
|
||||
return_value=httpx.Response(200, json={"TranscriptionJob": {"TranscriptionJobStatus": "COMPLETED"}})
|
||||
)
|
||||
route = upstream.post(TRANSCRIBE_UPSTREAM).mock(return_value=httpx.Response(200, json=_owned_job("user-a")))
|
||||
response = transcribe_client.post(
|
||||
"/transcribe",
|
||||
json={"TranscriptionJobName": "litellm-job-1"},
|
||||
|
|
@ -5338,21 +5356,76 @@ class TestTranscribeProxyRoute:
|
|||
},
|
||||
)
|
||||
|
||||
assert (response.status_code, response.json()) == (
|
||||
200,
|
||||
{"TranscriptionJob": {"TranscriptionJobStatus": "COMPLETED"}},
|
||||
)
|
||||
assert (response.status_code, response.json()) == (200, _owned_job("user-a"))
|
||||
assert [call.request.headers["x-amz-target"] for call in route.calls] == ["Transcribe.GetTranscriptionJob"] * 2
|
||||
sent = route.calls.last.request
|
||||
assert sent.headers["x-amz-target"] == "Transcribe.GetTranscriptionJob"
|
||||
assert "Credential=test-access-key/" in sent.headers["authorization"]
|
||||
assert "sk-virtual" not in sent.headers["authorization"]
|
||||
|
||||
@pytest.mark.parametrize("operation", ["GetTranscriptionJob", "DeleteTranscriptionJob"])
|
||||
@pytest.mark.parametrize("owner", ["user-b", None])
|
||||
def test_jobs_started_by_others_are_not_reachable(
|
||||
self, transcribe_client: TestClient, operation: str, owner: str | None
|
||||
) -> None:
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
route = upstream.post(TRANSCRIBE_UPSTREAM).mock(return_value=httpx.Response(200, json=_owned_job(owner)))
|
||||
response = transcribe_client.post(
|
||||
f"/transcribe/{operation}", json={"TranscriptionJobName": "litellm-job-1"}
|
||||
)
|
||||
|
||||
assert response.status_code == 404
|
||||
assert [call.request.headers["x-amz-target"] for call in route.calls] == ["Transcribe.GetTranscriptionJob"]
|
||||
|
||||
def test_the_owner_may_delete_the_job(self, transcribe_client: TestClient) -> None:
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
route = upstream.post(TRANSCRIBE_UPSTREAM)
|
||||
route.side_effect = [httpx.Response(200, json=_owned_job("user-a")), httpx.Response(200, json={})]
|
||||
response = transcribe_client.post(
|
||||
"/transcribe/DeleteTranscriptionJob", json={"TranscriptionJobName": "litellm-job-1"}
|
||||
)
|
||||
|
||||
assert (response.status_code, response.json()) == (200, {})
|
||||
assert [call.request.headers["x-amz-target"] for call in route.calls] == [
|
||||
"Transcribe.GetTranscriptionJob",
|
||||
"Transcribe.DeleteTranscriptionJob",
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("operation", ["ListTranscriptionJobs", "ListVocabularies", "DeleteVocabulary"])
|
||||
def test_account_wide_operations_need_a_proxy_admin(self, transcribe_client: TestClient, operation: str) -> None:
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
route = upstream.post(TRANSCRIBE_UPSTREAM)
|
||||
response = transcribe_client.post(f"/transcribe/{operation}", json={})
|
||||
|
||||
assert response.status_code == 403
|
||||
assert operation in response.json()["detail"]
|
||||
assert not route.called
|
||||
|
||||
def test_a_proxy_admin_reaches_account_wide_operations(
|
||||
self, transcribe_client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
monkeypatch.setitem(
|
||||
app.dependency_overrides,
|
||||
user_api_key_auth,
|
||||
lambda: UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
upstream.post(TRANSCRIBE_UPSTREAM).mock(
|
||||
return_value=httpx.Response(200, json={"TranscriptionJobSummaries": []})
|
||||
)
|
||||
response = transcribe_client.post("/transcribe/ListTranscriptionJobs", json={})
|
||||
|
||||
assert (response.status_code, response.json()) == (200, {"TranscriptionJobSummaries": []})
|
||||
|
||||
def test_upstream_error_status_and_body_are_returned(self, transcribe_client: TestClient) -> None:
|
||||
aws_error = {"__type": "BadRequestException", "Message": "The requested job couldn't be found."}
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
upstream.post(TRANSCRIBE_UPSTREAM).mock(return_value=httpx.Response(400, json=aws_error))
|
||||
response = transcribe_client.post(
|
||||
"/transcribe/GetTranscriptionJob", json={"TranscriptionJobName": "missing"}
|
||||
"/transcribe/StartTranscriptionJob",
|
||||
json={**dict(self.START_JOB_BODY), "TranscriptionJobName": "missing"},
|
||||
)
|
||||
|
||||
assert (response.status_code, response.json()) == (400, aws_error)
|
||||
|
|
@ -5386,7 +5459,7 @@ class TestTranscribeProxyRoute:
|
|||
with respx.mock(assert_all_called=False) as upstream:
|
||||
route = upstream.post(TRANSCRIBE_UPSTREAM)
|
||||
response = transcribe_client.post(
|
||||
"/transcribe/ListTranscriptionJobs", content=raw_body, headers={"Content-Type": "application/json"}
|
||||
"/transcribe/GetTranscriptionJob", content=raw_body, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
|
|
@ -5399,7 +5472,7 @@ class TestTranscribeProxyRoute:
|
|||
monkeypatch.delenv(name, raising=False)
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
route = upstream.post(TRANSCRIBE_UPSTREAM)
|
||||
response = transcribe_client.post("/transcribe/ListTranscriptionJobs", json={})
|
||||
response = transcribe_client.post("/transcribe/GetTranscriptionJob", json={})
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "AWS region" in response.json()["detail"]
|
||||
|
|
@ -5495,9 +5568,7 @@ class TestVertexAILiveWebsocketPassthrough:
|
|||
]
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
|
||||
monkeypatch.setattr(
|
||||
passthrough_module.passthrough_endpoint_router, "default_vertex_config", None
|
||||
)
|
||||
monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None)
|
||||
self._clear_vertex_env(monkeypatch)
|
||||
websocket = self._websocket()
|
||||
ensure_token = AsyncMock(return_value=("token-abc", "proj-db"))
|
||||
|
|
@ -5639,9 +5710,7 @@ class TestVertexAILiveWebsocketPassthrough:
|
|||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr(
|
||||
passthrough_module.passthrough_endpoint_router, "default_vertex_config", None
|
||||
)
|
||||
monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None)
|
||||
self._clear_vertex_env(monkeypatch)
|
||||
websocket = self._websocket()
|
||||
ensure_token = AsyncMock(side_effect=Exception("Unable to find your credentials"))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue