fix(provider): preserve think tags in streamed reasoning

This commit is contained in:
fdc310
2026-08-06 01:50:20 +08:00
parent d6f01adb8a
commit a1c2c975cb
2 changed files with 74 additions and 5 deletions
@@ -1135,6 +1135,8 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
role = 'assistant'
tool_call_state: dict[int, dict[str, typing.Any]] = {}
think_state = _ThinkStripState() if remove_think else None
reasoning_started = False
reasoning_closed = False
try:
response = await acompletion(**args)
@@ -1165,18 +1167,36 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
if 'role' in delta and delta['role']:
role = delta['role']
delta_content = delta.get('content', '')
reasoning_content = delta.get('reasoning_content', '')
delta_content = delta.get('content') or ''
reasoning_content = delta.get('reasoning_content') or ''
provider_fields = dict(delta.get('provider_specific_fields') or {})
# Handle reasoning_content based on remove_think flag
if reasoning_content:
provider_fields['reasoning_content'] = reasoning_content
if remove_think:
delta_content = None
delta_content = delta_content or None
else:
# Use reasoning_content as the displayed content
delta_content = reasoning_content
# Stream explicit markers so downstream adapters and
# the debug page see the same format as non-streaming
# responses.
if not reasoning_started:
delta_content = '<think>\n'
reasoning_started = True
else:
delta_content = ''
delta_content += reasoning_content
if delta.get('content'):
delta_content += f'\n</think>\n{delta.get("content")}'
reasoning_closed = True
elif delta_content and not remove_think and reasoning_started and not reasoning_closed:
delta_content = f'\n</think>\n{delta_content}'
reasoning_closed = True
if finish_reason and not remove_think and reasoning_started and not reasoning_closed:
delta_content = f'{delta_content}\n</think>\n'
reasoning_closed = True
if think_state is not None and delta_content:
delta_content = think_state.feed(delta_content)
@@ -1205,6 +1225,13 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
yield provider_message.MessageChunk(**chunk_data)
chunk_idx += 1
if reasoning_started and not reasoning_closed:
yield provider_message.MessageChunk(
role=role,
content='\n</think>\n',
is_final=True,
)
if think_state is not None:
pending_content = think_state.flush()
if pending_content:
@@ -714,3 +714,45 @@ async def test_stream_reasoning_round_trip_with_hidden_display(monkeypatch):
assert emitted is not None
assert emitted.content == 'answer'
assert emitted.provider_specific_fields == {'reasoning_content': 'private '}
@pytest.mark.asyncio
async def test_stream_reasoning_content_is_wrapped_for_display(monkeypatch):
request = _requester('deepseek')
request._build_completion_args = AsyncMock(return_value={})
async def chunks():
yield SimpleNamespace(
choices=[
SimpleNamespace(
delta=_Dumpable({'role': 'assistant', 'reasoning_content': 'private '}),
finish_reason=None,
)
],
usage=None,
)
yield SimpleNamespace(
choices=[
SimpleNamespace(
delta=_Dumpable({'content': 'answer'}),
finish_reason='stop',
)
],
usage=None,
)
monkeypatch.setattr(litellmchat, 'acompletion', AsyncMock(return_value=chunks()))
accumulator = _StreamAccumulator(remove_think=False)
emitted: provider_message.MessageChunk | None = None
async for chunk in request.invoke_llm_stream(
None,
_runtime_model(request),
[],
remove_think=False,
):
emitted = accumulator.add(chunk) or emitted
assert emitted is not None
assert emitted.content == '<think>\nprivate \n</think>\nanswer'
assert emitted.provider_specific_fields == {'reasoning_content': 'private '}