fix(llm): Add async streaming support to ChatNVIDIA provider patch (#1504)

* feat(llm): Add async streaming support to ChatNVIDIA provider

Enables stream_async() to work with ChatNVIDIA/NIM models by
implementing async streaming decorator and _agenerate method. Prior to
this fix, stream_async() would fail with NIM engine configurations.

* fix: ensure stream_async background task completes before exit (#1508)

Wrap the returned iterator to await the background generation task in a
finally block, preventing "Task was destroyed but it is pending"
warning.
Add overloaded type signatures to provide accurate return types based on
the include_generation_metadata parameter.

Author: Pouyanpi

nemoguardrails/llm/providers/_langchain_nvidia_ai_endpoints_patch.py

@@ -17,12 +17,18 @@
 from functools import wraps
 from typing import Any, List, Optional
 
-from langchain_core.callbacks.manager import CallbackManagerForLLMRun
-from langchain_core.language_models.chat_models import generate_from_stream
+from langchain_core.callbacks.manager import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
+from langchain_core.language_models.chat_models import (
+    agenerate_from_stream,
+    generate_from_stream,
+)
 from langchain_core.messages import BaseMessage
 from langchain_core.outputs import ChatResult
 from langchain_nvidia_ai_endpoints import ChatNVIDIA as ChatNVIDIAOriginal
-from pydantic.v1 import Field
+from pydantic import Field
 
 log = logging.getLogger(__name__)
 
@@ -49,6 +55,28 @@ def wrapper(
     return wrapper
 
 
+def async_stream_decorator(func):  # pragma: no cover
+    @wraps(func)
+    async def wrapper(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        stream: Optional[bool] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        should_stream = stream if stream is not None else self.streaming
+        if should_stream:
+            stream_iter = self._astream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return await agenerate_from_stream(stream_iter)
+        else:
+            return await func(self, messages, stop, run_manager, **kwargs)
+
+    return wrapper
+
+
 # NOTE: this needs to have the same name as the original class,
 #   otherwise, there's a check inside `langchain-nvidia-ai-endpoints` that will fail.
 class ChatNVIDIA(ChatNVIDIAOriginal):
@@ -65,6 +93,21 @@ def _generate(
         **kwargs: Any,
     ) -> ChatResult:
         return super()._generate(
+            messages=messages,
+            stop=stop,
+            run_manager=run_manager,
+            **kwargs,
+        )
+
+    @async_stream_decorator
+    async def _agenerate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        return await super()._agenerate(
             messages=messages, stop=stop, run_manager=run_manager, **kwargs
         )
 

nemoguardrails/rails/llm/llmrails.py

@@ -30,11 +30,13 @@
     Callable,
     Dict,
     List,
+    Literal,
     Optional,
     Tuple,
     Type,
     Union,
     cast,
+    overload,
 )
 
 from langchain_core.language_models import BaseChatModel
@@ -1255,15 +1257,39 @@ def _validate_streaming_with_output_rails(self) -> None:
                 "generate_async() instead of stream_async()."
             )
 
+    @overload
     def stream_async(
         self,
         prompt: Optional[str] = None,
         messages: Optional[List[dict]] = None,
         options: Optional[Union[dict, GenerationOptions]] = None,
         state: Optional[Union[dict, State]] = None,
-        include_generation_metadata: Optional[bool] = False,
+        include_generation_metadata: Literal[False] = False,
         generator: Optional[AsyncIterator[str]] = None,
     ) -> AsyncIterator[str]:
+        ...
+
+    @overload
+    def stream_async(
+        self,
+        prompt: Optional[str] = None,
+        messages: Optional[List[dict]] = None,
+        options: Optional[Union[dict, GenerationOptions]] = None,
+        state: Optional[Union[dict, State]] = None,
+        include_generation_metadata: Literal[True] = ...,
+        generator: Optional[AsyncIterator[str]] = None,
+    ) -> AsyncIterator[Union[str, dict]]:
+        ...
+
+    def stream_async(
+        self,
+        prompt: Optional[str] = None,
+        messages: Optional[List[dict]] = None,
+        options: Optional[Union[dict, GenerationOptions]] = None,
+        state: Optional[Union[dict, State]] = None,
+        include_generation_metadata: Optional[bool] = False,
+        generator: Optional[AsyncIterator[str]] = None,
+    ) -> AsyncIterator[Union[str, dict]]:
         """Simplified interface for getting directly the streamed tokens from the LLM."""
 
         self._validate_streaming_with_output_rails()
@@ -1328,15 +1354,24 @@ def task_done_callback(task):
             self.config.rails.output.streaming
             and self.config.rails.output.streaming.enabled
         ):
-            # returns an async generator
-            return self._run_output_rails_in_streaming(
+            base_iterator = self._run_output_rails_in_streaming(
                 streaming_handler=streaming_handler,
                 output_rails_streaming_config=self.config.rails.output.streaming,
                 messages=messages,
                 prompt=prompt,
             )
         else:
-            return streaming_handler
+            base_iterator = streaming_handler
+
+        async def wrapped_iterator():
+            try:
+                async for chunk in base_iterator:
+                    if chunk is not None:
+                        yield chunk
+            finally:
+                await task
+
+        return wrapped_iterator()
 
     def generate(
         self,