diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py index 23a69bea8..e927a3baf 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py @@ -16,7 +16,11 @@ import langbot_plugin.api.entities.builtin.provider.message as provider_message class _ThinkStripState: - """Stateful filter that drops think blocks across chunks.""" + """Drop explicit think blocks, buffering only possible split tag prefixes. + + Untagged text must stream immediately: an orphan closing tag cannot identify + earlier text as reasoning without delaying every ordinary answer until EOF. + """ _THINK_OPEN = '' _THINK_CLOSE = '' @@ -31,7 +35,6 @@ class _ThinkStripState: self._open_tags = tuple(open_tag for open_tag, _close_tag in self._pairs) self._buf = '' self._close_tag: str | None = None - self._pending_initial = True def feed(self, chunk: str) -> str: """Feed a streaming delta and return user-visible content.""" @@ -64,7 +67,6 @@ class _ThinkStripState: if close_idx != -1: self._close_tag = None self._buf = '' - self._pending_initial = False return self._process_visible_text(text[close_idx + len(close_tag) :]) self._buf = self._close_prefix(text, close_tag) @@ -75,37 +77,14 @@ class _ThinkStripState: index = 0 while index < len(text): - if self._pending_initial: - open_idx, open_tag, close_tag = self._find_next_open(text, index) - orphan_close_idx, orphan_close_tag = self._find_next_close(text, index) - - if orphan_close_idx != -1 and (open_idx == -1 or orphan_close_idx < open_idx): - self._pending_initial = False - index = orphan_close_idx + len(orphan_close_tag) - continue - - if open_idx == -1: - self._buf = text[index:] - return ''.join(out) - - if open_idx > index: - self._pending_initial = False - out.append(text[index:open_idx]) - index = open_idx - continue - open_idx, open_tag, close_tag = self._find_next_open(text, index) if open_idx == -1: emit_end = self._visible_emit_end(text, index) out.append(text[index:emit_end]) - if emit_end > index: - self._pending_initial = False self._buf = text[emit_end:] return ''.join(out) out.append(text[index:open_idx]) - if open_idx > index: - self._pending_initial = False body_start = open_idx + len(open_tag) close_idx = text.find(close_tag, body_start) if close_idx == -1: @@ -113,7 +92,6 @@ class _ThinkStripState: self._buf = self._close_prefix(text[body_start:], close_tag) return ''.join(out) - self._pending_initial = False index = close_idx + len(close_tag) self._buf = '' @@ -131,16 +109,6 @@ class _ThinkStripState: best_close = close_tag return best_idx, best_open, best_close - def _find_next_close(self, text: str, start: int) -> tuple[int, str]: - best_idx = -1 - best_close = '' - for _open_tag, close_tag in self._pairs: - idx = text.find(close_tag, start) - if idx != -1 and (best_idx == -1 or idx < best_idx): - best_idx = idx - best_close = close_tag - return best_idx, best_close - def _visible_emit_end(self, text: str, start: int) -> int: visible = text[start:] limit = min(len(visible), max(len(open_tag) for open_tag in self._open_tags) - 1) diff --git a/tests/unit_tests/provider/test_litellmchat.py b/tests/unit_tests/provider/test_litellmchat.py index 4c089816d..f73fac2e4 100644 --- a/tests/unit_tests/provider/test_litellmchat.py +++ b/tests/unit_tests/provider/test_litellmchat.py @@ -379,43 +379,33 @@ class TestInvokeLLMStreamUsage: assert ''.join(chunk.content or '' for chunk in collected) == 'Visible answer' @pytest.mark.asyncio - async def test_stream_removes_initial_orphan_think_close(self): - """Initial reasoning content without an open tag is removed until .""" - import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query - import langbot_plugin.api.entities.builtin.provider.message as provider_message - + @pytest.mark.parametrize('separate_reasoning', [False, True]) + async def test_remove_think_preserves_incremental_plain_text(self, separate_reasoning): + """Visible text must arrive before the provider produces the next delta.""" mock_ap = Mock() - mock_ap.tool_mgr = Mock() mock_ap.tool_mgr.generate_tools_for_openai = AsyncMock(return_value=None) requester = litellmchat.LiteLLMRequester(ap=mock_ap, config={}) - model = MockRuntimeModel('minimax-m3', 'test-api-key') + model = MockRuntimeModel('gpt-4o', 'test-api-key') + consumed = [] - chunks = [ - self._make_chunk(content='hidden reasoning'), - self._make_chunk(content=' still hidden', ''), + ('CRETIRE_REASONING_BEGINk', 'CRETIRE_REASONING_ENDk'), + ], + ) + def test_split_tags_and_repeated_blocks(self, opening, closing): + state = litellmchat._ThinkStripState() + assert state.feed('Before ') == 'Before ' + for char in opening + 'hidden' + closing: + assert state.feed(char) == '' + assert state.feed('Answer ') == 'Answer ' + assert state.feed(opening + 'more hidden' + closing + 'end') == 'end' + assert state.flush() == '' + + def test_unclosed_think_is_discarded(self): + state = litellmchat._ThinkStripState() + assert state.feed('hidden') == '' + assert state.flush() == '' + + def test_partial_non_tag_is_released(self): + state = litellmchat._ThinkStripState() + assert state.feed('Answer ') == '' + assert state.feed(' <') == ' ' + assert state.flush() == '<'