Fix ordinary streaming output buffering with stripthink enabled (#1311)

src/python/txtai/pipeline/llm/generation.py

@@ -173,12 +173,15 @@ class Generation:
         """
 
         # Consume "thinking" tokens
-        text, buffer = None, ""
+        text, buffer = "", ""
         for chunk in results:
             buffer += chunk
             text = self.cleanthink(buffer)
             if text != buffer:
                 break
+            prefix = buffer.lstrip()
+            if not any(marker.startswith(prefix) or prefix.startswith(marker) for marker in ("<think>", "<|")):
+                break
 
         # Yield remaining tokens
         yield from text.lstrip()

test/python/testpipeline/testllm/testllm.py

@@ -94,6 +94,11 @@ class TestLLM(unittest.TestCase):
             # Test always keeping as prompt text
             self.assertEqual(type(generator.format([message], "prompt")[0]), str)
 
+    def testEmptyStream(self):
+        """Return no output when the upstream stream is empty."""
+        model, _ = self.streammodel([])
+        self.assertEqual(list(model("question", stream=True, stripthink=True)), [])
+
     def testExternal(self):
         """
         Test externally loaded model
@@ -124,6 +129,28 @@ class TestLLM(unittest.TestCase):
         generation = Generation()
         self.assertRaises(NotImplementedError, generation.stream, None, None, None, None)
 
+    def testPartialThinkingPrefixes(self):
+        """Handle split thinking prefixes, ordinary tags, and whitespace."""
+        for chunks, expected in [
+            (["<", "th", "ink>", "reason", "</think>", "answer"], "answer"),
+            (["<", "table>", "answer"], "<table>answer"),
+            (["  ", "plain", " answer"], "plain answer"),
+            (["  ", "\n"], ""),
+            (["<", "t"], "<t"),
+            (["<|start|>assistant<|channel|>analysis<|message|>", "reason", "<|channel|>final<|message|>answer"], "answer"),
+        ]:
+            with self.subTest(chunks=chunks):
+                model, _ = self.streammodel(chunks)
+                self.assertEqual("".join(model("question", stream=True, stripthink=True)), expected)
+
+    def testPlainAnswerIsIncremental(self):
+        """Yield ordinary text before consuming subsequent chunks."""
+        model, consumed = self.streammodel(["  blue", " sky"])
+        result = model("question", stream=True, stripthink=True)
+        self.assertEqual(next(result), "b")
+        self.assertEqual(consumed, ["  blue"])
+        self.assertEqual("".join(result), "lue sky")
+
     def testStop(self):
         """
         Test stop strings
@@ -180,6 +207,14 @@ class TestLLM(unittest.TestCase):
             self.assertEqual("".join(model("Hello, how are", stripthink=True, stream=True)), "you")
             self.assertEqual("".join(model("Hello, how are", stripthink=False, stream=True)), "".join(list(method())))
 
+    def testStripthinkDisabled(self):
+        """Preserve chunks and consume them on demand when stripping is disabled."""
+        model, consumed = self.streammodel(["<think>reason</think>", "answer"])
+        result = model("question", stream=True, stripthink=False)
+        self.assertEqual(next(result), "<think>reason</think>")
+        self.assertEqual(consumed, ["<think>reason</think>"])
+        self.assertEqual(list(result), ["answer"])
+
     def testVision(self):
         """
         Test vision LLM
@@ -191,3 +226,16 @@ class TestLLM(unittest.TestCase):
         )
 
         self.assertIsNotNone(result)
+
+    def streammodel(self, chunks):
+        """Create an LLM pipeline with a controlled upstream stream."""
+        consumed = []
+
+        def execute(*_args, **_kwargs):
+            for chunk in chunks:
+                consumed.append(chunk)
+                yield chunk
+
+        model = LLM("test", method="txtai.pipeline.Generation")
+        model.generator.execute = execute
+        return model, consumed