From 2383e9498595d0862ed7914d787ea27c5bf48acf Mon Sep 17 00:00:00 2001 From: Safzan Pirani Date: Wed, 22 Apr 2026 23:53:29 +0530 Subject: [PATCH] Fix Kimi tool-call continuation reasoning --- backend/open_webui/test/util/test_misc.py | 158 ++++++++++++++++++++++ backend/open_webui/utils/middleware.py | 15 +- backend/open_webui/utils/misc.py | 113 ++++++++++++++++ 3 files changed, 283 insertions(+), 3 deletions(-) create mode 100644 backend/open_webui/test/util/test_misc.py diff --git a/backend/open_webui/test/util/test_misc.py b/backend/open_webui/test/util/test_misc.py new file mode 100644 index 0000000000..4ac05f1662 --- /dev/null +++ b/backend/open_webui/test/util/test_misc.py @@ -0,0 +1,158 @@ +from open_webui.utils.misc import add_reasoning_content_to_tool_messages + + +class TestAddReasoningContentToToolMessages: + def test_adds_reasoning_content_for_kimi_tool_messages(self): + messages = [ + { + 'role': 'assistant', + 'content': 'Need to search.', + 'tool_calls': [ + { + 'id': 'call_123', + 'type': 'function', + 'function': {'name': 'search_web', 'arguments': '{"query":"Rakhi Kaag"}'}, + } + ], + }, + {'role': 'tool', 'tool_call_id': 'call_123', 'content': 'Search results'}, + ] + output = [ + { + 'type': 'reasoning', + 'summary': [{'type': 'output_text', 'text': 'Need to search.'}], + }, + { + 'type': 'function_call', + 'call_id': 'call_123', + 'name': 'search_web', + 'arguments': '{"query":"Rakhi Kaag"}', + }, + {'type': 'function_call_output', 'call_id': 'call_123', 'output': []}, + ] + + adjusted = add_reasoning_content_to_tool_messages(messages, output, 'kimi-k2.6') + + assert adjusted[0]['reasoning_content'] == 'Need to search.' + assert adjusted[1] == messages[1] + + def test_scopes_reasoning_to_matching_assistant_message(self): + messages = [ + { + 'role': 'assistant', + 'content': 'Need to search the web.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': {'name': 'search_web', 'arguments': '{"query":"Rakhi Kaag"}'}, + } + ], + }, + {'role': 'tool', 'tool_call_id': 'call_1', 'content': 'First search results'}, + { + 'role': 'assistant', + 'content': 'Need to search notes.', + 'tool_calls': [ + { + 'id': 'call_2', + 'type': 'function', + 'function': {'name': 'search_notes', 'arguments': '{"query":"Rakhi Kaag"}'}, + } + ], + }, + ] + output = [ + { + 'type': 'reasoning', + 'summary': [{'type': 'output_text', 'text': 'Need to search the web.'}], + }, + { + 'type': 'function_call', + 'call_id': 'call_1', + 'name': 'search_web', + 'arguments': '{"query":"Rakhi Kaag"}', + }, + {'type': 'function_call_output', 'call_id': 'call_1', 'output': []}, + { + 'type': 'reasoning', + 'summary': [{'type': 'output_text', 'text': 'Need to search notes.'}], + }, + { + 'type': 'function_call', + 'call_id': 'call_2', + 'name': 'search_notes', + 'arguments': '{"query":"Rakhi Kaag"}', + }, + {'type': 'function_call_output', 'call_id': 'call_2', 'output': []}, + ] + + adjusted = add_reasoning_content_to_tool_messages(messages, output, 'kimi-k2.6') + + assert adjusted[0]['reasoning_content'] == 'Need to search the web.' + assert adjusted[2]['reasoning_content'] == 'Need to search notes.' + + def test_leaves_non_kimi_models_unchanged(self): + messages = [ + { + 'role': 'assistant', + 'content': 'Need to search.', + 'tool_calls': [ + { + 'id': 'call_123', + 'type': 'function', + 'function': {'name': 'search_web', 'arguments': '{"query":"Rakhi Kaag"}'}, + } + ], + } + ] + output = [ + { + 'type': 'reasoning', + 'summary': [{'type': 'output_text', 'text': 'Need to search.'}], + }, + { + 'type': 'function_call', + 'call_id': 'call_123', + 'name': 'search_web', + 'arguments': '{"query":"Rakhi Kaag"}', + }, + {'type': 'function_call_output', 'call_id': 'call_123', 'output': []}, + ] + + adjusted = add_reasoning_content_to_tool_messages(messages, output, 'gpt-4.1') + + assert adjusted == messages + + def test_keeps_existing_reasoning_content(self): + messages = [ + { + 'role': 'assistant', + 'content': 'Need to search.', + 'reasoning_content': 'Existing reasoning', + 'tool_calls': [ + { + 'id': 'call_123', + 'type': 'function', + 'function': {'name': 'search_web', 'arguments': '{"query":"Rakhi Kaag"}'}, + } + ], + } + ] + output = [ + { + 'type': 'reasoning', + 'summary': [{'type': 'output_text', 'text': 'Need to search.'}], + }, + { + 'type': 'function_call', + 'call_id': 'call_123', + 'name': 'search_web', + 'arguments': '{"query":"Rakhi Kaag"}', + }, + {'type': 'function_call_output', 'call_id': 'call_123', 'output': []}, + ] + + adjusted = add_reasoning_content_to_tool_messages(messages, output, 'moonshot-v1') + + assert adjusted[0]['reasoning_content'] == 'Existing reasoning' diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 8d3b6dd267..7e927308e8 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -98,6 +98,7 @@ from open_webui.utils.misc import ( convert_logit_bias_input_to_json, get_content_from_message, convert_output_to_messages, + add_reasoning_content_to_tool_messages, strip_empty_content_blocks, ) from open_webui.utils.tools import ( @@ -4660,12 +4661,17 @@ async def streaming_chat_response_handler(response, ctx): if ENABLE_RESPONSES_API_STATEFUL and last_response_id: system_message = get_system_message(form_data['messages']) + tool_messages = add_reasoning_content_to_tool_messages( + convert_output_to_messages(output, raw=True), output, model_id + ) new_form_data['messages'] = ( [system_message] if system_message else [] - ) + convert_output_to_messages(output, raw=True) + ) + tool_messages new_form_data['previous_response_id'] = last_response_id else: - tool_messages = convert_output_to_messages(output, raw=True) + tool_messages = add_reasoning_content_to_tool_messages( + convert_output_to_messages(output, raw=True), output, model_id + ) # Chat Completions providers don't support multimodal # tool messages. Extract images into a user message. @@ -4875,13 +4881,16 @@ async def streaming_chat_response_handler(response, ctx): ) try: + tool_messages = add_reasoning_content_to_tool_messages( + convert_output_to_messages(output, raw=True), output, model_id + ) new_form_data = { **form_data, 'model': model_id, 'stream': True, 'messages': [ *form_data['messages'], - *convert_output_to_messages(output, raw=True), + *tool_messages, ], } diff --git a/backend/open_webui/utils/misc.py b/backend/open_webui/utils/misc.py index 670a94b512..c3d28b898e 100644 --- a/backend/open_webui/utils/misc.py +++ b/backend/open_webui/utils/misc.py @@ -129,6 +129,119 @@ def get_content_from_message(message: dict) -> Optional[str]: return None +def model_requires_reasoning_content_for_tool_calls(model_id: Optional[str]) -> bool: + """ + Some OpenAI-compatible thinking models require assistant tool-call messages + to preserve top-level reasoning_content on continuation turns. + """ + if not model_id: + return False + + model_name = model_id.lower() + return 'kimi' in model_name or 'moonshot' in model_name + + +def _extract_reasoning_content_from_item(item: dict) -> str: + reasoning_parts = [] + + source_list = item.get('summary', []) or item.get('content', []) + for part in source_list: + text = part.get('text', '') + if text: + reasoning_parts.append(text) + + return ''.join(reasoning_parts).strip() + + +def get_assistant_reasoning_content_from_output(output: list[dict]) -> list[str]: + """ + Mirror convert_output_to_messages() assistant-message boundaries so + reasoning_content can be reattached to the matching assistant tool-call + message instead of being applied globally. + """ + assistant_reasoning = [] + pending_reasoning = [] + pending_tool_calls = [] + pending_content = [] + + def flush_pending(): + nonlocal pending_reasoning, pending_tool_calls, pending_content + if pending_content or pending_tool_calls: + assistant_reasoning.append(''.join(pending_reasoning).strip()) + pending_reasoning = [] + pending_tool_calls = [] + pending_content = [] + + for item in output: + item_type = item.get('type', '') + + if item_type == 'message': + content_parts = item.get('content', []) + text = '' + for part in content_parts: + if part.get('type') == 'output_text': + text += part.get('text', '') + if text: + pending_content.append(text) + + elif item_type == 'function_call': + pending_tool_calls.append(item) + + elif item_type == 'function_call_output': + flush_pending() + + elif item_type == 'reasoning': + reasoning_text = _extract_reasoning_content_from_item(item) + if reasoning_text: + pending_reasoning.append(reasoning_text) + + elif item_type == 'open_webui:code_interpreter': + if item.get('code') or item.get('output'): + pending_content.append('code_interpreter') + + elif item_type.startswith('open_webui:'): + pass + + flush_pending() + return assistant_reasoning + + +def add_reasoning_content_to_tool_messages( + messages: list[dict], output: list[dict], model_id: Optional[str] +) -> list[dict]: + """ + Re-add top-level reasoning_content to assistant tool-call messages for + providers that require it during the post-tool continuation request. + """ + if not model_requires_reasoning_content_for_tool_calls(model_id): + return messages + + assistant_reasoning = get_assistant_reasoning_content_from_output(output) + if not any(assistant_reasoning): + return messages + + adjusted_messages = [] + assistant_index = 0 + for message in messages: + reasoning_content = None + if message.get('role') == 'assistant': + if assistant_index < len(assistant_reasoning): + reasoning_content = assistant_reasoning[assistant_index] + assistant_index += 1 + + if ( + message.get('role') == 'assistant' + and message.get('tool_calls') + and reasoning_content + and not message.get('reasoning_content') + ): + message = {**message, 'reasoning_content': reasoning_content} + + adjusted_messages.append(message) + + return adjusted_messages + + def convert_output_to_messages(output: list, raw: bool = False) -> list[dict]: """ Convert OR-aligned output items to OpenAI Chat Completion-format messages.