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 hiddenVisible answer', finish_reason='stop'),
- ]
+ async def provider_stream():
+ if separate_reasoning:
+ reasoning_chunk = self._make_chunk()
+ reasoning_chunk.choices[0].delta.model_dump.return_value['reasoning_content'] = 'hidden reasoning'
+ yield reasoning_chunk
+ for text in ['Hello', ' world']:
+ consumed.append(text)
+ yield self._make_chunk(content=text)
+ yield self._make_chunk(content=None, finish_reason='stop')
- async def _aiter(*args, **kwargs):
- for c in chunks:
- yield c
-
- query = Mock(spec=pipeline_query.Query)
- query.variables = {}
- messages = [provider_message.Message(role='user', content='Hi')]
-
- with patch.object(litellmchat, 'acompletion', new=AsyncMock(side_effect=lambda **kw: _aiter())):
- collected = [
- chunk
- async for chunk in requester.invoke_llm_stream(
- query=query,
- model=model,
- messages=messages,
- remove_think=True,
- )
- ]
-
- assert ''.join(chunk.content or '' for chunk in collected) == 'Visible answer'
+ with patch.object(litellmchat, 'acompletion', new=AsyncMock(side_effect=lambda **kw: provider_stream())):
+ received = []
+ async for chunk in requester.invoke_llm_stream(query=None, model=model, messages=[], remove_think=True):
+ if chunk.content:
+ received.append(chunk.content)
+ assert received == consumed
+ assert not chunk.is_final
+ assert received == ['Hello', ' world']
@pytest.mark.asyncio
async def test_stream_removes_non_leading_think_content(self):
@@ -1540,3 +1530,33 @@ class TestScanModels:
if __name__ == '__main__':
pytest.main([__file__, '-v'])
+
+
+class TestStreamingThinkFilter:
+ @pytest.mark.parametrize(
+ 'opening,closing',
+ [
+ ('', ''),
+ ('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() == '<'