bhetherman/open-webui
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek committed 2026-06-17 02:52:35 +02:00
1 parent e8d55c0a8b
commit 5cdcdbaeec
53 files changed
+7829 -7954
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/open_webui/routers/tasks.py
+75 -99
View File
@@ -16,6 +16,7 @@ from open_webui.config import (
DEFAULT_VOICE_MODE_PROMPT_TEMPLATE,
)
from open_webui.constants import ERROR_MESSAGES, TASKS
from open_webui.models.config import Config
from open_webui.routers.pipelines import process_pipeline_inlet_filter
from open_webui.utils.auth import get_admin_user, get_verified_user
from open_webui.utils.chat import generate_chat_completion
@@ -36,6 +37,35 @@ log = logging.getLogger(__name__)
router = APIRouter()
TASK_CONFIG_KEYS = {
'TASK_MODEL': 'task.model.default',
'TASK_MODEL_EXTERNAL': 'task.model.external',
'TITLE_GENERATION_PROMPT_TEMPLATE': 'task.title.prompt_template',
'IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE': 'task.image.prompt_template',
'ENABLE_AUTOCOMPLETE_GENERATION': 'task.autocomplete.enable',
'AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH': 'task.autocomplete.input_max_length',
'TAGS_GENERATION_PROMPT_TEMPLATE': 'task.tags.prompt_template',
'FOLLOW_UP_GENERATION_PROMPT_TEMPLATE': 'task.follow_up.prompt_template',
'ENABLE_FOLLOW_UP_GENERATION': 'task.follow_up.enable',
'ENABLE_TAGS_GENERATION': 'task.tags.enable',
'ENABLE_TITLE_GENERATION': 'task.title.enable',
'ENABLE_SEARCH_QUERY_GENERATION': 'task.query.search.enable',
'ENABLE_RETRIEVAL_QUERY_GENERATION': 'task.query.retrieval.enable',
'QUERY_GENERATION_PROMPT_TEMPLATE': 'task.query.prompt_template',
'TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE': 'task.tools.prompt_template',
'ENABLE_VOICE_MODE_PROMPT': 'task.voice.prompt.enable',
'VOICE_MODE_PROMPT_TEMPLATE': 'task.voice.prompt_template',
}
async def get_config_values(key_map: dict[str, str]) -> dict:
values = await Config.get_many(*key_map.values())
return {field: values[storage_key] for field, storage_key in key_map.items() if storage_key in values}
def config_updates(data: dict, key_map: dict[str, str]) -> dict:
return {key_map[field]: value for field, value in data.items() if field in key_map}
##################################
#
@@ -59,25 +89,7 @@ async def check_active_chats(request: Request, form_data: ActiveChatsForm, user=
@router.get('/config')
async def get_task_config(request: Request, user=Depends(get_verified_user)):
return {
'TASK_MODEL': request.app.state.config.TASK_MODEL,
'TASK_MODEL_EXTERNAL': request.app.state.config.TASK_MODEL_EXTERNAL,
'TITLE_GENERATION_PROMPT_TEMPLATE': request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
'IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE': request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE,
'ENABLE_AUTOCOMPLETE_GENERATION': request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION,
'AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH': request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH,
'TAGS_GENERATION_PROMPT_TEMPLATE': request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE,
'FOLLOW_UP_GENERATION_PROMPT_TEMPLATE': request.app.state.config.FOLLOW_UP_GENERATION_PROMPT_TEMPLATE,
'ENABLE_FOLLOW_UP_GENERATION': request.app.state.config.ENABLE_FOLLOW_UP_GENERATION,
'ENABLE_TAGS_GENERATION': request.app.state.config.ENABLE_TAGS_GENERATION,
'ENABLE_TITLE_GENERATION': request.app.state.config.ENABLE_TITLE_GENERATION,
'ENABLE_SEARCH_QUERY_GENERATION': request.app.state.config.ENABLE_SEARCH_QUERY_GENERATION,
'ENABLE_RETRIEVAL_QUERY_GENERATION': request.app.state.config.ENABLE_RETRIEVAL_QUERY_GENERATION,
'QUERY_GENERATION_PROMPT_TEMPLATE': request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE,
'TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE': request.app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
'ENABLE_VOICE_MODE_PROMPT': request.app.state.config.ENABLE_VOICE_MODE_PROMPT,
'VOICE_MODE_PROMPT_TEMPLATE': request.app.state.config.VOICE_MODE_PROMPT_TEMPLATE,
}
return await get_config_values(TASK_CONFIG_KEYS)
class TaskConfigForm(BaseModel):
@@ -102,56 +114,13 @@ class TaskConfigForm(BaseModel):
@router.post('/config/update')
async def update_task_config(request: Request, form_data: TaskConfigForm, user=Depends(get_admin_user)):
request.app.state.config.TASK_MODEL = form_data.TASK_MODEL
request.app.state.config.TASK_MODEL_EXTERNAL = form_data.TASK_MODEL_EXTERNAL
request.app.state.config.ENABLE_TITLE_GENERATION = form_data.ENABLE_TITLE_GENERATION
request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE = form_data.TITLE_GENERATION_PROMPT_TEMPLATE
request.app.state.config.ENABLE_FOLLOW_UP_GENERATION = form_data.ENABLE_FOLLOW_UP_GENERATION
request.app.state.config.FOLLOW_UP_GENERATION_PROMPT_TEMPLATE = form_data.FOLLOW_UP_GENERATION_PROMPT_TEMPLATE
request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE = form_data.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE
request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION = form_data.ENABLE_AUTOCOMPLETE_GENERATION
request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH = (
form_data.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH
)
request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE = form_data.TAGS_GENERATION_PROMPT_TEMPLATE
request.app.state.config.ENABLE_TAGS_GENERATION = form_data.ENABLE_TAGS_GENERATION
request.app.state.config.ENABLE_SEARCH_QUERY_GENERATION = form_data.ENABLE_SEARCH_QUERY_GENERATION
request.app.state.config.ENABLE_RETRIEVAL_QUERY_GENERATION = form_data.ENABLE_RETRIEVAL_QUERY_GENERATION
request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE = form_data.QUERY_GENERATION_PROMPT_TEMPLATE
request.app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE = form_data.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
request.app.state.config.ENABLE_VOICE_MODE_PROMPT = form_data.ENABLE_VOICE_MODE_PROMPT
request.app.state.config.VOICE_MODE_PROMPT_TEMPLATE = form_data.VOICE_MODE_PROMPT_TEMPLATE
return {
'TASK_MODEL': request.app.state.config.TASK_MODEL,
'TASK_MODEL_EXTERNAL': request.app.state.config.TASK_MODEL_EXTERNAL,
'ENABLE_TITLE_GENERATION': request.app.state.config.ENABLE_TITLE_GENERATION,
'TITLE_GENERATION_PROMPT_TEMPLATE': request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
'IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE': request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE,
'ENABLE_AUTOCOMPLETE_GENERATION': request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION,
'AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH': request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH,
'TAGS_GENERATION_PROMPT_TEMPLATE': request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE,
'ENABLE_TAGS_GENERATION': request.app.state.config.ENABLE_TAGS_GENERATION,
'ENABLE_FOLLOW_UP_GENERATION': request.app.state.config.ENABLE_FOLLOW_UP_GENERATION,
'FOLLOW_UP_GENERATION_PROMPT_TEMPLATE': request.app.state.config.FOLLOW_UP_GENERATION_PROMPT_TEMPLATE,
'ENABLE_SEARCH_QUERY_GENERATION': request.app.state.config.ENABLE_SEARCH_QUERY_GENERATION,
'ENABLE_RETRIEVAL_QUERY_GENERATION': request.app.state.config.ENABLE_RETRIEVAL_QUERY_GENERATION,
'QUERY_GENERATION_PROMPT_TEMPLATE': request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE,
'TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE': request.app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
'ENABLE_VOICE_MODE_PROMPT': request.app.state.config.ENABLE_VOICE_MODE_PROMPT,
'VOICE_MODE_PROMPT_TEMPLATE': request.app.state.config.VOICE_MODE_PROMPT_TEMPLATE,
}
await Config.upsert(config_updates(form_data.model_dump(), TASK_CONFIG_KEYS))
return await get_config_values(TASK_CONFIG_KEYS)
@router.post('/title/completions')
async def generate_title(request: Request, form_data: dict, user=Depends(get_verified_user)):
if not request.app.state.config.ENABLE_TITLE_GENERATION:
if not await Config.get('task.title.enable'):
return JSONResponse(
status_code=status.HTTP_200_OK,
content={'detail': 'Title generation is disabled'},
@@ -181,15 +150,16 @@ async def generate_title(request: Request, form_data: dict, user=Depends(get_ver
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
request.app.state.config.TASK_MODEL,
request.app.state.config.TASK_MODEL_EXTERNAL,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
log.debug(f'generating chat title using model {task_model_id} for user {user.email} ')
if request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE != '':
template = request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE
title_template = await Config.get('task.title.prompt_template')
if title_template != '':
template = title_template
else:
template = DEFAULT_TITLE_GENERATION_PROMPT_TEMPLATE
@@ -234,7 +204,7 @@ async def generate_title(request: Request, form_data: dict, user=Depends(get_ver
@router.post('/follow_up/completions')
async def generate_follow_ups(request: Request, form_data: dict, user=Depends(get_verified_user)):
if not request.app.state.config.ENABLE_FOLLOW_UP_GENERATION:
if not await Config.get('task.follow_up.enable'):
return JSONResponse(
status_code=status.HTTP_200_OK,
content={'detail': 'Follow-up generation is disabled'},
@@ -259,15 +229,16 @@ async def generate_follow_ups(request: Request, form_data: dict, user=Depends(ge
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
request.app.state.config.TASK_MODEL,
request.app.state.config.TASK_MODEL_EXTERNAL,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
log.debug(f'generating chat title using model {task_model_id} for user {user.email} ')
if request.app.state.config.FOLLOW_UP_GENERATION_PROMPT_TEMPLATE != '':
template = request.app.state.config.FOLLOW_UP_GENERATION_PROMPT_TEMPLATE
follow_up_template = await Config.get('task.follow_up.prompt_template')
if follow_up_template != '':
template = follow_up_template
else:
template = DEFAULT_FOLLOW_UP_GENERATION_PROMPT_TEMPLATE
@@ -303,7 +274,7 @@ async def generate_follow_ups(request: Request, form_data: dict, user=Depends(ge
@router.post('/tags/completions')
async def generate_chat_tags(request: Request, form_data: dict, user=Depends(get_verified_user)):
if not request.app.state.config.ENABLE_TAGS_GENERATION:
if not await Config.get('task.tags.enable'):
return JSONResponse(
status_code=status.HTTP_200_OK,
content={'detail': 'Tags generation is disabled'},
@@ -328,15 +299,16 @@ async def generate_chat_tags(request: Request, form_data: dict, user=Depends(get
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
request.app.state.config.TASK_MODEL,
request.app.state.config.TASK_MODEL_EXTERNAL,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
log.debug(f'generating chat tags using model {task_model_id} for user {user.email} ')
if request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE != '':
template = request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE
tags_template = await Config.get('task.tags.prompt_template')
if tags_template != '':
template = tags_template
else:
template = DEFAULT_TAGS_GENERATION_PROMPT_TEMPLATE
@@ -391,15 +363,16 @@ async def generate_image_prompt(request: Request, form_data: dict, user=Depends(
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
request.app.state.config.TASK_MODEL,
request.app.state.config.TASK_MODEL_EXTERNAL,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
log.debug(f'generating image prompt using model {task_model_id} for user {user.email} ')
if request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE != '':
template = request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE
image_prompt_template = await Config.get('task.image.prompt_template')
if image_prompt_template != '':
template = image_prompt_template
else:
template = DEFAULT_IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE
@@ -437,13 +410,13 @@ async def generate_image_prompt(request: Request, form_data: dict, user=Depends(
async def generate_queries(request: Request, form_data: dict, user=Depends(get_verified_user)):
type = form_data.get('type')
if type == 'web_search':
if not request.app.state.config.ENABLE_SEARCH_QUERY_GENERATION:
if not await Config.get('task.query.search.enable'):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.FEATURE_DISABLED('Search query generation'),
)
elif type == 'retrieval':
if not request.app.state.config.ENABLE_RETRIEVAL_QUERY_GENERATION:
if not await Config.get('task.query.retrieval.enable'):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.FEATURE_DISABLED('Query generation'),
@@ -472,15 +445,16 @@ async def generate_queries(request: Request, form_data: dict, user=Depends(get_v
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
request.app.state.config.TASK_MODEL,
request.app.state.config.TASK_MODEL_EXTERNAL,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
log.debug(f'generating {type} queries using model {task_model_id} for user {user.email}')
if (request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE).strip() != '':
template = request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE
query_template = await Config.get('task.query.prompt_template')
if query_template.strip() != '':
template = query_template
else:
template = DEFAULT_QUERY_GENERATION_PROMPT_TEMPLATE
@@ -515,7 +489,7 @@ async def generate_queries(request: Request, form_data: dict, user=Depends(get_v
@router.post('/auto/completions')
async def generate_autocompletion(request: Request, form_data: dict, user=Depends(get_verified_user)):
if not request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION:
if not await Config.get('task.autocomplete.enable'):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.FEATURE_DISABLED('Autocompletion generation'),
@@ -525,11 +499,12 @@ async def generate_autocompletion(request: Request, form_data: dict, user=Depend
prompt = form_data.get('prompt')
messages = form_data.get('messages')
if request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH > 0:
if len(prompt) > request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH:
autocomplete_input_max_length = await Config.get('task.autocomplete.input_max_length')
if autocomplete_input_max_length > 0:
if len(prompt) > autocomplete_input_max_length:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.INPUT_TOO_LONG(request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH),
detail=ERROR_MESSAGES.INPUT_TOO_LONG(autocomplete_input_max_length),
)
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
@@ -551,15 +526,16 @@ async def generate_autocompletion(request: Request, form_data: dict, user=Depend
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
request.app.state.config.TASK_MODEL,
request.app.state.config.TASK_MODEL_EXTERNAL,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
log.debug(f'generating autocompletion using model {task_model_id} for user {user.email}')
if (request.app.state.config.AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE).strip() != '':
template = request.app.state.config.AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE
autocomplete_template = await Config.get('task.autocomplete.prompt_template')
if autocomplete_template.strip() != '':
template = autocomplete_template
else:
template = DEFAULT_AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE
@@ -614,8 +590,8 @@ async def generate_emoji(request: Request, form_data: dict, user=Depends(get_ver
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
request.app.state.config.TASK_MODEL,
request.app.state.config.TASK_MODEL_EXTERNAL,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)