xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-10-21 16:05:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(provider): preserve streaming when removing thinking content

Browse Source
This commit is contained in:
Hyu
2026-10-12 13:30:17 +08:00
parent 8ae182dbd7
commit 09a84f6642
2 changed files with 57 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py
+5 -37
View File
@@ -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>'
_THINK_CLOSE = '</think>'
@@ -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)
tests/unit_tests/provider/test_litellmchat.py
+52 -32
View File
@@ -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 </think>."""
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</thi'),
self._make_chunk(content='nk>Visible 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',
[
('<think>', '</think>'),
('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('<think>hidden') == ''
assert state.flush() == ''
def test_partial_non_tag_is_released(self):
state = litellmchat._ThinkStripState()
assert state.feed('Answer <thi') == 'Answer '
assert state.feed('ng>') == '<thing>'
assert state.feed(' <') == ' '
assert state.flush() == '<'