Pass OutputRailsStreamingConfig into _run_output_rails_in_streaming as mandatory argument as all 2 calls check for None

Author: tgasser-nv

nemoguardrails/rails/llm/llmrails.py

@@ -78,7 +78,11 @@
 from nemoguardrails.logging.verbose import set_verbose
 from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop
 from nemoguardrails.rails.llm.buffer import get_buffer_strategy
-from nemoguardrails.rails.llm.config import EmbeddingSearchProvider, RailsConfig
+from nemoguardrails.rails.llm.config import (
+    EmbeddingSearchProvider,
+    OutputRailsStreamingConfig,
+    RailsConfig,
+)
 from nemoguardrails.rails.llm.options import (
     GenerationLog,
     GenerationOptions,
@@ -1289,6 +1293,7 @@ def stream_async(
             ):
                 return self._run_output_rails_in_streaming(
                     streaming_handler=generator,
+                    output_rails_streaming_config=self.config.rails.output.streaming,
                     messages=messages,
                     prompt=prompt,
                 )
@@ -1344,6 +1349,7 @@ def task_done_callback(task):
             # returns an async generator
             return self._run_output_rails_in_streaming(
                 streaming_handler=streaming_handler,
+                output_rails_streaming_config=self.config.rails.output.streaming,
                 messages=messages,
                 prompt=prompt,
             )
@@ -1572,6 +1578,7 @@ def __setstate__(self, state):
     async def _run_output_rails_in_streaming(
         self,
         streaming_handler: AsyncIterator[str],
+        output_rails_streaming_config: OutputRailsStreamingConfig,
         prompt: Optional[str] = None,
         messages: Optional[List[dict]] = None,
         stream_first: Optional[bool] = None,
@@ -1674,9 +1681,6 @@ def _prepare_params(
                 **action_params,
             }
 
-        output_rails_streaming_config = self.config.rails.output.streaming
-        if output_rails_streaming_config is None:
-            raise ValueError("Output rails streaming config is not available")
         buffer_strategy = get_buffer_strategy(output_rails_streaming_config)
         output_rails_flows_id = self.config.rails.output.flows
         stream_first = stream_first or output_rails_streaming_config.stream_first