From 53a8f131f305a3beb87376d5dbc05889883fb950 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Tue, 13 Oct 2026 00:42:34 +0400 Subject: [PATCH] refac --- backend/open_webui/routers/tasks.py | 4 ++- backend/open_webui/utils/middleware.py | 31 ++++++++++------ backend/open_webui/utils/task.py | 49 ++++++++++++++++++++++++-- 3 files changed, 70 insertions(+), 14 deletions(-) diff --git a/backend/open_webui/routers/tasks.py b/backend/open_webui/routers/tasks.py index c2505b0ee0..921809a12f 100644 --- a/backend/open_webui/routers/tasks.py +++ b/backend/open_webui/routers/tasks.py @@ -248,7 +248,9 @@ async def generate_follow_ups(request: Request, form_data: dict, user=Depends(ge else: template = DEFAULT_FOLLOW_UP_GENERATION_PROMPT_TEMPLATE - content = await follow_up_generation_template(template, form_data['messages'], user) + content = await follow_up_generation_template( + template, form_data['messages'], user, use_default=follow_up_template == '' + ) payload = { 'model': task_model_id, diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index fe967019a8..57ac2e7517 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -140,6 +140,7 @@ from open_webui.utils.skills import ( ) from open_webui.utils.task import ( get_task_model_id, + parse_follow_ups, rag_template, tools_function_calling_generation_template, ) @@ -3892,10 +3893,11 @@ async def background_tasks_handler(ctx): for message in message_list: content = message.get('content', '') if isinstance(content, list): - for item in content: - if item.get('type') == 'text': - content = item['text'] - break + content = '\n'.join( + item['text'] + for item in content + if isinstance(item, dict) and item.get('type') == 'text' and isinstance(item.get('text'), str) + ) if isinstance(content, str): content = re.sub( @@ -3917,16 +3919,27 @@ async def background_tasks_handler(ctx): message = get_last_user_message_item(form_data.get('messages', [])) messages = form_data.get('messages', []) if message: - message['model'] = form_data.get('model') + message = {**message, 'model': form_data.get('model')} if message and 'model' in message: if tasks and messages: if TASKS.FOLLOW_UP_GENERATION in tasks and tasks[TASKS.FOLLOW_UP_GENERATION]: + follow_up_messages = list(messages) + assistant_message = ctx.get('assistant_message') + if not is_saved_chat_id(metadata.get('chat_id')) and assistant_message: + if follow_up_messages[-1].get('role') == 'assistant' and ( + follow_up_messages[-1].get('id') == metadata['message_id'] + or get_content_from_message(follow_up_messages[-1]) + == get_content_from_message(assistant_message) + ): + follow_up_messages.pop() + follow_up_messages.append({**assistant_message, 'role': 'assistant', 'id': metadata['message_id']}) + res = await generate_follow_ups( request, { 'model': message['model'], - 'messages': messages, + 'messages': follow_up_messages, 'message_id': metadata['message_id'], 'chat_id': metadata['chat_id'], }, @@ -3943,12 +3956,8 @@ async def background_tasks_handler(ctx): else: follow_ups_string = '' - follow_ups_string = follow_ups_string[ - follow_ups_string.find('{') : follow_ups_string.rfind('}') + 1 - ] - try: - follow_ups = JSONCodec.loads(follow_ups_string).get('follow_ups', []) + follow_ups = parse_follow_ups(follow_ups_string) await event_emitter( { 'type': 'chat:message:follow_ups', diff --git a/backend/open_webui/utils/task.py b/backend/open_webui/utils/task.py index 5c6f50142a..b5cee1dcb3 100644 --- a/backend/open_webui/utils/task.py +++ b/backend/open_webui/utils/task.py @@ -6,7 +6,8 @@ from datetime import datetime from typing import Any, Optional from open_webui.config import DEFAULT_RAG_TEMPLATE -from open_webui.utils.misc import get_last_user_message, get_messages_content +from open_webui.utils.json_codec import JSONCodec +from open_webui.utils.misc import get_content_from_message, get_last_user_message, get_messages_content log = logging.getLogger(__name__) @@ -311,7 +312,31 @@ async def title_generation_template(template: str, messages: list[dict], user: O return template -async def follow_up_generation_template(template: str, messages: list[dict], user: Optional[Any] = None) -> str: +async def follow_up_generation_template( + template: str, messages: list[dict], user: Optional[Any] = None, *, use_default: bool = False +) -> str: + if use_default: + start = max(0, len(messages) - 4) + first_user = next((i for i, message in enumerate(messages) if message.get('role') == 'user'), None) + messages = ([messages[first_user]] if first_user is not None and first_user < start else []) + messages[start:] + + normalized = [] + for message in messages: + content = message.get('content') + if isinstance(content, list): + content = '\n'.join( + item['text'] + for item in content + if isinstance(item, dict) and item.get('type') == 'text' and isinstance(item.get('text'), str) + ) + else: + content = get_content_from_message(message) or '' + if use_default and len(content) > 1000: + # truncate_content adds three characters for the omission marker. + content = truncate_content(content, 997) + normalized.append({**message, 'content': content}) + messages = normalized + prompt = get_last_user_message(messages) template = replace_prompt_variable(template, prompt) template = replace_messages_variable(template, messages) @@ -320,6 +345,26 @@ async def follow_up_generation_template(template: str, messages: list[dict], use return template +def parse_follow_ups(content: str, max_count: int = 5) -> list[str]: + if not isinstance(content, str): + return [] + content = re.sub(r'^```(?:json)?\s*|\s*```$', '', content.strip()) + try: + data = JSONCodec.loads(content) + except (ValueError, TypeError): + return [] + suggestions = data.get('follow_ups') if isinstance(data, dict) else None + if not isinstance(suggestions, list) or any(not isinstance(item, str) for item in suggestions): + return [] + + result = [] + for item in suggestions: + item = item.strip() + if item and len(item) <= 400 and item not in result: + result.append(item) + return result[:max_count] + + async def tags_generation_template(template: str, messages: list[dict], user: Optional[Any] = None) -> str: prompt = get_last_user_message(messages) template = replace_prompt_variable(template, prompt)