diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index 651a2215052..748203b22d9 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -16,7 +16,9 @@ The handler is intentionally a thin orchestration layer: 3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS) vs self-hosted (any host:port, optional TLS) modes. 4. Stream the audio through Riva's ``streaming_response_generator`` and - aggregate ``is_final`` results into a single transcript. + aggregate ``is_final`` results into a single transcript, or, when + ``riva_offline`` is set, send the whole buffer through the unary + ``offline_recognize`` RPC for NIMs that only ship offline engines. 5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed on ``_hidden_params`` so cost calculation works. @@ -26,13 +28,19 @@ without the optional STT extras installed. import asyncio import inspect -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol +import httpx + +import litellm +from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_name, process_audio_file, ) +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.llms.nvidia_riva.audio_transcription.audio_utils import ( resample_to_riva_pcm, ) @@ -45,6 +53,7 @@ from litellm.llms.nvidia_riva.common_utils import ( NvidiaRivaException, grpc_error_to_litellm_exception, ) +from litellm.secret_managers.main import str_to_bool from litellm.types.utils import FileTypes, TranscriptionResponse from litellm.utils import convert_to_model_response_object @@ -67,10 +76,20 @@ class _RivaAuth(Protocol): """Opaque ``riva.client.Auth`` handle.""" +class _RecognizeFuture(Protocol): + def result(self, timeout: float | None = None) -> object: ... + + def done(self) -> bool: ... + + def cancel(self) -> bool: ... + + class _AsrService(Protocol): @property def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ... + def offline_recognize(self, audio_bytes: bytes, config: object, future: bool) -> _RecognizeFuture: ... + class _EndpointingConfig(Protocol): """Opaque ``EndpointingConfig`` protobuf message.""" @@ -128,6 +147,38 @@ def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding: return riva_asr_module.AudioEncoding +def _parse_riva_offline(value: object, drop_params: bool) -> bool: + match value: + case None: + return False + case bool(): + return value + case str() if (parsed := str_to_bool(value)) is not None: + return parsed + case _ if drop_params: + return False + case _: + raise NvidiaRivaException( + status_code=400, + message=f"riva_offline must be true or false, got {value!r}. Set drop_params to ignore it.", + ) + + +def _deadline_seconds(timeout: float | str | httpx.Timeout, drop_params: bool) -> float: + try: + limit: Final = CompletionTimeout.as_limit(timeout) + except ValueError: + if drop_params: + return COMPLETION_HTTP_FALLBACK_SECONDS + raise NvidiaRivaException( + status_code=400, + message=f"timeout must be a number of seconds, got {timeout!r}. Set drop_params to ignore it.", + ) from None + if not isinstance(limit, httpx.Timeout): + return limit + return COMPLETION_HTTP_FALLBACK_SECONDS if limit.read is None else limit.read + + class NvidiaRivaAudioTranscription: """Sync + async entry point for Riva ASR.""" @@ -138,7 +189,7 @@ class NvidiaRivaAudioTranscription: optional_params: dict, litellm_params: dict, model_response: TranscriptionResponse, - timeout: float, + timeout: float | str | httpx.Timeout | None, logging_obj: "LiteLLMLoggingObj", api_key: str | None, api_base: str | None, @@ -183,7 +234,7 @@ class NvidiaRivaAudioTranscription: optional_params: dict, litellm_params: dict, model_response: TranscriptionResponse, - timeout: float, + timeout: float | str | httpx.Timeout | None, logging_obj: "LiteLLMLoggingObj", api_key: str | None, api_base: str | None, @@ -213,7 +264,7 @@ class NvidiaRivaAudioTranscription: optional_params: dict, litellm_params: dict, model_response: TranscriptionResponse, - timeout: float, + timeout: float | str | httpx.Timeout | None, logging_obj: "LiteLLMLoggingObj", api_key: str | None, api_base: str | None, @@ -259,6 +310,10 @@ class NvidiaRivaAudioTranscription: response_format: Final = request_payload.get("response_format") or "json" timestamp_granularities: Final = request_payload.get("timestamp_granularities") + drop_params: Final = ( + "drop_params" in litellm_params and litellm_params["drop_params"] is True + ) or litellm.drop_params is True + riva_offline: Final = _parse_riva_offline(optional_params.get("riva_offline"), drop_params=drop_params) riva_module, riva_asr_module = _import_riva() auth_obj: Final = self._construct_auth( @@ -272,10 +327,6 @@ class NvidiaRivaAudioTranscription: riva_asr_module=riva_asr_module, recognition_config_dict=recognition_config_dict, ) - streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig( - config=recognition_config, interim_results=False - ) - logging_obj.pre_call( input=None, api_key=api_key, @@ -286,24 +337,31 @@ class NvidiaRivaAudioTranscription: "recognition_config": recognition_config_dict, "nvcf_function_id_set": bool(optional_params.get("nvcf_function_id")), "use_ssl": optional_params.get("use_ssl"), + "riva_offline": riva_offline, }, }, ) try: asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj) - audio_chunks: Final = self._iter_audio_chunks(resampled.pcm_bytes) - stream_kwargs: Final[dict[str, object]] = { - "audio_chunks": audio_chunks, - "streaming_config": streaming_config, - } - # Forward the deadline so the stream cannot block forever if the - # server stalls. Older riva-client versions do not accept a - # ``timeout`` kwarg, so pass it only when supported. - if timeout is not None and self._supports_timeout_kwarg(asr_service.streaming_response_generator): - stream_kwargs["timeout"] = float(timeout) - stream: Final = asr_service.streaming_response_generator(**stream_kwargs) - final_results: Final = self._collect_final_results(stream) + responses: Final = ( + self._offline_responses( + asr_service=asr_service, + pcm_bytes=resampled.pcm_bytes, + recognition_config=recognition_config, + timeout=None if timeout is None else _deadline_seconds(timeout, drop_params), + ) + if riva_offline + else self._streaming_responses( + asr_service=asr_service, + riva_asr_module=riva_asr_module, + pcm_bytes=resampled.pcm_bytes, + recognition_config=recognition_config, + timeout=timeout, + drop_params=drop_params, + ) + ) + final_results: Final = self._collect_final_results(responses, require_final=not riva_offline) except NvidiaRivaException: raise except Exception as e: @@ -409,6 +467,54 @@ class NvidiaRivaAudioTranscription: return config + def _streaming_responses( + self, + asr_service: _AsrService, + riva_asr_module: _RivaAsrModule, + pcm_bytes: bytes, + recognition_config: _RecognitionConfig, + timeout: float | str | httpx.Timeout | None, + drop_params: bool, + ) -> Iterable[object]: + streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig( + config=recognition_config, interim_results=False + ) + return asr_service.streaming_response_generator( + audio_chunks=self._iter_audio_chunks(pcm_bytes), + streaming_config=streaming_config, + **self._timeout_kwargs(asr_service.streaming_response_generator, timeout, drop_params), + ) + + def _offline_responses( + self, + asr_service: _AsrService, + pcm_bytes: bytes, + recognition_config: _RecognitionConfig, + timeout: float | None, + ) -> Iterable[object]: + pending: Final = asr_service.offline_recognize(audio_bytes=pcm_bytes, config=recognition_config, future=True) + try: + return (pending.result(timeout=timeout),) + except Exception as e: + if pending.done(): + raise + pending.cancel() + raise NvidiaRivaException( + status_code=408, + message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds", + ) from e + + @classmethod + def _timeout_kwargs( + cls, rpc: Callable[..., object], timeout: float | str | httpx.Timeout | None, drop_params: bool + ) -> Mapping[str, float]: + # Forward the deadline so the call cannot block forever if the + # server stalls. Older riva-client versions do not accept a + # ``timeout`` kwarg, so pass it only when supported. + if timeout is None or not cls._supports_timeout_kwarg(rpc): + return MappingProxyType({}) + return MappingProxyType({"timeout": _deadline_seconds(timeout, drop_params)}) + @staticmethod def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool: try: @@ -429,9 +535,9 @@ class NvidiaRivaAudioTranscription: yield chunk @staticmethod - def _collect_final_results(stream) -> list[dict[str, object]]: + def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> list[dict[str, object]]: """ - Walk the gRPC stream, ignore empty / non-final chunks, and return a + Walk the gRPC responses, ignore empty / non-final chunks, and return a list of normalized final-result dicts. Matching the user's note: the ``id`` blocks with no ``results`` are streaming heartbeats and must be skipped. @@ -440,7 +546,7 @@ class NvidiaRivaAudioTranscription: for response in stream: results = getattr(response, "results", None) or [] for result in results: - if not getattr(result, "is_final", False): + if require_final and not getattr(result, "is_final", False): continue alternatives = getattr(result, "alternatives", None) or [] if not alternatives: diff --git a/litellm/llms/nvidia_riva/common_utils.py b/litellm/llms/nvidia_riva/common_utils.py index b654e725dc8..c671b14e742 100644 --- a/litellm/llms/nvidia_riva/common_utils.py +++ b/litellm/llms/nvidia_riva/common_utils.py @@ -41,6 +41,9 @@ _GRPC_STATUS_CODE_TO_HTTP: Final[dict] = { } +_GRPC_MESSAGE_TOO_LARGE: Final = "message larger than max" + + def _extract_grpc_status_name(error: object) -> str | None: """ Best-effort extraction of a gRPC StatusCode name from an arbitrary error. @@ -79,8 +82,16 @@ def grpc_error_to_litellm_exception(error: Exception) -> NvidiaRivaException: a ``NvidiaRivaException`` with an appropriate HTTP-equivalent status code. """ status_name: Final = _extract_grpc_status_name(error) - http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500) - detail: Final = _extract_grpc_details(error) or str(error) + if status_name == "RESOURCE_EXHAUSTED" and _GRPC_MESSAGE_TOO_LARGE in detail: + return NvidiaRivaException( + status_code=413, + message=( + f"NVIDIA Riva gRPC error ({status_name}): {detail}. The audio is too large for one offline " + "request to this Riva server: send shorter audio, or drop riva_offline on models that support streaming." + ), + ) + + http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500) message = f"NVIDIA Riva gRPC error ({status_name}): {detail}" if status_name else f"NVIDIA Riva error: {detail}" return NvidiaRivaException(status_code=http_status, message=message) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 68cd8420166..7a186d84aa5 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -446,6 +446,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # the request away from the admin's pinned configuration. "nvcf_function_id", "use_ssl", + "riva_offline", # Per-deployment opt-in that hands the whole call to the Rust core. It is a # deployment decision, not a request one: the Rust path uses its own client # rather than the one the deployment configured, and reports no post_call, diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py index 7ecc0b47d9f..dbb5cde6060 100644 --- a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py @@ -12,11 +12,12 @@ import io from types import SimpleNamespace from unittest.mock import MagicMock +import httpx import numpy as np import pytest import soundfile as sf - +import litellm from litellm.llms.nvidia_riva.audio_transcription import handler as handler_mod from litellm.llms.nvidia_riva.audio_transcription.handler import ( NvidiaRivaAudioTranscription, @@ -93,6 +94,24 @@ def mock_riva(monkeypatch): LINEAR_PCM = "LINEAR_PCM" streaming_responses_holder = {"value": []} + offline_response_holder = {"value": None} + offline_error_holder = {"value": None} + offline_future_holder = {"value": None} + streaming_calls = [] + offline_calls = {} + + class FakeRecognizeFuture: + def result(self, timeout: float | None = None) -> object: + offline_calls["timeout"] = timeout + if offline_error_holder["value"] is not None: + raise offline_error_holder["value"] + return offline_response_holder["value"] + + def done(self) -> bool: + return True + + def cancel(self) -> bool: + return False class FakeASRService: def __init__(self, auth): @@ -100,9 +119,15 @@ def mock_riva(monkeypatch): def streaming_response_generator(self, audio_chunks, streaming_config): # Drain audio_chunks generator so we exercise the chunking path. - list(audio_chunks) + streaming_calls.append(list(audio_chunks)) yield from streaming_responses_holder["value"] + def offline_recognize(self, audio_bytes: bytes, config: object, future: bool = False) -> object: + offline_calls["audio_bytes"] = audio_bytes + offline_calls["config"] = config + pending = offline_future_holder["value"] or FakeRecognizeFuture() + return pending if future else pending.result() + fake_riva_client = SimpleNamespace( Auth=FakeAuth, ASRService=FakeASRService, @@ -120,6 +145,11 @@ def mock_riva(monkeypatch): return SimpleNamespace( auth_calls=auth_calls, responses=streaming_responses_holder, + offline_response=offline_response_holder, + offline_error=offline_error_holder, + offline_future=offline_future_holder, + streaming_calls=streaming_calls, + offline_calls=offline_calls, client=fake_riva_client, ) @@ -203,6 +233,339 @@ def test_sync_path_aggregates_only_final_results(mock_riva, logging_obj): ) +def test_riva_offline_uses_unary_recognize_and_keeps_results_without_is_final( + mock_riva: SimpleNamespace, logging_obj: MagicMock +) -> None: + """ + Offline-only NIMs (Parakeet TDT, Whisper) reject StreamingRecognize. With + ``riva_offline`` the handler must call ``offline_recognize`` with the whole + resampled buffer and a bare ``RecognitionConfig``, and must keep every + result even though ``RecognizeResponse`` results carry no ``is_final``. + """ + mock_riva.offline_response["value"] = _fake_response( + results=[ + SimpleNamespace(alternatives=[_fake_alternative("Hello,", words=[_fake_word("Hello,", 0, 320)])]), + SimpleNamespace(alternatives=[_fake_alternative(" world.", words=[_fake_word("world.", 480, 870)])]), + ] + ) + + impl = NvidiaRivaAudioTranscription() + response: TranscriptionResponse = impl.audio_transcriptions( + model="nvidia/parakeet-0.6b-tdt", + audio_file=_make_wav_bytes(), + optional_params={ + "language_code": "multi", + "riva_offline": True, + "enable_word_time_offsets": True, + "response_format": "verbose_json", + "timestamp_granularities": ["word"], + }, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert response.text == "Hello, world." + assert response["words"][1]["end"] == pytest.approx(0.87) + assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0, abs=0.05) + assert mock_riva.streaming_calls == [] + assert len(mock_riva.offline_calls["audio_bytes"]) == pytest.approx(32000, abs=1600) + assert mock_riva.offline_calls["config"]._kwargs["language_code"] == "multi" + assert logging_obj.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"]["riva_offline"] is True + + +def test_default_mode_streams_and_never_calls_offline_recognize( + mock_riva: SimpleNamespace, logging_obj: MagicMock +) -> None: + mock_riva.responses["value"] = [ + _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + ] + + impl = NvidiaRivaAudioTranscription() + response = impl.audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US"}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + assert response.text == "streamed" + assert mock_riva.offline_calls == {} + assert len(mock_riva.streaming_calls) == 1 + assert logging_obj.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"]["riva_offline"] is False + + + +@pytest.mark.parametrize(("riva_offline", "expect_offline"), [("false", False), ("False", False), ("true", True)]) +def test_riva_offline_string_from_multipart_form_is_parsed( + mock_riva: SimpleNamespace, logging_obj: MagicMock, riva_offline: str, expect_offline: bool +) -> None: + mock_riva.responses["value"] = [ + _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + ] + mock_riva.offline_response["value"] = _fake_response( + results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] + ) + + response = NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": riva_offline}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert response.text == ("offline" if expect_offline else "streamed") + assert (mock_riva.offline_calls != {}) is expect_offline + + +@pytest.mark.parametrize("riva_offline", ["yes", "", 1]) +def test_riva_offline_malformed_value_is_rejected_before_any_rpc( + mock_riva: SimpleNamespace, logging_obj: MagicMock, riva_offline: object +) -> None: + with pytest.raises(NvidiaRivaException) as excinfo: + NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": riva_offline}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert excinfo.value.status_code == 400 + assert "riva_offline" in excinfo.value.message + assert mock_riva.offline_calls == {} + assert mock_riva.streaming_calls == [] + + +def test_riva_offline_malformed_value_is_dropped_under_drop_params( + mock_riva: SimpleNamespace, logging_obj: MagicMock +) -> None: + mock_riva.responses["value"] = [ + _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + ] + + response = NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": "yes"}, + litellm_params={"drop_params": True}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert response.text == "streamed" + assert mock_riva.offline_calls == {} + + +def test_deployment_riva_offline_reaches_the_handler_through_transcription(mock_riva: SimpleNamespace) -> None: + mock_riva.offline_response["value"] = _fake_response( + results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] + ) + + response = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + riva_offline=True, + ) + + assert response.text == "offline" + assert mock_riva.streaming_calls == [] + assert mock_riva.offline_calls != {} + + +@pytest.mark.parametrize( + ("timeout", "waited_for"), + [("7.5", 7.5), (httpx.Timeout(30.0, read=7.5), 7.5), (httpx.Timeout(5.0, read=None), 600.0), (None, None)], +) +def test_offline_recognition_waits_the_timeout_in_seconds_whatever_its_form( + mock_riva: SimpleNamespace, timeout: str | httpx.Timeout | None, waited_for: float | None +) -> None: + mock_riva.offline_response["value"] = _fake_response( + results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] + ) + + response = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + riva_offline=True, + timeout=timeout, + ) + + assert response.text == "offline" + assert mock_riva.offline_calls["timeout"] == waited_for + + +@pytest.mark.parametrize("timeout", ["abc", "30s"]) +def test_malformed_timeout_on_offline_call_is_a_400_before_any_rpc( + mock_riva: SimpleNamespace, timeout: str +) -> None: + with pytest.raises(NvidiaRivaException) as excinfo: + litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + riva_offline=True, + timeout=timeout, + ) + + assert excinfo.value.status_code == 400 + assert "timeout" in excinfo.value.message + assert "audio_bytes" not in mock_riva.offline_calls + + +def test_malformed_timeout_on_offline_call_waits_the_default_under_drop_params( + mock_riva: SimpleNamespace, +) -> None: + mock_riva.offline_response["value"] = _fake_response( + results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] + ) + + response = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + riva_offline=True, + timeout="abc", + drop_params=True, + ) + + assert response.text == "offline" + assert mock_riva.offline_calls["timeout"] == 600.0 + + +def test_malformed_timeout_still_streams_when_the_client_takes_no_timeout(mock_riva: SimpleNamespace) -> None: + mock_riva.responses["value"] = [ + _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + ] + + response = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-ctc-1.1b", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + timeout="abc", + ) + + assert response.text == "streamed" + assert len(mock_riva.streaming_calls) == 1 + + +def test_oversized_offline_audio_is_a_non_retryable_413(mock_riva: SimpleNamespace, logging_obj: MagicMock) -> None: + class FakeGrpcError(Exception): + def code(self) -> SimpleNamespace: + return SimpleNamespace(name="RESOURCE_EXHAUSTED") + + def details(self) -> str: + return "SERVER: Received message larger than max (4348825 vs. 4194304)" + + mock_riva.offline_error["value"] = FakeGrpcError("rpc fail") + + with pytest.raises(NvidiaRivaException) as excinfo: + NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": True}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert excinfo.value.status_code == 413 + assert litellm.utils.should_retry(excinfo.value.status_code) is False + + +def test_stalled_offline_recognition_is_cancelled_and_times_out_as_408( + mock_riva: SimpleNamespace, logging_obj: MagicMock +) -> None: + class FutureTimeoutError(Exception): + pass + + class StalledRecognizeFuture: + def __init__(self) -> None: + self.waited_for: float | None = None + self.cancelled = False + + def result(self, timeout: float | None = None) -> object: + self.waited_for = timeout + raise FutureTimeoutError() + + def done(self) -> bool: + return self.cancelled + + def cancel(self) -> bool: + self.cancelled = True + return True + + stalled = StalledRecognizeFuture() + mock_riva.offline_future["value"] = stalled + + with pytest.raises(NvidiaRivaException) as excinfo: + NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": True}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=7.5, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert excinfo.value.status_code == 408 + assert stalled.waited_for == 7.5 + assert stalled.cancelled is True + +def test_riva_offline_grpc_error_is_wrapped(mock_riva: SimpleNamespace, logging_obj: MagicMock) -> None: + class FakeGrpcError(Exception): + def code(self) -> SimpleNamespace: + return SimpleNamespace(name="INVALID_ARGUMENT") + + def details(self) -> str: + return "Unavailable model requested" + + mock_riva.offline_error["value"] = FakeGrpcError("rpc fail") + + impl = NvidiaRivaAudioTranscription() + with pytest.raises(NvidiaRivaException) as excinfo: + impl.audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "multi", "riva_offline": True}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + assert excinfo.value.status_code == 400 + assert "INVALID_ARGUMENT" in excinfo.value.message + + def test_auth_nvcf_defaults_use_ssl_and_attaches_function_id(mock_riva, logging_obj): mock_riva.responses["value"] = [ _fake_response( @@ -343,9 +706,10 @@ def test_async_path_uses_to_thread(mock_riva, logging_obj): ) +@pytest.mark.parametrize("timeout", [12.5, "12.5", httpx.Timeout(60.0, read=12.5)]) def test_timeout_is_forwarded_to_streaming_generator_when_supported( - mock_riva, logging_obj -): + mock_riva: SimpleNamespace, logging_obj: MagicMock, timeout: float | str | httpx.Timeout +) -> None: """ Without a deadline the gRPC stream can block forever on a stalled Riva server. The handler must forward the call-level ``timeout`` to @@ -374,7 +738,7 @@ def test_timeout_is_forwarded_to_streaming_generator_when_supported( optional_params={"language_code": "en-US"}, litellm_params={}, model_response=TranscriptionResponse(), - timeout=12.5, + timeout=timeout, logging_obj=logging_obj, api_key=None, api_base="localhost:50051", diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 7597847ef6d..ee1ae43fd3a 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3902,6 +3902,38 @@ class TestIsRequestBodySafeBlocksNVCFFunctionOverride: ) +def test_riva_offline_in_request_body_is_rejected() -> None: + with pytest.raises(ValueError, match="riva_offline"): + is_request_body_safe( + request_body={"model": "nvidia_riva/parakeet", "riva_offline": True}, + general_settings={}, + llm_router=None, + model="nvidia_riva/parakeet", + ) + + +def test_admin_opt_in_proxy_wide_allows_riva_offline() -> None: + assert ( + is_request_body_safe( + request_body={"model": "nvidia_riva/parakeet", "riva_offline": True}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="nvidia_riva/parakeet", + ) + is True + ) + + +def test_api_base_override_keeps_the_deployments_riva_offline_mode() -> None: + from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + + out: Final = get_dynamic_litellm_params( + litellm_params={"model": "nvidia_riva/nvidia/parakeet-0.6b-tdt", "riva_offline": True}, + request_kwargs={"api_base": "caller-riva:50051"}, + ) + assert (out["api_base"], out["riva_offline"]) == ("caller-riva:50051", True) + + class TestIsRequestBodySafeBlocksRivaUseSsl: """``use_ssl`` is rejected as a request-body param unless the admin opted in proxy-wide or per-deployment."""