Merged #1410. Patch coverage 100% -> 72%

Author: tgasser-nv

nemoguardrails/context.py

@@ -14,14 +14,11 @@
 # limitations under the License.
 
 import contextvars
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
-
-from nemoguardrails.logging.explain import LLMCallInfo
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 if TYPE_CHECKING:
     from nemoguardrails.logging.explain import ExplainInfo
-    from nemoguardrails.logging.stats import LLMStats
-    from nemoguardrails.rails.llm.options import GenerationOptions
+    from nemoguardrails.rails.llm.options import GenerationOptions, LLMStats
     from nemoguardrails.streaming import StreamingHandler
 
 streaming_handler_var: contextvars.ContextVar[
@@ -34,9 +31,9 @@
 ] = contextvars.ContextVar("explain_info", default=None)
 
 # The current LLM call.
-llm_call_info_var: contextvars.ContextVar[
-    Optional[LLMCallInfo]
-] = contextvars.ContextVar("llm_call_info", default=None)
+llm_call_info_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "llm_call_info", default=None
+)
 
 # All the generation options applicable to the current context.
 generation_options_var: contextvars.ContextVar[
@@ -50,9 +47,9 @@
 
 # The raw LLM request that comes from the user.
 # This is used in passthrough mode.
-raw_llm_request: contextvars.ContextVar[
-    Optional[Union[str, List[Dict[str, Any]]]]
-] = contextvars.ContextVar("raw_llm_request", default=None)
+raw_llm_request: contextvars.ContextVar[Optional[Any]] = contextvars.ContextVar(
+    "raw_llm_request", default=None
+)
 
 reasoning_trace_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
     "reasoning_trace", default=None

nemoguardrails/rails/llm/buffer.py

@@ -114,9 +114,7 @@ def format_chunks(self, chunks: List[str]) -> str:
         ...
 
     @abstractmethod
-    async def process_stream(
-        self, streaming_handler
-    ) -> AsyncGenerator[ChunkBatch, None]:
+    async def process_stream(self, streaming_handler):
         """Process streaming chunks and yield chunk batches.
 
         This is the main method that concrete buffer strategies must implement.
@@ -141,8 +139,7 @@ async def process_stream(
             ...     print(f"Processing: {context_formatted}")
             ...     print(f"User: {user_formatted}")
         """
-        raise NotImplementedError  # pragma: no cover
-        yield
+        yield ChunkBatch([], [])  # pragma: no cover
 
     async def __call__(self, streaming_handler):
         """Callable interface that delegates to process_stream.

nemoguardrails/rails/llm/config.py

@@ -1750,7 +1750,7 @@ def streaming_supported(self):
             # if we have output rails streaming enabled
             # we keep it in case it was needed when we have
             # support per rails
-            if self.rails.output.streaming.enabled:
+            if self.rails.output.streaming and self.rails.output.streaming.enabled:
                 return True
             return False