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 dd85d723fb5..b977cf3ccc1 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 @@ -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, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 685c19062bb..a0866d11374 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index df539d687d7..c7f4f5c16f5 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -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, 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 2d0d65db805..481533fd7d4 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 @@ -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: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d854ee39ff4..ac601079d80 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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,