refac: async db
This commit is contained in:
Timothy Jaeryang Baek
@@ -26,8 +26,8 @@ from open_webui.constants import ERROR_MESSAGES
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from pydantic import BaseModel, HttpUrl
|
||||
from open_webui.internal.db import get_session
|
||||
from sqlalchemy.orm import Session
|
||||
from open_webui.internal.db import get_async_session
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
@@ -42,13 +42,13 @@ router = APIRouter()
|
||||
|
||||
|
||||
@router.get('/', response_model=list[FunctionResponse])
|
||||
async def get_functions(user=Depends(get_verified_user), db: Session = Depends(get_session)):
|
||||
return Functions.get_functions(db=db)
|
||||
async def get_functions(user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)):
|
||||
return await Functions.get_functions(db=db)
|
||||
|
||||
|
||||
@router.get('/list', response_model=list[FunctionUserResponse])
|
||||
async def get_function_list(user=Depends(get_admin_user), db: Session = Depends(get_session)):
|
||||
return Functions.get_function_list(db=db)
|
||||
async def get_function_list(user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
|
||||
return await Functions.get_function_list(db=db)
|
||||
|
||||
|
||||
############################
|
||||
@@ -60,9 +60,9 @@ async def get_function_list(user=Depends(get_admin_user), db: Session = Depends(
|
||||
async def get_functions(
|
||||
include_valves: bool = False,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
return Functions.get_functions(include_valves=include_valves, db=db)
|
||||
return await Functions.get_functions(include_valves=include_valves, db=db)
|
||||
|
||||
|
||||
############################
|
||||
@@ -145,12 +145,12 @@ async def sync_functions(
|
||||
request: Request,
|
||||
form_data: SyncFunctionsForm,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
try:
|
||||
for function in form_data.functions:
|
||||
function.content = replace_imports(function.content)
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(
|
||||
function_module, function_type, frontmatter = await load_function_module_by_id(
|
||||
function.id,
|
||||
content=function.content,
|
||||
)
|
||||
@@ -163,7 +163,7 @@ async def sync_functions(
|
||||
log.exception(f'Error validating valves for function {function.id}: {e}')
|
||||
raise e
|
||||
|
||||
return Functions.sync_functions(user.id, form_data.functions, db=db)
|
||||
return await Functions.sync_functions(user.id, form_data.functions, db=db)
|
||||
except Exception as e:
|
||||
log.exception(f'Failed to load a function: {e}')
|
||||
raise HTTPException(
|
||||
@@ -182,7 +182,7 @@ async def create_new_function(
|
||||
request: Request,
|
||||
form_data: FunctionForm,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
if not form_data.id.isidentifier():
|
||||
raise HTTPException(
|
||||
@@ -192,11 +192,11 @@ async def create_new_function(
|
||||
|
||||
form_data.id = form_data.id.lower()
|
||||
|
||||
function = Functions.get_function_by_id(form_data.id, db=db)
|
||||
function = await Functions.get_function_by_id(form_data.id, db=db)
|
||||
if function is None:
|
||||
try:
|
||||
form_data.content = replace_imports(form_data.content)
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(
|
||||
function_module, function_type, frontmatter = await load_function_module_by_id(
|
||||
form_data.id,
|
||||
content=form_data.content,
|
||||
)
|
||||
@@ -205,13 +205,13 @@ async def create_new_function(
|
||||
FUNCTIONS = request.app.state.FUNCTIONS
|
||||
FUNCTIONS[form_data.id] = function_module
|
||||
|
||||
function = Functions.insert_new_function(user.id, function_type, form_data, db=db)
|
||||
function = await Functions.insert_new_function(user.id, function_type, form_data, db=db)
|
||||
|
||||
function_cache_dir = CACHE_DIR / 'functions' / form_data.id
|
||||
function_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if function_type == 'filter' and getattr(function_module, 'toggle', None):
|
||||
Functions.update_function_metadata_by_id(form_data.id, {'toggle': True}, db=db)
|
||||
await Functions.update_function_metadata_by_id(form_data.id, {'toggle': True}, db=db)
|
||||
|
||||
if function:
|
||||
return function
|
||||
@@ -239,8 +239,8 @@ async def create_new_function(
|
||||
|
||||
|
||||
@router.get('/id/{id}', response_model=Optional[FunctionModel])
|
||||
async def get_function_by_id(id: str, user=Depends(get_admin_user), db: Session = Depends(get_session)):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
async def get_function_by_id(id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
|
||||
if function:
|
||||
return function
|
||||
@@ -257,10 +257,10 @@ async def get_function_by_id(id: str, user=Depends(get_admin_user), db: Session
|
||||
|
||||
|
||||
@router.post('/id/{id}/toggle', response_model=Optional[FunctionModel])
|
||||
async def toggle_function_by_id(id: str, user=Depends(get_admin_user), db: Session = Depends(get_session)):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
async def toggle_function_by_id(id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
if function:
|
||||
function = Functions.update_function_by_id(id, {'is_active': not function.is_active}, db=db)
|
||||
function = await Functions.update_function_by_id(id, {'is_active': not function.is_active}, db=db)
|
||||
|
||||
if function:
|
||||
return function
|
||||
@@ -282,10 +282,10 @@ async def toggle_function_by_id(id: str, user=Depends(get_admin_user), db: Sessi
|
||||
|
||||
|
||||
@router.post('/id/{id}/toggle/global', response_model=Optional[FunctionModel])
|
||||
async def toggle_global_by_id(id: str, user=Depends(get_admin_user), db: Session = Depends(get_session)):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
async def toggle_global_by_id(id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
if function:
|
||||
function = Functions.update_function_by_id(id, {'is_global': not function.is_global}, db=db)
|
||||
function = await Functions.update_function_by_id(id, {'is_global': not function.is_global}, db=db)
|
||||
|
||||
if function:
|
||||
return function
|
||||
@@ -312,11 +312,11 @@ async def update_function_by_id(
|
||||
id: str,
|
||||
form_data: FunctionForm,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
try:
|
||||
form_data.content = replace_imports(form_data.content)
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(id, content=form_data.content)
|
||||
function_module, function_type, frontmatter = await load_function_module_by_id(id, content=form_data.content)
|
||||
form_data.meta.manifest = frontmatter
|
||||
|
||||
FUNCTIONS = request.app.state.FUNCTIONS
|
||||
@@ -325,10 +325,10 @@ async def update_function_by_id(
|
||||
updated = {**form_data.model_dump(exclude={'id'}), 'type': function_type}
|
||||
log.debug(updated)
|
||||
|
||||
function = Functions.update_function_by_id(id, updated, db=db)
|
||||
function = await Functions.update_function_by_id(id, updated, db=db)
|
||||
|
||||
if function_type == 'filter' and getattr(function_module, 'toggle', None):
|
||||
Functions.update_function_metadata_by_id(id, {'toggle': True}, db=db)
|
||||
await Functions.update_function_metadata_by_id(id, {'toggle': True}, db=db)
|
||||
|
||||
if function:
|
||||
return function
|
||||
@@ -355,9 +355,9 @@ async def delete_function_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
result = Functions.delete_function_by_id(id, db=db)
|
||||
result = await Functions.delete_function_by_id(id, db=db)
|
||||
|
||||
if result:
|
||||
FUNCTIONS = request.app.state.FUNCTIONS
|
||||
@@ -373,11 +373,11 @@ async def delete_function_by_id(
|
||||
|
||||
|
||||
@router.get('/id/{id}/valves', response_model=Optional[dict])
|
||||
async def get_function_valves_by_id(id: str, user=Depends(get_admin_user), db: Session = Depends(get_session)):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
async def get_function_valves_by_id(id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
if function:
|
||||
try:
|
||||
valves = Functions.get_function_valves_by_id(id, db=db)
|
||||
valves = await Functions.get_function_valves_by_id(id, db=db)
|
||||
return valves
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
@@ -401,11 +401,11 @@ async def get_function_valves_spec_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
if function:
|
||||
function_module, function_type, frontmatter = get_function_module_from_cache(request, id)
|
||||
function_module, function_type, frontmatter = await get_function_module_from_cache(request, id)
|
||||
|
||||
if hasattr(function_module, 'Valves'):
|
||||
Valves = function_module.Valves
|
||||
@@ -432,11 +432,11 @@ async def update_function_valves_by_id(
|
||||
id: str,
|
||||
form_data: dict,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
if function:
|
||||
function_module, function_type, frontmatter = get_function_module_from_cache(request, id)
|
||||
function_module, function_type, frontmatter = await get_function_module_from_cache(request, id)
|
||||
|
||||
if hasattr(function_module, 'Valves'):
|
||||
Valves = function_module.Valves
|
||||
@@ -446,7 +446,7 @@ async def update_function_valves_by_id(
|
||||
valves = Valves(**form_data)
|
||||
|
||||
valves_dict = valves.model_dump(exclude_unset=True)
|
||||
Functions.update_function_valves_by_id(id, valves_dict, db=db)
|
||||
await Functions.update_function_valves_by_id(id, valves_dict, db=db)
|
||||
return valves_dict
|
||||
except Exception as e:
|
||||
log.exception(f'Error updating function values by id {id}: {e}')
|
||||
@@ -473,11 +473,11 @@ async def update_function_valves_by_id(
|
||||
|
||||
|
||||
@router.get('/id/{id}/valves/user', response_model=Optional[dict])
|
||||
async def get_function_user_valves_by_id(id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
async def get_function_user_valves_by_id(id: str, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)):
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
if function:
|
||||
try:
|
||||
user_valves = Functions.get_user_valves_by_id_and_user_id(id, user.id, db=db)
|
||||
user_valves = await Functions.get_user_valves_by_id_and_user_id(id, user.id, db=db)
|
||||
return user_valves
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
@@ -496,11 +496,11 @@ async def get_function_user_valves_spec_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
if function:
|
||||
function_module, function_type, frontmatter = get_function_module_from_cache(request, id)
|
||||
function_module, function_type, frontmatter = await get_function_module_from_cache(request, id)
|
||||
|
||||
if hasattr(function_module, 'UserValves'):
|
||||
UserValves = function_module.UserValves
|
||||
@@ -522,12 +522,12 @@ async def update_function_user_valves_by_id(
|
||||
id: str,
|
||||
form_data: dict,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
function = Functions.get_function_by_id(id, db=db)
|
||||
function = await Functions.get_function_by_id(id, db=db)
|
||||
|
||||
if function:
|
||||
function_module, function_type, frontmatter = get_function_module_from_cache(request, id)
|
||||
function_module, function_type, frontmatter = await get_function_module_from_cache(request, id)
|
||||
|
||||
if hasattr(function_module, 'UserValves'):
|
||||
UserValves = function_module.UserValves
|
||||
@@ -536,7 +536,7 @@ async def update_function_user_valves_by_id(
|
||||
form_data = {k: v for k, v in form_data.items() if v is not None}
|
||||
user_valves = UserValves(**form_data)
|
||||
user_valves_dict = user_valves.model_dump(exclude_unset=True)
|
||||
Functions.update_user_valves_by_id_and_user_id(id, user.id, user_valves_dict, db=db)
|
||||
await Functions.update_user_valves_by_id_and_user_id(id, user.id, user_valves_dict, db=db)
|
||||
return user_valves_dict
|
||||
except Exception as e:
|
||||
log.exception(f'Error updating function user valves by id {id}: {e}')
|
||||
|
||||
Reference in New Issue
Block a user