diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py index 002027f0b..4267f0190 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py @@ -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 = '\n' + reasoning_started = True + else: + delta_content = '' + delta_content += reasoning_content + if delta.get('content'): + delta_content += f'\n\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\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\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\n', + is_final=True, + ) + if think_state is not None: pending_content = think_state.flush() if pending_content: diff --git a/tests/unit_tests/provider/test_reasoning_control.py b/tests/unit_tests/provider/test_reasoning_control.py index 35056a86c..2f4b24dc8 100644 --- a/tests/unit_tests/provider/test_reasoning_control.py +++ b/tests/unit_tests/provider/test_reasoning_control.py @@ -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 == '\nprivate \n\nanswer' + assert emitted.provider_specific_fields == {'reasoning_content': 'private '}