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
2026-08-11 01:15:05 -06:00
parent be4afd7545
commit f0bfcd4097
7 changed files with 133 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/routers/tasks.py
+55 -72
View File
@@ -20,6 +20,7 @@ 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
from open_webui.utils.payload import apply_params_to_form_data
from open_webui.utils.task import (
autocomplete_generation_template,
emoji_generation_template,
@@ -40,6 +41,7 @@ router = APIRouter()
TASK_CONFIG_KEYS = {
'TASK_MODEL': 'task.model.default',
'TASK_MODEL_EXTERNAL': 'task.model.external',
'TASK_MODEL_PARAMS': 'task.model.params',
'TITLE_GENERATION_PROMPT_TEMPLATE': 'task.title.prompt_template',
'IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE': 'task.image.prompt_template',
'ENABLE_AUTOCOMPLETE_GENERATION': 'task.autocomplete.enable',
@@ -68,6 +70,34 @@ 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}
def apply_task_model_params(payload: dict, models: dict, task_model_id: str, params: dict | None = None) -> dict:
model = models.get(payload.get('model')) or models.get(task_model_id)
if not model or (not params and not payload.get('params')):
return payload
return apply_params_to_form_data(payload, model, params or None)
async def get_task_model_generation_config(default_model_id: str, models) -> tuple[str, dict]:
config = await Config.get_many(
'task.model.default',
'task.model.external',
'task.model.params',
)
params = config.get('task.model.params') or {}
if not isinstance(params, dict):
params = {}
return (
get_task_model_id(
default_model_id,
config.get('task.model.default'),
config.get('task.model.external'),
models,
),
{key: value for key, value in params.items() if value is not None and value != ''},
)
##################################
#
# Task Endpoints
@@ -83,6 +113,7 @@ async def get_task_config(request: Request, user=Depends(get_verified_user)):
class TaskConfigForm(BaseModel):
TASK_MODEL: Optional[str]
TASK_MODEL_EXTERNAL: Optional[str]
TASK_MODEL_PARAMS: dict | None = None
ENABLE_TITLE_GENERATION: bool
TITLE_GENERATION_PROMPT_TEMPLATE: str
IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE: str
@@ -135,14 +166,7 @@ async def generate_title(request: Request, form_data: dict, user=Depends(get_ver
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(),
)
# Check if the user has a custom task model
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
task_model_id, task_model_params = await get_task_model_generation_config(model_id, models)
log.debug('generating chat title using model %s for user %s ', task_model_id, user.email)
@@ -153,20 +177,14 @@ async def generate_title(request: Request, form_data: dict, user=Depends(get_ver
template = DEFAULT_TITLE_GENERATION_PROMPT_TEMPLATE
content = await title_generation_template(template, form_data['messages'], user)
max_tokens = models[task_model_id].get('info', {}).get('params', {}).get('max_tokens', 1000)
task_model_params = task_model_params or {
'max_tokens': models[task_model_id].get('info', {}).get('params', {}).get('max_tokens', 1000)
}
payload = {
'model': task_model_id,
'messages': [{'role': 'user', 'content': content}],
'stream': False,
**(
{'max_tokens': max_tokens}
if models[task_model_id].get('owned_by') == 'ollama'
else {
'max_completion_tokens': max_tokens,
}
),
'metadata': {
**(request.state.metadata if hasattr(request.state, 'metadata') else {}),
'task': str(TASKS.TITLE_GENERATION),
@@ -181,6 +199,8 @@ async def generate_title(request: Request, form_data: dict, user=Depends(get_ver
except Exception as e:
raise e
payload = apply_task_model_params(payload, models, task_model_id, task_model_params)
try:
return await generate_chat_completion(request, form_data=payload, user=user)
except Exception as e:
@@ -214,14 +234,7 @@ async def generate_follow_ups(request: Request, form_data: dict, user=Depends(ge
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(),
)
# Check if the user has a custom task model
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
task_model_id, task_model_params = await get_task_model_generation_config(model_id, models)
log.debug('generating chat title using model %s for user %s ', task_model_id, user.email)
@@ -251,6 +264,8 @@ async def generate_follow_ups(request: Request, form_data: dict, user=Depends(ge
except Exception as e:
raise e
payload = apply_task_model_params(payload, models, task_model_id, task_model_params)
try:
return await generate_chat_completion(request, form_data=payload, user=user)
except Exception as e:
@@ -284,14 +299,7 @@ async def generate_chat_tags(request: Request, form_data: dict, user=Depends(get
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(),
)
# Check if the user has a custom task model
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
task_model_id, task_model_params = await get_task_model_generation_config(model_id, models)
log.debug('generating chat tags using model %s for user %s ', task_model_id, user.email)
@@ -321,6 +329,8 @@ async def generate_chat_tags(request: Request, form_data: dict, user=Depends(get
except Exception as e:
raise e
payload = apply_task_model_params(payload, models, task_model_id, task_model_params)
try:
return await generate_chat_completion(request, form_data=payload, user=user)
except Exception as e:
@@ -348,14 +358,7 @@ async def generate_image_prompt(request: Request, form_data: dict, user=Depends(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(),
)
# Check if the user has a custom task model
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
task_model_id, task_model_params = await get_task_model_generation_config(model_id, models)
log.debug('generating image prompt using model %s for user %s ', task_model_id, user.email)
@@ -385,6 +388,8 @@ async def generate_image_prompt(request: Request, form_data: dict, user=Depends(
except Exception as e:
raise e
payload = apply_task_model_params(payload, models, task_model_id, task_model_params)
try:
return await generate_chat_completion(request, form_data=payload, user=user)
except Exception as e:
@@ -430,14 +435,7 @@ async def generate_queries(request: Request, form_data: dict, user=Depends(get_v
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(),
)
# Check if the user has a custom task model
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
task_model_id, task_model_params = await get_task_model_generation_config(model_id, models)
log.debug('generating %s queries using model %s for user %s', type, task_model_id, user.email)
@@ -467,6 +465,8 @@ async def generate_queries(request: Request, form_data: dict, user=Depends(get_v
except Exception as e:
raise e
payload = apply_task_model_params(payload, models, task_model_id, task_model_params)
try:
return await generate_chat_completion(request, form_data=payload, user=user)
except Exception as e:
@@ -511,14 +511,7 @@ async def generate_autocompletion(request: Request, form_data: dict, user=Depend
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(),
)
# Check if the user has a custom task model
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
task_model_id, task_model_params = await get_task_model_generation_config(model_id, models)
log.debug('generating autocompletion using model %s for user %s', task_model_id, user.email)
@@ -548,6 +541,8 @@ async def generate_autocompletion(request: Request, form_data: dict, user=Depend
except Exception as e:
raise e
payload = apply_task_model_params(payload, models, task_model_id, task_model_params)
try:
return await generate_chat_completion(request, form_data=payload, user=user)
except Exception as e:
@@ -575,14 +570,7 @@ async def generate_emoji(request: Request, form_data: dict, user=Depends(get_ver
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(),
)
# Check if the user has a custom task model
# If the user has a custom task model, use that model
task_model_id = get_task_model_id(
model_id,
await Config.get('task.model.default'),
await Config.get('task.model.external'),
models,
)
task_model_id, _ = await get_task_model_generation_config(model_id, models)
log.debug('generating emoji using model %s for user %s ', task_model_id, user.email)
@@ -594,13 +582,6 @@ async def generate_emoji(request: Request, form_data: dict, user=Depends(get_ver
'model': task_model_id,
'messages': [{'role': 'user', 'content': content}],
'stream': False,
**(
{'max_tokens': 4}
if models[task_model_id].get('owned_by') == 'ollama'
else {
'max_completion_tokens': 4,
}
),
'metadata': {
**(request.state.metadata if hasattr(request.state, 'metadata') else {}),
'task': str(TASKS.EMOJI_GENERATION),
@@ -615,6 +596,8 @@ async def generate_emoji(request: Request, form_data: dict, user=Depends(get_ver
except Exception as e:
raise e
payload = apply_task_model_params(payload, models, task_model_id, {'max_tokens': 4})
try:
return await generate_chat_completion(request, form_data=payload, user=user)
except Exception as e: