mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
feat(nvidia_riva)!: add riva_offline mode for offline-only ASR models (#41021)
* feat(nvidia_riva): add riva_offline mode for offline-only ASR models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(nvidia_riva): assert offline vs streaming via fake service call tracking Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(ui): regenerate schema.d.ts after user soft_budget docstring removal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(nvidia_riva): parse riva_offline strings, map oversize offline audio to 413 A multipart riva_offline=false no longer turns offline mode on, a malformed value drops under drop_params and answers 400 without it, and a Riva RESOURCE_EXHAUSTED for a message over the server's max now returns a non-retryable 413 instead of a retried 429 * fix(nvidia_riva): apply the request timeout to offline recognition and keep riva_offline on an api_base override Offline Recognize now runs as a gRPC future waited on with the request timeout, so a stalled Riva server returns a 408 and the call is cancelled. riva_offline is a mode flag, not a credential, so an api_base override no longer clears it. The 413 hint for audio over the gRPC message limit is reworded, and the new tests carry type annotations. * fix(nvidia_riva): wait the request timeout in seconds whatever its form and cap the riva client to versions that import A timeout from os.environ, a multipart form field, or an httpx.Timeout reached the gRPC future as a string or object, so result() raised TypeError and every offline call came back as a false 408. The handler now turns the timeout into float seconds with CompletionTimeout.as_limit, taking the read limit of an httpx.Timeout, for both the offline future and the streaming deadline The stt-nvidia-riva extra caps nvidia-riva-client below 2.17, since every later release either pins an older grpcio or ships stubs that need a newer one than the grpcio==1.78.0 pin, so a fresh install failed at import * fix(nvidia_riva): parse the timeout only where the call uses it and drop the riva client cap A malformed timeout no longer fails streaming calls on clients that take no timeout. On offline calls it is a 400, or the default under drop_params, and an httpx.Timeout with no read limit waits the default instead of forever. The nvidia-riva-client cap made installs next to a newer client unresolvable, so the extra goes back to main's specifier. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
4aec80f5ee
commit
d7c3588efd
5 changed files with 545 additions and 31 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue