diff --git a/optillm/plugins/streaming_hints/retriever.py b/optillm/plugins/streaming_hints/retriever.py
index f62634d7..c7a4448c 100644
--- a/optillm/plugins/streaming_hints/retriever.py
+++ b/optillm/plugins/streaming_hints/retriever.py
@@ -54,6 +54,7 @@ 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]:
"""
@@ -61,6 +62,7 @@ def feed(self, delta: str) -> Optional[Injection]:
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
diff --git a/optillm/plugins/streaming_hints/streaming_hints.py b/optillm/plugins/streaming_hints/streaming_hints.py
index 96e18942..e2799a71 100644
--- a/optillm/plugins/streaming_hints/streaming_hints.py
+++ b/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
diff --git a/tests/test_streaming_hints.py b/tests/test_streaming_hints.py
index f1c7b3a2..85215ac0 100644
--- a/tests/test_streaming_hints.py
+++ b/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.\nCheck the units of every quantity." 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 "" 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 "Check the units of every quantity." not in calls[1]["messages"][1]["content"]
+ assert "3.14 after conversion.\n" 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."),