scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

fix(proxy): price a Transcribe job from the parsed start response instead of the relayed httpx body

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-18 02:48:09 +00:00
parent 2decf3761a
commit 3f824793e5
5 changed files with 69 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py
View file

@ -344,9 +344,9 @@ def transcribe_max_job_cost(cost_per_second: float) -> float:
return transcription_job_cost(TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS, cost_per_second)
def started_transcription_job(response_body: str) -> TranscriptionJobRecord | None:
def started_transcription_job(response_body: Mapping[str, object] | None) -> TranscriptionJobRecord | None:
try:
return _TranscriptionJobResponse.model_validate_json(response_body).TranscriptionJob
return _TranscriptionJobResponse.model_validate(response_body).TranscriptionJob
except ValidationError:
return None
@ -603,6 +603,7 @@ class TranscribePassthroughLoggingHandler:
def schedule_priced_job_logging(
self,
httpx_response: httpx.Response,
response_body: Mapping[str, object] | None,
logging_obj: LiteLLMLoggingObj,
url_route: str,
result: str,
@ -616,6 +617,7 @@ class TranscribePassthroughLoggingHandler:
task: Final = asyncio.create_task(
self._price_then_log(
httpx_response=httpx_response,
started_job=started_transcription_job(response_body),
logging_obj=logging_obj,
url_route=url_route,
result=result,
@ -634,6 +636,7 @@ class TranscribePassthroughLoggingHandler:
async def _price_then_log(
self,
httpx_response: httpx.Response,
started_job: TranscriptionJobRecord | None,
logging_obj: LiteLLMLoggingObj,
url_route: str,
result: str,
@ -654,7 +657,7 @@ class TranscribePassthroughLoggingHandler:
job_name if isinstance(job_name, str) else "",
aws_region_name,
cost_per_second,
started_transcription_job(httpx_response.text),
started_job,
)
payload: Final = self.transcribe_passthrough_handler(
httpx_response=httpx_response,

17
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -2668,17 +2668,22 @@ def _should_buffer_passthrough_response(response: httpx.Response) -> bool:
"""
Decide from the response headers whether the body must be read into memory.
JSON bodies (and upstream errors) stay buffered: spend logging, guardrails and
managed-id rewriting inspect them, and they are small in practice. Everything
else (jsonl batch results, octet-stream files, ...) is relayed to the client
chunk by chunk so a large body is never resident in full (LIT-4009). A missing
content-type is buffered because the body cannot be classified.
JSON bodies (including the AWS JSON protocol media types) and upstream errors
stay buffered: spend logging, guardrails and managed-id rewriting inspect them,
and they are small in practice. Everything else (jsonl batch results,
octet-stream files, ...) is relayed to the client chunk by chunk so a large
body is never resident in full (LIT-4009). A missing content-type is buffered
because the body cannot be classified.
"""
if response.status_code >= 400:
return True
content_type_header: Final[str] = response.headers.get("content-type", "")
media_type: Final = content_type_header.split(";")[0].strip().lower()
return media_type in ("", "application/json") or media_type.endswith("+json")
return (
media_type in ("", "application/json")
or media_type.endswith("+json")
or media_type.startswith("application/x-amz-json")
)
async def _relay_passthrough_response_bytes(

1
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -374,6 +374,7 @@ class PassThroughEndpointLogging:
):
self.transcribe_passthrough_logging_handler.schedule_priced_job_logging(
httpx_response=httpx_response,
response_body=response_body if isinstance(response_body, dict) else None,
logging_obj=logging_obj,
url_route=url_route,
result=result,

59
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py
View file

@ -40,13 +40,26 @@ from litellm.proxy.pass_through_endpoints.success_handler import (
COST_PER_SECOND = 0.0001
def _make_response(operation: str, text: str = '{"TranscriptionJob": {}}') -> httpx.Response:
def _make_response(operation: str) -> httpx.Response:
request = httpx.Request(
"POST",
"https://transcribe.us-west-2.amazonaws.com/",
headers={"X-Amz-Target": f"Transcribe.{operation}"},
)
return httpx.Response(200, request=request, text=text)
return httpx.Response(200, request=request, text='{"TranscriptionJob": {}}')
async def _relayed_response(operation: str, body: bytes) -> httpx.Response:
response = httpx.Response(
200,
request=_make_response(operation).request,
headers={"content-type": "application/x-amz-json-1.1"},
stream=httpx.ByteStream(body),
)
async for _ in response.aiter_bytes():
pass
await response.aclose()
return response
def _make_logging_obj() -> MagicMock:
@ -332,7 +345,7 @@ class TestPriceTranscriptionJob:
get_job, seen = _missing_job("BadRequestException")
media_seconds, measured = _media_probe(17.577)
started = started_transcription_job(
'{"TranscriptionJob": {"Media": {"MediaFileUri": "s3://b/started.wav"}, "CreationTime": 5.0}}'
{"TranscriptionJob": {"Media": {"MediaFileUri": "s3://b/started.wav"}, "CreationTime": 5.0}}
)
cost = await price_transcription_job(
@ -450,16 +463,22 @@ class TestMediaFileSeconds:
class TestStartedTranscriptionJob:
def test_reads_the_media_and_creation_time_from_the_start_response(self):
started = started_transcription_job(
'{"TranscriptionJob": {"TranscriptionJobName": "j", "Media": {"MediaFileUri": "s3://b/a.wav"},'
' "CreationTime": 1.5, "TranscriptionJobStatus": "IN_PROGRESS"}}'
{
"TranscriptionJob": {
"TranscriptionJobName": "j",
"Media": {"MediaFileUri": "s3://b/a.wav"},
"CreationTime": 1.5,
"TranscriptionJobStatus": "IN_PROGRESS",
}
}
)
assert started == TranscriptionJobRecord(
TranscriptionJobStatus="IN_PROGRESS", CreationTime=1.5, Media={"MediaFileUri": "s3://b/a.wav"}
)
@pytest.mark.parametrize("body", ["not json", "[]", '{"TranscriptionJob": {"CreationTime": "soon"}}'])
def test_unreadable_start_response_yields_no_record(self, body: str):
@pytest.mark.parametrize("body", [None, {"Message": "throttled"}, {"TranscriptionJob": {"CreationTime": "soon"}}])
def test_unreadable_start_response_yields_no_record(self, body: dict[str, object] | None):
assert started_transcription_job(body) is None
@ -744,6 +763,7 @@ class TestStartTranscriptionJobIsLoggedAtJobCost:
logging_obj = _make_logging_obj()
task = handler.schedule_priced_job_logging(
httpx_response=_make_response("StartTranscriptionJob"),
response_body={"TranscriptionJob": {}},
logging_obj=logging_obj,
url_route="https://transcribe.us-west-2.amazonaws.com/",
result='{"TranscriptionJob": {}}',
@ -778,6 +798,7 @@ class TestStartTranscriptionJobIsLoggedAtJobCost:
monkeypatch.delitem(litellm.model_cost, "transcribe/StartTranscriptionJob")
await TranscribePassthroughLoggingHandler(job_pricer=job_pricer).schedule_priced_job_logging(
httpx_response=_make_response("StartTranscriptionJob"),
response_body={"TranscriptionJob": {}},
logging_obj=_make_logging_obj(),
url_route="https://transcribe.us-west-2.amazonaws.com/",
result='{"TranscriptionJob": {}}',
@ -828,29 +849,34 @@ class TestStartTranscriptionJobIsLoggedAtJobCost:
assert [entry["response_cost"] for entry in immediate] == [0.0]
@pytest.mark.asyncio
async def test_pass_through_success_handler_gives_the_pricer_the_started_job_from_the_response(self):
async def test_pass_through_success_handler_prices_a_relayed_start_response_from_its_parsed_body(self):
started_jobs: list[TranscriptionJobRecord | None] = []
logged_costs: list[object] = []
async def job_pricer(
job_name: str, aws_region_name: str, cost_per_second: float, started_job: TranscriptionJobRecord | None
) -> float:
started_jobs.append(started_job)
return 0.0
return 18 * COST_PER_SECOND
async def log_dispatch(**kwargs: object) -> None:
pass
logged_costs.append(kwargs["response_cost"])
start_response = (
'{"TranscriptionJob": {"TranscriptionJobName": "litellm-job-1", "TranscriptionJobStatus": "IN_PROGRESS",'
' "Media": {"MediaFileUri": "s3://b/started.wav"}, "CreationTime": 5.0}}'
)
start_response = {
"TranscriptionJob": {
"TranscriptionJobName": "litellm-job-1",
"TranscriptionJobStatus": "IN_PROGRESS",
"Media": {"MediaFileUri": "s3://b/started.wav"},
"CreationTime": 5.0,
}
}
logging = PassThroughEndpointLogging(
TranscribePassthroughLoggingHandler(job_pricer=job_pricer), log_dispatch=log_dispatch
)
await logging.pass_through_async_success_handler(
httpx_response=_make_response("StartTranscriptionJob", text=start_response),
response_body=json.loads(start_response),
httpx_response=await _relayed_response("StartTranscriptionJob", json.dumps(start_response).encode()),
response_body=start_response,
logging_obj=_make_logging_obj(),
url_route="https://transcribe.us-west-2.amazonaws.com/",
result="",
@ -868,6 +894,7 @@ class TestStartTranscriptionJobIsLoggedAtJobCost:
TranscriptionJobStatus="IN_PROGRESS", CreationTime=5.0, Media={"MediaFileUri": "s3://b/started.wav"}
)
]
assert logged_costs == [pytest.approx(18 * COST_PER_SECOND)]
class TestIsTranscribeRoute:

14
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -4401,12 +4401,14 @@ async def test_pass_through_request_relays_non_json_body_without_buffering():
@pytest.mark.asyncio
async def test_pass_through_request_json_response_stays_buffered_for_logging():
@pytest.mark.parametrize("content_type", ["application/json", "application/x-amz-json-1.1"])
async def test_pass_through_request_json_response_stays_buffered_for_logging(content_type: str):
"""
JSON responses (content-type application/json) must keep the buffered
behavior: spend logging and guardrails inspect the parsed body, so the
handler reads the full upstream body and passes the parsed dict to the
success handler.
JSON responses (content-type application/json, and the AWS JSON protocol
media types AWS services such as Amazon Transcribe answer with) must keep
the buffered behavior: spend logging and guardrails inspect the parsed body,
so the handler reads the full upstream body and passes the parsed dict to
the success handler instead of handing it a relayed, already closed response.
"""
from fastapi.responses import StreamingResponse
@ -4417,7 +4419,7 @@ async def test_pass_through_request_json_response_stays_buffered_for_logging():
fake_client, cleanup = _inject_fake_passthrough_client(
_FakeUpstreamTransport(
status_code=200,
headers={"content-type": "application/json"},
headers={"content-type": content_type},
stream=upstream_stream,
),
timeout=312.0,