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."),