fix(streaming_hints): close a finished uncertain sentence

optillm/plugins/streaming_hints/retriever.py

@@ -54,13 +54,15 @@ def reset(self, problem: str) -> None:
         self.pending = ""
         self.chars_since_injection = 0
         self.used: List[str] = []
+        self._last_delta_length = 0
 
     def feed(self, delta: str) -> Optional[Injection]:
         """
         Consume the next piece of reasoning. Returns the first injection found
         in it, or None. On an injection the text after the offset is treated
         as discarded and is not kept in the retriever's state.
         """
+        self._last_delta_length = len(delta)
         if len(self.used) >= self.max_hints:
             self.text += delta
             return None
@@ -84,6 +86,31 @@ def feed(self, delta: str) -> Optional[Injection]:
         self.pending += delta[consumed:]
         return None
 
+    def flush(self) -> Optional[Injection]:
+        """Close the trailing sentence of a finished segment.
+
+        feed() ends a sentence only on punctuation followed by whitespace, or
+        on a newline, so a later chunk can still continue it. A segment that
+        stops has no later chunk, and a reply ending in "Wait, check the
+        units." would otherwise stay pending and never be an injection point.
+        The offset is the end of the segment passed to the last feed().
+        """
+        if len(self.used) >= self.max_hints:
+            return None
+        sentence = self.pending
+        if not sentence.strip():
+            self.pending = ""
+            return None
+        self.pending = ""
+        self.text += sentence
+        self.chars_since_injection += len(sentence)
+        if not self._is_injection_point(sentence):
+            return None
+        injection = self._retrieve(self._last_delta_length)
+        if injection:
+            self.chars_since_injection = 0
+        return injection
+
     def _is_injection_point(self, sentence: str) -> bool:
         if self.chars_since_injection < self.min_gap_chars:
             return False

optillm/plugins/streaming_hints/streaming_hints.py

@@ -99,6 +99,11 @@ def run_streaming_hints(
         truncated = choice.finish_reason == "length" or used_tokens >= segment_tokens
 
         injection = retriever.feed(text)
+        # Punctuation at the end of a stopped segment has no trailing
+        # whitespace, so feed() leaves that sentence pending. A truncated
+        # segment is still open: the next segment may continue it.
+        if injection is None and choice.finish_reason == "stop" and not truncated:
+            injection = retriever.flush()
         if injection:
             kept = text[: injection.offset]
             output += kept

tests/test_streaming_hints.py

@@ -82,6 +82,24 @@ def test_retriever_handles_sentences_split_across_feeds():
     assert injection.offset == len("its here.")
 
 
+def test_flush_injects_the_uncertain_sentence_that_ends_the_reply():
+    # feed() requires whitespace after .?! so a later chunk can continue the
+    # sentence. A finished reply has no later chunk, and this is how a model
+    # actually stops.
+    retriever = StreamingHintRetriever(units_bank(), threshold=0.1, min_gap_chars=20)
+    retriever.reset(PROBLEM)
+    text = "Speed is 80 km/h. Wait, the answer must be in meters per second."
+
+    assert retriever.feed(text) is None
+    injection = retriever.flush()
+
+    assert injection is not None
+    assert injection.hint.id == "units"
+    assert injection.offset == len(text)
+    assert text[: injection.offset].endswith("meters per second.")
+    assert injection.hint.uses == 1
+
+
 def test_retriever_respects_max_hints_and_threshold():
     retriever = StreamingHintRetriever(units_bank(), threshold=0.99, min_gap_chars=0)
     retriever.reset(PROBLEM)
@@ -123,6 +141,72 @@ def test_run_injects_hint_and_continues_from_cut():
     assert "discarded" not in continuation
 
 
+def test_run_injects_when_the_uncertain_sentence_ends_the_segment():
+    # Longer than the default 200-character gap, and the uncertain sentence
+    # ends the segment with no whitespace after its period.
+    text = (
+        "The distance is 120 km and the time is 1.5 hours. " * 5
+        + "Wait, the question asks for meters per second."
+    )
+    client, calls = fake_client(
+        response(text, tokens=30),
+        response(" About 22.2 m/s.", tokens=8),
+    )
+
+    answer, tokens = run_streaming_hints(
+        "system", PROBLEM, client, "model",
+        request_config={"hint_threshold": 0.1, "show_hints": True},
+        memory_bank=units_bank(),
+    )
+
+    assert tokens == 38
+    assert len(calls) == 2
+    continuation = calls[1]["messages"][1]["content"]
+    assert "meters per second.\n<hint>Check the units of every quantity.</hint>" in continuation
+    assert answer.endswith("About 22.2 m/s.")
+
+
+def test_run_only_flushes_a_normal_stop():
+    text = "The distance is 120 km and the time is 1.5 hours. " * 5 + "Wait, check the units."
+    client, calls = fake_client(response(text, reason="content_filter", tokens=30))
+
+    answer, tokens = run_streaming_hints(
+        "system", PROBLEM, client, "model",
+        request_config={"hint_threshold": 0.1, "show_hints": True},
+        memory_bank=units_bank(),
+    )
+
+    assert answer == text
+    assert tokens == 30
+    assert len(calls) == 1
+    assert "<hint>" not in answer
+
+
+def test_run_keeps_a_decimal_split_across_a_truncated_segment():
+    first = "The distance is 120 km and the time is 1.5 hours. " * 5 + "Wait, perhaps the answer is 3."
+    client, calls = fake_client(
+        response(first, reason="length", tokens=30),
+        response("14 after conversion.", tokens=8),
+        response(" Final answer.", tokens=3),
+    )
+
+    answer, tokens = run_streaming_hints(
+        "system", PROBLEM, client, "model",
+        request_config={
+            "hint_threshold": 0.1,
+            "segment_tokens": 30,
+            "show_hints": True,
+        },
+        memory_bank=units_bank(),
+    )
+
+    assert tokens == 41
+    assert len(calls) == 3
+    assert "<hint>Check the units of every quantity.</hint>" not in calls[1]["messages"][1]["content"]
+    assert "3.14 after conversion.\n<hint>" in calls[2]["messages"][1]["content"]
+    assert answer.endswith("Final answer.")
+
+
 def test_run_show_hints_returns_reasoning_with_hint_tags():
     client, _ = fake_client(
         response("Speed is 80. Wait, the units are meters per second. Tail."),