From 42b4cfc5bfa2031d48bafe11cf321f4bdd98eaa5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 4 Nov 2026 18:07:31 -0800 Subject: [PATCH] fix(nvidia_riva): send the api key exactly once by calling riva over grpc directly (#50140) * fix(nvidia_riva): send the api key once per call by talking to Riva over grpc directly Drops nvidia-riva-client from the stt-nvidia-riva extra. Depending on its version it sent the bearer token zero times on plaintext or twice on TLS, so token-checking Riva deployments rejected every request (#50059). The provider now declares the few Riva messages it uses, opens its own grpc channel and puts authorization and function-id on each call exactly once * test(nvidia_riva): run the provider against a real grpc server built from nvidia's own definitions * fix(nvidia_riva): validate endpointing settings with pydantic before building the riva message * test(nvidia_riva): check the key is sent once over tls, and drop a comment that restated the code --- .../audio_transcription/handler.py | 328 +---- .../audio_transcription/riva_asr_schema.py | 199 +++ .../audio_transcription/riva_grpc.py | 113 ++ pyproject.toml | 7 +- tests/integration/_support/riva.py | 16 +- .../providers/test_nvidia_riva_wire.py | 9 + .../audio_transcription/test_handler.py | 1253 ++++++++--------- .../test_riva_asr_schema.py | 41 + uv.lock | 10 +- 9 files changed, 985 insertions(+), 991 deletions(-) create mode 100644 litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py create mode 100644 litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py create mode 100644 tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index cafeece0a4e..d2ef33b89dd 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -13,24 +13,26 @@ The handler is intentionally a thin orchestration layer: 2. Build ``RecognitionConfig`` / ``StreamingRecognitionConfig`` protobufs from the structured dict produced by :class:`NvidiaRivaAudioTranscriptionConfig`. -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, or, when - ``riva_offline`` is set, send the whole buffer through the unary - ``offline_recognize`` RPC for NIMs that only ship offline engines. +3. Open a gRPC channel honoring NVCF (TLS by default) vs self-hosted (any + host:port, optional TLS) modes, and attach the ``authorization`` and NVCF + ``function-id`` metadata to each call exactly once. +4. Stream the audio through Riva's ``StreamingRecognize`` RPC and aggregate + ``is_final`` results into a single transcript, or, when ``riva_offline`` + is set, send the whole buffer through the unary ``Recognize`` RPC for NIMs + that only ship offline engines. 5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed on ``_hidden_params`` so cost calculation works. -``riva-client`` is imported lazily so ``litellm`` core remains usable -without the optional STT extras installed. +``grpcio`` and ``protobuf`` are imported lazily so ``litellm`` core remains +usable without the optional STT extras installed. """ +from __future__ import annotations + import asyncio -import inspect -from collections.abc import Callable, Iterable, Mapping, Sequence -from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol +from collections.abc import Iterable, Sequence +from contextlib import closing +from typing import TYPE_CHECKING, Final import httpx @@ -64,89 +66,19 @@ if TYPE_CHECKING: Logging as LiteLLMLoggingObj, ) -# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches -# NVIDIA's recommended chunk size for streaming ASR — small enough for -# responsive endpointing, large enough to keep per-RPC overhead low. -_DEFAULT_CHUNK_SAMPLES: Final = 1600 -_DEFAULT_CHUNK_BYTES: Final = _DEFAULT_CHUNK_SAMPLES * 2 # int16 = 2 bytes/sample - - _RIVA_INSTALL_HINT = "NVIDIA Riva client is not installed. Install with `pip install 'litellm[stt-nvidia-riva]'`." -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.""" - - -class _EndpointingConfigField(Protocol): - CopyFrom: Callable[[_EndpointingConfig], None] - - -class _RecognitionConfig(Protocol): - @property - def endpointing_config(self) -> _EndpointingConfigField: ... - - -class _StreamingRecognitionConfig(Protocol): - """Opaque ``StreamingRecognitionConfig`` protobuf message.""" - - -class _AudioEncoding(Protocol): - @property - def LINEAR_PCM(self) -> object: ... - - -class _RivaClientModule(Protocol): - """The ``riva.client`` entry points this handler calls.""" - - @property - def Auth(self) -> Callable[..., _RivaAuth]: ... - - @property - def ASRService(self) -> Callable[[_RivaAuth], _AsrService]: ... - - -class _RivaAsrModule(Protocol): - """The protobuf constructors this handler calls, from whichever module exposes them.""" - - @property - def AudioEncoding(self) -> _AudioEncoding: ... - - @property - def RecognitionConfig(self) -> Callable[..., _RecognitionConfig]: ... - - @property - def StreamingRecognitionConfig(self) -> Callable[..., _StreamingRecognitionConfig]: ... - - @property - def EndpointingConfig(self) -> Callable[..., _EndpointingConfig]: ... - - -def _auth_factory(riva_module: _RivaClientModule) -> Callable[..., _RivaAuth]: - return riva_module.Auth - - -def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding: - return riva_asr_module.AudioEncoding +def _use_ssl(optional_params: dict) -> bool: + """ + NVCF endpoints are TLS-only, so ``nvcf_function_id`` defaults ``use_ssl`` + to True. Self-hosted Riva defaults to plaintext but honors an explicit + ``use_ssl``, since TLS-terminating ingresses in front of it are common. + """ + use_ssl_override: Final = optional_params.get("use_ssl") + if use_ssl_override is not None: + return bool(use_ssl_override) + return bool(optional_params.get("nvcf_function_id")) def _parse_riva_offline(value: object, drop_params: bool) -> bool: @@ -192,7 +124,7 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, atranscription: bool = False, @@ -237,13 +169,11 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, provider_config: NvidiaRivaAudioTranscriptionConfig | None = None, ) -> TranscriptionResponse: - # ``riva-client`` exposes a sync streaming generator, so we offload - # the blocking call to a worker thread to keep the event loop free. return await asyncio.to_thread( self._run_sync, model=model, @@ -267,7 +197,7 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, provider_config: NvidiaRivaAudioTranscriptionConfig, @@ -317,18 +247,12 @@ class NvidiaRivaAudioTranscription: ) 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( - riva_module=riva_module, - api_base=api_base, - api_key=api_key, - optional_params=optional_params, - ) + try: + from litellm.llms.nvidia_riva.audio_transcription import riva_grpc + except ImportError as e: + raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - recognition_config: Final = self._build_recognition_config_proto( - riva_asr_module=riva_asr_module, - recognition_config_dict=recognition_config_dict, - ) + recognition_config: Final = riva_grpc.recognition_config(recognition_config_dict) logging_obj.pre_call( input=None, api_key=api_key, @@ -344,26 +268,18 @@ class NvidiaRivaAudioTranscription: }, ) + metadata: Final = riva_grpc.call_metadata(api_key, optional_params.get("nvcf_function_id")) try: - asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj) - 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), + deadline: Final = None if timeout is None else _deadline_seconds(timeout, drop_params) + with closing(riva_grpc.open_channel(api_base, _use_ssl(optional_params))) as channel: + responses: Final = ( + riva_grpc.recognize(channel, recognition_config, resampled.pcm_bytes, metadata, deadline) + if riva_offline + else riva_grpc.streaming_recognize( + channel, recognition_config, resampled.pcm_bytes, metadata, deadline + ) ) - 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) + final_results: Final = self._collect_final_results(responses, require_final=not riva_offline) except NvidiaRivaException: raise except Exception as e: @@ -400,142 +316,6 @@ class NvidiaRivaAudioTranscription: return final_response - def _construct_auth( - self, - riva_module: _RivaClientModule, - api_base: str, - api_key: str | None, - optional_params: dict, - ) -> _RivaAuth: - """ - Build a ``riva.client.Auth`` object. - - - When ``nvcf_function_id`` is provided we attach the NVCF - ``function-id`` and bearer ``authorization`` metadata, and default - ``use_ssl`` to True (NVCF endpoints are TLS-only). - - Otherwise (self-hosted) we default ``use_ssl`` to False but still - honor an explicit override — self-hosted Riva behind an ingress - with TLS termination is a real deployment topology. - """ - nvcf_function_id: Final = optional_params.get("nvcf_function_id") - use_ssl_override: Final = optional_params.get("use_ssl") - use_ssl: Final = bool(use_ssl_override) if use_ssl_override is not None else bool(nvcf_function_id) - - metadata: Final[list[tuple[str, str]]] = [] - if nvcf_function_id: - metadata.append(("function-id", str(nvcf_function_id))) - if api_key: - metadata.append(("authorization", f"Bearer {api_key}")) - - try: - return _auth_factory(riva_module)(uri=api_base, use_ssl=use_ssl, metadata_args=metadata) - except TypeError: - # Older riva-client signatures used positional-only args. - return _auth_factory(riva_module)(None, use_ssl, api_base, metadata) - - def _build_recognition_config_proto( - self, riva_asr_module: _RivaAsrModule, recognition_config_dict: dict[str, Any] - ) -> _RecognitionConfig: - encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper() - encoding_enum: Final[object] = getattr( - _audio_encoding(riva_asr_module), - encoding_name, - _audio_encoding(riva_asr_module).LINEAR_PCM, - ) - - config: Final[_RecognitionConfig] = riva_asr_module.RecognitionConfig( - encoding=encoding_enum, - sample_rate_hertz=int(recognition_config_dict["sample_rate_hertz"]), - language_code=recognition_config_dict["language_code"], - audio_channel_count=int(recognition_config_dict["audio_channel_count"]), - enable_automatic_punctuation=bool(recognition_config_dict.get("enable_automatic_punctuation", True)), - enable_word_time_offsets=bool(recognition_config_dict.get("enable_word_time_offsets", False)), - max_alternatives=int(recognition_config_dict.get("max_alternatives", 1)), - model=recognition_config_dict.get("model", "") or "", - verbatim_transcripts=bool(recognition_config_dict.get("verbatim_transcripts", False)), - profanity_filter=bool(recognition_config_dict.get("profanity_filter", False)), - ) - - endpointing: Final = recognition_config_dict.get("endpointing_config") - if isinstance(endpointing, dict) and endpointing: - try: - ep: Final[_EndpointingConfig] = riva_asr_module.EndpointingConfig(**endpointing) - config.endpointing_config.CopyFrom(ep) - except Exception: - # If the user supplied an unknown EndpointingConfig field - # (older Riva server), fall back to Riva's defaults rather - # than failing the whole request. - pass - - 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: - sig: Final = inspect.signature(callable_obj) - except (TypeError, ValueError): - return False - params: Final = sig.parameters - if "timeout" in params: - return True - return any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) - - @staticmethod - def _iter_audio_chunks(pcm_bytes: bytes): - for offset in range(0, len(pcm_bytes), _DEFAULT_CHUNK_BYTES): - chunk = pcm_bytes[offset : offset + _DEFAULT_CHUNK_BYTES] - if not chunk: - continue - yield chunk - @staticmethod def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> Sequence[RivaFinalResult]: """ @@ -570,27 +350,3 @@ class NvidiaRivaAudioTranscription: } ) return final_results - - -def _import_riva() -> tuple[_RivaClientModule, _RivaAsrModule]: - """ - Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``. - - We try the SDK first (preferred) and fall back to importing the proto - module separately when the SDK packaging changes between versions. - """ - try: - import riva.client as riva_client - except ImportError as e: - raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - - riva_asr_module = riva_client - if not hasattr(riva_asr_module, "RecognitionConfig"): - try: - from riva.client.proto import riva_asr_pb2 - - riva_asr_module = riva_asr_pb2 - except ImportError as e: - raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - - return riva_client, riva_asr_module diff --git a/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py b/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py new file mode 100644 index 00000000000..13af5820f07 --- /dev/null +++ b/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py @@ -0,0 +1,199 @@ +""" +The part of NVIDIA Riva's speech recognition gRPC schema that LiteLLM sends and reads. + +Names and numbers match ``riva/proto/riva_asr.proto`` and ``riva/proto/riva_audio.proto`` in +https://github.com/nvidia-riva/common. Fields LiteLLM never sets or reads are left out, and the +wire format skips them. The messages live in their own descriptor pool so they never collide +with nvidia-riva-client's copies when both are imported in one process. +""" + +from dataclasses import dataclass +from enum import Enum, IntEnum +from typing import Final + +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory +from google.protobuf.message import Message + +_PACKAGE: Final = "nvidia.riva.asr" +_SERVICE: Final = f"/{_PACKAGE}.RivaSpeechRecognition" +RECOGNIZE_METHOD: Final = f"{_SERVICE}/Recognize" +STREAMING_RECOGNIZE_METHOD: Final = f"{_SERVICE}/StreamingRecognize" + + +class _Kind(Enum): + FLOAT = "TYPE_FLOAT" + INT32 = "TYPE_INT32" + BOOL = "TYPE_BOOL" + STRING = "TYPE_STRING" + MESSAGE = "TYPE_MESSAGE" + BYTES = "TYPE_BYTES" + ENUM = "TYPE_ENUM" + + +class AudioEncoding(IntEnum): + ENCODING_UNSPECIFIED = 0 + LINEAR_PCM = 1 + + +LINEAR_PCM: Final = AudioEncoding.LINEAR_PCM.value + + +@dataclass(frozen=True, slots=True) +class _Field: + name: str + number: int + kind: _Kind + type_name: str = "" + repeated: bool = False + optional: bool = False + oneof: str = "" + + +_MESSAGES: Final[tuple[tuple[str, tuple[_Field, ...]], ...]] = ( + ( + "RecognitionConfig", + ( + _Field("encoding", 1, _Kind.ENUM, "AudioEncoding"), + _Field("sample_rate_hertz", 2, _Kind.INT32), + _Field("language_code", 3, _Kind.STRING), + _Field("max_alternatives", 4, _Kind.INT32), + _Field("profanity_filter", 5, _Kind.BOOL), + _Field("audio_channel_count", 7, _Kind.INT32), + _Field("enable_word_time_offsets", 8, _Kind.BOOL), + _Field("enable_automatic_punctuation", 11, _Kind.BOOL), + _Field("model", 13, _Kind.STRING), + _Field("verbatim_transcripts", 14, _Kind.BOOL), + _Field("endpointing_config", 25, _Kind.MESSAGE, "EndpointingConfig", optional=True), + ), + ), + ( + "EndpointingConfig", + ( + _Field("start_history", 1, _Kind.INT32, optional=True), + _Field("start_threshold", 2, _Kind.FLOAT, optional=True), + _Field("stop_history", 3, _Kind.INT32, optional=True), + _Field("stop_threshold", 4, _Kind.FLOAT, optional=True), + _Field("stop_history_eou", 5, _Kind.INT32, optional=True), + _Field("stop_threshold_eou", 6, _Kind.FLOAT, optional=True), + ), + ), + ( + "StreamingRecognitionConfig", + ( + _Field("config", 1, _Kind.MESSAGE, "RecognitionConfig"), + _Field("interim_results", 2, _Kind.BOOL), + ), + ), + ( + "StreamingRecognizeRequest", + ( + _Field("streaming_config", 1, _Kind.MESSAGE, "StreamingRecognitionConfig", oneof="streaming_request"), + _Field("audio_content", 2, _Kind.BYTES, oneof="streaming_request"), + ), + ), + ( + "WordInfo", + ( + _Field("start_time", 1, _Kind.INT32), + _Field("end_time", 2, _Kind.INT32), + _Field("word", 3, _Kind.STRING), + ), + ), + ( + "SpeechRecognitionAlternative", + ( + _Field("transcript", 1, _Kind.STRING), + _Field("words", 3, _Kind.MESSAGE, "WordInfo", repeated=True), + ), + ), + ( + "StreamingRecognitionResult", + ( + _Field("alternatives", 1, _Kind.MESSAGE, "SpeechRecognitionAlternative", repeated=True), + _Field("is_final", 2, _Kind.BOOL), + ), + ), + ( + "StreamingRecognizeResponse", + (_Field("results", 1, _Kind.MESSAGE, "StreamingRecognitionResult", repeated=True),), + ), + ( + "RecognizeRequest", + ( + _Field("config", 1, _Kind.MESSAGE, "RecognitionConfig"), + _Field("audio", 2, _Kind.BYTES), + ), + ), + ( + "SpeechRecognitionResult", + (_Field("alternatives", 1, _Kind.MESSAGE, "SpeechRecognitionAlternative", repeated=True),), + ), + ( + "RecognizeResponse", + (_Field("results", 1, _Kind.MESSAGE, "SpeechRecognitionResult", repeated=True),), + ), +) + + +def _oneof_names(fields: tuple[_Field, ...]) -> tuple[str, ...]: + declared: Final = tuple(dict.fromkeys(field.oneof for field in fields if field.oneof)) + synthetic: Final = tuple(f"_{field.name}" for field in fields if field.optional) + return declared + synthetic + + +def _field_proto(field: _Field, oneofs: tuple[str, ...]) -> descriptor_pb2.FieldDescriptorProto: + oneof: Final = f"_{field.name}" if field.optional else field.oneof + return descriptor_pb2.FieldDescriptorProto( + name=field.name, + number=field.number, + type=descriptor_pb2.FieldDescriptorProto.Type.Value(field.kind.value), + label=( + descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED + if field.repeated + else descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL + ), + type_name=f".{_PACKAGE}.{field.type_name}" if field.type_name else None, + oneof_index=oneofs.index(oneof) if oneof else None, + proto3_optional=field.optional or None, + ) + + +def _message_proto(name: str, fields: tuple[_Field, ...]) -> descriptor_pb2.DescriptorProto: + oneofs: Final = _oneof_names(fields) + return descriptor_pb2.DescriptorProto( + name=name, + field=[_field_proto(field, oneofs) for field in fields], + oneof_decl=[descriptor_pb2.OneofDescriptorProto(name=oneof) for oneof in oneofs], + ) + + +_FILE: Final = descriptor_pb2.FileDescriptorProto( + name="litellm/nvidia_riva/riva_asr.proto", + package=_PACKAGE, + syntax="proto3", + enum_type=[ + descriptor_pb2.EnumDescriptorProto( + name="AudioEncoding", + value=[ + descriptor_pb2.EnumValueDescriptorProto(name=encoding.name, number=encoding.value) + for encoding in AudioEncoding + ], + ) + ], + message_type=[_message_proto(name, fields) for name, fields in _MESSAGES], +) +_POOL: Final = descriptor_pool.DescriptorPool() +_POOL.AddSerializedFile(_FILE.SerializeToString()) + + +def _message_class(name: str) -> type[Message]: + return message_factory.GetMessageClass(_POOL.FindMessageTypeByName(f"{_PACKAGE}.{name}")) + + +RecognitionConfig: Final = _message_class("RecognitionConfig") +EndpointingConfig: Final = _message_class("EndpointingConfig") +StreamingRecognitionConfig: Final = _message_class("StreamingRecognitionConfig") +StreamingRecognizeRequest: Final = _message_class("StreamingRecognizeRequest") +StreamingRecognizeResponse: Final = _message_class("StreamingRecognizeResponse") +RecognizeRequest: Final = _message_class("RecognizeRequest") +RecognizeResponse: Final = _message_class("RecognizeResponse") diff --git a/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py b/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py new file mode 100644 index 00000000000..8367853673f --- /dev/null +++ b/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py @@ -0,0 +1,113 @@ +""" +gRPC calls to NVIDIA Riva's speech recognition service. + +The ``authorization`` and NVCF ``function-id`` metadata ride on each call, never on the channel, so every +request carries them exactly once over both plaintext and TLS. +""" + +from collections.abc import Iterator, Mapping, Sequence +from typing import Any, Final, TypeAlias + +import grpc +from google.protobuf.message import Message +from pydantic import TypeAdapter + +from litellm.llms.nvidia_riva.audio_transcription import riva_asr_schema as schema +from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException + +# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches +# NVIDIA's recommended chunk size for streaming ASR, small enough for +# responsive endpointing, large enough to keep per-RPC overhead low. +_CHUNK_SAMPLES: Final = 1600 +_CHUNK_BYTES: Final = _CHUNK_SAMPLES * 2 + +CallMetadata: TypeAlias = tuple[tuple[str, str], ...] + +_ENDPOINTING_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +def call_metadata(api_key: str | None, nvcf_function_id: object) -> CallMetadata: + function_id: Final = (("function-id", str(nvcf_function_id)),) if nvcf_function_id else () + authorization: Final = (("authorization", f"Bearer {api_key}"),) if api_key else () + return function_id + authorization + + +def open_channel(api_base: str, use_ssl: bool) -> grpc.Channel: + if use_ssl: + return grpc.secure_channel(api_base, grpc.ssl_channel_credentials()) + return grpc.insecure_channel(api_base) + + +def recognition_config(config: Mapping[str, Any]) -> Message: + endpointing: Final = _endpointing_config(config.get("endpointing_config")) + return schema.RecognitionConfig( + encoding=schema.LINEAR_PCM, + sample_rate_hertz=int(config["sample_rate_hertz"]), + language_code=config["language_code"], + audio_channel_count=int(config["audio_channel_count"]), + enable_automatic_punctuation=bool(config.get("enable_automatic_punctuation", True)), + enable_word_time_offsets=bool(config.get("enable_word_time_offsets", False)), + max_alternatives=int(config.get("max_alternatives", 1)), + model=config.get("model", "") or "", + verbatim_transcripts=bool(config.get("verbatim_transcripts", False)), + profanity_filter=bool(config.get("profanity_filter", False)), + **({} if endpointing is None else {"endpointing_config": endpointing}), + ) + + +def _endpointing_config(endpointing: object) -> Message | None: + try: + fields: Final = _ENDPOINTING_FIELDS.validate_python(endpointing) + return schema.EndpointingConfig(**fields) if fields else None + except (ValueError, TypeError): + # A missing, unknown or mistyped field falls back to Riva's default + # endpointing rather than failing the whole request. + return None + + +def _streaming_requests(config: Message, pcm_bytes: bytes) -> Iterator[Message]: + yield schema.StreamingRecognizeRequest( + streaming_config=schema.StreamingRecognitionConfig(config=config, interim_results=False) + ) + for offset in range(0, len(pcm_bytes), _CHUNK_BYTES): + yield schema.StreamingRecognizeRequest(audio_content=pcm_bytes[offset : offset + _CHUNK_BYTES]) + + +def streaming_recognize( + channel: grpc.Channel, + config: Message, + pcm_bytes: bytes, + metadata: CallMetadata, + timeout: float | None, +) -> Iterator[Message]: + rpc: Final = channel.stream_stream( + schema.STREAMING_RECOGNIZE_METHOD, + request_serializer=schema.StreamingRecognizeRequest.SerializeToString, + response_deserializer=schema.StreamingRecognizeResponse.FromString, + ) + return rpc(_streaming_requests(config, pcm_bytes), metadata=metadata, timeout=timeout) + + +def recognize( + channel: grpc.Channel, + config: Message, + pcm_bytes: bytes, + metadata: CallMetadata, + timeout: float | None, +) -> Sequence[Message]: + rpc: Final = channel.unary_unary( + schema.RECOGNIZE_METHOD, + request_serializer=schema.RecognizeRequest.SerializeToString, + response_deserializer=schema.RecognizeResponse.FromString, + ) + pending: Final = rpc.future(schema.RecognizeRequest(config=config, audio=pcm_bytes), metadata=metadata) + 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 diff --git a/pyproject.toml b/pyproject.toml index f2774e87b7b..2cc2d0f70a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,7 +155,8 @@ stt-vertex-chirp = [ stt-nvidia-riva = [ # NVIDIA Riva STT provider (gRPC). These are imported lazily inside the # provider handler so litellm core remains usable without them. - "nvidia-riva-client>=2.15.0", + "grpcio>=1.67.1,<2.0", + "protobuf>=5.26.1,<7.0", "soundfile>=0.12.1", "audioread>=3.0.1", "numpy>=1.26.0", @@ -241,6 +242,7 @@ dev = [ "pytest-timeout==2.4.0", "vcrpy==8.2.1", "pytest-recording==0.13.4", + "nvidia-riva-client==2.16.0", ] e2e-dev = [ "playwright==1.61.0", @@ -275,6 +277,9 @@ ci = [ # protobuf, Pillow is a compiled C extension). "tenacity==8.5.0", "google-generativeai==0.8.6", + # The Riva tests serve NVIDIA's own message definitions, which checks + # LiteLLM's hand-written copy of them on every run. + "nvidia-riva-client==2.16.0", "Pillow==12.3.0; python_version == '3.10'", "Pillow==13.0.0; python_version >= '3.11'", # Azure batch E2E tests still import psycopg2 directly. diff --git a/tests/integration/_support/riva.py b/tests/integration/_support/riva.py index 8552bc2197d..438d877386b 100644 --- a/tests/integration/_support/riva.py +++ b/tests/integration/_support/riva.py @@ -39,6 +39,7 @@ class RivaCall: rpc: Rpc audio: bytes cancelled: bool + authorization: tuple[str, ...] = () @dataclass(frozen=True, slots=True) @@ -68,6 +69,10 @@ def pcm_bytes(seconds: float) -> int: return int(seconds * SAMPLE_RATE_HZ) * 2 +def _authorization(context: grpc.ServicerContext) -> tuple[str, ...]: + return tuple(str(value) for key, value in context.invocation_metadata() if key == "authorization") + + def carries_clip(audio: bytes, seconds: float) -> bool: sent: Final = _clip_samples(seconds) if len(audio) != len(sent) * 2: @@ -100,9 +105,9 @@ class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): released: Final = threading.Event() context.add_callback(released.set) released.wait(timeout=_STALL_LIMIT_SECONDS) - self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=not context.is_active())) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, not context.is_active(), _authorization(context))) return riva_asr_pb2.RecognizeResponse() - self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=False)) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, False, _authorization(context))) return riva_asr_pb2.RecognizeResponse( results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[self._alternative(request.config)])] ) @@ -111,13 +116,12 @@ class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: if self._offline_only: - self._calls.put(RivaCall(Rpc.STREAMING, b"", cancelled=False)) + self._calls.put(RivaCall(Rpc.STREAMING, b"", False, _authorization(context))) context.abort(grpc.StatusCode.INVALID_ARGUMENT, OFFLINE_ONLY_REFUSAL) requests: Final = tuple(request_iterator) config: Final = requests[0].streaming_config.config if requests else riva_asr_pb2.RecognitionConfig() - self._calls.put( - RivaCall(Rpc.STREAMING, b"".join(request.audio_content for request in requests), cancelled=False) - ) + audio: Final = b"".join(request.audio_content for request in requests) + self._calls.put(RivaCall(Rpc.STREAMING, audio, False, _authorization(context))) yield riva_asr_pb2.StreamingRecognizeResponse( results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[self._alternative(config)], is_final=True)] ) diff --git a/tests/integration/providers/test_nvidia_riva_wire.py b/tests/integration/providers/test_nvidia_riva_wire.py index ef36c449761..56d54a87600 100644 --- a/tests/integration/providers/test_nvidia_riva_wire.py +++ b/tests/integration/providers/test_nvidia_riva_wire.py @@ -135,3 +135,12 @@ def test_a_stalled_offline_recognize_answers_408_at_the_deployment_timeout(gatew assert timed_out.status_code == 408, timed_out.text calls: Final = eventually(server.drain, lambda observed: len(observed) > 0, seconds=30) assert all(call.rpc is Rpc.RECOGNIZE and call.cancelled for call in calls), calls + + +@pytest.mark.parametrize("riva_offline", (False, True), ids=("streaming", "offline")) +def test_the_deployment_api_key_reaches_riva_as_exactly_one_bearer_token(gateway: Gateway, riva_offline: bool) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, api_key="riva-deployment-key", riva_offline=riva_offline) + served: Final = _transcribe(gateway, scenario.key(), model) + assert served.status_code == 200, served.text + assert [call.authorization for call in server.drain()] == [("Bearer riva-deployment-key",)] 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 dbb5cde6060..40ee4d61ac4 100644 --- a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py @@ -1,780 +1,641 @@ -""" -End-to-end-ish tests for NvidiaRivaAudioTranscription. - -We mock ``riva.client`` so the test does not need the real gRPC SDK or a -running Riva server. The mock also lets us assert how Auth metadata is -constructed (NVCF vs self-hosted) and how the streaming generator output -is aggregated. -""" - import asyncio import io -from types import SimpleNamespace -from unittest.mock import MagicMock +import ipaddress +import threading +import uuid +import wave +from array import array +from collections.abc import Iterable, Iterator +from concurrent import futures +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timezone +from queue import SimpleQueue +from typing import Final +import grpc import httpx -import numpy as np import pytest -import soundfile as sf +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID +from riva.client.proto import riva_asr_pb2, riva_asr_pb2_grpc, riva_audio_pb2 import litellm -from litellm.llms.nvidia_riva.audio_transcription import handler as handler_mod -from litellm.llms.nvidia_riva.audio_transcription.handler import ( - NvidiaRivaAudioTranscription, -) +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.nvidia_riva.audio_transcription.handler import NvidiaRivaAudioTranscription from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException from litellm.types.utils import TranscriptionResponse +_SAMPLE_RATE_HZ: Final = 16000 +_SERVER_WAIT_SECONDS: Final = 30 +_STALL_LIMIT_SECONDS: Final = 60 +_NO_DEADLINE_SECONDS: Final = 10**9 +_STALL_TIMEOUT_SECONDS: Final = 0.5 +_API_KEY: Final = "nvapi-test" -def _make_wav_bytes(seconds: float = 1.0, sample_rate: int = 16000) -> bytes: - n = int(sample_rate * seconds) - samples = (0.05 * np.sin(np.linspace(0, 2 * np.pi * 220 * seconds, n))).astype( - np.float32 + +@dataclass(frozen=True, slots=True) +class _Script: + streaming: tuple[riva_asr_pb2.StreamingRecognizeResponse, ...] = () + offline: riva_asr_pb2.RecognizeResponse = field(default_factory=riva_asr_pb2.RecognizeResponse) + abort: tuple[grpc.StatusCode, str] | None = None + stall: bool = False + + +@dataclass(frozen=True, slots=True) +class _Call: + rpc: str + metadata: tuple[tuple[str, str], ...] + config: riva_asr_pb2.RecognitionConfig + audio: bytes + deadline_seconds: float | None + cancelled: bool + + def header(self, key: str) -> tuple[str, ...]: + return tuple(value for name, value in self.metadata if name == key) + + +@dataclass(frozen=True, slots=True) +class _Riva: + address: str + calls: SimpleQueue[_Call] + + def only_call(self) -> _Call: + call: Final = self.calls.get(timeout=_SERVER_WAIT_SECONDS) + assert self.calls.empty(), "Riva received more than one call" + return call + + +class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): + def __init__(self, script: _Script, calls: SimpleQueue[_Call]) -> None: + self._script: Final = script + self._calls: Final = calls + + def _settle( + self, context: grpc.ServicerContext, rpc: str, config: riva_asr_pb2.RecognitionConfig, audio: bytes + ) -> None: + remaining: Final = context.time_remaining() + released: Final = threading.Event() + if self._script.stall and context.add_callback(released.set): + released.wait(timeout=_STALL_LIMIT_SECONDS) + self._calls.put( + _Call( + rpc=rpc, + metadata=tuple((item.key, item.value) for item in context.invocation_metadata()), + config=config, + audio=audio, + deadline_seconds=None if remaining is None or remaining > _NO_DEADLINE_SECONDS else remaining, + cancelled=not context.is_active(), + ) + ) + if self._script.abort is not None: + context.abort(*self._script.abort) + + def Recognize( + self, request: riva_asr_pb2.RecognizeRequest, context: grpc.ServicerContext + ) -> riva_asr_pb2.RecognizeResponse: + self._settle(context, "Recognize", request.config, request.audio) + return self._script.offline + + def StreamingRecognize( + self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext + ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: + first, *rest = tuple(request_iterator) + self._settle( + context, "StreamingRecognize", first.streaming_config.config, b"".join(r.audio_content for r in rest) + ) + yield from self._script.streaming + + +@dataclass(frozen=True, slots=True) +class _Tls: + key_pem: bytes + cert_pem: bytes + + +def _self_signed_loopback_tls() -> _Tls: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]) + cert: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime(2000, 1, 1, tzinfo=timezone.utc)) + .not_valid_after(datetime(2099, 1, 1, tzinfo=timezone.utc)) + .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True) + .add_extension( + x509.SubjectAlternativeName([x509.DNSName("localhost"), x509.IPAddress(ipaddress.ip_address("127.0.0.1"))]), + critical=False, + ) + .sign(key, hashes.SHA256()) ) - buf = io.BytesIO() - sf.write(buf, samples, sample_rate, format="WAV", subtype="PCM_16") - return buf.getvalue() - - -def _fake_word(word: str, start_ms: int, end_ms: int): - return SimpleNamespace(word=word, start_time=start_ms, end_time=end_ms) - - -def _fake_alternative(transcript: str, words=None): - return SimpleNamespace(transcript=transcript, words=words or []) - - -def _fake_result(is_final: bool, alternatives): - return SimpleNamespace(is_final=is_final, alternatives=alternatives) - - -def _fake_response(results): - return SimpleNamespace(results=results) - - -@pytest.fixture -def mock_riva(monkeypatch): - """ - Stand-ins for the bits of ``riva.client`` the handler touches: - - ``Auth`` (constructor) - - ``ASRService`` with ``streaming_response_generator`` - - ``RecognitionConfig``, ``StreamingRecognitionConfig``, ``EndpointingConfig`` - - ``AudioEncoding`` namespace with ``LINEAR_PCM`` - """ - auth_calls = {} - - class FakeAuth: - def __init__(self, *args, **kwargs): - # Support both keyword and positional Auth constructors. - if kwargs: - auth_calls["uri"] = kwargs.get("uri") - auth_calls["use_ssl"] = kwargs.get("use_ssl") - auth_calls["metadata_args"] = kwargs.get("metadata_args") - else: - # positional: (None, use_ssl, uri, metadata) - auth_calls["use_ssl"] = args[1] if len(args) > 1 else None - auth_calls["uri"] = args[2] if len(args) > 2 else None - auth_calls["metadata_args"] = args[3] if len(args) > 3 else None - - class FakeRecognitionConfig: - def __init__(self, **kwargs): - self._kwargs = kwargs - self.endpointing_config = SimpleNamespace(CopyFrom=lambda _: None) - - class FakeStreamingRecognitionConfig: - def __init__(self, config, interim_results): - self.config = config - self.interim_results = interim_results - - class FakeEndpointingConfig: - def __init__(self, **kwargs): - self._kwargs = kwargs - - class FakeAudioEncoding: - 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): - self.auth = auth - - def streaming_response_generator(self, audio_chunks, streaming_config): - # Drain audio_chunks generator so we exercise the chunking path. - 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, - RecognitionConfig=FakeRecognitionConfig, - StreamingRecognitionConfig=FakeStreamingRecognitionConfig, - EndpointingConfig=FakeEndpointingConfig, - AudioEncoding=FakeAudioEncoding, - ) - - def fake_import_riva(): - return fake_riva_client, fake_riva_client - - monkeypatch.setattr(handler_mod, "_import_riva", fake_import_riva) - - 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, + return _Tls( + key_pem=key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ), + cert_pem=cert.public_bytes(serialization.Encoding.PEM), ) -@pytest.fixture -def logging_obj(): - return MagicMock() +@contextmanager +def _riva_server(script: _Script, max_receive_bytes: int | None = None, tls: _Tls | None = None) -> Iterator[_Riva]: + calls: Final[SimpleQueue[_Call]] = SimpleQueue() + options: Final = () if max_receive_bytes is None else (("grpc.max_receive_message_length", max_receive_bytes),) + server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=4), options=options) + riva_asr_pb2_grpc.add_RivaSpeechRecognitionServicer_to_server(_Asr(script, calls), server) + port: Final = ( + server.add_insecure_port("127.0.0.1:0") + if tls is None + else server.add_secure_port("127.0.0.1:0", grpc.ssl_server_credentials(((tls.key_pem, tls.cert_pem),))) + ) + server.start() + try: + yield _Riva(f"127.0.0.1:{port}", calls) + finally: + server.stop(grace=None).wait() -def test_sync_path_aggregates_only_final_results(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - # Empty heartbeat chunk: ignore. - _fake_response(results=[]), - # Interim chunk (not final): ignore. - _fake_response( - results=[ - _fake_result( - is_final=False, alternatives=[_fake_alternative("partial...")] - ) - ] - ), - # Two final chunks aggregated. - _fake_response( - results=[ - _fake_result( - is_final=True, - alternatives=[ - _fake_alternative( - "Hello,", - words=[_fake_word("Hello,", 0, 320)], - ) - ], - ) - ] - ), - _fake_response( - results=[ - _fake_result( - is_final=True, - alternatives=[ - _fake_alternative( - " world.", - words=[_fake_word("world.", 480, 870)], - ) - ], - ) - ] - ), - ] +def _samples(seconds: float) -> array[int]: + return array("h", (((index * 37) % 2001 - 1000) * 16 for index in range(int(seconds * _SAMPLE_RATE_HZ)))) - impl = NvidiaRivaAudioTranscription() - response: TranscriptionResponse = impl.audio_transcriptions( - model="nvidia/parakeet-ctc-1_1b-asr", - audio_file=_make_wav_bytes(), - optional_params={ - "language_code": "en-US", - "enable_word_time_offsets": True, - "response_format": "verbose_json", - "timestamp_granularities": ["word"], - }, - litellm_params={}, + +def _wav(seconds: float = 1.0) -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(_SAMPLE_RATE_HZ) + writer.writeframes(_samples(seconds).tobytes()) + return buffer.getvalue() + + +def _carries_clip(audio: bytes, seconds: float) -> bool: + sent: Final = _samples(seconds) + if len(audio) != len(sent) * 2: + return False + return all(abs(left - right) <= 1 for left, right in zip(sent, array("h", audio), strict=True)) + + +def _alternative( + transcript: str, words: tuple[tuple[str, int, int], ...] = () +) -> riva_asr_pb2.SpeechRecognitionAlternative: + return riva_asr_pb2.SpeechRecognitionAlternative( + transcript=transcript, + words=[riva_asr_pb2.WordInfo(word=word, start_time=start, end_time=end) for word, start, end in words], + ) + + +def _streamed(transcript: str, is_final: bool = True) -> riva_asr_pb2.StreamingRecognizeResponse: + return riva_asr_pb2.StreamingRecognizeResponse( + results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[_alternative(transcript)], is_final=is_final)] + ) + + +def _offline(*transcripts: str) -> riva_asr_pb2.RecognizeResponse: + return riva_asr_pb2.RecognizeResponse( + results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative(text)]) for text in transcripts] + ) + + +def _answers(transcript: str) -> _Script: + return _Script(streaming=(_streamed(transcript),), offline=_offline(transcript)) + + +def _logging() -> Logging: + return Logging( + model="m", + messages=[], + stream=False, + call_type="transcription", + start_time=datetime(2026, 1, 1), + litellm_call_id=str(uuid.uuid4()), + function_id="riva-test", + ) + + +def _transcribe( + riva: _Riva, + optional_params: dict[str, object] | None = None, + *, + api_key: str | None = None, + timeout: float | str | httpx.Timeout | None = 60, + litellm_params: dict[str, object] | None = None, +) -> TranscriptionResponse: + return NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_wav(), + optional_params={"language_code": "en-US", **(optional_params or {})}, + litellm_params=litellm_params or {}, model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", - ) - - assert response.text == "Hello, world." - # duration is propagated from the resampler. - assert response._hidden_params["audio_transcription_duration"] == pytest.approx( - 1.0, abs=0.05 - ) - # word timestamps converted from ms to seconds. - words = response["words"] - assert words[0]["start"] == pytest.approx(0.0) - assert words[1]["end"] == pytest.approx(0.87) - assert ( - logging_obj.pre_call.call_args.kwargs["additional_args"]["atranscription"] - is False + timeout=timeout, + logging_obj=_logging(), + api_key=api_key, + api_base=riva.address, ) -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( +def test_streaming_aggregates_only_final_results_and_their_word_timings() -> None: + heartbeat: Final = riva_asr_pb2.StreamingRecognizeResponse() + hello: Final = riva_asr_pb2.StreamingRecognizeResponse( results=[ - SimpleNamespace(alternatives=[_fake_alternative("Hello,", words=[_fake_word("Hello,", 0, 320)])]), - SimpleNamespace(alternatives=[_fake_alternative(" world.", words=[_fake_word("world.", 480, 870)])]), + riva_asr_pb2.StreamingRecognitionResult( + alternatives=[_alternative("Hello,", (("Hello,", 0, 320),))], is_final=True + ) ] ) - - 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", + world: Final = riva_asr_pb2.StreamingRecognizeResponse( + results=[ + riva_asr_pb2.StreamingRecognitionResult( + alternatives=[_alternative(" world.", (("world.", 480, 870),))], is_final=True + ) + ] ) + with _riva_server(_Script(streaming=(heartbeat, _streamed("partial...", is_final=False), hello, world))) as riva: + response: Final = _transcribe( + riva, + {"enable_word_time_offsets": True, "response_format": "verbose_json", "timestamp_granularities": ["word"]}, + ) + call: Final = riva.only_call() 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")])]) + assert response["words"] == [ + {"word": "Hello,", "start": 0.0, "end": 0.32}, + {"word": "world.", "start": 0.48, "end": 0.87}, ] + assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0) + assert call.rpc == "StreamingRecognize" + assert _carries_clip(call.audio, 1.0), "streamed chunks do not add up to the resampled clip" - 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", + +def test_riva_offline_uses_recognize_and_keeps_results_without_is_final() -> None: + offline: Final = riva_asr_pb2.RecognizeResponse( + results=[ + riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative("Hello,", (("Hello,", 0, 320),))]), + riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative(" world.", (("world.", 480, 870),))]), + ] ) - 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", + with _riva_server(_Script(offline=offline)) as riva: + response: Final = _transcribe( + riva, + { + "language_code": "multi", + "riva_offline": True, + "enable_word_time_offsets": True, + "response_format": "verbose_json", + "timestamp_granularities": ["word"], + }, ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 400 - assert "riva_offline" in excinfo.value.message - assert mock_riva.offline_calls == {} - assert mock_riva.streaming_calls == [] + assert response.text == "Hello, world." + assert response["words"][1] == {"word": "world.", "start": 0.48, "end": 0.87} + assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0) + assert call.rpc == "Recognize" + assert call.config.language_code == "multi" + assert _carries_clip(call.audio, 1.0), "offline request does not carry the whole resampled clip" -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")])]) - ] +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_recognition_config_reaches_riva_as_nvidia_decodes_it(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe( + riva, + { + "riva_offline": riva_offline, + "enable_word_time_offsets": True, + "riva_model_name": "parakeet", + "sample_rate_hertz": 8000, + "profanity_filter": True, + }, + ) + call: Final = riva.only_call() - 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 call.config == riva_asr_pb2.RecognitionConfig( + encoding=riva_audio_pb2.AudioEncoding.LINEAR_PCM, + sample_rate_hertz=_SAMPLE_RATE_HZ, + language_code="en-US", + max_alternatives=1, + profanity_filter=True, + audio_channel_count=1, + enable_word_time_offsets=True, + enable_automatic_punctuation=True, + model="parakeet", ) + +def test_default_mode_streams_and_never_calls_recognize() -> None: + with _riva_server(_answers("streamed")) as riva: + response: Final = _transcribe(riva) + call: Final = riva.only_call() + 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 != {} + assert call.rpc == "StreamingRecognize" @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)], + ("riva_offline", "rpc"), [("false", "StreamingRecognize"), ("False", "StreamingRecognize"), ("true", "Recognize")] ) -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")])] - ) +def test_riva_offline_string_from_multipart_form_is_parsed(riva_offline: str, rpc: str) -> None: + with _riva_server(_Script(streaming=(_streamed("streamed"),), offline=_offline("offline"))) as riva: + response: Final = _transcribe(riva, {"riva_offline": riva_offline}) + call: Final = riva.only_call() - 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 call.rpc == rpc + assert response.text == ("offline" if rpc == "Recognize" else "streamed") + + +@pytest.mark.parametrize("riva_offline", ["yes", "", 1]) +def test_riva_offline_malformed_value_is_rejected_before_any_rpc(riva_offline: object) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": riva_offline}) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 400 + assert "riva_offline" in excinfo.value.message + + +def test_riva_offline_malformed_value_is_dropped_under_drop_params() -> None: + with _riva_server(_Script(streaming=(_streamed("streamed"),), offline=_offline("offline"))) as riva: + response: Final = _transcribe(riva, {"riva_offline": "yes"}, litellm_params={"drop_params": True}) + call: Final = riva.only_call() + + assert response.text == "streamed" + assert call.rpc == "StreamingRecognize" + + +def test_deployment_riva_offline_reaches_the_handler_through_transcription() -> None: + with _riva_server(_Script(offline=_offline("offline"))) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=True, + ) + call: Final = riva.only_call() assert response.text == "offline" - assert mock_riva.offline_calls["timeout"] == waited_for + assert call.rpc == "Recognize" -@pytest.mark.parametrize("timeout", ["abc", "30s"]) -def test_malformed_timeout_on_offline_call_is_a_400_before_any_rpc( - mock_riva: SimpleNamespace, timeout: str +@pytest.mark.parametrize( + ("timeout", "deadline_seconds"), + [ + (12.5, 12.5), + ("12.5", 12.5), + (httpx.Timeout(60.0, read=12.5), 12.5), + (httpx.Timeout(5.0, read=None), 600.0), + (None, None), + ], +) +def test_streaming_deadline_is_the_timeout_in_seconds_whatever_its_form( + timeout: float | str | httpx.Timeout | None, deadline_seconds: float | None ) -> 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, + with _riva_server(_answers("streamed")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-ctc-1.1b", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, timeout=timeout, ) + call: Final = riva.only_call() + + assert response.text == "streamed" + assert call.deadline_seconds == (None if deadline_seconds is None else pytest.approx(deadline_seconds, rel=0.25)) + + +@pytest.mark.parametrize("riva_offline", [False, True]) +@pytest.mark.parametrize("timeout", ["abc", "30s"]) +def test_malformed_timeout_is_a_400_before_any_rpc(riva_offline: bool, timeout: str) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=riva_offline, + timeout=timeout, + ) + + assert riva.calls.empty() 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", +def test_malformed_streaming_timeout_falls_back_to_the_default_deadline_under_drop_params() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-ctc-1.1b", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + timeout="abc", + drop_params=True, ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 413 + assert response.text == "ok" + assert call.rpc == "StreamingRecognize" + assert call.deadline_seconds == pytest.approx(600.0, rel=0.25) + + +def test_malformed_offline_timeout_is_ignored_under_drop_params() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=True, + timeout="abc", + drop_params=True, + ) + call: Final = riva.only_call() + + assert response.text == "ok" + assert call.rpc == "Recognize" + + +@pytest.mark.parametrize( + "timeout", [_STALL_TIMEOUT_SECONDS, str(_STALL_TIMEOUT_SECONDS), httpx.Timeout(30.0, read=_STALL_TIMEOUT_SECONDS)] +) +def test_stalled_offline_recognition_is_cancelled_and_times_out_as_408(timeout: float | str | httpx.Timeout) -> None: + with _riva_server(_Script(stall=True)) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": True}, timeout=timeout) + call: Final = riva.only_call() + + assert excinfo.value.status_code == 408 + assert f"within {_STALL_TIMEOUT_SECONDS} seconds" in excinfo.value.message + assert call.cancelled is True + + +def test_stalled_streaming_recognition_hits_its_deadline_as_504() -> None: + with _riva_server(_Script(stall=True)) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, timeout=_STALL_TIMEOUT_SECONDS) + call: Final = riva.only_call() + + assert excinfo.value.status_code == 504, excinfo.value.message + assert "DEADLINE_EXCEEDED" in excinfo.value.message + assert call.cancelled is True + + +def test_oversized_offline_audio_is_a_non_retryable_413() -> None: + with _riva_server(_answers("ok"), max_receive_bytes=4096) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": True}) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 413, excinfo.value.message 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 +@pytest.mark.parametrize( + ("riva_offline", "code", "status_code"), + [ + (False, grpc.StatusCode.UNAUTHENTICATED, 401), + (True, grpc.StatusCode.INVALID_ARGUMENT, 400), + (False, grpc.StatusCode.RESOURCE_EXHAUSTED, 429), + ], +) +def test_grpc_error_is_wrapped_as_nvidia_riva_exception( + riva_offline: bool, code: grpc.StatusCode, status_code: int ) -> None: - class FutureTimeoutError(Exception): - pass + with _riva_server(_Script(abort=(code, "Riva said no"))) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": riva_offline}, api_key=_API_KEY) - class StalledRecognizeFuture: - def __init__(self) -> None: - self.waited_for: float | None = None - self.cancelled = False + assert excinfo.value.status_code == status_code + assert code.name in excinfo.value.message + assert "Riva said no" in excinfo.value.message - def result(self, timeout: float | None = None) -> object: - self.waited_for = timeout - raise FutureTimeoutError() - def done(self) -> bool: - return self.cancelled +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_self_hosted_api_key_is_sent_as_one_bearer_token(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"riva_offline": riva_offline}, api_key=_API_KEY) + call: Final = riva.only_call() - def cancel(self) -> bool: - self.cancelled = True - return True + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + assert call.header("function-id") == () - stalled = StalledRecognizeFuture() - mock_riva.offline_future["value"] = stalled +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_nvcf_function_id_and_api_key_are_each_sent_once(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe( + riva, {"riva_offline": riva_offline, "nvcf_function_id": "abc-123", "use_ssl": False}, api_key=_API_KEY + ) + call: Final = riva.only_call() + + assert call.header("function-id") == ("abc-123",) + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + + +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_no_api_key_sends_no_authorization(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"riva_offline": riva_offline}) + call: Final = riva.only_call() + + assert call.header("authorization") == () + assert call.header("function-id") == () + + +@pytest.mark.parametrize("riva_offline", [False, True]) +@pytest.mark.parametrize( + ("optional_params", "function_ids"), + [({"use_ssl": True}, ()), ({"nvcf_function_id": "abc-123"}, ("abc-123",))], + ids=["use_ssl", "nvcf_default_tls"], +) +def test_tls_api_key_is_sent_as_one_bearer_token( + monkeypatch: pytest.MonkeyPatch, + optional_params: dict[str, object], + function_ids: tuple[str, ...], + riva_offline: bool, +) -> None: + tls: Final = _self_signed_loopback_tls() + system_roots: Final = grpc.ssl_channel_credentials + monkeypatch.setattr(grpc, "ssl_channel_credentials", lambda: system_roots(root_certificates=tls.cert_pem)) + with _riva_server(_answers("ok"), tls=tls) as riva: + response: Final = _transcribe(riva, {**optional_params, "riva_offline": riva_offline}, api_key=_API_KEY) + call: Final = riva.only_call() + + assert response.text == "ok" + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + assert call.header("function-id") == function_ids + + +@pytest.mark.parametrize("optional_params", [{"use_ssl": True}, {"nvcf_function_id": "abc-123"}]) +def test_tls_never_falls_back_to_plaintext(optional_params: dict[str, object]) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, optional_params, api_key=_API_KEY, timeout=_SERVER_WAIT_SECONDS) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 503, excinfo.value.message + + +def test_explicit_endpointing_config_keeps_zero_thresholds() -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"endpointing_config": {"stop_threshold": 0.0, "stop_history": 800}}) + call: Final = riva.only_call() + + assert call.config.HasField("endpointing_config") + assert call.config.endpointing_config.HasField("stop_threshold") + assert call.config.endpointing_config == riva_asr_pb2.EndpointingConfig(stop_threshold=0.0, stop_history=800) + + +def test_unknown_endpointing_field_falls_back_to_riva_defaults() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = _transcribe(riva, {"endpointing_config": {"stop_threshold": 0.5, "not_a_field": 1}}) + call: Final = riva.only_call() + + assert response.text == "ok" + assert not call.config.HasField("endpointing_config") + + +def test_missing_api_base_raises_clear_error() -> 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": 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( - results=[ - _fake_result( - is_final=True, - alternatives=[_fake_alternative("ok")], - ) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={ - "nvcf_function_id": "abc-123", - "language_code": "en-US", - }, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", - ) - - assert mock_riva.auth_calls["uri"] == "grpc.nvcf.nvidia.com:443" - assert mock_riva.auth_calls["use_ssl"] is True - metadata = dict(mock_riva.auth_calls["metadata_args"]) - assert metadata["function-id"] == "abc-123" - assert metadata["authorization"] == "Bearer nvapi-xxx" - - -def test_auth_self_hosted_defaults_no_ssl_and_no_function_id(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - 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 mock_riva.auth_calls["uri"] == "localhost:50051" - assert mock_riva.auth_calls["use_ssl"] is False - metadata = dict(mock_riva.auth_calls["metadata_args"]) - # No function-id, no authorization metadata. - assert "function-id" not in metadata - assert "authorization" not in metadata - - -def test_explicit_use_ssl_override_wins(mock_riva, logging_obj): - """ - Self-hosted Riva behind an ingress with TLS termination is a real - deployment topology. ``use_ssl=True`` must be honored even without an - NVCF function id. - """ - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"use_ssl": True, "language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="riva.internal.company.com:443", - ) - - assert mock_riva.auth_calls["use_ssl"] is True - - -def test_missing_api_base_raises_clear_error(mock_riva, logging_obj): - impl = NvidiaRivaAudioTranscription() - with pytest.raises(NvidiaRivaException) as excinfo: - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), + audio_file=_wav(), optional_params={}, litellm_params={}, model_response=TranscriptionResponse(), timeout=60, - logging_obj=logging_obj, + logging_obj=_logging(), api_key=None, api_base=None, ) + + assert excinfo.value.status_code == 400 assert "api_base" in excinfo.value.message -def test_async_path_uses_to_thread(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result( - is_final=True, alternatives=[_fake_alternative("async ok")] - ) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - response = asyncio.run( - impl.async_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 == "async ok" - assert ( - logging_obj.pre_call.call_args.kwargs["additional_args"]["atranscription"] - is True - ) - - -@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: 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 - ``streaming_response_generator`` whenever the installed riva-client - accepts a ``timeout`` kwarg. - """ - captured_kwargs = {} - - def streaming_with_timeout(self, audio_chunks, streaming_config, timeout=None): - captured_kwargs["timeout"] = timeout - list(audio_chunks) - yield from [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] +def test_async_path_transcribes() -> None: + with _riva_server(_answers("async ok")) as riva: + response: Final = asyncio.run( + NvidiaRivaAudioTranscription().async_audio_transcriptions( + model="m", + audio_file=_wav(), + optional_params={"language_code": "en-US"}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=_logging(), + api_key=_API_KEY, + api_base=riva.address, ) - ] - - mock_riva.client.ASRService.streaming_response_generator = streaming_with_timeout - - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=timeout, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - assert captured_kwargs["timeout"] == pytest.approx(12.5) - - -def test_grpc_error_is_wrapped_as_nvidia_riva_exception(mock_riva, logging_obj): - class FakeGrpcError(Exception): - def code(self): - return SimpleNamespace(name="UNAUTHENTICATED") - - def details(self): - return "bad token" - - def raising_streaming_response_generator(self, audio_chunks, streaming_config): - list(audio_chunks) - raise FakeGrpcError("rpc fail") - - mock_riva.client.ASRService.streaming_response_generator = ( - raising_streaming_response_generator - ) - - impl = NvidiaRivaAudioTranscription() - with pytest.raises(NvidiaRivaException) as excinfo: - 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="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 401 - assert "UNAUTHENTICATED" in excinfo.value.message + assert response.text == "async ok" + assert call.header("authorization") == (f"Bearer {_API_KEY}",) diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py new file mode 100644 index 00000000000..124e3c28291 --- /dev/null +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py @@ -0,0 +1,41 @@ +from collections.abc import Iterator +from typing import Final + +from google.protobuf.descriptor import FieldDescriptor +from riva.client.proto import riva_asr_pb2, riva_audio_pb2 + +from litellm.llms.nvidia_riva.audio_transcription import riva_asr_schema + + +def _declared_fields() -> Iterator[tuple[str, FieldDescriptor]]: + for message_name, _ in riva_asr_schema._MESSAGES: + message: Final = riva_asr_schema._POOL.FindMessageTypeByName(f"{riva_asr_schema._PACKAGE}.{message_name}") + yield from ((message_name, field) for field in message.fields) + + +def _wire_shape(field: FieldDescriptor | None) -> tuple[object, ...] | None: + if field is None: + return None + return ( + field.number, + field.type, + field.label, + field.has_presence, + field.message_type.full_name if field.message_type else None, + field.enum_type.name if field.enum_type else None, + field.containing_oneof.name if field.containing_oneof else None, + ) + + +def _nvidia_field(message_name: str, field_name: str) -> FieldDescriptor | None: + return riva_asr_pb2.DESCRIPTOR.message_types_by_name[message_name].fields_by_name.get(field_name) + + +def test_every_declared_field_matches_nvidias_riva_schema() -> None: + declared: Final = tuple(_declared_fields()) + + assert len(declared) == sum(len(fields) for _, fields in riva_asr_schema._MESSAGES) + assert {f"{message}.{field.name}": _wire_shape(field) for message, field in declared} == { + f"{message}.{field.name}": _wire_shape(_nvidia_field(message, field.name)) for message, field in declared + } + assert riva_asr_schema.LINEAR_PCM == riva_audio_pb2.AudioEncoding.LINEAR_PCM diff --git a/uv.lock b/uv.lock index 0c796b677e2..4f36ef20bfb 100644 --- a/uv.lock +++ b/uv.lock @@ -5134,9 +5134,10 @@ semantic-router = [ ] stt-nvidia-riva = [ { name = "audioread" }, + { name = "grpcio" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "nvidia-riva-client" }, + { name = "protobuf" }, { name = "soundfile" }, ] stt-vertex-chirp = [ @@ -5176,6 +5177,7 @@ ci = [ { name = "logfire" }, { name = "lunary", version = "1.4.36", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "lunary", version = "1.4.37", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "nvidia-riva-client" }, { name = "pillow", version = "12.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pillow", version = "13.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "psutil" }, @@ -5200,6 +5202,7 @@ dev = [ { name = "mypy" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "nvidia-riva-client" }, { name = "openapi-core" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, @@ -5300,6 +5303,7 @@ requires-dist = [ { name = "granian", marker = "python_full_version < '3.15' and extra == 'proxy'", specifier = ">=2.7.4,<3.0" }, { name = "grpcio", marker = "extra == 'grpc'", specifier = "==1.84.1" }, { name = "grpcio", marker = "extra == 'proxy-runtime'", specifier = "==1.84.1" }, + { name = "grpcio", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.67.1,<2.0" }, { name = "gunicorn", marker = "extra == 'proxy'", specifier = ">=23.0.0,<24.0" }, { name = "hiredis", marker = "python_full_version >= '3.15' and extra == 'proxy'", specifier = ">=3.5.0,<4.0" }, { name = "hiredis", marker = "python_full_version < '3.15' and extra == 'proxy'", specifier = ">=3.0.0,<4.0" }, @@ -5323,7 +5327,6 @@ requires-dist = [ { name = "mlflow", marker = "extra == 'mlflow'", specifier = ">=3.17.2,<4.0" }, { name = "numpy", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.26.0" }, { name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" }, - { name = "nvidia-riva-client", marker = "extra == 'stt-nvidia-riva'", specifier = ">=2.15.0" }, { name = "openai", specifier = ">=2.20.0,<4.0.0" }, { name = "opentelemetry-api", marker = "extra == 'proxy-runtime'", specifier = "==1.45.1" }, { name = "opentelemetry-exporter-otlp", marker = "extra == 'proxy-runtime'", specifier = "==1.45.1" }, @@ -5335,6 +5338,7 @@ requires-dist = [ { name = "polars", marker = "extra == 'proxy'", specifier = ">=1.38.1,<2.0" }, { name = "prisma", marker = "extra == 'extra-proxy'", specifier = ">=0.11.0,<1.0" }, { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = ">=0.20.0,<1.0" }, + { name = "protobuf", marker = "extra == 'stt-nvidia-riva'", specifier = ">=5.26.1,<7.0" }, { name = "psycopg", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, { name = "psycopg-binary", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, { name = "pydantic", marker = "python_full_version < '3.14'", specifier = ">=2.11.0,<3.0.0" }, @@ -5406,6 +5410,7 @@ ci = [ { name = "logfire", specifier = "==5.2.0" }, { name = "lunary", marker = "python_full_version == '3.10.*'", specifier = "==1.4.36" }, { name = "lunary", marker = "python_full_version >= '3.11'", specifier = "==1.4.37" }, + { name = "nvidia-riva-client", specifier = "==2.16.0" }, { name = "pillow", marker = "python_full_version == '3.10.*'", specifier = "==12.3.0" }, { name = "pillow", marker = "python_full_version >= '3.11'", specifier = "==13.0.0" }, { name = "psutil", specifier = "==7.2.2" }, @@ -5429,6 +5434,7 @@ dev = [ { name = "langfuse", specifier = ">=4.7,<5.0" }, { name = "mypy", specifier = "==1.20.1" }, { name = "numpy", specifier = ">=1.26.0,<3.0" }, + { name = "nvidia-riva-client", specifier = "==2.16.0" }, { name = "openapi-core", specifier = "==0.22.0" }, { name = "opentelemetry-api", specifier = "==1.45.1" }, { name = "opentelemetry-exporter-otlp", specifier = "==1.45.1" },