scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-10-14 04:08:46 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

refac

Browse source
This commit is contained in:
Timothy Jaeryang Baek 2026-10-13 00:42:34 +04:00
parent 7ae0996745
commit 53a8f131f3
3 changed files with 70 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
backend/open_webui/routers/tasks.py
View file

@ -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,

31
backend/open_webui/utils/middleware.py
View file

@ -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',

49
backend/open_webui/utils/task.py
View file

@ -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)