Classic298
243a39dc9d
`request.app.state.MODELS` is a `RedisDict` when Redis is configured. Unpacking it with `{**pool}` makes Python call `keys()` and then `__getitem__` once per key, which is one HKEYS plus one HGET per model, issued sequentially through a synchronous client. At 200 models that is 201 blocking Redis round trips per call.
`RedisDict.items()` is a single HGETALL, so `dict(pool.items())` fetches the same data in one round trip. `utils/chat.py:184` already does exactly this and carries a comment explaining why; these ten call sites were missed.
They are on the direct-connection branch of the task endpoints (title, tags, follow-up, autocomplete, query generation and the rest), of `chat_completed`, and of context compaction, so they run for background tasks fired on ordinary chat turns.
Behaviour is unchanged. The merged mapping is identical, the explicitly added direct model still overrides any pool entry with the same id, and when Redis is not configured the pool is a plain dict where `dict(d.items())` and `{**d}` are equivalent.
It also closes a race. `RedisDict.set` writes with HSET and then HDELs the stale keys, so a key returned by HKEYS could be deleted before its HGET arrived, raising `KeyError` out of the dict literal and failing the request mid model refresh. The old path could likewise observe a mix of pre- and post-refresh entries. HGETALL is atomic, so the caller now always sees one coherent snapshot.
678 lines
24 KiB
Python
678 lines
24 KiB
Python
import logging
|
|
import re
|
|
from typing import Optional
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
|
|
from fastapi.responses import JSONResponse, RedirectResponse
|
|
from open_webui.config import (
|
|
DEFAULT_AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE,
|
|
DEFAULT_EMOJI_GENERATION_PROMPT_TEMPLATE,
|
|
DEFAULT_FOLLOW_UP_GENERATION_PROMPT_TEMPLATE,
|
|
DEFAULT_IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE,
|
|
DEFAULT_MOA_GENERATION_PROMPT_TEMPLATE,
|
|
DEFAULT_QUERY_GENERATION_PROMPT_TEMPLATE,
|
|
DEFAULT_TAGS_GENERATION_PROMPT_TEMPLATE,
|
|
DEFAULT_TITLE_GENERATION_PROMPT_TEMPLATE,
|
|
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
|
|
from open_webui.utils.task import (
|
|
autocomplete_generation_template,
|
|
emoji_generation_template,
|
|
follow_up_generation_template,
|
|
get_task_model_id,
|
|
image_prompt_generation_template,
|
|
moa_response_generation_template,
|
|
query_generation_template,
|
|
tags_generation_template,
|
|
title_generation_template,
|
|
)
|
|
from pydantic import BaseModel
|
|
|
|
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',
|
|
'AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE': 'task.autocomplete.prompt_template',
|
|
'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}
|
|
|
|
|
|
##################################
|
|
#
|
|
# Task Endpoints
|
|
#
|
|
##################################
|
|
|
|
|
|
@router.get('/config')
|
|
async def get_task_config(request: Request, user=Depends(get_verified_user)):
|
|
return await get_config_values(TASK_CONFIG_KEYS)
|
|
|
|
|
|
class TaskConfigForm(BaseModel):
|
|
TASK_MODEL: Optional[str]
|
|
TASK_MODEL_EXTERNAL: Optional[str]
|
|
ENABLE_TITLE_GENERATION: bool
|
|
TITLE_GENERATION_PROMPT_TEMPLATE: str
|
|
IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE: str
|
|
ENABLE_AUTOCOMPLETE_GENERATION: bool
|
|
AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH: int
|
|
AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE: str
|
|
TAGS_GENERATION_PROMPT_TEMPLATE: str
|
|
FOLLOW_UP_GENERATION_PROMPT_TEMPLATE: str
|
|
ENABLE_FOLLOW_UP_GENERATION: bool
|
|
ENABLE_TAGS_GENERATION: bool
|
|
ENABLE_SEARCH_QUERY_GENERATION: bool
|
|
ENABLE_RETRIEVAL_QUERY_GENERATION: bool
|
|
QUERY_GENERATION_PROMPT_TEMPLATE: str
|
|
TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE: str
|
|
ENABLE_VOICE_MODE_PROMPT: bool
|
|
VOICE_MODE_PROMPT_TEMPLATE: Optional[str]
|
|
|
|
|
|
@router.post('/config/update')
|
|
async def update_task_config(request: Request, form_data: TaskConfigForm, user=Depends(get_admin_user)):
|
|
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 await Config.get('task.title.enable'):
|
|
return JSONResponse(
|
|
status_code=status.HTTP_200_OK,
|
|
content={'detail': 'Title generation is disabled'},
|
|
)
|
|
|
|
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
|
|
models = {
|
|
**dict(request.app.state.MODELS.items()),
|
|
request.state.model['id']: request.state.model,
|
|
}
|
|
else:
|
|
models = request.app.state.MODELS
|
|
|
|
model_id = form_data['model']
|
|
if not model_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail='No model specified for title generation. Please ensure a model is selected for this chat.',
|
|
)
|
|
if model_id not in models:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
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,
|
|
)
|
|
|
|
log.debug(f'generating chat title using model {task_model_id} for user {user.email} ')
|
|
|
|
title_template = await Config.get('task.title.prompt_template')
|
|
if title_template != '':
|
|
template = title_template
|
|
else:
|
|
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)
|
|
|
|
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),
|
|
'task_body': form_data,
|
|
'chat_id': form_data.get('chat_id', None),
|
|
},
|
|
}
|
|
|
|
# Process the payload through the pipeline
|
|
try:
|
|
payload = await process_pipeline_inlet_filter(request, payload, user, models)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
try:
|
|
return await generate_chat_completion(request, form_data=payload, user=user)
|
|
except Exception as e:
|
|
log.error('Exception occurred', exc_info=True)
|
|
return JSONResponse(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
content={'detail': 'An internal error has occurred.'},
|
|
)
|
|
|
|
|
|
@router.post('/follow_up/completions')
|
|
async def generate_follow_ups(request: Request, form_data: dict, user=Depends(get_verified_user)):
|
|
if not await Config.get('task.follow_up.enable'):
|
|
return JSONResponse(
|
|
status_code=status.HTTP_200_OK,
|
|
content={'detail': 'Follow-up generation is disabled'},
|
|
)
|
|
|
|
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
|
|
models = {
|
|
**dict(request.app.state.MODELS.items()),
|
|
request.state.model['id']: request.state.model,
|
|
}
|
|
else:
|
|
models = request.app.state.MODELS
|
|
|
|
model_id = form_data['model']
|
|
if model_id not in models:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
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,
|
|
)
|
|
|
|
log.debug(f'generating chat title using model {task_model_id} for user {user.email} ')
|
|
|
|
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
|
|
|
|
content = await follow_up_generation_template(template, form_data['messages'], user)
|
|
|
|
payload = {
|
|
'model': task_model_id,
|
|
'messages': [{'role': 'user', 'content': content}],
|
|
'stream': False,
|
|
'metadata': {
|
|
**(request.state.metadata if hasattr(request.state, 'metadata') else {}),
|
|
'task': str(TASKS.FOLLOW_UP_GENERATION),
|
|
'task_body': form_data,
|
|
'chat_id': form_data.get('chat_id', None),
|
|
},
|
|
}
|
|
|
|
# Process the payload through the pipeline
|
|
try:
|
|
payload = await process_pipeline_inlet_filter(request, payload, user, models)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
try:
|
|
return await generate_chat_completion(request, form_data=payload, user=user)
|
|
except Exception as e:
|
|
log.error('Exception occurred', exc_info=True)
|
|
return JSONResponse(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
content={'detail': 'An internal error has occurred.'},
|
|
)
|
|
|
|
|
|
@router.post('/tags/completions')
|
|
async def generate_chat_tags(request: Request, form_data: dict, user=Depends(get_verified_user)):
|
|
if not await Config.get('task.tags.enable'):
|
|
return JSONResponse(
|
|
status_code=status.HTTP_200_OK,
|
|
content={'detail': 'Tags generation is disabled'},
|
|
)
|
|
|
|
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
|
|
models = {
|
|
**dict(request.app.state.MODELS.items()),
|
|
request.state.model['id']: request.state.model,
|
|
}
|
|
else:
|
|
models = request.app.state.MODELS
|
|
|
|
model_id = form_data['model']
|
|
if model_id not in models:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
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,
|
|
)
|
|
|
|
log.debug(f'generating chat tags using model {task_model_id} for user {user.email} ')
|
|
|
|
tags_template = await Config.get('task.tags.prompt_template')
|
|
if tags_template != '':
|
|
template = tags_template
|
|
else:
|
|
template = DEFAULT_TAGS_GENERATION_PROMPT_TEMPLATE
|
|
|
|
content = await tags_generation_template(template, form_data['messages'], user)
|
|
|
|
payload = {
|
|
'model': task_model_id,
|
|
'messages': [{'role': 'user', 'content': content}],
|
|
'stream': False,
|
|
'metadata': {
|
|
**(request.state.metadata if hasattr(request.state, 'metadata') else {}),
|
|
'task': str(TASKS.TAGS_GENERATION),
|
|
'task_body': form_data,
|
|
'chat_id': form_data.get('chat_id', None),
|
|
},
|
|
}
|
|
|
|
# Process the payload through the pipeline
|
|
try:
|
|
payload = await process_pipeline_inlet_filter(request, payload, user, models)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
try:
|
|
return await generate_chat_completion(request, form_data=payload, user=user)
|
|
except Exception as e:
|
|
log.error(f'Error generating chat completion: {e}')
|
|
return JSONResponse(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
content={'detail': 'An internal error has occurred.'},
|
|
)
|
|
|
|
|
|
@router.post('/image_prompt/completions')
|
|
async def generate_image_prompt(request: Request, form_data: dict, user=Depends(get_verified_user)):
|
|
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
|
|
models = {
|
|
**dict(request.app.state.MODELS.items()),
|
|
request.state.model['id']: request.state.model,
|
|
}
|
|
else:
|
|
models = request.app.state.MODELS
|
|
|
|
model_id = form_data['model']
|
|
if model_id not in models:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
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,
|
|
)
|
|
|
|
log.debug(f'generating image prompt using model {task_model_id} for user {user.email} ')
|
|
|
|
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
|
|
|
|
content = await image_prompt_generation_template(template, form_data['messages'], user)
|
|
|
|
payload = {
|
|
'model': task_model_id,
|
|
'messages': [{'role': 'user', 'content': content}],
|
|
'stream': False,
|
|
'metadata': {
|
|
**(request.state.metadata if hasattr(request.state, 'metadata') else {}),
|
|
'task': str(TASKS.IMAGE_PROMPT_GENERATION),
|
|
'task_body': form_data,
|
|
'chat_id': form_data.get('chat_id', None),
|
|
},
|
|
}
|
|
|
|
# Process the payload through the pipeline
|
|
try:
|
|
payload = await process_pipeline_inlet_filter(request, payload, user, models)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
try:
|
|
return await generate_chat_completion(request, form_data=payload, user=user)
|
|
except Exception as e:
|
|
log.error('Exception occurred', exc_info=True)
|
|
return JSONResponse(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
content={'detail': 'An internal error has occurred.'},
|
|
)
|
|
|
|
|
|
@router.post('/queries/completions')
|
|
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 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 await Config.get('task.query.retrieval.enable'):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail=ERROR_MESSAGES.FEATURE_DISABLED('Query generation'),
|
|
)
|
|
|
|
if getattr(request.state, 'cached_queries', None):
|
|
log.info(f'Reusing cached queries: {request.state.cached_queries}')
|
|
return request.state.cached_queries
|
|
|
|
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
|
|
models = {
|
|
**dict(request.app.state.MODELS.items()),
|
|
request.state.model['id']: request.state.model,
|
|
}
|
|
else:
|
|
models = request.app.state.MODELS
|
|
|
|
model_id = form_data['model']
|
|
if model_id not in models:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
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,
|
|
)
|
|
|
|
log.debug(f'generating {type} queries using model {task_model_id} for user {user.email}')
|
|
|
|
query_template = await Config.get('task.query.prompt_template')
|
|
if query_template.strip() != '':
|
|
template = query_template
|
|
else:
|
|
template = DEFAULT_QUERY_GENERATION_PROMPT_TEMPLATE
|
|
|
|
content = await query_generation_template(template, form_data['messages'], user)
|
|
|
|
payload = {
|
|
'model': task_model_id,
|
|
'messages': [{'role': 'user', 'content': content}],
|
|
'stream': False,
|
|
'metadata': {
|
|
**(request.state.metadata if hasattr(request.state, 'metadata') else {}),
|
|
'task': str(TASKS.QUERY_GENERATION),
|
|
'task_body': form_data,
|
|
'chat_id': form_data.get('chat_id', None),
|
|
},
|
|
}
|
|
|
|
# Process the payload through the pipeline
|
|
try:
|
|
payload = await process_pipeline_inlet_filter(request, payload, user, models)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
try:
|
|
return await generate_chat_completion(request, form_data=payload, user=user)
|
|
except Exception as e:
|
|
return JSONResponse(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
content={'detail': str(e)},
|
|
)
|
|
|
|
|
|
@router.post('/auto/completions')
|
|
async def generate_autocompletion(request: Request, form_data: dict, user=Depends(get_verified_user)):
|
|
if not await Config.get('task.autocomplete.enable'):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail=ERROR_MESSAGES.FEATURE_DISABLED('Autocompletion generation'),
|
|
)
|
|
|
|
type = form_data.get('type')
|
|
prompt = form_data.get('prompt')
|
|
messages = form_data.get('messages')
|
|
|
|
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(autocomplete_input_max_length),
|
|
)
|
|
|
|
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
|
|
models = {
|
|
**dict(request.app.state.MODELS.items()),
|
|
request.state.model['id']: request.state.model,
|
|
}
|
|
else:
|
|
models = request.app.state.MODELS
|
|
|
|
model_id = form_data['model']
|
|
if model_id not in models:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
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,
|
|
)
|
|
|
|
log.debug(f'generating autocompletion using model {task_model_id} for user {user.email}')
|
|
|
|
autocomplete_template = await Config.get('task.autocomplete.prompt_template')
|
|
if autocomplete_template.strip() != '':
|
|
template = autocomplete_template
|
|
else:
|
|
template = DEFAULT_AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE
|
|
|
|
content = await autocomplete_generation_template(template, prompt, messages, type, user)
|
|
|
|
payload = {
|
|
'model': task_model_id,
|
|
'messages': [{'role': 'user', 'content': content}],
|
|
'stream': False,
|
|
'metadata': {
|
|
**(request.state.metadata if hasattr(request.state, 'metadata') else {}),
|
|
'task': str(TASKS.AUTOCOMPLETE_GENERATION),
|
|
'task_body': form_data,
|
|
'chat_id': form_data.get('chat_id', None),
|
|
},
|
|
}
|
|
|
|
# Process the payload through the pipeline
|
|
try:
|
|
payload = await process_pipeline_inlet_filter(request, payload, user, models)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
try:
|
|
return await generate_chat_completion(request, form_data=payload, user=user)
|
|
except Exception as e:
|
|
log.error(f'Error generating chat completion: {e}')
|
|
return JSONResponse(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
content={'detail': 'An internal error has occurred.'},
|
|
)
|
|
|
|
|
|
@router.post('/emoji/completions')
|
|
async def generate_emoji(request: Request, form_data: dict, user=Depends(get_verified_user)):
|
|
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
|
|
models = {
|
|
**dict(request.app.state.MODELS.items()),
|
|
request.state.model['id']: request.state.model,
|
|
}
|
|
else:
|
|
models = request.app.state.MODELS
|
|
|
|
model_id = form_data['model']
|
|
if model_id not in models:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
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,
|
|
)
|
|
|
|
log.debug(f'generating emoji using model {task_model_id} for user {user.email} ')
|
|
|
|
template = DEFAULT_EMOJI_GENERATION_PROMPT_TEMPLATE
|
|
|
|
content = await emoji_generation_template(template, form_data['prompt'], user)
|
|
|
|
payload = {
|
|
'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),
|
|
'task_body': form_data,
|
|
'chat_id': form_data.get('chat_id', None),
|
|
},
|
|
}
|
|
|
|
# Process the payload through the pipeline
|
|
try:
|
|
payload = await process_pipeline_inlet_filter(request, payload, user, models)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
try:
|
|
return await generate_chat_completion(request, form_data=payload, user=user)
|
|
except Exception as e:
|
|
return JSONResponse(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
content={'detail': str(e)},
|
|
)
|
|
|
|
|
|
@router.post('/moa/completions')
|
|
async def generate_moa_response(request: Request, form_data: dict, user=Depends(get_verified_user)):
|
|
if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
|
|
models = {
|
|
**dict(request.app.state.MODELS.items()),
|
|
request.state.model['id']: request.state.model,
|
|
}
|
|
else:
|
|
models = request.app.state.MODELS
|
|
|
|
model_id = form_data['model']
|
|
|
|
if model_id not in models:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(),
|
|
)
|
|
|
|
template = DEFAULT_MOA_GENERATION_PROMPT_TEMPLATE
|
|
|
|
content = moa_response_generation_template(
|
|
template,
|
|
form_data['prompt'],
|
|
form_data['responses'],
|
|
)
|
|
|
|
payload = {
|
|
'model': model_id,
|
|
'messages': [{'role': 'user', 'content': content}],
|
|
'stream': form_data.get('stream', False),
|
|
'metadata': {
|
|
**(request.state.metadata if hasattr(request.state, 'metadata') else {}),
|
|
'chat_id': form_data.get('chat_id', None),
|
|
'task': str(TASKS.MOA_RESPONSE_GENERATION),
|
|
'task_body': form_data,
|
|
},
|
|
}
|
|
|
|
# Process the payload through the pipeline
|
|
try:
|
|
payload = await process_pipeline_inlet_filter(request, payload, user, models)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
try:
|
|
return await generate_chat_completion(request, form_data=payload, user=user)
|
|
except Exception as e:
|
|
return JSONResponse(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
content={'detail': str(e)},
|
|
)
|