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

refac: async db

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-04-12 14:22:11 -05:00
parent b618d84065
commit 27169124f2
74 changed files with 4831 additions and 4479 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/models/functions.py
+102 -101
View File
@@ -2,8 +2,9 @@ import logging
import time
from typing import Optional
from sqlalchemy.orm import Session, defer
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, update
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.users import Users, UserModel, UserResponse
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Boolean, Column, String, Text, Index
@@ -107,12 +108,12 @@ class FunctionValves(BaseModel):
class FunctionsTable:
def insert_new_function(
async def insert_new_function(
self,
user_id: str,
type: str,
form_data: FunctionForm,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[FunctionModel]:
function = FunctionModel(
**{
@@ -125,11 +126,11 @@ class FunctionsTable:
)
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
result = Function(**function.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return FunctionModel.model_validate(result)
else:
@@ -138,17 +139,18 @@ class FunctionsTable:
log.exception(f'Error creating a new function: {e}')
return None
def sync_functions(
async def sync_functions(
self,
user_id: str,
functions: list[FunctionWithValvesModel],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[FunctionWithValvesModel]:
# Synchronize functions for a user by updating existing ones, inserting new ones, and removing those that are no longer present.
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Get existing functions
existing_functions = db.query(Function).all()
result = await db.execute(select(Function))
existing_functions = result.scalars().all()
existing_ids = {func.id for func in existing_functions}
# Prepare a set of new function IDs
@@ -157,12 +159,12 @@ class FunctionsTable:
# Update or insert functions
for func in functions:
if func.id in existing_ids:
db.query(Function).filter_by(id=func.id).update(
{
await db.execute(
update(Function).filter_by(id=func.id).values(
**func.model_dump(),
'user_id': user_id,
'updated_at': int(time.time()),
}
user_id=user_id,
updated_at=int(time.time()),
)
)
else:
new_func = Function(
@@ -177,24 +179,25 @@ class FunctionsTable:
# Remove functions that are no longer present
for func in existing_functions:
if func.id not in new_function_ids:
db.delete(func)
await db.delete(func)
db.commit()
await db.commit()
return [FunctionModel.model_validate(func) for func in db.query(Function).all()]
result = await db.execute(select(Function))
return [FunctionModel.model_validate(func) for func in result.scalars().all()]
except Exception as e:
log.exception(f'Error syncing functions for user {user_id}: {e}')
return []
def get_function_by_id(self, id: str, db: Optional[Session] = None) -> Optional[FunctionModel]:
async def get_function_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[FunctionModel]:
try:
with get_db_context(db) as db:
function = db.get(Function, id)
return FunctionModel.model_validate(function)
async with get_async_db_context(db) as db:
function = await db.get(Function, id)
return FunctionModel.model_validate(function) if function else None
except Exception:
return None
def get_functions_by_ids(self, ids: list[str], db: Optional[Session] = None) -> list[FunctionModel]:
async def get_functions_by_ids(self, ids: list[str], db: Optional[AsyncSession] = None) -> list[FunctionModel]:
"""
Batch fetch multiple functions by their IDs in a single query.
Returns functions in the same order as the input IDs (None entries filtered out).
@@ -202,8 +205,9 @@ class FunctionsTable:
if not ids:
return []
try:
with get_db_context(db) as db:
functions = db.query(Function).filter(Function.id.in_(ids)).all()
async with get_async_db_context(db) as db:
result = await db.execute(select(Function).filter(Function.id.in_(ids)))
functions = result.scalars().all()
# Create a dict for O(1) lookup
func_dict = {f.id: FunctionModel.model_validate(f) for f in functions}
# Return in original order, filtering out any not found
@@ -211,27 +215,31 @@ class FunctionsTable:
except Exception:
return []
def get_functions(
self, active_only=False, include_valves=False, db: Optional[Session] = None
async def get_functions(
self, active_only=False, include_valves=False, db: Optional[AsyncSession] = None
) -> list[FunctionModel | FunctionWithValvesModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
if active_only:
functions = db.query(Function).filter_by(is_active=True).all()
result = await db.execute(select(Function).filter_by(is_active=True))
else:
functions = db.query(Function).all()
result = await db.execute(select(Function))
functions = result.scalars().all()
if include_valves:
return [FunctionWithValvesModel.model_validate(function) for function in functions]
else:
return [FunctionModel.model_validate(function) for function in functions]
def get_function_list(self, db: Optional[Session] = None) -> list[FunctionUserResponse]:
with get_db_context(db) as db:
functions = db.query(Function).options(defer(Function.content)).order_by(Function.updated_at.desc()).all()
async def get_function_list(self, db: Optional[AsyncSession] = None) -> list[FunctionUserResponse]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(Function).order_by(Function.updated_at.desc())
)
functions = result.scalars().all()
user_ids = list(set(func.user_id for func in functions))
users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else []
users = await Users.get_users_by_user_ids(user_ids, db=db) if user_ids else []
users_dict = {user.id: user for user in users}
return [
@@ -253,42 +261,34 @@ class FunctionsTable:
for func in functions
]
def get_functions_by_type(self, type: str, active_only=False, db: Optional[Session] = None) -> list[FunctionModel]:
with get_db_context(db) as db:
async def get_functions_by_type(self, type: str, active_only=False, db: Optional[AsyncSession] = None) -> list[FunctionModel]:
async with get_async_db_context(db) as db:
if active_only:
return [
FunctionModel.model_validate(function)
for function in db.query(Function).filter_by(type=type, is_active=True).all()
]
result = await db.execute(select(Function).filter_by(type=type, is_active=True))
else:
return [
FunctionModel.model_validate(function) for function in db.query(Function).filter_by(type=type).all()
]
result = await db.execute(select(Function).filter_by(type=type))
return [FunctionModel.model_validate(function) for function in result.scalars().all()]
def get_global_filter_functions(self, db: Optional[Session] = None) -> list[FunctionModel]:
with get_db_context(db) as db:
return [
FunctionModel.model_validate(function)
for function in db.query(Function).filter_by(type='filter', is_active=True, is_global=True).all()
]
async def get_global_filter_functions(self, db: Optional[AsyncSession] = None) -> list[FunctionModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Function).filter_by(type='filter', is_active=True, is_global=True))
return [FunctionModel.model_validate(function) for function in result.scalars().all()]
def get_global_action_functions(self, db: Optional[Session] = None) -> list[FunctionModel]:
with get_db_context(db) as db:
return [
FunctionModel.model_validate(function)
for function in db.query(Function).filter_by(type='action', is_active=True, is_global=True).all()
]
async def get_global_action_functions(self, db: Optional[AsyncSession] = None) -> list[FunctionModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Function).filter_by(type='action', is_active=True, is_global=True))
return [FunctionModel.model_validate(function) for function in result.scalars().all()]
def get_function_valves_by_id(self, id: str, db: Optional[Session] = None) -> Optional[dict]:
with get_db_context(db) as db:
async def get_function_valves_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[dict]:
async with get_async_db_context(db) as db:
try:
function = db.get(Function, id)
function = await db.get(Function, id)
return function.valves if function.valves else {}
except Exception as e:
log.exception(f'Error getting function valves by id {id}: {e}')
return None
def get_function_valves_by_ids(self, ids: list[str], db: Optional[Session] = None) -> dict[str, dict]:
async def get_function_valves_by_ids(self, ids: list[str], db: Optional[AsyncSession] = None) -> dict[str, dict]:
"""
Batch fetch valves for multiple functions in a single query.
Returns a dict mapping function_id -> valves dict.
@@ -297,33 +297,34 @@ class FunctionsTable:
if not ids:
return {}
try:
with get_db_context(db) as db:
functions = db.query(Function.id, Function.valves).filter(Function.id.in_(ids)).all()
async with get_async_db_context(db) as db:
result = await db.execute(select(Function.id, Function.valves).filter(Function.id.in_(ids)))
functions = result.all()
return {f.id: (f.valves if f.valves else {}) for f in functions}
except Exception as e:
log.exception(f'Error batch-fetching function valves: {e}')
return {}
def update_function_valves_by_id(
self, id: str, valves: dict, db: Optional[Session] = None
async def update_function_valves_by_id(
self, id: str, valves: dict, db: Optional[AsyncSession] = None
) -> Optional[FunctionValves]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
try:
function = db.get(Function, id)
function = await db.get(Function, id)
function.valves = valves
function.updated_at = int(time.time())
db.commit()
db.refresh(function)
await db.commit()
await db.refresh(function)
return FunctionModel.model_validate(function)
except Exception:
return None
def update_function_metadata_by_id(
self, id: str, metadata: dict, db: Optional[Session] = None
async def update_function_metadata_by_id(
self, id: str, metadata: dict, db: Optional[AsyncSession] = None
) -> Optional[FunctionModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
try:
function = db.get(Function, id)
function = await db.get(Function, id)
if function:
if function.meta:
@@ -332,8 +333,8 @@ class FunctionsTable:
function.meta = metadata
function.updated_at = int(time.time())
db.commit()
db.refresh(function)
await db.commit()
await db.refresh(function)
return FunctionModel.model_validate(function)
else:
return None
@@ -341,9 +342,9 @@ class FunctionsTable:
log.exception(f'Error updating function metadata by id {id}: {e}')
return None
def get_user_valves_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> Optional[dict]:
async def get_user_valves_by_id_and_user_id(self, id: str, user_id: str, db: Optional[AsyncSession] = None) -> Optional[dict]:
try:
user = Users.get_user_by_id(user_id, db=db)
user = await Users.get_user_by_id(user_id, db=db)
user_settings = user.settings.model_dump() if user.settings else {}
# Check if user has "functions" and "valves" settings
@@ -357,11 +358,11 @@ class FunctionsTable:
log.exception(f'Error getting user values by id {id} and user id {user_id}')
return None
def update_user_valves_by_id_and_user_id(
self, id: str, user_id: str, valves: dict, db: Optional[Session] = None
async def update_user_valves_by_id_and_user_id(
self, id: str, user_id: str, valves: dict, db: Optional[AsyncSession] = None
) -> Optional[dict]:
try:
user = Users.get_user_by_id(user_id, db=db)
user = await Users.get_user_by_id(user_id, db=db)
user_settings = user.settings.model_dump() if user.settings else {}
# Check if user has "functions" and "valves" settings
@@ -373,47 +374,47 @@ class FunctionsTable:
user_settings['functions']['valves'][id] = valves
# Update the user settings in the database
Users.update_user_by_id(user_id, {'settings': user_settings}, db=db)
await Users.update_user_by_id(user_id, {'settings': user_settings}, db=db)
return user_settings['functions']['valves'][id]
except Exception as e:
log.exception(f'Error updating user valves by id {id} and user_id {user_id}: {e}')
return None
def update_function_by_id(self, id: str, updated: dict, db: Optional[Session] = None) -> Optional[FunctionModel]:
with get_db_context(db) as db:
async def update_function_by_id(self, id: str, updated: dict, db: Optional[AsyncSession] = None) -> Optional[FunctionModel]:
async with get_async_db_context(db) as db:
try:
db.query(Function).filter_by(id=id).update(
{
await db.execute(
update(Function).filter_by(id=id).values(
**updated,
'updated_at': int(time.time()),
}
updated_at=int(time.time()),
)
)
db.commit()
function = db.get(Function, id)
await db.commit()
function = await db.get(Function, id)
return FunctionModel.model_validate(function) if function else None
except Exception:
return None
def deactivate_all_functions(self, db: Optional[Session] = None) -> Optional[bool]:
with get_db_context(db) as db:
async def deactivate_all_functions(self, db: Optional[AsyncSession] = None) -> Optional[bool]:
async with get_async_db_context(db) as db:
try:
db.query(Function).update(
{
'is_active': False,
'updated_at': int(time.time()),
}
await db.execute(
update(Function).values(
is_active=False,
updated_at=int(time.time()),
)
)
db.commit()
await db.commit()
return True
except Exception:
return None
def delete_function_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def delete_function_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
db.query(Function).filter_by(id=id).delete()
db.commit()
await db.execute(delete(Function).filter_by(id=id))
await db.commit()
return True
except Exception: