fix(streaming): handle None stop tokens in streaming handler (#1685)

Author: Pouyanpi

nemoguardrails/streaming.py

@@ -16,7 +16,7 @@
 import asyncio
 import logging
 import warnings
-from typing import Any, AsyncIterator, Dict, Optional, Union
+from typing import Any, AsyncIterator, Dict, List, Optional, Union
 
 from nemoguardrails.utils import new_uuid
 
@@ -83,12 +83,19 @@ def __init__(
         # If set, the chunk will be piped to the specified handler rather than added to the queue or printed
         self.pipe_to = None
 
-        # The stop chunks
-        self.stop = []
+        self._stop = []
 
         self.include_metadata = include_metadata
         self.current_metadata = {}
 
+    @property
+    def stop(self) -> List[str]:
+        return self._stop
+
+    @stop.setter
+    def stop(self, value: Optional[List[str]]) -> None:
+        self._stop = value or []
+
     def set_pattern(self, prefix: Optional[str] = None, suffix: Optional[str] = None):
         """Sets the pattern that is expected.
 

tests/test_actions_llm_utils.py

@@ -29,6 +29,7 @@
     _infer_provider_from_module,
     _store_reasoning_traces,
     _store_tool_calls,
+    _stream_llm_call,
     llm_call,
 )
 from nemoguardrails.context import reasoning_trace_var, tool_calls_var
@@ -712,3 +713,48 @@ async def test_llm_call_does_not_mutate_llm_params(self):
         original_params = {"max_tokens": 100}
         await llm_call(mock_llm, "prompt", stop=["User:"], llm_params=original_params)
         assert original_params == {"max_tokens": 100}
+
+
+async def _empty_astream(*args, **kwargs):
+    return
+    yield
+
+
+class _FakeLLM:
+    def __init__(self, stop=None, kwargs=None):
+        self.stop = stop
+        if kwargs is not None:
+            self.kwargs = kwargs
+        self.astream = _empty_astream
+
+
+class TestStreamLlmCallStopCoercion:
+    @pytest.mark.asyncio
+    async def test_llm_stop_attr_none_coerced_to_list(self):
+        from nemoguardrails.streaming import StreamingHandler
+
+        llm = _FakeLLM(stop=None)
+        handler = StreamingHandler()
+        await _stream_llm_call(llm, "prompt", handler)
+
+        assert handler.stop == []
+
+    @pytest.mark.asyncio
+    async def test_llm_kwargs_stop_none_coerced_to_list(self):
+        from nemoguardrails.streaming import StreamingHandler
+
+        llm = _FakeLLM(kwargs={"stop": None})
+        handler = StreamingHandler()
+        await _stream_llm_call(llm, "prompt", handler)
+
+        assert handler.stop == []
+
+    @pytest.mark.asyncio
+    async def test_llm_with_valid_stop_preserved(self):
+        from nemoguardrails.streaming import StreamingHandler
+
+        llm = _FakeLLM(stop=["User:"])
+        handler = StreamingHandler()
+        await _stream_llm_call(llm, "prompt", handler)
+
+        assert handler.stop == ["User:"]

tests/test_streaming_handler.py

@@ -435,6 +435,40 @@ async def test_multiple_stop_tokens():
         await consumer.cancel()
 
 
+@pytest.mark.asyncio
+async def test_stop_none_does_not_raise():
+    handler = StreamingHandler()
+    consumer = StreamingConsumer(handler)
+
+    try:
+        handler.stop = None
+        await handler.push_chunk("Hello world")
+        await handler.push_chunk(END_OF_STREAM)
+
+        chunks = await consumer.get_chunks()
+        assert chunks == ["Hello world"]
+        assert handler.completion == "Hello world"
+    finally:
+        await consumer.cancel()
+
+
+@pytest.mark.asyncio
+async def test_stop_none_with_pattern():
+    await _test_pattern_case(
+        prefix='Bot message: "',
+        suffix='"',
+        stop=None,
+        chunks=[
+            "Bot",
+            " message: ",
+            '"',
+            "This is a message",
+            '"',
+        ],
+        final_chunks=["This is a message"],
+    )
+
+
 @pytest.mark.asyncio
 async def test_enable_print_functionality():
     """Test the enable_print functionality."""