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/access_grants.py
+101 -89
View File
@@ -3,8 +3,9 @@ import time
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db_context
from sqlalchemy import select, delete
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, Text, UniqueConstraint, or_, and_
@@ -281,20 +282,20 @@ def grants_to_access_control(grants: list) -> Optional[dict]:
class AccessGrantsTable:
def grant_access(
async def grant_access(
self,
resource_type: str,
resource_id: str,
principal_type: str,
principal_id: str,
permission: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[AccessGrantModel]:
"""Add a single access grant. Idempotent (ignores duplicates)."""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Check for existing grant
existing = (
db.query(AccessGrant)
result = await db.execute(
select(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
@@ -302,8 +303,8 @@ class AccessGrantsTable:
principal_id=principal_id,
permission=permission,
)
.first()
)
existing = result.scalars().first()
if existing:
return AccessGrantModel.model_validate(existing)
@@ -317,23 +318,23 @@ class AccessGrantsTable:
created_at=int(time.time()),
)
db.add(grant)
db.commit()
db.refresh(grant)
await db.commit()
await db.refresh(grant)
return AccessGrantModel.model_validate(grant)
def revoke_access(
async def revoke_access(
self,
resource_type: str,
resource_id: str,
principal_type: str,
principal_id: str,
permission: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
"""Remove a single access grant."""
with get_db_context(db) as db:
deleted = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
delete(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
@@ -341,47 +342,47 @@ class AccessGrantsTable:
principal_id=principal_id,
permission=permission,
)
.delete()
)
db.commit()
return deleted > 0
await db.commit()
return result.rowcount > 0
def revoke_all_access(
async def revoke_all_access(
self,
resource_type: str,
resource_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> int:
"""Remove all access grants for a resource."""
with get_db_context(db) as db:
deleted = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
delete(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
.delete()
)
db.commit()
return deleted
await db.commit()
return result.rowcount
def set_access_control(
async def set_access_control(
self,
resource_type: str,
resource_id: str,
access_control: Optional[dict],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[AccessGrantModel]:
"""
Replace all grants for a resource from an access_control JSON dict.
This is the primary bridge for backward compat with the frontend.
"""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Delete all existing grants for this resource
db.query(AccessGrant).filter_by(
resource_type=resource_type,
resource_id=resource_id,
).delete()
await db.execute(
delete(AccessGrant).filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
)
# Convert JSON to grant dicts
grant_dicts = access_control_to_grants(resource_type, resource_id, access_control)
@@ -397,25 +398,27 @@ class AccessGrantsTable:
db.add(grant)
results.append(grant)
db.commit()
await db.commit()
return [AccessGrantModel.model_validate(g) for g in results]
def set_access_grants(
async def set_access_grants(
self,
resource_type: str,
resource_id: str,
access_grants: Optional[list],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[AccessGrantModel]:
"""
Replace all grants for a resource from a direct access_grants list.
"""
with get_db_context(db) as db:
db.query(AccessGrant).filter_by(
resource_type=resource_type,
resource_id=resource_id,
).delete()
async with get_async_db_context(db) as db:
await db.execute(
delete(AccessGrant).filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
)
normalized_grants = normalize_access_grants(access_grants)
@@ -433,80 +436,80 @@ class AccessGrantsTable:
db.add(grant)
results.append(grant)
db.commit()
await db.commit()
return [AccessGrantModel.model_validate(g) for g in results]
def get_access_control(
async def get_access_control(
self,
resource_type: str,
resource_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[dict]:
"""
Reconstruct the old-style access_control JSON dict from grants.
For backward compat with the frontend.
"""
with get_db_context(db) as db:
grants = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
.all()
)
grants = result.scalars().all()
grant_models = [AccessGrantModel.model_validate(g) for g in grants]
return grants_to_access_control(grant_models)
def get_grants_by_resource(
async def get_grants_by_resource(
self,
resource_type: str,
resource_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[AccessGrantModel]:
"""Get all grants for a specific resource."""
with get_db_context(db) as db:
grants = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
.all()
)
grants = result.scalars().all()
return [AccessGrantModel.model_validate(g) for g in grants]
def get_grants_by_resources(
async def get_grants_by_resources(
self,
resource_type: str,
resource_ids: list[str],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, list[AccessGrantModel]]:
"""Batch-fetch grants for multiple resources. Returns {resource_id: [grants]}."""
if not resource_ids:
return {}
with get_db_context(db) as db:
grants = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AccessGrant)
.filter(
AccessGrant.resource_type == resource_type,
AccessGrant.resource_id.in_(resource_ids),
)
.all()
)
result: dict[str, list[AccessGrantModel]] = {rid: [] for rid in resource_ids}
grants = result.scalars().all()
result_dict: dict[str, list[AccessGrantModel]] = {rid: [] for rid in resource_ids}
for g in grants:
result[g.resource_id].append(AccessGrantModel.model_validate(g))
return result
result_dict[g.resource_id].append(AccessGrantModel.model_validate(g))
return result_dict
def has_access(
async def has_access(
self,
user_id: str,
resource_type: str,
resource_id: str,
permission: str = 'read',
user_group_ids: Optional[set[str]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
"""
Check if a user has the specified permission on a resource.
@@ -516,7 +519,7 @@ class AccessGrantsTable:
- There's a grant for the specific user with the requested permission
- There's a grant for any of the user's groups with the requested permission
"""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Build conditions for matching grants
conditions = [
# Public access
@@ -535,7 +538,7 @@ class AccessGrantsTable:
if user_group_ids is None:
from open_webui.models.groups import Groups
user_groups = Groups.get_groups_by_member_id(user_id, db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
if user_group_ids:
@@ -546,26 +549,27 @@ class AccessGrantsTable:
)
)
exists = (
db.query(AccessGrant)
result = await db.execute(
select(AccessGrant)
.filter(
AccessGrant.resource_type == resource_type,
AccessGrant.resource_id == resource_id,
AccessGrant.permission == permission,
or_(*conditions),
)
.first()
.limit(1)
)
return exists is not None
grant = result.scalars().first()
return grant is not None
def get_accessible_resource_ids(
async def get_accessible_resource_ids(
self,
user_id: str,
resource_type: str,
resource_ids: list[str],
permission: str = 'read',
user_group_ids: Optional[set[str]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> set[str]:
"""
Batch check: return the subset of resource_ids that the user can access.
@@ -575,7 +579,7 @@ class AccessGrantsTable:
if not resource_ids:
return set()
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
conditions = [
and_(
AccessGrant.principal_type == 'user',
@@ -590,7 +594,7 @@ class AccessGrantsTable:
if user_group_ids is None:
from open_webui.models.groups import Groups
user_groups = Groups.get_groups_by_member_id(user_id, db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
if user_group_ids:
@@ -601,8 +605,8 @@ class AccessGrantsTable:
)
)
rows = (
db.query(AccessGrant.resource_id)
result = await db.execute(
select(AccessGrant.resource_id)
.filter(
AccessGrant.resource_type == resource_type,
AccessGrant.resource_id.in_(resource_ids),
@@ -610,16 +614,16 @@ class AccessGrantsTable:
or_(*conditions),
)
.distinct()
.all()
)
rows = result.all()
return {row[0] for row in rows}
def get_users_with_access(
async def get_users_with_access(
self,
resource_type: str,
resource_id: str,
permission: str = 'read',
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list:
"""
Get all users who have the specified permission on a resource.
@@ -628,21 +632,21 @@ class AccessGrantsTable:
from open_webui.models.users import Users, UserModel
from open_webui.models.groups import Groups
with get_db_context(db) as db:
grants = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
permission=permission,
)
.all()
)
grants = result.scalars().all()
# Check for public access
for grant in grants:
if grant.principal_type == 'user' and grant.principal_id == '*':
result = Users.get_users(filter={'roles': ['!pending']}, db=db)
result = await Users.get_users(filter={'roles': ['!pending']}, db=db)
return result.get('users', [])
user_ids_with_access = set()
@@ -651,14 +655,14 @@ class AccessGrantsTable:
if grant.principal_type == 'user':
user_ids_with_access.add(grant.principal_id)
elif grant.principal_type == 'group':
group_user_ids = Groups.get_group_user_ids_by_id(grant.principal_id, db=db)
group_user_ids = await Groups.get_group_user_ids_by_id(grant.principal_id, db=db)
if group_user_ids:
user_ids_with_access.update(group_user_ids)
if not user_ids_with_access:
return []
return Users.get_users_by_user_ids(list(user_ids_with_access), db=db)
return await Users.get_users_by_user_ids(list(user_ids_with_access), db=db)
def has_permission_filter(
self,
@@ -673,6 +677,10 @@ class AccessGrantsTable:
Apply access control filtering to a SQLAlchemy query by JOINing with access_grant.
This replaces the old JSON-column-based filtering with a proper relational JOIN.
Note: This method builds SQLAlchemy expressions and does NOT perform I/O itself,
so it remains synchronous. The caller is responsible for executing the query
asynchronously with `await db.execute(...)`.
"""
group_ids = filter.get('group_ids', [])
user_id = filter.get('user_id')
@@ -718,7 +726,7 @@ class AccessGrantsTable:
# LEFT JOIN access_grant and filter
# We use a subquery approach to avoid duplicates from multiple matching grants
from sqlalchemy import exists as sa_exists, select
from sqlalchemy import exists as sa_exists
grant_exists = (
select(AccessGrant.id)
@@ -776,11 +784,15 @@ class AccessGrantsTable:
"""
Filter for items where user has read BUT NOT write access.
Public items are NOT considered read_only.
Note: This method builds SQLAlchemy expressions and does NOT perform I/O itself,
so it remains synchronous. The caller is responsible for executing the query
asynchronously with `await db.execute(...)`.
"""
group_ids = filter.get('group_ids', [])
user_id = filter.get('user_id')
from sqlalchemy import exists as sa_exists, select
from sqlalchemy import exists as sa_exists
# Has read grant (not public)
read_grant_exists = (
backend/open_webui/models/auths.py
+40 -38
View File
@@ -2,8 +2,9 @@ import logging
import uuid
from typing import Optional
from sqlalchemy.orm import Session
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 User, UserModel, UserProfileImageResponse, Users
from open_webui.utils.validate import validate_profile_image_url
from pydantic import BaseModel, field_validator
@@ -88,7 +89,7 @@ class AddUserForm(SignupForm):
class AuthsTable:
def insert_new_auth(
async def insert_new_auth(
self,
email: str,
password: str,
@@ -96,9 +97,9 @@ class AuthsTable:
profile_image_url: str = '/user.png',
role: str = 'pending',
oauth: Optional[dict] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[UserModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
log.info('insert_new_auth')
id = str(uuid.uuid4())
@@ -107,28 +108,29 @@ class AuthsTable:
result = Auth(**auth.model_dump())
db.add(result)
user = Users.insert_new_user(id, name, email, profile_image_url, role, oauth=oauth, db=db)
user = await Users.insert_new_user(id, name, email, profile_image_url, role, oauth=oauth, db=db)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result and user:
return user
else:
return None
def authenticate_user(
self, email: str, verify_password: callable, db: Optional[Session] = None
async def authenticate_user(
self, email: str, verify_password: callable, db: Optional[AsyncSession] = None
) -> Optional[UserModel]:
log.info(f'authenticate_user: {email}')
user = Users.get_user_by_email(email, db=db)
user = await Users.get_user_by_email(email, db=db)
if not user:
return None
try:
with get_db_context(db) as db:
auth = db.query(Auth).filter_by(id=user.id, active=True).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Auth).filter_by(id=user.id, active=True))
auth = result.scalars().first()
if auth:
if verify_password(auth.password):
return user
@@ -139,66 +141,66 @@ class AuthsTable:
except Exception:
return None
def authenticate_user_by_api_key(self, api_key: str, db: Optional[Session] = None) -> Optional[UserModel]:
async def authenticate_user_by_api_key(self, api_key: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
log.info(f'authenticate_user_by_api_key')
# if no api_key, return None
if not api_key:
return None
try:
user = Users.get_user_by_api_key(api_key, db=db)
user = await Users.get_user_by_api_key(api_key, db=db)
return user if user else None
except Exception:
return False
def authenticate_user_by_email(self, email: str, db: Optional[Session] = None) -> Optional[UserModel]:
async def authenticate_user_by_email(self, email: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
log.info(f'authenticate_user_by_email: {email}')
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Single JOIN query instead of two separate queries
result = (
db.query(Auth, User)
result = await db.execute(
select(Auth, User)
.join(User, Auth.id == User.id)
.filter(Auth.email == email, Auth.active == True)
.first()
)
if result:
_, user = result
row = result.first()
if row:
_, user = row
return UserModel.model_validate(user)
return None
except Exception:
return None
def update_user_password_by_id(self, id: str, new_password: str, db: Optional[Session] = None) -> bool:
async def update_user_password_by_id(self, id: str, new_password: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
result = db.query(Auth).filter_by(id=id).update({'password': new_password})
db.commit()
return True if result == 1 else False
async with get_async_db_context(db) as db:
result = await db.execute(update(Auth).filter_by(id=id).values(password=new_password))
await db.commit()
return True if result.rowcount == 1 else False
except Exception:
return False
def update_email_by_id(self, id: str, email: str, db: Optional[Session] = None) -> bool:
async def update_email_by_id(self, id: str, email: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
result = db.query(Auth).filter_by(id=id).update({'email': email})
db.commit()
if result == 1:
Users.update_user_by_id(id, {'email': email}, db=db)
async with get_async_db_context(db) as db:
result = await db.execute(update(Auth).filter_by(id=id).values(email=email))
await db.commit()
if result.rowcount == 1:
await Users.update_user_by_id(id, {'email': email}, db=db)
return True
return False
except Exception:
return False
def delete_auth_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_auth_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Delete User
result = Users.delete_user_by_id(id, db=db)
result = await Users.delete_user_by_id(id, db=db)
if result:
db.query(Auth).filter_by(id=id).delete()
db.commit()
await db.execute(delete(Auth).filter_by(id=id))
await db.commit()
return True
else:
backend/open_webui/models/automations.py
+85 -75
View File
@@ -4,10 +4,10 @@ from typing import Optional
from uuid import uuid4
from pydantic import BaseModel, ConfigDict
from sqlalchemy import Column, Text, JSON, Boolean, BigInteger, Index, select, or_, func, cast, String
from sqlalchemy.orm import Session
from sqlalchemy import Column, Text, JSON, Boolean, BigInteger, Index, select, or_, func, cast, String, delete, update
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_db, get_db_context
from open_webui.internal.db import Base, get_async_db_context
log = logging.getLogger(__name__)
@@ -118,14 +118,14 @@ class AutomationListResponse(BaseModel):
class AutomationTable:
def insert(
async def insert(
self,
user_id: str,
form: AutomationForm,
next_run_at: int,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> AutomationModel:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
now = int(time.time_ns())
row = Automation(
id=str(uuid4()),
@@ -139,35 +139,38 @@ class AutomationTable:
updated_at=now,
)
db.add(row)
db.commit()
db.refresh(row)
await db.commit()
await db.refresh(row)
return AutomationModel.model_validate(row)
def count_by_user(self, user_id: str, db: Optional[Session] = None) -> int:
with get_db_context(db) as db:
return db.query(Automation).filter_by(user_id=user_id).count()
async def count_by_user(self, user_id: str, db: Optional[AsyncSession] = None) -> int:
async with get_async_db_context(db) as db:
result = await db.execute(
select(func.count()).select_from(Automation).filter_by(user_id=user_id)
)
return result.scalar()
def get_by_id(self, id: str, db: Optional[Session] = None) -> Optional[AutomationModel]:
with get_db_context(db) as db:
row = db.get(Automation, id)
async def get_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[AutomationModel]:
async with get_async_db_context(db) as db:
row = await db.get(Automation, id)
return AutomationModel.model_validate(row) if row else None
def search_automations(
async def search_automations(
self,
user_id: str,
query: Optional[str] = None,
status: Optional[str] = None,
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> 'AutomationListResponse':
with get_db_context(db) as db:
q = db.query(Automation).filter_by(user_id=user_id)
async with get_async_db_context(db) as db:
stmt = select(Automation).filter_by(user_id=user_id)
if query:
search = f'%{query}%'
# Search in name and prompt inside JSON data
q = q.filter(
stmt = stmt.filter(
or_(
Automation.name.ilike(search),
cast(Automation.data, String).ilike(search),
@@ -175,34 +178,39 @@ class AutomationTable:
)
if status == 'active':
q = q.filter(Automation.is_active == True)
stmt = stmt.filter(Automation.is_active == True)
elif status == 'paused':
q = q.filter(Automation.is_active == False)
stmt = stmt.filter(Automation.is_active == False)
q = q.order_by(Automation.created_at.desc())
stmt = stmt.order_by(Automation.created_at.desc())
total = q.count()
# Get total count
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
q = q.offset(skip)
stmt = stmt.offset(skip)
if limit:
q = q.limit(limit)
stmt = stmt.limit(limit)
rows = q.all()
result = await db.execute(stmt)
rows = result.scalars().all()
return AutomationListResponse(
items=[AutomationModel.model_validate(r) for r in rows],
total=total,
)
def update_by_id(
async def update_by_id(
self,
id: str,
form: AutomationForm,
next_run_at: int,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[AutomationModel]:
with get_db_context(db) as db:
row = db.get(Automation, id)
async with get_async_db_context(db) as db:
row = await db.get(Automation, id)
if not row:
return None
row.name = form.name
@@ -212,37 +220,37 @@ class AutomationTable:
row.is_active = form.is_active
row.next_run_at = next_run_at
row.updated_at = int(time.time_ns())
db.commit()
db.refresh(row)
await db.commit()
await db.refresh(row)
return AutomationModel.model_validate(row)
def toggle(
async def toggle(
self,
id: str,
next_run_at: Optional[int],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[AutomationModel]:
with get_db_context(db) as db:
row = db.get(Automation, id)
async with get_async_db_context(db) as db:
row = await db.get(Automation, id)
if not row:
return None
row.is_active = not row.is_active
row.next_run_at = next_run_at if row.is_active else None
row.updated_at = int(time.time_ns())
db.commit()
db.refresh(row)
await db.commit()
await db.refresh(row)
return AutomationModel.model_validate(row)
def delete(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
row = db.get(Automation, id)
async def delete(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
row = await db.get(Automation, id)
if not row:
return False
db.delete(row)
db.commit()
await db.delete(row)
await db.commit()
return True
def claim_due(self, now_ns: int, limit: int = 10, db: Optional[Session] = None) -> list[AutomationModel]:
async def claim_due(self, now_ns: int, limit: int = 10, db: Optional[AsyncSession] = None) -> list[AutomationModel]:
"""
Atomically claim due automations for execution.
@@ -250,7 +258,7 @@ class AutomationTable:
double-claimed. On PostgreSQL, uses FOR UPDATE SKIP LOCKED
for zero-contention distributed work claiming.
"""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
stmt = (
select(Automation)
.where(
@@ -264,7 +272,8 @@ class AutomationTable:
if db.bind.dialect.name == 'postgresql':
stmt = stmt.with_for_update(skip_locked=True)
rows = db.execute(stmt).scalars().all()
result = await db.execute(stmt)
rows = result.scalars().all()
from open_webui.utils.automations import next_run_ns
@@ -272,7 +281,7 @@ class AutomationTable:
row.last_run_at = now_ns
row.next_run_at = next_run_ns(row.data.get('rrule', ''))
db.commit()
await db.commit()
return [AutomationModel.model_validate(r) for r in rows]
@@ -283,15 +292,15 @@ class AutomationTable:
class AutomationRunTable:
def insert(
async def insert(
self,
automation_id: str,
status: str,
chat_id: Optional[str] = None,
error: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> AutomationRunModel:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
row = AutomationRun(
id=str(uuid4()),
automation_id=automation_id,
@@ -301,30 +310,31 @@ class AutomationRunTable:
created_at=int(time.time_ns()),
)
db.add(row)
db.commit()
db.refresh(row)
await db.commit()
await db.refresh(row)
return AutomationRunModel.model_validate(row)
def get_latest(self, automation_id: str, db: Optional[Session] = None) -> Optional[AutomationRunModel]:
with get_db_context(db) as db:
row = (
db.query(AutomationRun)
async def get_latest(self, automation_id: str, db: Optional[AsyncSession] = None) -> Optional[AutomationRunModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(AutomationRun)
.filter_by(automation_id=automation_id)
.order_by(AutomationRun.created_at.desc())
.first()
.limit(1)
)
row = result.scalars().first()
return AutomationRunModel.model_validate(row) if row else None
def get_latest_batch(
self, automation_ids: list[str], db: Optional[Session] = None
async def get_latest_batch(
self, automation_ids: list[str], db: Optional[AsyncSession] = None
) -> dict[str, AutomationRunModel]:
"""Fetch the latest run for each automation in a single query."""
if not automation_ids:
return {}
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Subquery: max created_at per automation_id
subq = (
db.query(
select(
AutomationRun.automation_id,
func.max(AutomationRun.created_at).label('max_created'),
)
@@ -332,43 +342,43 @@ class AutomationRunTable:
.group_by(AutomationRun.automation_id)
.subquery()
)
rows = (
db.query(AutomationRun)
result = await db.execute(
select(AutomationRun)
.join(
subq,
(AutomationRun.automation_id == subq.c.automation_id)
& (AutomationRun.created_at == subq.c.max_created),
)
.all()
)
rows = result.scalars().all()
return {
row.automation_id: AutomationRunModel.model_validate(row)
for row in rows
}
def get_by_automation(
async def get_by_automation(
self,
automation_id: str,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[AutomationRunModel]:
with get_db_context(db) as db:
rows = (
db.query(AutomationRun)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AutomationRun)
.filter_by(automation_id=automation_id)
.order_by(AutomationRun.created_at.desc())
.offset(skip)
.limit(limit)
.all()
)
rows = result.scalars().all()
return [AutomationRunModel.model_validate(r) for r in rows]
def delete_by_automation(self, automation_id: str, db: Optional[Session] = None) -> int:
with get_db_context(db) as db:
count = db.query(AutomationRun).filter_by(automation_id=automation_id).delete()
db.commit()
return count
async def delete_by_automation(self, automation_id: str, db: Optional[AsyncSession] = None) -> int:
async with get_async_db_context(db) as db:
result = await db.execute(delete(AutomationRun).filter_by(automation_id=automation_id))
await db.commit()
return result.rowcount
Automations = AutomationTable()
backend/open_webui/models/channels.py
+244 -242
View File
@@ -4,8 +4,9 @@ import time
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, update, func, case, or_, and_
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.groups import Groups
from open_webui.models.access_grants import (
AccessGrantModel,
@@ -25,11 +26,7 @@ from sqlalchemy import (
Text,
JSON,
UniqueConstraint,
case,
cast,
)
from sqlalchemy import or_, func, select, and_, text
from sqlalchemy.sql import exists
####################
# Channel DB Schema
@@ -249,22 +246,22 @@ class ChannelWebhookForm(BaseModel):
class ChannelTable:
def _get_access_grants(self, channel_id: str, db: Optional[Session] = None) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource('channel', channel_id, db=db)
async def _get_access_grants(self, channel_id: str, db: Optional[AsyncSession] = None) -> list[AccessGrantModel]:
return await AccessGrants.get_grants_by_resource('channel', channel_id, db=db)
def _to_channel_model(
async def _to_channel_model(
self,
channel: Channel,
access_grants: Optional[list[AccessGrantModel]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> ChannelModel:
channel_data = ChannelModel.model_validate(channel).model_dump(exclude={'access_grants'})
channel_data['access_grants'] = (
access_grants if access_grants is not None else self._get_access_grants(channel_data['id'], db=db)
access_grants if access_grants is not None else await self._get_access_grants(channel_data['id'], db=db)
)
return ChannelModel.model_validate(channel_data)
def _collect_unique_user_ids(
async def _collect_unique_user_ids(
self,
invited_by: str,
user_ids: Optional[list[str]] = None,
@@ -281,7 +278,8 @@ class ChannelTable:
users.add(invited_by)
for group_id in group_ids or []:
users.update(Groups.get_group_user_ids_by_id(group_id))
group_user_ids = await Groups.get_group_user_ids_by_id(group_id)
users.update(group_user_ids)
return users
@@ -321,10 +319,20 @@ class ChannelTable:
return memberships
def insert_new_channel(
self, form_data: CreateChannelForm, user_id: str, db: Optional[Session] = None
def _has_permission(self, db, query, filter: dict, permission: str = 'read'):
return AccessGrants.has_permission_filter(
db=db,
query=query,
DocumentModel=Channel,
filter=filter,
resource_type='channel',
permission=permission,
)
async def insert_new_channel(
self, form_data: CreateChannelForm, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[ChannelModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
channel = ChannelModel(
**{
**form_data.model_dump(exclude={'access_grants'}),
@@ -340,7 +348,7 @@ class ChannelTable:
new_channel = Channel(**channel.model_dump(exclude={'access_grants'}))
if form_data.type in ['group', 'dm']:
users = self._collect_unique_user_ids(
users = await self._collect_unique_user_ids(
invited_by=user_id,
user_ids=form_data.user_ids,
group_ids=form_data.group_ids,
@@ -353,17 +361,18 @@ class ChannelTable:
db.add_all(memberships)
db.add(new_channel)
db.commit()
AccessGrants.set_access_grants('channel', new_channel.id, form_data.access_grants, db=db)
return self._to_channel_model(new_channel, db=db)
await db.commit()
await AccessGrants.set_access_grants('channel', new_channel.id, form_data.access_grants, db=db)
return await self._to_channel_model(new_channel, db=db)
def get_channels(self, db: Optional[Session] = None) -> list[ChannelModel]:
with get_db_context(db) as db:
channels = db.query(Channel).all()
async def get_channels(self, db: Optional[AsyncSession] = None) -> list[ChannelModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Channel))
channels = result.scalars().all()
channel_ids = [channel.id for channel in channels]
grants_map = AccessGrants.get_grants_by_resources('channel', channel_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('channel', channel_ids, db=db)
return [
self._to_channel_model(
await self._to_channel_model(
channel,
access_grants=grants_map.get(channel.id, []),
db=db,
@@ -371,22 +380,12 @@ class ChannelTable:
for channel in channels
]
def _has_permission(self, db, query, filter: dict, permission: str = 'read'):
return AccessGrants.has_permission_filter(
db=db,
query=query,
DocumentModel=Channel,
filter=filter,
resource_type='channel',
permission=permission,
)
async def get_channels_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[ChannelModel]:
async with get_async_db_context(db) as db:
user_group_ids = [group.id for group in await Groups.get_groups_by_member_id(user_id, db=db)]
def get_channels_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[ChannelModel]:
with get_db_context(db) as db:
user_group_ids = [group.id for group in Groups.get_groups_by_member_id(user_id, db=db)]
membership_channels = (
db.query(Channel)
result = await db.execute(
select(Channel)
.join(ChannelMember, Channel.id == ChannelMember.channel_id)
.filter(
Channel.deleted_at.is_(None),
@@ -395,10 +394,10 @@ class ChannelTable:
ChannelMember.user_id == user_id,
ChannelMember.is_active.is_(True),
)
.all()
)
membership_channels = result.scalars().all()
query = db.query(Channel).filter(
stmt = select(Channel).filter(
Channel.deleted_at.is_(None),
Channel.archived_at.is_(None),
or_(
@@ -407,17 +406,18 @@ class ChannelTable:
and_(Channel.type != 'group', Channel.type != 'dm'),
),
)
query = self._has_permission(db, query, {'user_id': user_id, 'group_ids': user_group_ids})
stmt = self._has_permission(db, stmt, {'user_id': user_id, 'group_ids': user_group_ids})
standard_channels = query.all()
result = await db.execute(stmt)
standard_channels = result.scalars().all()
all_channels = membership_channels + standard_channels
all_channels = list(membership_channels) + list(standard_channels)
channel_ids = [c.id for c in all_channels]
grants_map = AccessGrants.get_grants_by_resources('channel', channel_ids, db=db)
return [self._to_channel_model(c, access_grants=grants_map.get(c.id, []), db=db) for c in all_channels]
grants_map = await AccessGrants.get_grants_by_resources('channel', channel_ids, db=db)
return [await self._to_channel_model(c, access_grants=grants_map.get(c.id, []), db=db) for c in all_channels]
def get_dm_channel_by_user_ids(self, user_ids: list[str], db: Optional[Session] = None) -> Optional[ChannelModel]:
with get_db_context(db) as db:
async def get_dm_channel_by_user_ids(self, user_ids: list[str], db: Optional[AsyncSession] = None) -> Optional[ChannelModel]:
async with get_async_db_context(db) as db:
# Ensure uniqueness in case a list with duplicates is passed
unique_user_ids = list(set(user_ids))
@@ -429,7 +429,7 @@ class ChannelTable:
)
subquery = (
db.query(ChannelMember.channel_id)
select(ChannelMember.channel_id)
.group_by(ChannelMember.channel_id)
# 1. Channel must have exactly len(user_ids) members
.having(func.count(ChannelMember.user_id) == len(unique_user_ids))
@@ -438,33 +438,34 @@ class ChannelTable:
.subquery()
)
channel = (
db.query(Channel)
result = await db.execute(
select(Channel)
.filter(
Channel.id.in_(subquery),
Channel.id.in_(select(subquery.c.channel_id)),
Channel.type == 'dm',
)
.first()
.limit(1)
)
channel = result.scalars().first()
return self._to_channel_model(channel, db=db) if channel else None
return await self._to_channel_model(channel, db=db) if channel else None
def add_members_to_channel(
async def add_members_to_channel(
self,
channel_id: str,
invited_by: str,
user_ids: Optional[list[str]] = None,
group_ids: Optional[list[str]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[ChannelMemberModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# 1. Collect all user_ids including groups + inviter
requested_users = self._collect_unique_user_ids(invited_by, user_ids, group_ids)
requested_users = await self._collect_unique_user_ids(invited_by, user_ids, group_ids)
existing_users = {
row.user_id
for row in db.query(ChannelMember.user_id).filter(ChannelMember.channel_id == channel_id).all()
}
result = await db.execute(
select(ChannelMember.user_id).filter(ChannelMember.channel_id == channel_id)
)
existing_users = {row[0] for row in result.all()}
new_user_ids = requested_users - existing_users
if not new_user_ids:
@@ -473,59 +474,54 @@ class ChannelTable:
new_memberships = self._create_membership_models(channel_id, invited_by, new_user_ids)
db.add_all(new_memberships)
db.commit()
await db.commit()
return [ChannelMemberModel.model_validate(membership) for membership in new_memberships]
def remove_members_from_channel(
async def remove_members_from_channel(
self,
channel_id: str,
user_ids: list[str],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> int:
with get_db_context(db) as db:
result = (
db.query(ChannelMember)
.filter(
async with get_async_db_context(db) as db:
result = await db.execute(
delete(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id.in_(user_ids),
)
.delete(synchronize_session=False)
)
db.commit()
return result # number of rows deleted
await db.commit()
return result.rowcount # number of rows deleted
def is_user_channel_manager(self, channel_id: str, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
# Check if the user is the creator of the channel
# or has a 'manager' role in ChannelMember
channel = db.query(Channel).filter(Channel.id == channel_id).first()
async def is_user_channel_manager(self, channel_id: str, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(select(Channel).filter(Channel.id == channel_id))
channel = result.scalars().first()
if channel and channel.user_id == user_id:
return True
membership = (
db.query(ChannelMember)
.filter(
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id == user_id,
ChannelMember.is_active.is_(True),
ChannelMember.role == 'manager',
)
.first()
)
membership = result.scalars().first()
return membership is not None
def join_channel(self, channel_id: str, user_id: str, db: Optional[Session] = None) -> Optional[ChannelMemberModel]:
with get_db_context(db) as db:
async def join_channel(self, channel_id: str, user_id: str, db: Optional[AsyncSession] = None) -> Optional[ChannelMemberModel]:
async with get_async_db_context(db) as db:
# Check if the membership already exists
existing_membership = (
db.query(ChannelMember)
.filter(
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id == user_id,
)
.first()
)
existing_membership = result.scalars().first()
if existing_membership:
return ChannelMemberModel.model_validate(existing_membership)
@@ -549,19 +545,18 @@ class ChannelTable:
new_membership = ChannelMember(**channel_member.model_dump())
db.add(new_membership)
db.commit()
await db.commit()
return channel_member
def leave_channel(self, channel_id: str, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
membership = (
db.query(ChannelMember)
.filter(
async def leave_channel(self, channel_id: str, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id == user_id,
)
.first()
)
membership = result.scalars().first()
if not membership:
return False
@@ -570,126 +565,127 @@ class ChannelTable:
membership.left_at = int(time.time_ns())
membership.updated_at = int(time.time_ns())
db.commit()
await db.commit()
return True
def get_member_by_channel_and_user_id(
self, channel_id: str, user_id: str, db: Optional[Session] = None
async def get_member_by_channel_and_user_id(
self, channel_id: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[ChannelMemberModel]:
with get_db_context(db) as db:
membership = (
db.query(ChannelMember)
.filter(
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id == user_id,
)
.first()
)
membership = result.scalars().first()
return ChannelMemberModel.model_validate(membership) if membership else None
def get_members_by_channel_id(self, channel_id: str, db: Optional[Session] = None) -> list[ChannelMemberModel]:
with get_db_context(db) as db:
memberships = db.query(ChannelMember).filter(ChannelMember.channel_id == channel_id).all()
async def get_members_by_channel_id(self, channel_id: str, db: Optional[AsyncSession] = None) -> list[ChannelMemberModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelMember).filter(ChannelMember.channel_id == channel_id)
)
memberships = result.scalars().all()
return [ChannelMemberModel.model_validate(membership) for membership in memberships]
def pin_channel(
async def pin_channel(
self,
channel_id: str,
user_id: str,
is_pinned: bool,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
with get_db_context(db) as db:
membership = (
db.query(ChannelMember)
.filter(
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id == user_id,
)
.first()
)
membership = result.scalars().first()
if not membership:
return False
membership.is_channel_pinned = is_pinned
membership.updated_at = int(time.time_ns())
db.commit()
await db.commit()
return True
def update_member_last_read_at(self, channel_id: str, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
membership = (
db.query(ChannelMember)
.filter(
async def update_member_last_read_at(self, channel_id: str, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id == user_id,
)
.first()
)
membership = result.scalars().first()
if not membership:
return False
membership.last_read_at = int(time.time_ns())
membership.updated_at = int(time.time_ns())
db.commit()
await db.commit()
return True
def update_member_active_status(
async def update_member_active_status(
self,
channel_id: str,
user_id: str,
is_active: bool,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
with get_db_context(db) as db:
membership = (
db.query(ChannelMember)
.filter(
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id == user_id,
)
.first()
)
membership = result.scalars().first()
if not membership:
return False
membership.is_active = is_active
membership.updated_at = int(time.time_ns())
db.commit()
await db.commit()
return True
def is_user_channel_member(self, channel_id: str, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
membership = (
db.query(ChannelMember)
.filter(
async def is_user_channel_member(self, channel_id: str, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel_id,
ChannelMember.user_id == user_id,
ChannelMember.is_active.is_(True),
)
.first()
).limit(1)
)
membership = result.scalars().first()
return membership is not None
def get_channel_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ChannelModel]:
async def get_channel_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ChannelModel]:
try:
with get_db_context(db) as db:
channel = db.query(Channel).filter(Channel.id == id).first()
return self._to_channel_model(channel, db=db) if channel else None
async with get_async_db_context(db) as db:
result = await db.execute(select(Channel).filter(Channel.id == id))
channel = result.scalars().first()
return await self._to_channel_model(channel, db=db) if channel else None
except Exception:
return None
def get_channels_by_file_id(self, file_id: str, db: Optional[Session] = None) -> list[ChannelModel]:
with get_db_context(db) as db:
channel_files = db.query(ChannelFile).filter(ChannelFile.file_id == file_id).all()
async def get_channels_by_file_id(self, file_id: str, db: Optional[AsyncSession] = None) -> list[ChannelModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(ChannelFile).filter(ChannelFile.file_id == file_id))
channel_files = result.scalars().all()
channel_ids = [cf.channel_id for cf in channel_files]
channels = db.query(Channel).filter(Channel.id.in_(channel_ids)).all()
grants_map = AccessGrants.get_grants_by_resources('channel', channel_ids, db=db)
result = await db.execute(select(Channel).filter(Channel.id.in_(channel_ids)))
channels = result.scalars().all()
grants_map = await AccessGrants.get_grants_by_resources('channel', channel_ids, db=db)
return [
self._to_channel_model(
await self._to_channel_model(
channel,
access_grants=grants_map.get(channel.id, []),
db=db,
@@ -697,123 +693,123 @@ class ChannelTable:
for channel in channels
]
def get_channels_by_file_id_and_user_id(
self, file_id: str, user_id: str, db: Optional[Session] = None
async def get_channels_by_file_id_and_user_id(
self, file_id: str, user_id: str, db: Optional[AsyncSession] = None
) -> list[ChannelModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# 1. Determine which channels have this file
channel_file_rows = db.query(ChannelFile).filter(ChannelFile.file_id == file_id).all()
result = await db.execute(select(ChannelFile).filter(ChannelFile.file_id == file_id))
channel_file_rows = result.scalars().all()
channel_ids = [row.channel_id for row in channel_file_rows]
if not channel_ids:
return []
# 2. Load all channel rows that still exist
channels = (
db.query(Channel)
.filter(
result = await db.execute(
select(Channel).filter(
Channel.id.in_(channel_ids),
Channel.deleted_at.is_(None),
Channel.archived_at.is_(None),
)
.all()
)
channels = result.scalars().all()
if not channels:
return []
# Preload user's group membership
user_group_ids = [g.id for g in Groups.get_groups_by_member_id(user_id, db=db)]
user_group_ids = [g.id for g in await Groups.get_groups_by_member_id(user_id, db=db)]
allowed_channels = []
for channel in channels:
# --- Case A: group or dm => user must be an active member ---
if channel.type in ['group', 'dm']:
membership = (
db.query(ChannelMember)
.filter(
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == channel.id,
ChannelMember.user_id == user_id,
ChannelMember.is_active.is_(True),
)
.first()
).limit(1)
)
membership = result.scalars().first()
if membership:
allowed_channels.append(self._to_channel_model(channel, db=db))
allowed_channels.append(await self._to_channel_model(channel, db=db))
continue
# --- Case B: standard channel => rely on ACL permissions ---
query = db.query(Channel).filter(Channel.id == channel.id)
stmt = select(Channel).filter(Channel.id == channel.id)
query = self._has_permission(
stmt = self._has_permission(
db,
query,
stmt,
{'user_id': user_id, 'group_ids': user_group_ids},
permission='read',
)
allowed = query.first()
result = await db.execute(stmt)
allowed = result.scalars().first()
if allowed:
allowed_channels.append(self._to_channel_model(allowed, db=db))
allowed_channels.append(await self._to_channel_model(allowed, db=db))
return allowed_channels
def get_channel_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[Session] = None
async def get_channel_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[ChannelModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Fetch the channel
channel: Channel = (
db.query(Channel)
.filter(
result = await db.execute(
select(Channel).filter(
Channel.id == id,
Channel.deleted_at.is_(None),
Channel.archived_at.is_(None),
)
.first()
)
channel = result.scalars().first()
if not channel:
return None
# If the channel is a group or dm, read access requires membership (active)
if channel.type in ['group', 'dm']:
membership = (
db.query(ChannelMember)
.filter(
result = await db.execute(
select(ChannelMember).filter(
ChannelMember.channel_id == id,
ChannelMember.user_id == user_id,
ChannelMember.is_active.is_(True),
)
.first()
).limit(1)
)
membership = result.scalars().first()
if membership:
return self._to_channel_model(channel, db=db)
return await self._to_channel_model(channel, db=db)
else:
return None
# For channels that are NOT group/dm, fall back to ACL-based read access
query = db.query(Channel).filter(Channel.id == id)
stmt = select(Channel).filter(Channel.id == id)
# Determine user groups
user_group_ids = [group.id for group in Groups.get_groups_by_member_id(user_id, db=db)]
user_group_ids = [group.id for group in await Groups.get_groups_by_member_id(user_id, db=db)]
# Apply ACL rules
query = self._has_permission(
stmt = self._has_permission(
db,
query,
stmt,
{'user_id': user_id, 'group_ids': user_group_ids},
permission='read',
)
channel_allowed = query.first()
return self._to_channel_model(channel_allowed, db=db) if channel_allowed else None
result = await db.execute(stmt)
channel_allowed = result.scalars().first()
return await self._to_channel_model(channel_allowed, db=db) if channel_allowed else None
def update_channel_by_id(
self, id: str, form_data: ChannelForm, db: Optional[Session] = None
async def update_channel_by_id(
self, id: str, form_data: ChannelForm, db: Optional[AsyncSession] = None
) -> Optional[ChannelModel]:
with get_db_context(db) as db:
channel = db.query(Channel).filter(Channel.id == id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Channel).filter(Channel.id == id))
channel = result.scalars().first()
if not channel:
return None
@@ -825,16 +821,16 @@ class ChannelTable:
channel.meta = form_data.meta
if form_data.access_grants is not None:
AccessGrants.set_access_grants('channel', id, form_data.access_grants, db=db)
await AccessGrants.set_access_grants('channel', id, form_data.access_grants, db=db)
channel.updated_at = int(time.time_ns())
db.commit()
return self._to_channel_model(channel, db=db) if channel else None
await db.commit()
return await self._to_channel_model(channel, db=db) if channel else None
def add_file_to_channel_by_id(
self, channel_id: str, file_id: str, user_id: str, db: Optional[Session] = None
async def add_file_to_channel_by_id(
self, channel_id: str, file_id: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[ChannelFileModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
channel_file = ChannelFileModel(
**{
'id': str(uuid.uuid4()),
@@ -849,8 +845,8 @@ class ChannelTable:
try:
result = ChannelFile(**channel_file.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return ChannelFileModel.model_validate(result)
else:
@@ -858,55 +854,58 @@ class ChannelTable:
except Exception:
return None
def set_file_message_id_in_channel_by_id(
async def set_file_message_id_in_channel_by_id(
self,
channel_id: str,
file_id: str,
message_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
try:
with get_db_context(db) as db:
channel_file = db.query(ChannelFile).filter_by(channel_id=channel_id, file_id=file_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelFile).filter_by(channel_id=channel_id, file_id=file_id)
)
channel_file = result.scalars().first()
if not channel_file:
return False
channel_file.message_id = message_id
channel_file.updated_at = int(time.time())
db.commit()
await db.commit()
return True
except Exception:
return False
def remove_file_from_channel_by_id(self, channel_id: str, file_id: str, db: Optional[Session] = None) -> bool:
async def remove_file_from_channel_by_id(self, channel_id: str, file_id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
db.query(ChannelFile).filter_by(channel_id=channel_id, file_id=file_id).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(ChannelFile).filter_by(channel_id=channel_id, file_id=file_id))
await db.commit()
return True
except Exception:
return False
def delete_channel_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
AccessGrants.revoke_all_access('channel', id, db=db)
db.query(Channel).filter(Channel.id == id).delete()
db.commit()
async def delete_channel_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
await AccessGrants.revoke_all_access('channel', id, db=db)
await db.execute(delete(Channel).filter(Channel.id == id))
await db.commit()
return True
####################
# Webhook Methods
####################
def insert_webhook(
async def insert_webhook(
self,
channel_id: str,
user_id: str,
form_data: ChannelWebhookForm,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[ChannelWebhookModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
webhook = ChannelWebhookModel(
id=str(uuid.uuid4()),
channel_id=channel_id,
@@ -919,63 +918,66 @@ class ChannelTable:
updated_at=int(time.time_ns()),
)
db.add(ChannelWebhook(**webhook.model_dump()))
db.commit()
await db.commit()
return webhook
def get_webhooks_by_channel_id(self, channel_id: str, db: Optional[Session] = None) -> list[ChannelWebhookModel]:
with get_db_context(db) as db:
webhooks = db.query(ChannelWebhook).filter(ChannelWebhook.channel_id == channel_id).all()
async def get_webhooks_by_channel_id(self, channel_id: str, db: Optional[AsyncSession] = None) -> list[ChannelWebhookModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(ChannelWebhook).filter(ChannelWebhook.channel_id == channel_id))
webhooks = result.scalars().all()
return [ChannelWebhookModel.model_validate(w) for w in webhooks]
def get_webhook_by_id(self, webhook_id: str, db: Optional[Session] = None) -> Optional[ChannelWebhookModel]:
with get_db_context(db) as db:
webhook = db.query(ChannelWebhook).filter(ChannelWebhook.id == webhook_id).first()
async def get_webhook_by_id(self, webhook_id: str, db: Optional[AsyncSession] = None) -> Optional[ChannelWebhookModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(ChannelWebhook).filter(ChannelWebhook.id == webhook_id))
webhook = result.scalars().first()
return ChannelWebhookModel.model_validate(webhook) if webhook else None
def get_webhook_by_id_and_token(
self, webhook_id: str, token: str, db: Optional[Session] = None
async def get_webhook_by_id_and_token(
self, webhook_id: str, token: str, db: Optional[AsyncSession] = None
) -> Optional[ChannelWebhookModel]:
with get_db_context(db) as db:
webhook = (
db.query(ChannelWebhook)
.filter(
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChannelWebhook).filter(
ChannelWebhook.id == webhook_id,
ChannelWebhook.token == token,
)
.first()
)
webhook = result.scalars().first()
return ChannelWebhookModel.model_validate(webhook) if webhook else None
def update_webhook_by_id(
async def update_webhook_by_id(
self,
webhook_id: str,
form_data: ChannelWebhookForm,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[ChannelWebhookModel]:
with get_db_context(db) as db:
webhook = db.query(ChannelWebhook).filter(ChannelWebhook.id == webhook_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(ChannelWebhook).filter(ChannelWebhook.id == webhook_id))
webhook = result.scalars().first()
if not webhook:
return None
webhook.name = form_data.name
webhook.profile_image_url = form_data.profile_image_url
webhook.updated_at = int(time.time_ns())
db.commit()
await db.commit()
return ChannelWebhookModel.model_validate(webhook)
def update_webhook_last_used_at(self, webhook_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
webhook = db.query(ChannelWebhook).filter(ChannelWebhook.id == webhook_id).first()
async def update_webhook_last_used_at(self, webhook_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(select(ChannelWebhook).filter(ChannelWebhook.id == webhook_id))
webhook = result.scalars().first()
if not webhook:
return False
webhook.last_used_at = int(time.time_ns())
db.commit()
await db.commit()
return True
def delete_webhook_by_id(self, webhook_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
result = db.query(ChannelWebhook).filter(ChannelWebhook.id == webhook_id).delete()
db.commit()
return result > 0
async def delete_webhook_by_id(self, webhook_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(delete(ChannelWebhook).filter(ChannelWebhook.id == webhook_id))
await db.commit()
return result.rowcount > 0
Channels = ChannelTable()
backend/open_webui/models/chat_messages.py
+180 -151
View File
@@ -3,8 +3,9 @@ import time
import uuid
from typing import Any, Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db_context
from sqlalchemy import select, delete, func, cast, Integer
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from open_webui.utils.response import normalize_usage
from pydantic import BaseModel, ConfigDict
@@ -16,7 +17,6 @@ from sqlalchemy import (
Text,
JSON,
Index,
func,
)
####################
@@ -129,23 +129,23 @@ class ChatMessageModel(BaseModel):
class ChatMessageTable:
def upsert_message(
async def upsert_message(
self,
message_id: str,
chat_id: str,
user_id: str,
data: dict,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[ChatMessageModel]:
"""Insert or update a chat message."""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
now = int(time.time())
timestamp = data.get('timestamp', now)
# Use composite ID: {chat_id}-{message_id}
composite_id = f'{chat_id}-{message_id}'
existing = db.get(ChatMessage, composite_id)
existing = await db.get(ChatMessage, composite_id)
if existing:
# Update existing
if 'role' in data:
@@ -178,8 +178,8 @@ class ChatMessageTable:
# from accidentally clearing the primary response's token counts
existing.usage = {**(existing.usage or {}), **usage}
existing.updated_at = now
db.commit()
db.refresh(existing)
await db.commit()
await db.refresh(existing)
return ChatMessageModel.model_validate(existing)
else:
# Insert new
@@ -205,143 +205,155 @@ class ChatMessageTable:
updated_at=now,
)
db.add(message)
db.commit()
db.refresh(message)
await db.commit()
await db.refresh(message)
return ChatMessageModel.model_validate(message)
def get_message_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ChatMessageModel]:
with get_db_context(db) as db:
message = db.get(ChatMessage, id)
async def get_message_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ChatMessageModel]:
async with get_async_db_context(db) as db:
message = await db.get(ChatMessage, id)
return ChatMessageModel.model_validate(message) if message else None
def get_messages_by_chat_id(self, chat_id: str, db: Optional[Session] = None) -> list[ChatMessageModel]:
with get_db_context(db) as db:
messages = db.query(ChatMessage).filter_by(chat_id=chat_id).order_by(ChatMessage.created_at.asc()).all()
async def get_messages_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> list[ChatMessageModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChatMessage).filter_by(chat_id=chat_id).order_by(ChatMessage.created_at.asc())
)
messages = result.scalars().all()
return [ChatMessageModel.model_validate(message) for message in messages]
def get_messages_by_user_id(
async def get_messages_by_user_id(
self,
user_id: str,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[ChatMessageModel]:
with get_db_context(db) as db:
messages = (
db.query(ChatMessage)
async with get_async_db_context(db) as db:
result = await db.execute(
select(ChatMessage)
.filter_by(user_id=user_id)
.order_by(ChatMessage.created_at.desc())
.offset(skip)
.limit(limit)
.all()
)
messages = result.scalars().all()
return [ChatMessageModel.model_validate(message) for message in messages]
def get_messages_by_model_id(
async def get_messages_by_model_id(
self,
model_id: str,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
skip: int = 0,
limit: int = 100,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[ChatMessageModel]:
with get_db_context(db) as db:
query = db.query(ChatMessage).filter_by(model_id=model_id)
async with get_async_db_context(db) as db:
stmt = select(ChatMessage).filter_by(model_id=model_id)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
messages = query.order_by(ChatMessage.created_at.desc()).offset(skip).limit(limit).all()
stmt = stmt.filter(ChatMessage.created_at <= end_date)
stmt = stmt.order_by(ChatMessage.created_at.desc()).offset(skip).limit(limit)
result = await db.execute(stmt)
messages = result.scalars().all()
return [ChatMessageModel.model_validate(message) for message in messages]
def get_chat_ids_by_model_id(
async def get_chat_ids_by_model_id(
self,
model_id: str,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[str]:
"""Get distinct chat_ids that used a specific model."""
with get_db_context(db) as db:
query = db.query(
ChatMessage.chat_id,
func.max(ChatMessage.created_at).label('last_message_at'),
).filter(ChatMessage.model_id == model_id)
async with get_async_db_context(db) as db:
stmt = (
select(
ChatMessage.chat_id,
func.max(ChatMessage.created_at).label('last_message_at'),
)
.filter(ChatMessage.model_id == model_id)
)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
stmt = stmt.filter(ChatMessage.created_at <= end_date)
# Group by chat_id and order by most recent message in each chat
# Secondary sort on chat_id ensures deterministic pagination
# (prevents duplicates across pages when timestamps tie)
chat_ids = (
query.group_by(ChatMessage.chat_id)
stmt = (
stmt.group_by(ChatMessage.chat_id)
.order_by(func.max(ChatMessage.created_at).desc(), ChatMessage.chat_id)
.offset(skip)
.limit(limit)
.all()
)
result = await db.execute(stmt)
chat_ids = result.all()
return [chat_id for chat_id, _ in chat_ids]
def delete_messages_by_chat_id(self, chat_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
db.query(ChatMessage).filter_by(chat_id=chat_id).delete()
db.commit()
async def delete_messages_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
await db.execute(delete(ChatMessage).filter_by(chat_id=chat_id))
await db.commit()
return True
# Analytics methods
def get_message_count_by_model(
async def get_message_count_by_model(
self,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
group_id: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, int]:
with get_db_context(db) as db:
from sqlalchemy import func
async with get_async_db_context(db) as db:
from open_webui.models.groups import GroupMember
query = db.query(ChatMessage.model_id, func.count(ChatMessage.id).label('count')).filter(
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
~ChatMessage.user_id.like('shared-%'),
stmt = (
select(ChatMessage.model_id, func.count(ChatMessage.id).label('count'))
.filter(
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
~ChatMessage.user_id.like('shared-%'),
)
)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
stmt = stmt.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
stmt = stmt.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.model_id).all()
return {row.model_id: row.count for row in results}
stmt = stmt.group_by(ChatMessage.model_id)
result = await db.execute(stmt)
return {row.model_id: row.count for row in result.all()}
def get_token_usage_by_model(
async def get_token_usage_by_model(
self,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
group_id: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, dict]:
"""Aggregate token usage by model using database-level aggregation."""
with get_db_context(db) as db:
from sqlalchemy import func, cast, Integer
async with get_async_db_context(db) as db:
from open_webui.models.groups import GroupMember
dialect = db.bind.dialect.name
# We need the dialect to determine JSON extraction syntax
# For async sessions, access via get_bind()
bind = await db.connection()
dialect = bind.dialect.name
if dialect == 'sqlite':
input_tokens = cast(func.json_extract(ChatMessage.usage, '$.input_tokens'), Integer)
output_tokens = cast(func.json_extract(ChatMessage.usage, '$.output_tokens'), Integer)
elif dialect == 'postgresql':
# Use json_extract_path_text for PostgreSQL JSON columns
input_tokens = cast(
func.json_extract_path_text(ChatMessage.usage, 'input_tokens'),
Integer,
@@ -353,27 +365,31 @@ class ChatMessageTable:
else:
raise NotImplementedError(f'Unsupported dialect: {dialect}')
query = db.query(
ChatMessage.model_id,
func.coalesce(func.sum(input_tokens), 0).label('input_tokens'),
func.coalesce(func.sum(output_tokens), 0).label('output_tokens'),
func.count(ChatMessage.id).label('message_count'),
).filter(
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
ChatMessage.usage.isnot(None),
~ChatMessage.user_id.like('shared-%'),
stmt = (
select(
ChatMessage.model_id,
func.coalesce(func.sum(input_tokens), 0).label('input_tokens'),
func.coalesce(func.sum(output_tokens), 0).label('output_tokens'),
func.count(ChatMessage.id).label('message_count'),
)
.filter(
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
ChatMessage.usage.isnot(None),
~ChatMessage.user_id.like('shared-%'),
)
)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
stmt = stmt.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
stmt = stmt.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.model_id).all()
stmt = stmt.group_by(ChatMessage.model_id)
result = await db.execute(stmt)
return {
row.model_id: {
@@ -382,28 +398,27 @@ class ChatMessageTable:
'total_tokens': row.input_tokens + row.output_tokens,
'message_count': row.message_count,
}
for row in results
for row in result.all()
}
def get_token_usage_by_user(
async def get_token_usage_by_user(
self,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
group_id: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, dict]:
"""Aggregate token usage by user using database-level aggregation."""
with get_db_context(db) as db:
from sqlalchemy import func, cast, Integer
async with get_async_db_context(db) as db:
from open_webui.models.groups import GroupMember
dialect = db.bind.dialect.name
bind = await db.connection()
dialect = bind.dialect.name
if dialect == 'sqlite':
input_tokens = cast(func.json_extract(ChatMessage.usage, '$.input_tokens'), Integer)
output_tokens = cast(func.json_extract(ChatMessage.usage, '$.output_tokens'), Integer)
elif dialect == 'postgresql':
# Use json_extract_path_text for PostgreSQL JSON columns
input_tokens = cast(
func.json_extract_path_text(ChatMessage.usage, 'input_tokens'),
Integer,
@@ -415,27 +430,31 @@ class ChatMessageTable:
else:
raise NotImplementedError(f'Unsupported dialect: {dialect}')
query = db.query(
ChatMessage.user_id,
func.coalesce(func.sum(input_tokens), 0).label('input_tokens'),
func.coalesce(func.sum(output_tokens), 0).label('output_tokens'),
func.count(ChatMessage.id).label('message_count'),
).filter(
ChatMessage.role == 'assistant',
ChatMessage.user_id.isnot(None),
ChatMessage.usage.isnot(None),
~ChatMessage.user_id.like('shared-%'),
stmt = (
select(
ChatMessage.user_id,
func.coalesce(func.sum(input_tokens), 0).label('input_tokens'),
func.coalesce(func.sum(output_tokens), 0).label('output_tokens'),
func.count(ChatMessage.id).label('message_count'),
)
.filter(
ChatMessage.role == 'assistant',
ChatMessage.user_id.isnot(None),
ChatMessage.usage.isnot(None),
~ChatMessage.user_id.like('shared-%'),
)
)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
stmt = stmt.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
stmt = stmt.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.user_id).all()
stmt = stmt.group_by(ChatMessage.user_id)
result = await db.execute(stmt)
return {
row.user_id: {
@@ -444,88 +463,94 @@ class ChatMessageTable:
'total_tokens': row.input_tokens + row.output_tokens,
'message_count': row.message_count,
}
for row in results
for row in result.all()
}
def get_message_count_by_user(
async def get_message_count_by_user(
self,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
group_id: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, int]:
with get_db_context(db) as db:
from sqlalchemy import func
async with get_async_db_context(db) as db:
from open_webui.models.groups import GroupMember
query = db.query(ChatMessage.user_id, func.count(ChatMessage.id).label('count')).filter(
~ChatMessage.user_id.like('shared-%')
stmt = (
select(ChatMessage.user_id, func.count(ChatMessage.id).label('count'))
.filter(~ChatMessage.user_id.like('shared-%'))
)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
stmt = stmt.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
stmt = stmt.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.user_id).all()
return {row.user_id: row.count for row in results}
stmt = stmt.group_by(ChatMessage.user_id)
result = await db.execute(stmt)
return {row.user_id: row.count for row in result.all()}
def get_message_count_by_chat(
async def get_message_count_by_chat(
self,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
group_id: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, int]:
with get_db_context(db) as db:
from sqlalchemy import func
async with get_async_db_context(db) as db:
from open_webui.models.groups import GroupMember
query = db.query(ChatMessage.chat_id, func.count(ChatMessage.id).label('count')).filter(
~ChatMessage.user_id.like('shared-%')
stmt = (
select(ChatMessage.chat_id, func.count(ChatMessage.id).label('count'))
.filter(~ChatMessage.user_id.like('shared-%'))
)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
stmt = stmt.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
stmt = stmt.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.chat_id).all()
return {row.chat_id: row.count for row in results}
stmt = stmt.group_by(ChatMessage.chat_id)
result = await db.execute(stmt)
return {row.chat_id: row.count for row in result.all()}
def get_daily_message_counts_by_model(
async def get_daily_message_counts_by_model(
self,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
group_id: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, dict[str, int]]:
"""Get message counts grouped by day and model."""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
from datetime import datetime, timedelta
from open_webui.models.groups import GroupMember
query = db.query(ChatMessage.created_at, ChatMessage.model_id).filter(
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
~ChatMessage.user_id.like('shared-%'),
stmt = (
select(ChatMessage.created_at, ChatMessage.model_id)
.filter(
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
~ChatMessage.user_id.like('shared-%'),
)
)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
stmt = stmt.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
stmt = stmt.filter(ChatMessage.user_id.in_(group_users))
results = query.all()
result = await db.execute(stmt)
results = result.all()
# Group by date -> model -> count
daily_counts: dict[str, dict[str, int]] = {}
@@ -547,28 +572,32 @@ class ChatMessageTable:
return daily_counts
def get_hourly_message_counts_by_model(
async def get_hourly_message_counts_by_model(
self,
start_date: Optional[int] = None,
end_date: Optional[int] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, dict[str, int]]:
"""Get message counts grouped by hour and model."""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
from datetime import datetime, timedelta
query = db.query(ChatMessage.created_at, ChatMessage.model_id).filter(
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
~ChatMessage.user_id.like('shared-%'),
stmt = (
select(ChatMessage.created_at, ChatMessage.model_id)
.filter(
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
~ChatMessage.user_id.like('shared-%'),
)
)
if start_date:
query = query.filter(ChatMessage.created_at >= start_date)
stmt = stmt.filter(ChatMessage.created_at >= start_date)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
stmt = stmt.filter(ChatMessage.created_at <= end_date)
results = query.all()
result = await db.execute(stmt)
results = result.all()
# Group by hour -> model -> count
hourly_counts: dict[str, dict[str, int]] = {}
backend/open_webui/models/chats.py
+428 -425
View File
File diff suppressed because it is too large Load Diff
backend/open_webui/models/feedbacks.py
+119 -113
View File
@@ -3,9 +3,10 @@ import time
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from open_webui.models.users import User
from sqlalchemy import select, delete, func
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.users import User, UserModel
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, Text, JSON, Boolean
@@ -139,10 +140,10 @@ class ModelHistoryResponse(BaseModel):
class FeedbackTable:
def insert_new_feedback(
self, user_id: str, form_data: FeedbackForm, db: Optional[Session] = None
async def insert_new_feedback(
self, user_id: str, form_data: FeedbackForm, db: Optional[AsyncSession] = None
) -> Optional[FeedbackModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
id = str(uuid.uuid4())
feedback = FeedbackModel(
**{
@@ -157,8 +158,8 @@ class FeedbackTable:
try:
result = Feedback(**feedback.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return FeedbackModel.model_validate(result)
else:
@@ -167,97 +168,101 @@ class FeedbackTable:
log.exception(f'Error creating a new feedback: {e}')
return None
def get_feedback_by_id(self, id: str, db: Optional[Session] = None) -> Optional[FeedbackModel]:
async def get_feedback_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[FeedbackModel]:
try:
with get_db_context(db) as db:
feedback = db.query(Feedback).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Feedback).filter_by(id=id))
feedback = result.scalars().first()
if not feedback:
return None
return FeedbackModel.model_validate(feedback)
except Exception:
return None
def get_feedback_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[Session] = None
async def get_feedback_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[FeedbackModel]:
try:
with get_db_context(db) as db:
feedback = db.query(Feedback).filter_by(id=id, user_id=user_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Feedback).filter_by(id=id, user_id=user_id))
feedback = result.scalars().first()
if not feedback:
return None
return FeedbackModel.model_validate(feedback)
except Exception:
return None
def get_feedbacks_by_chat_id(self, chat_id: str, db: Optional[Session] = None) -> list[FeedbackModel]:
async def get_feedbacks_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> list[FeedbackModel]:
"""Get all feedbacks for a specific chat."""
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# meta.chat_id stores the chat reference
feedbacks = (
db.query(Feedback)
result = await db.execute(
select(Feedback)
.filter(Feedback.meta['chat_id'].as_string() == chat_id)
.order_by(Feedback.created_at.desc())
.all()
)
feedbacks = result.scalars().all()
return [FeedbackModel.model_validate(fb) for fb in feedbacks]
except Exception:
return []
def get_feedback_items(
async def get_feedback_items(
self,
filter: dict = {},
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> FeedbackListResponse:
with get_db_context(db) as db:
query = db.query(Feedback, User).join(User, Feedback.user_id == User.id)
async with get_async_db_context(db) as db:
stmt = select(Feedback, User).join(User, Feedback.user_id == User.id)
if filter:
# Apply model_id filter (exact match)
model_id = filter.get('model_id')
if model_id:
query = query.filter(Feedback.data['model_id'].as_string() == model_id)
stmt = stmt.filter(Feedback.data['model_id'].as_string() == model_id)
order_by = filter.get('order_by')
direction = filter.get('direction')
if order_by == 'username':
if direction == 'asc':
query = query.order_by(User.name.asc())
stmt = stmt.order_by(User.name.asc())
else:
query = query.order_by(User.name.desc())
stmt = stmt.order_by(User.name.desc())
elif order_by == 'model_id':
# it's stored in feedback.data['model_id']
if direction == 'asc':
query = query.order_by(Feedback.data['model_id'].as_string().asc())
stmt = stmt.order_by(Feedback.data['model_id'].as_string().asc())
else:
query = query.order_by(Feedback.data['model_id'].as_string().desc())
stmt = stmt.order_by(Feedback.data['model_id'].as_string().desc())
elif order_by == 'rating':
# it's stored in feedback.data['rating']
if direction == 'asc':
query = query.order_by(Feedback.data['rating'].as_string().asc())
stmt = stmt.order_by(Feedback.data['rating'].as_string().asc())
else:
query = query.order_by(Feedback.data['rating'].as_string().desc())
stmt = stmt.order_by(Feedback.data['rating'].as_string().desc())
elif order_by == 'updated_at':
if direction == 'asc':
query = query.order_by(Feedback.updated_at.asc())
stmt = stmt.order_by(Feedback.updated_at.asc())
else:
query = query.order_by(Feedback.updated_at.desc())
stmt = stmt.order_by(Feedback.updated_at.desc())
else:
query = query.order_by(Feedback.created_at.desc())
stmt = stmt.order_by(Feedback.created_at.desc())
# Count BEFORE pagination
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit:
query = query.limit(limit)
stmt = stmt.limit(limit)
items = query.all()
result = await db.execute(stmt)
items = result.all()
feedbacks = []
for feedback, user in items:
@@ -267,15 +272,17 @@ class FeedbackTable:
return FeedbackListResponse(items=feedbacks, total=total)
def get_all_feedbacks(self, db: Optional[Session] = None) -> list[FeedbackModel]:
with get_db_context(db) as db:
return [
FeedbackModel.model_validate(feedback)
for feedback in db.query(Feedback).order_by(Feedback.updated_at.desc()).all()
]
async def get_all_feedbacks(self, db: Optional[AsyncSession] = None) -> list[FeedbackModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Feedback).order_by(Feedback.updated_at.desc()))
return [FeedbackModel.model_validate(feedback) for feedback in result.scalars().all()]
def get_all_feedback_ids(self, db: Optional[Session] = None) -> list[FeedbackIdResponse]:
with get_db_context(db) as db:
async def get_all_feedback_ids(self, db: Optional[AsyncSession] = None) -> list[FeedbackIdResponse]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(Feedback.id, Feedback.user_id, Feedback.created_at, Feedback.updated_at)
.order_by(Feedback.updated_at.desc())
)
return [
FeedbackIdResponse(
id=row.id,
@@ -283,36 +290,28 @@ class FeedbackTable:
created_at=row.created_at,
updated_at=row.updated_at,
)
for row in db.query(
Feedback.id,
Feedback.user_id,
Feedback.created_at,
Feedback.updated_at,
)
.order_by(Feedback.updated_at.desc())
.all()
for row in result.all()
]
def get_distinct_model_ids(self, db: Optional[Session] = None) -> list[str]:
async def get_distinct_model_ids(self, db: Optional[AsyncSession] = None) -> list[str]:
"""Get distinct model_ids from feedback data for filter dropdowns."""
with get_db_context(db) as db:
rows = (
db.query(Feedback.data['model_id'].as_string())
async with get_async_db_context(db) as db:
result = await db.execute(
select(Feedback.data['model_id'].as_string())
.filter(Feedback.data['model_id'].as_string().isnot(None))
.distinct()
.all()
)
rows = result.all()
return sorted([row[0] for row in rows if row[0]])
def get_feedbacks_for_leaderboard(self, db: Optional[Session] = None) -> list[LeaderboardFeedbackData]:
async def get_feedbacks_for_leaderboard(self, db: Optional[AsyncSession] = None) -> list[LeaderboardFeedbackData]:
"""Fetch only id and data for leaderboard computation (excludes snapshot/meta)."""
with get_db_context(db) as db:
return [
LeaderboardFeedbackData(id=row.id, data=row.data) for row in db.query(Feedback.id, Feedback.data).all()
]
async with get_async_db_context(db) as db:
result = await db.execute(select(Feedback.id, Feedback.data))
return [LeaderboardFeedbackData(id=row.id, data=row.data) for row in result.all()]
def get_model_evaluation_history(
self, model_id: str, days: int = 30, db: Optional[Session] = None
async def get_model_evaluation_history(
self, model_id: str, days: int = 30, db: Optional[AsyncSession] = None
) -> list[ModelHistoryEntry]:
"""
Get daily wins/losses for a specific model over the past N days.
@@ -322,13 +321,16 @@ class FeedbackTable:
from datetime import datetime, timedelta
from collections import defaultdict
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
if days == 0:
# All time - no cutoff
rows = db.query(Feedback.created_at, Feedback.data).all()
result = await db.execute(select(Feedback.created_at, Feedback.data))
else:
cutoff = int(time.time()) - (days * 86400)
rows = db.query(Feedback.created_at, Feedback.data).filter(Feedback.created_at >= cutoff).all()
result = await db.execute(
select(Feedback.created_at, Feedback.data).filter(Feedback.created_at >= cutoff)
)
rows = result.all()
daily_counts = defaultdict(lambda: {'won': 0, 'lost': 0})
first_date = None
@@ -374,25 +376,26 @@ class FeedbackTable:
return result
def get_feedbacks_by_type(self, type: str, db: Optional[Session] = None) -> list[FeedbackModel]:
with get_db_context(db) as db:
return [
FeedbackModel.model_validate(feedback)
for feedback in db.query(Feedback).filter_by(type=type).order_by(Feedback.updated_at.desc()).all()
]
async def get_feedbacks_by_type(self, type: str, db: Optional[AsyncSession] = None) -> list[FeedbackModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(Feedback).filter_by(type=type).order_by(Feedback.updated_at.desc())
)
return [FeedbackModel.model_validate(feedback) for feedback in result.scalars().all()]
def get_feedbacks_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[FeedbackModel]:
with get_db_context(db) as db:
return [
FeedbackModel.model_validate(feedback)
for feedback in db.query(Feedback).filter_by(user_id=user_id).order_by(Feedback.updated_at.desc()).all()
]
async def get_feedbacks_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[FeedbackModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(Feedback).filter_by(user_id=user_id).order_by(Feedback.updated_at.desc())
)
return [FeedbackModel.model_validate(feedback) for feedback in result.scalars().all()]
def update_feedback_by_id(
self, id: str, form_data: FeedbackForm, db: Optional[Session] = None
async def update_feedback_by_id(
self, id: str, form_data: FeedbackForm, db: Optional[AsyncSession] = None
) -> Optional[FeedbackModel]:
with get_db_context(db) as db:
feedback = db.query(Feedback).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Feedback).filter_by(id=id))
feedback = result.scalars().first()
if not feedback:
return None
@@ -405,18 +408,19 @@ class FeedbackTable:
feedback.updated_at = int(time.time())
db.commit()
await db.commit()
return FeedbackModel.model_validate(feedback)
def update_feedback_by_id_and_user_id(
async def update_feedback_by_id_and_user_id(
self,
id: str,
user_id: str,
form_data: FeedbackForm,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[FeedbackModel]:
with get_db_context(db) as db:
feedback = db.query(Feedback).filter_by(id=id, user_id=user_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Feedback).filter_by(id=id, user_id=user_id))
feedback = result.scalars().first()
if not feedback:
return None
@@ -429,38 +433,40 @@ class FeedbackTable:
feedback.updated_at = int(time.time())
db.commit()
await db.commit()
return FeedbackModel.model_validate(feedback)
def delete_feedback_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
feedback = db.query(Feedback).filter_by(id=id).first()
async def delete_feedback_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(select(Feedback).filter_by(id=id))
feedback = result.scalars().first()
if not feedback:
return False
db.delete(feedback)
db.commit()
await db.delete(feedback)
await db.commit()
return True
def delete_feedback_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
feedback = db.query(Feedback).filter_by(id=id, user_id=user_id).first()
async def delete_feedback_by_id_and_user_id(self, id: str, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(select(Feedback).filter_by(id=id, user_id=user_id))
feedback = result.scalars().first()
if not feedback:
return False
db.delete(feedback)
db.commit()
await db.delete(feedback)
await db.commit()
return True
def delete_feedbacks_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
result = db.query(Feedback).filter_by(user_id=user_id).delete()
db.commit()
return result > 0
async def delete_feedbacks_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(delete(Feedback).filter_by(user_id=user_id))
await db.commit()
return result.rowcount > 0
def delete_all_feedbacks(self, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
result = db.query(Feedback).delete()
db.commit()
return result > 0
async def delete_all_feedbacks(self, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(delete(Feedback))
await db.commit()
return result.rowcount > 0
Feedbacks = FeedbackTable()
backend/open_webui/models/files.py
+102 -86
View File
@@ -2,8 +2,9 @@ import logging
import time
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, func
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.utils.misc import sanitize_metadata
from pydantic import BaseModel, ConfigDict, model_validator
from sqlalchemy import BigInteger, Column, String, Text, JSON
@@ -124,8 +125,8 @@ class FileUpdateForm(BaseModel):
class FilesTable:
def insert_new_file(self, user_id: str, form_data: FileForm, db: Optional[Session] = None) -> Optional[FileModel]:
with get_db_context(db) as db:
async def insert_new_file(self, user_id: str, form_data: FileForm, db: Optional[AsyncSession] = None) -> Optional[FileModel]:
async with get_async_db_context(db) as db:
file_data = form_data.model_dump()
# Sanitize meta to remove non-JSON-serializable objects
@@ -145,8 +146,8 @@ class FilesTable:
try:
result = File(**file.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return FileModel.model_validate(result)
else:
@@ -155,21 +156,22 @@ class FilesTable:
log.exception(f'Error inserting a new file: {e}')
return None
def get_file_by_id(self, id: str, db: Optional[Session] = None) -> Optional[FileModel]:
async def get_file_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[FileModel]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
try:
file = db.get(File, id)
return FileModel.model_validate(file)
file = await db.get(File, id)
return FileModel.model_validate(file) if file else None
except Exception:
return None
except Exception:
return None
def get_file_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> Optional[FileModel]:
with get_db_context(db) as db:
async def get_file_by_id_and_user_id(self, id: str, user_id: str, db: Optional[AsyncSession] = None) -> Optional[FileModel]:
async with get_async_db_context(db) as db:
try:
file = db.query(File).filter_by(id=id, user_id=user_id).first()
result = await db.execute(select(File).filter_by(id=id, user_id=user_id))
file = result.scalars().first()
if file:
return FileModel.model_validate(file)
else:
@@ -177,10 +179,12 @@ class FilesTable:
except Exception:
return None
def get_file_metadata_by_id(self, id: str, db: Optional[Session] = None) -> Optional[FileMetadataResponse]:
with get_db_context(db) as db:
async def get_file_metadata_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[FileMetadataResponse]:
async with get_async_db_context(db) as db:
try:
file = db.get(File, id)
file = await db.get(File, id)
if not file:
return None
return FileMetadataResponse(
id=file.id,
hash=file.hash,
@@ -191,12 +195,13 @@ class FilesTable:
except Exception:
return None
def get_files(self, db: Optional[Session] = None) -> list[FileModel]:
with get_db_context(db) as db:
return [FileModel.model_validate(file) for file in db.query(File).all()]
async def get_files(self, db: Optional[AsyncSession] = None) -> list[FileModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(File))
return [FileModel.model_validate(file) for file in result.scalars().all()]
def check_access_by_user_id(self, id, user_id, permission='write', db: Optional[Session] = None) -> bool:
file = self.get_file_by_id(id, db=db)
async def check_access_by_user_id(self, id, user_id, permission='write', db: Optional[AsyncSession] = None) -> bool:
file = await self.get_file_by_id(id, db=db)
if not file:
return False
if file.user_id == user_id:
@@ -204,50 +209,59 @@ class FilesTable:
# Implement additional access control logic here as needed
return False
def get_files_by_ids(self, ids: list[str], db: Optional[Session] = None) -> list[FileModel]:
with get_db_context(db) as db:
return [
FileModel.model_validate(file)
for file in db.query(File).filter(File.id.in_(ids)).order_by(File.updated_at.desc()).all()
]
async def get_files_by_ids(self, ids: list[str], db: Optional[AsyncSession] = None) -> list[FileModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(File).filter(File.id.in_(ids)).order_by(File.updated_at.desc())
)
return [FileModel.model_validate(file) for file in result.scalars().all()]
def get_file_metadatas_by_ids(self, ids: list[str], db: Optional[Session] = None) -> list[FileMetadataResponse]:
with get_db_context(db) as db:
return [
FileMetadataResponse(
id=file.id,
hash=file.hash,
meta=file.meta,
created_at=file.created_at,
updated_at=file.updated_at,
)
for file in db.query(File.id, File.hash, File.meta, File.created_at, File.updated_at)
async def get_file_metadatas_by_ids(self, ids: list[str], db: Optional[AsyncSession] = None) -> list[FileMetadataResponse]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(File.id, File.hash, File.meta, File.created_at, File.updated_at)
.filter(File.id.in_(ids))
.order_by(File.updated_at.desc())
.all()
)
return [
FileMetadataResponse(
id=row.id,
hash=row.hash,
meta=row.meta,
created_at=row.created_at,
updated_at=row.updated_at,
)
for row in result.all()
]
def get_files_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[FileModel]:
with get_db_context(db) as db:
return [FileModel.model_validate(file) for file in db.query(File).filter_by(user_id=user_id).all()]
async def get_files_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[FileModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(File).filter_by(user_id=user_id))
return [FileModel.model_validate(file) for file in result.scalars().all()]
def get_file_list(
async def get_file_list(
self,
user_id: Optional[str] = None,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> 'FileListResponse':
with get_db_context(db) as db:
query = db.query(File)
async with get_async_db_context(db) as db:
stmt = select(File)
if user_id:
query = query.filter_by(user_id=user_id)
stmt = stmt.filter_by(user_id=user_id)
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
result = await db.execute(
stmt.order_by(File.updated_at.desc(), File.id.desc()).offset(skip).limit(limit)
)
items = [
FileModelResponse.model_validate(file, from_attributes=True)
for file in query.order_by(File.updated_at.desc(), File.id.desc()).offset(skip).limit(limit).all()
for file in result.scalars().all()
]
return FileListResponse(items=items, total=total)
@@ -275,13 +289,13 @@ class FilesTable:
pattern = pattern.replace('?', '_')
return pattern
def search_files(
async def search_files(
self,
user_id: Optional[str] = None,
filename: str = '*',
skip: int = 0,
limit: int = 100,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[FileModel]:
"""
Search files with glob pattern matching, optional user filter, and pagination.
@@ -296,27 +310,28 @@ class FilesTable:
Returns:
List of matching FileModel objects, ordered by created_at descending.
"""
with get_db_context(db) as db:
query = db.query(File)
async with get_async_db_context(db) as db:
stmt = select(File)
if user_id:
query = query.filter_by(user_id=user_id)
stmt = stmt.filter_by(user_id=user_id)
pattern = self._glob_to_like_pattern(filename)
if pattern != '%':
query = query.filter(File.filename.ilike(pattern, escape='\\'))
stmt = stmt.filter(File.filename.ilike(pattern, escape='\\'))
return [
FileModel.model_validate(file)
for file in query.order_by(File.created_at.desc(), File.id.desc()).offset(skip).limit(limit).all()
]
result = await db.execute(
stmt.order_by(File.created_at.desc(), File.id.desc()).offset(skip).limit(limit)
)
return [FileModel.model_validate(file) for file in result.scalars().all()]
def update_file_by_id(
self, id: str, form_data: FileUpdateForm, db: Optional[Session] = None
async def update_file_by_id(
self, id: str, form_data: FileUpdateForm, db: Optional[AsyncSession] = None
) -> Optional[FileModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
try:
file = db.query(File).filter_by(id=id).first()
result = await db.execute(select(File).filter_by(id=id))
file = result.scalars().first()
if form_data.hash is not None:
file.hash = form_data.hash
@@ -328,63 +343,64 @@ class FilesTable:
file.meta = {**(file.meta if file.meta else {}), **form_data.meta}
file.updated_at = int(time.time())
db.commit()
await db.commit()
return FileModel.model_validate(file)
except Exception as e:
log.exception(f'Error updating file completely by id: {e}')
return None
def update_file_hash_by_id(self, id: str, hash: Optional[str], db: Optional[Session] = None) -> Optional[FileModel]:
with get_db_context(db) as db:
async def update_file_hash_by_id(self, id: str, hash: Optional[str], db: Optional[AsyncSession] = None) -> Optional[FileModel]:
async with get_async_db_context(db) as db:
try:
file = db.query(File).filter_by(id=id).first()
result = await db.execute(select(File).filter_by(id=id))
file = result.scalars().first()
file.hash = hash
file.updated_at = int(time.time())
db.commit()
await db.commit()
return FileModel.model_validate(file)
except Exception:
return None
def update_file_data_by_id(self, id: str, data: dict, db: Optional[Session] = None) -> Optional[FileModel]:
with get_db_context(db) as db:
async def update_file_data_by_id(self, id: str, data: dict, db: Optional[AsyncSession] = None) -> Optional[FileModel]:
async with get_async_db_context(db) as db:
try:
file = db.query(File).filter_by(id=id).first()
result = await db.execute(select(File).filter_by(id=id))
file = result.scalars().first()
file.data = {**(file.data if file.data else {}), **data}
file.updated_at = int(time.time())
db.commit()
await db.commit()
return FileModel.model_validate(file)
except Exception as e:
return None
def update_file_metadata_by_id(self, id: str, meta: dict, db: Optional[Session] = None) -> Optional[FileModel]:
with get_db_context(db) as db:
async def update_file_metadata_by_id(self, id: str, meta: dict, db: Optional[AsyncSession] = None) -> Optional[FileModel]:
async with get_async_db_context(db) as db:
try:
file = db.query(File).filter_by(id=id).first()
result = await db.execute(select(File).filter_by(id=id))
file = result.scalars().first()
file.meta = {**(file.meta if file.meta else {}), **meta}
file.updated_at = int(time.time())
db.commit()
await db.commit()
return FileModel.model_validate(file)
except Exception:
return None
return False
def delete_file_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def delete_file_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
db.query(File).filter_by(id=id).delete()
db.commit()
await db.execute(delete(File).filter_by(id=id))
await db.commit()
return True
except Exception:
return False
def delete_all_files(self, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def delete_all_files(self, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
db.query(File).delete()
db.commit()
await db.execute(delete(File))
await db.commit()
return True
except Exception:
backend/open_webui/models/folders.py
+86 -77
View File
@@ -6,10 +6,10 @@ import re
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, Text, JSON, Boolean, func
from sqlalchemy.orm import Session
from sqlalchemy import BigInteger, Column, Text, JSON, Boolean, func, select, delete
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from open_webui.internal.db import Base, JSONField, get_async_db_context
log = logging.getLogger(__name__)
@@ -85,14 +85,14 @@ class FolderUpdateForm(BaseModel):
class FolderTable:
def insert_new_folder(
async def insert_new_folder(
self,
user_id: str,
form_data: FolderForm,
parent_id: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[FolderModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
id = str(uuid.uuid4())
folder = FolderModel(
**{
@@ -107,8 +107,8 @@ class FolderTable:
try:
result = Folder(**folder.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return FolderModel.model_validate(result)
else:
@@ -117,12 +117,13 @@ class FolderTable:
log.exception(f'Error inserting a new folder: {e}')
return None
def get_folder_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[Session] = None
async def get_folder_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[FolderModel]:
try:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
folder = result.scalars().first()
if not folder:
return None
@@ -131,48 +132,50 @@ class FolderTable:
except Exception:
return None
def get_children_folders_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[Session] = None
async def get_children_folders_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[list[FolderModel]]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
folders = []
def get_children(folder):
children = self.get_folders_by_parent_id_and_user_id(folder.id, user_id, db=db)
async def get_children(folder):
children = await self.get_folders_by_parent_id_and_user_id(folder.id, user_id, db=db)
for child in children:
get_children(child)
await get_children(child)
folders.append(child)
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
folder = result.scalars().first()
if not folder:
return None
get_children(folder)
await get_children(folder)
return folders
except Exception:
return None
def get_folders_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[FolderModel]:
with get_db_context(db) as db:
return [FolderModel.model_validate(folder) for folder in db.query(Folder).filter_by(user_id=user_id).all()]
async def get_folders_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[FolderModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(user_id=user_id))
return [FolderModel.model_validate(folder) for folder in result.scalars().all()]
def get_folder_by_parent_id_and_user_id_and_name(
async def get_folder_by_parent_id_and_user_id_and_name(
self,
parent_id: Optional[str],
user_id: str,
name: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[FolderModel]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Check if folder exists
folder = (
db.query(Folder)
result = await db.execute(
select(Folder)
.filter_by(parent_id=parent_id, user_id=user_id)
.filter(Folder.name.ilike(name))
.first()
)
folder = result.scalars().first()
if not folder:
return None
@@ -182,25 +185,24 @@ class FolderTable:
log.error(f'get_folder_by_parent_id_and_user_id_and_name: {e}')
return None
def get_folders_by_parent_id_and_user_id(
self, parent_id: Optional[str], user_id: str, db: Optional[Session] = None
async def get_folders_by_parent_id_and_user_id(
self, parent_id: Optional[str], user_id: str, db: Optional[AsyncSession] = None
) -> list[FolderModel]:
with get_db_context(db) as db:
return [
FolderModel.model_validate(folder)
for folder in db.query(Folder).filter_by(parent_id=parent_id, user_id=user_id).all()
]
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(parent_id=parent_id, user_id=user_id))
return [FolderModel.model_validate(folder) for folder in result.scalars().all()]
def update_folder_parent_id_by_id_and_user_id(
async def update_folder_parent_id_by_id_and_user_id(
self,
id: str,
user_id: str,
parent_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[FolderModel]:
try:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
folder = result.scalars().first()
if not folder:
return None
@@ -208,38 +210,39 @@ class FolderTable:
folder.parent_id = parent_id
folder.updated_at = int(time.time())
db.commit()
await db.commit()
return FolderModel.model_validate(folder)
except Exception as e:
log.error(f'update_folder: {e}')
return
def update_folder_by_id_and_user_id(
async def update_folder_by_id_and_user_id(
self,
id: str,
user_id: str,
form_data: FolderUpdateForm,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[FolderModel]:
try:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
folder = result.scalars().first()
if not folder:
return None
form_data = form_data.model_dump(exclude_unset=True)
existing_folder = (
db.query(Folder)
existing_result = await db.execute(
select(Folder)
.filter_by(
name=form_data.get('name'),
parent_id=folder.parent_id,
user_id=user_id,
)
.first()
)
existing_folder = existing_result.scalars().first()
if existing_folder and existing_folder.id != id:
return None
@@ -258,19 +261,20 @@ class FolderTable:
}
folder.updated_at = int(time.time())
db.commit()
await db.commit()
return FolderModel.model_validate(folder)
except Exception as e:
log.error(f'update_folder: {e}')
return
def update_folder_is_expanded_by_id_and_user_id(
self, id: str, user_id: str, is_expanded: bool, db: Optional[Session] = None
async def update_folder_is_expanded_by_id_and_user_id(
self, id: str, user_id: str, is_expanded: bool, db: Optional[AsyncSession] = None
) -> Optional[FolderModel]:
try:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
folder = result.scalars().first()
if not folder:
return None
@@ -278,37 +282,39 @@ class FolderTable:
folder.is_expanded = is_expanded
folder.updated_at = int(time.time())
db.commit()
await db.commit()
return FolderModel.model_validate(folder)
except Exception as e:
log.error(f'update_folder: {e}')
return
def delete_folder_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> list[str]:
async def delete_folder_by_id_and_user_id(self, id: str, user_id: str, db: Optional[AsyncSession] = None) -> list[str]:
try:
folder_ids = []
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
folder = result.scalars().first()
if not folder:
return folder_ids
folder_ids.append(folder.id)
# Delete all children folders
def delete_children(folder):
folder_children = self.get_folders_by_parent_id_and_user_id(folder.id, user_id, db=db)
async def delete_children(folder):
folder_children = await self.get_folders_by_parent_id_and_user_id(folder.id, user_id, db=db)
for folder_child in folder_children:
delete_children(folder_child)
await delete_children(folder_child)
folder_ids.append(folder_child.id)
folder = db.query(Folder).filter_by(id=folder_child.id).first()
db.delete(folder)
db.commit()
child_result = await db.execute(select(Folder).filter_by(id=folder_child.id))
child_folder = child_result.scalars().first()
await db.delete(child_folder)
await db.commit()
delete_children(folder)
db.delete(folder)
db.commit()
await delete_children(folder)
await db.delete(folder)
await db.commit()
return folder_ids
except Exception as e:
log.error(f'delete_folder: {e}')
@@ -319,8 +325,8 @@ class FolderTable:
name = re.sub(r'[\s_]+', ' ', name)
return name.strip().lower()
def search_folders_by_names(
self, user_id: str, queries: list[str], db: Optional[Session] = None
async def search_folders_by_names(
self, user_id: str, queries: list[str], db: Optional[AsyncSession] = None
) -> list[FolderModel]:
"""
Search for folders for a user where the name matches any of the queries, treating _ and space as equivalent, case-insensitive.
@@ -330,16 +336,18 @@ class FolderTable:
return []
results = {}
with get_db_context(db) as db:
folders = db.query(Folder).filter_by(user_id=user_id).all()
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(user_id=user_id))
folders = result.scalars().all()
for folder in folders:
if self.normalize_folder_name(folder.name) in normalized_queries:
results[folder.id] = FolderModel.model_validate(folder)
# get children folders
children = self.get_children_folders_by_id_and_user_id(folder.id, user_id, db=db)
for child in children:
results[child.id] = child
children = await self.get_children_folders_by_id_and_user_id(folder.id, user_id, db=db)
if children:
for child in children:
results[child.id] = child
# Return the results as a list
if not results:
@@ -348,16 +356,17 @@ class FolderTable:
results = list(results.values())
return results
def search_folders_by_name_contains(
self, user_id: str, query: str, db: Optional[Session] = None
async def search_folders_by_name_contains(
self, user_id: str, query: str, db: Optional[AsyncSession] = None
) -> list[FolderModel]:
"""
Partial match: normalized name contains (as substring) the normalized query.
"""
normalized_query = self.normalize_folder_name(query)
results = []
with get_db_context(db) as db:
folders = db.query(Folder).filter_by(user_id=user_id).all()
async with get_async_db_context(db) as db:
result = await db.execute(select(Folder).filter_by(user_id=user_id))
folders = result.scalars().all()
for folder in folders:
norm_name = self.normalize_folder_name(folder.name)
if normalized_query in norm_name:
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:
backend/open_webui/models/groups.py
+163 -144
View File
@@ -4,8 +4,9 @@ import time
from typing import Optional
import uuid
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, update, func, and_, or_, cast, String
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.env import DEFAULT_GROUP_SHARE_PERMISSION
from open_webui.models.files import FileMetadataResponse
@@ -15,15 +16,9 @@ from pydantic import BaseModel, ConfigDict
from sqlalchemy import (
BigInteger,
Column,
String,
Text,
JSON,
and_,
func,
ForeignKey,
cast,
or_,
select,
)
log = logging.getLogger(__name__)
@@ -143,10 +138,10 @@ class GroupTable:
group_data['data']['config']['share'] = DEFAULT_GROUP_SHARE_PERMISSION
return group_data
def insert_new_group(
self, user_id: str, form_data: GroupForm, db: Optional[Session] = None
async def insert_new_group(
self, user_id: str, form_data: GroupForm, db: Optional[AsyncSession] = None
) -> Optional[GroupModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
group_data = self._ensure_default_share_config(form_data.model_dump(exclude_none=True))
group = GroupModel(
**{
@@ -161,8 +156,8 @@ class GroupTable:
try:
result = Group(**group.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return GroupModel.model_validate(result)
else:
@@ -171,18 +166,20 @@ class GroupTable:
except Exception:
return None
def get_all_groups(self, db: Optional[Session] = None) -> list[GroupModel]:
with get_db_context(db) as db:
groups = db.query(Group).order_by(Group.updated_at.desc()).all()
async def get_all_groups(self, db: Optional[AsyncSession] = None) -> list[GroupModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Group).order_by(Group.updated_at.desc()))
groups = result.scalars().all()
return [GroupModel.model_validate(group) for group in groups]
def get_group_by_name(self, name: str, db: Optional[Session] = None) -> Optional[GroupModel]:
with get_db_context(db) as db:
group = db.query(Group).filter(Group.name == name).first()
async def get_group_by_name(self, name: str, db: Optional[AsyncSession] = None) -> Optional[GroupModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Group).filter(Group.name == name))
group = result.scalars().first()
return GroupModel.model_validate(group) if group else None
def get_groups(self, filter, db: Optional[Session] = None) -> list[GroupResponse]:
with get_db_context(db) as db:
async def get_groups(self, filter, db: Optional[AsyncSession] = None) -> list[GroupResponse]:
async with get_async_db_context(db) as db:
member_count = (
select(func.count(GroupMember.user_id))
.where(GroupMember.group_id == Group.id)
@@ -190,11 +187,11 @@ class GroupTable:
.scalar_subquery()
.label('member_count')
)
query = db.query(Group, member_count)
stmt = select(Group, member_count)
if filter:
if 'query' in filter:
query = query.filter(Group.name.ilike(f'%{filter["query"]}%'))
stmt = stmt.filter(Group.name.ilike(f'%{filter["query"]}%'))
# When share filter is present, member check is handled in the share logic
if 'share' in filter:
@@ -218,20 +215,21 @@ class GroupTable:
json_share_lower == 'members',
Group.id.in_(member_groups_select),
)
query = query.filter(or_(anyone_can_share, members_only_and_is_member))
stmt = stmt.filter(or_(anyone_can_share, members_only_and_is_member))
else:
query = query.filter(anyone_can_share)
stmt = stmt.filter(anyone_can_share)
else:
query = query.filter(and_(Group.data.isnot(None), json_share_lower == 'false'))
stmt = stmt.filter(and_(Group.data.isnot(None), json_share_lower == 'false'))
else:
# Only apply member_id filter when share filter is NOT present
if 'member_id' in filter:
query = query.filter(
stmt = stmt.filter(
Group.id.in_(select(GroupMember.group_id).where(GroupMember.user_id == filter['member_id']))
)
results = query.order_by(Group.updated_at.desc()).all()
result = await db.execute(stmt.order_by(Group.updated_at.desc()))
rows = result.all()
return [
GroupResponse.model_validate(
@@ -240,32 +238,36 @@ class GroupTable:
'member_count': count or 0,
}
)
for group, count in results
for group, count in rows
]
def search_groups(
async def search_groups(
self,
filter: Optional[dict] = None,
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> GroupListResponse:
with get_db_context(db) as db:
query = db.query(Group)
async with get_async_db_context(db) as db:
stmt = select(Group)
if filter:
if 'query' in filter:
query = query.filter(Group.name.ilike(f'%{filter["query"]}%'))
stmt = stmt.filter(Group.name.ilike(f'%{filter["query"]}%'))
if 'member_id' in filter:
query = query.filter(
stmt = stmt.filter(
Group.id.in_(select(GroupMember.group_id).where(GroupMember.user_id == filter['member_id']))
)
if 'share' in filter:
share_value = filter['share']
query = query.filter(Group.data.op('->>')('share') == str(share_value))
stmt = stmt.filter(Group.data.op('->>') ('share') == str(share_value))
total = query.count()
# Get total count
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
member_count = (
select(func.count(GroupMember.user_id))
@@ -274,7 +276,14 @@ class GroupTable:
.scalar_subquery()
.label('member_count')
)
results = query.add_columns(member_count).order_by(Group.updated_at.desc()).offset(skip).limit(limit).all()
result = await db.execute(
select(Group, member_count)
.where(Group.id.in_(select(stmt.subquery().c.id)))
.order_by(Group.updated_at.desc())
.offset(skip)
.limit(limit)
)
rows = result.all()
return {
'items': [
@@ -284,65 +293,67 @@ class GroupTable:
'member_count': count or 0,
}
)
for group, count in results
for group, count in rows
],
'total': total,
}
def get_groups_by_member_id(self, user_id: str, db: Optional[Session] = None) -> list[GroupModel]:
with get_db_context(db) as db:
return [
GroupModel.model_validate(group)
for group in db.query(Group)
async def get_groups_by_member_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[GroupModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(Group)
.join(GroupMember, GroupMember.group_id == Group.id)
.filter(GroupMember.user_id == user_id)
.order_by(Group.updated_at.desc())
.all()
]
)
return [GroupModel.model_validate(group) for group in result.scalars().all()]
def get_groups_by_member_ids(
self, user_ids: list[str], db: Optional[Session] = None
async def get_groups_by_member_ids(
self, user_ids: list[str], db: Optional[AsyncSession] = None
) -> dict[str, list[GroupModel]]:
"""Fetch groups for multiple users in a single query to avoid N+1."""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Query GroupMember joined with Group, filtering by user_ids
results = (
db.query(GroupMember.user_id, Group)
result = await db.execute(
select(GroupMember.user_id, Group)
.join(Group, Group.id == GroupMember.group_id)
.filter(GroupMember.user_id.in_(user_ids))
.order_by(Group.updated_at.desc())
.all()
)
rows = result.all()
# Group groups by user_id
user_groups: dict[str, list[GroupModel]] = {uid: [] for uid in user_ids}
for user_id, group in results:
for user_id, group in rows:
user_groups[user_id].append(GroupModel.model_validate(group))
return user_groups
def get_group_by_id(self, id: str, db: Optional[Session] = None) -> Optional[GroupModel]:
async def get_group_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[GroupModel]:
try:
with get_db_context(db) as db:
group = db.query(Group).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Group).filter_by(id=id))
group = result.scalars().first()
return GroupModel.model_validate(group) if group else None
except Exception:
return None
def get_group_user_ids_by_id(self, id: str, db: Optional[Session] = None) -> list[str]:
with get_db_context(db) as db:
members = db.query(GroupMember.user_id).filter(GroupMember.group_id == id).all()
async def get_group_user_ids_by_id(self, id: str, db: Optional[AsyncSession] = None) -> list[str]:
async with get_async_db_context(db) as db:
result = await db.execute(select(GroupMember.user_id).filter(GroupMember.group_id == id))
members = result.all()
if not members:
return []
return [m[0] for m in members]
def get_group_user_ids_by_ids(self, group_ids: list[str], db: Optional[Session] = None) -> dict[str, list[str]]:
with get_db_context(db) as db:
members = (
db.query(GroupMember.group_id, GroupMember.user_id).filter(GroupMember.group_id.in_(group_ids)).all()
async def get_group_user_ids_by_ids(self, group_ids: list[str], db: Optional[AsyncSession] = None) -> dict[str, list[str]]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(GroupMember.group_id, GroupMember.user_id).filter(GroupMember.group_id.in_(group_ids))
)
members = result.all()
group_user_ids: dict[str, list[str]] = {group_id: [] for group_id in group_ids}
@@ -351,10 +362,10 @@ class GroupTable:
return group_user_ids
def set_group_user_ids_by_id(self, group_id: str, user_ids: list[str], db: Optional[Session] = None) -> None:
with get_db_context(db) as db:
async def set_group_user_ids_by_id(self, group_id: str, user_ids: list[str], db: Optional[AsyncSession] = None) -> None:
async with get_async_db_context(db) as db:
# Delete existing members
db.query(GroupMember).filter(GroupMember.group_id == group_id).delete()
await db.execute(delete(GroupMember).filter(GroupMember.group_id == group_id))
# Insert new members
now = int(time.time())
@@ -370,101 +381,106 @@ class GroupTable:
]
db.add_all(new_members)
db.commit()
await db.commit()
def get_group_member_count_by_id(self, id: str, db: Optional[Session] = None) -> int:
with get_db_context(db) as db:
count = db.query(func.count(GroupMember.user_id)).filter(GroupMember.group_id == id).scalar()
async def get_group_member_count_by_id(self, id: str, db: Optional[AsyncSession] = None) -> int:
async with get_async_db_context(db) as db:
result = await db.execute(select(func.count(GroupMember.user_id)).filter(GroupMember.group_id == id))
count = result.scalar()
return count if count else 0
def get_group_member_counts_by_ids(self, ids: list[str], db: Optional[Session] = None) -> dict[str, int]:
async def get_group_member_counts_by_ids(self, ids: list[str], db: Optional[AsyncSession] = None) -> dict[str, int]:
if not ids:
return {}
with get_db_context(db) as db:
rows = (
db.query(GroupMember.group_id, func.count(GroupMember.user_id))
async with get_async_db_context(db) as db:
result = await db.execute(
select(GroupMember.group_id, func.count(GroupMember.user_id))
.filter(GroupMember.group_id.in_(ids))
.group_by(GroupMember.group_id)
.all()
)
rows = result.all()
return {group_id: count for group_id, count in rows}
def update_group_by_id(
async def update_group_by_id(
self,
id: str,
form_data: GroupUpdateForm,
overwrite: bool = False,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[GroupModel]:
try:
with get_db_context(db) as db:
db.query(Group).filter_by(id=id).update(
{
async with get_async_db_context(db) as db:
await db.execute(
update(Group).filter_by(id=id).values(
**form_data.model_dump(exclude_none=True),
'updated_at': int(time.time()),
}
updated_at=int(time.time()),
)
)
db.commit()
return self.get_group_by_id(id=id, db=db)
await db.commit()
return await self.get_group_by_id(id=id, db=db)
except Exception as e:
log.exception(e)
return None
def delete_group_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_group_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
db.query(Group).filter_by(id=id).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(Group).filter_by(id=id))
await db.commit()
return True
except Exception:
return False
def delete_all_groups(self, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def delete_all_groups(self, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
db.query(Group).delete()
db.commit()
await db.execute(delete(Group))
await db.commit()
return True
except Exception:
return False
def remove_user_from_all_groups(self, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def remove_user_from_all_groups(self, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
# Find all groups the user belongs to
groups = (
db.query(Group)
result = await db.execute(
select(Group)
.join(GroupMember, GroupMember.group_id == Group.id)
.filter(GroupMember.user_id == user_id)
.all()
)
groups = result.scalars().all()
# Remove the user from each group
for group in groups:
db.query(GroupMember).filter(
GroupMember.group_id == group.id, GroupMember.user_id == user_id
).delete()
await db.execute(
delete(GroupMember).filter(
GroupMember.group_id == group.id, GroupMember.user_id == user_id
)
)
db.query(Group).filter_by(id=group.id).update({'updated_at': int(time.time())})
await db.execute(
update(Group).filter_by(id=group.id).values(updated_at=int(time.time()))
)
db.commit()
await db.commit()
return True
except Exception:
db.rollback()
await db.rollback()
return False
def create_groups_by_group_names(
self, user_id: str, group_names: list[str], db: Optional[Session] = None
async def create_groups_by_group_names(
self, user_id: str, group_names: list[str], db: Optional[AsyncSession] = None
) -> list[GroupModel]:
# check for existing groups
existing_groups = self.get_all_groups(db=db)
existing_groups = await self.get_all_groups(db=db)
existing_group_names = {group.name for group in existing_groups}
new_groups = []
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
for group_name in group_names:
if group_name not in existing_group_names:
new_group = GroupModel(
@@ -483,31 +499,31 @@ class GroupTable:
try:
result = Group(**new_group.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
new_groups.append(GroupModel.model_validate(result))
except Exception as e:
log.exception(e)
continue
return new_groups
def sync_groups_by_group_names(self, user_id: str, group_names: list[str], db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def sync_groups_by_group_names(self, user_id: str, group_names: list[str], db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
now = int(time.time())
# 1. Groups that SHOULD contain the user
target_groups = db.query(Group).filter(Group.name.in_(group_names)).all()
result = await db.execute(select(Group).filter(Group.name.in_(group_names)))
target_groups = result.scalars().all()
target_group_ids = {g.id for g in target_groups}
# 2. Groups the user is CURRENTLY in
existing_group_ids = {
g.id
for g in db.query(Group)
result = await db.execute(
select(Group)
.join(GroupMember, GroupMember.group_id == Group.id)
.filter(GroupMember.user_id == user_id)
.all()
}
)
existing_group_ids = {g.id for g in result.scalars().all()}
# 3. Determine adds + removals
groups_to_add = target_group_ids - existing_group_ids
@@ -515,13 +531,15 @@ class GroupTable:
# 4. Remove in one bulk delete
if groups_to_remove:
db.query(GroupMember).filter(
GroupMember.user_id == user_id,
GroupMember.group_id.in_(groups_to_remove),
).delete(synchronize_session=False)
await db.execute(
delete(GroupMember).filter(
GroupMember.user_id == user_id,
GroupMember.group_id.in_(groups_to_remove),
)
)
db.query(Group).filter(Group.id.in_(groups_to_remove)).update(
{'updated_at': now}, synchronize_session=False
await db.execute(
update(Group).filter(Group.id.in_(groups_to_remove)).values(updated_at=now)
)
# 5. Bulk insert missing memberships
@@ -537,27 +555,28 @@ class GroupTable:
)
if groups_to_add:
db.query(Group).filter(Group.id.in_(groups_to_add)).update(
{'updated_at': now}, synchronize_session=False
await db.execute(
update(Group).filter(Group.id.in_(groups_to_add)).values(updated_at=now)
)
db.commit()
await db.commit()
return True
except Exception as e:
log.exception(e)
db.rollback()
await db.rollback()
return False
def add_users_to_group(
async def add_users_to_group(
self,
id: str,
user_ids: Optional[list[str]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[GroupModel]:
try:
with get_db_context(db) as db:
group = db.query(Group).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Group).filter_by(id=id))
group = result.scalars().first()
if not group:
return None
@@ -574,15 +593,14 @@ class GroupTable:
updated_at=now,
)
)
db.flush() # Detect unique constraint violation early
await db.flush() # Detect unique constraint violation early
except Exception:
db.rollback() # Clear failed INSERT
db.begin() # Start a new transaction
await db.rollback() # Clear failed INSERT
continue # Duplicate → ignore
group.updated_at = now
db.commit()
db.refresh(group)
await db.commit()
await db.refresh(group)
return GroupModel.model_validate(group)
@@ -590,15 +608,16 @@ class GroupTable:
log.exception(e)
return None
def remove_users_from_group(
async def remove_users_from_group(
self,
id: str,
user_ids: Optional[list[str]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[GroupModel]:
try:
with get_db_context(db) as db:
group = db.query(Group).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Group).filter_by(id=id))
group = result.scalars().first()
if not group:
return None
@@ -606,15 +625,15 @@ class GroupTable:
return GroupModel.model_validate(group)
# Remove users from group_member in batch
db.query(GroupMember).filter(GroupMember.group_id == id, GroupMember.user_id.in_(user_ids)).delete(
synchronize_session=False
await db.execute(
delete(GroupMember).filter(GroupMember.group_id == id, GroupMember.user_id.in_(user_ids))
)
# Update group timestamp
group.updated_at = int(time.time())
db.commit()
db.refresh(group)
await db.commit()
await db.refresh(group)
return GroupModel.model_validate(group)
except Exception as e:
backend/open_webui/models/knowledge.py
+180 -163
View File
@@ -4,8 +4,9 @@ import time
from typing import Optional
import uuid
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, update, or_, func
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.files import (
File,
@@ -27,7 +28,6 @@ from sqlalchemy import (
Text,
JSON,
UniqueConstraint,
or_,
)
log = logging.getLogger(__name__)
@@ -134,25 +134,25 @@ class KnowledgeFileListResponse(BaseModel):
class KnowledgeTable:
def _get_access_grants(self, knowledge_id: str, db: Optional[Session] = None) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource('knowledge', knowledge_id, db=db)
async def _get_access_grants(self, knowledge_id: str, db: Optional[AsyncSession] = None) -> list[AccessGrantModel]:
return await AccessGrants.get_grants_by_resource('knowledge', knowledge_id, db=db)
def _to_knowledge_model(
async def _to_knowledge_model(
self,
knowledge: Knowledge,
access_grants: Optional[list[AccessGrantModel]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> KnowledgeModel:
knowledge_data = KnowledgeModel.model_validate(knowledge).model_dump(exclude={'access_grants'})
knowledge_data['access_grants'] = (
access_grants if access_grants is not None else self._get_access_grants(knowledge_data['id'], db=db)
access_grants if access_grants is not None else await self._get_access_grants(knowledge_data['id'], db=db)
)
return KnowledgeModel.model_validate(knowledge_data)
def insert_new_knowledge(
self, user_id: str, form_data: KnowledgeForm, db: Optional[Session] = None
async def insert_new_knowledge(
self, user_id: str, form_data: KnowledgeForm, db: Optional[AsyncSession] = None
) -> Optional[KnowledgeModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
knowledge = KnowledgeModel(
**{
**form_data.model_dump(exclude={'access_grants'}),
@@ -167,27 +167,28 @@ class KnowledgeTable:
try:
result = Knowledge(**knowledge.model_dump(exclude={'access_grants'}))
db.add(result)
db.commit()
db.refresh(result)
AccessGrants.set_access_grants('knowledge', result.id, form_data.access_grants, db=db)
await db.commit()
await db.refresh(result)
await AccessGrants.set_access_grants('knowledge', result.id, form_data.access_grants, db=db)
if result:
return self._to_knowledge_model(result, db=db)
return await self._to_knowledge_model(result, db=db)
else:
return None
except Exception:
return None
def get_knowledge_bases(
self, skip: int = 0, limit: int = 30, db: Optional[Session] = None
async def get_knowledge_bases(
self, skip: int = 0, limit: int = 30, db: Optional[AsyncSession] = None
) -> list[KnowledgeUserModel]:
with get_db_context(db) as db:
all_knowledge = db.query(Knowledge).order_by(Knowledge.updated_at.desc()).all()
async with get_async_db_context(db) as db:
result = await db.execute(select(Knowledge).order_by(Knowledge.updated_at.desc()))
all_knowledge = result.scalars().all()
user_ids = list(set(knowledge.user_id for knowledge in all_knowledge))
knowledge_ids = [knowledge.id for knowledge in all_knowledge]
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}
grants_map = AccessGrants.get_grants_by_resources('knowledge', knowledge_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('knowledge', knowledge_ids, db=db)
knowledge_bases = []
for knowledge in all_knowledge:
@@ -195,33 +196,33 @@ class KnowledgeTable:
knowledge_bases.append(
KnowledgeUserModel.model_validate(
{
**self._to_knowledge_model(
**(await self._to_knowledge_model(
knowledge,
access_grants=grants_map.get(knowledge.id, []),
db=db,
).model_dump(),
)).model_dump(),
'user': user.model_dump() if user else None,
}
)
)
return knowledge_bases
def search_knowledge_bases(
async def search_knowledge_bases(
self,
user_id: str,
filter: dict,
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> KnowledgeListResponse:
try:
with get_db_context(db) as db:
query = db.query(Knowledge, User).outerjoin(User, User.id == Knowledge.user_id)
async with get_async_db_context(db) as db:
stmt = select(Knowledge, User).outerjoin(User, User.id == Knowledge.user_id)
if filter:
query_key = filter.get('query')
if query_key:
query = query.filter(
stmt = stmt.filter(
or_(
Knowledge.name.ilike(f'%{query_key}%'),
Knowledge.description.ilike(f'%{query_key}%'),
@@ -233,42 +234,46 @@ class KnowledgeTable:
view_option = filter.get('view_option')
if view_option == 'created':
query = query.filter(Knowledge.user_id == user_id)
stmt = stmt.filter(Knowledge.user_id == user_id)
elif view_option == 'shared':
query = query.filter(Knowledge.user_id != user_id)
stmt = stmt.filter(Knowledge.user_id != user_id)
query = AccessGrants.has_permission_filter(
stmt = AccessGrants.has_permission_filter(
db=db,
query=query,
query=stmt,
DocumentModel=Knowledge,
filter=filter,
resource_type='knowledge',
permission='read',
)
query = query.order_by(Knowledge.updated_at.desc(), Knowledge.id.asc())
stmt = stmt.order_by(Knowledge.updated_at.desc(), Knowledge.id.asc())
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit:
query = query.limit(limit)
stmt = stmt.limit(limit)
items = query.all()
result = await db.execute(stmt)
items = result.all()
knowledge_ids = [kb.id for kb, _ in items]
grants_map = AccessGrants.get_grants_by_resources('knowledge', knowledge_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('knowledge', knowledge_ids, db=db)
knowledge_bases = []
for knowledge_base, user in items:
knowledge_bases.append(
KnowledgeUserModel.model_validate(
{
**self._to_knowledge_model(
**(await self._to_knowledge_model(
knowledge_base,
access_grants=grants_map.get(knowledge_base.id, []),
db=db,
).model_dump(),
)).model_dump(),
'user': (UserModel.model_validate(user).model_dump() if user else None),
}
)
@@ -279,28 +284,27 @@ class KnowledgeTable:
print(e)
return KnowledgeListResponse(items=[], total=0)
def search_knowledge_files(
self, filter: dict, skip: int = 0, limit: int = 30, db: Optional[Session] = None
async def search_knowledge_files(
self, filter: dict, skip: int = 0, limit: int = 30, db: Optional[AsyncSession] = None
) -> KnowledgeFileListResponse:
"""
Scalable version: search files across all knowledge bases the user has
READ access to, without loading all KBs or using large IN() lists.
"""
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Base query: join Knowledge → KnowledgeFile → File
query = (
db.query(File, User, Knowledge)
stmt = (
select(File, User, Knowledge)
.join(KnowledgeFile, File.id == KnowledgeFile.file_id)
.join(Knowledge, KnowledgeFile.knowledge_id == Knowledge.id)
.outerjoin(User, User.id == KnowledgeFile.user_id)
)
# Apply access-control directly to the joined query
# This makes the database handle filtering, even with 10k+ KBs
query = AccessGrants.has_permission_filter(
stmt = AccessGrants.has_permission_filter(
db=db,
query=query,
query=stmt,
DocumentModel=Knowledge,
filter=filter,
resource_type='knowledge',
@@ -311,20 +315,24 @@ class KnowledgeTable:
if filter:
q = filter.get('query')
if q:
query = query.filter(File.filename.ilike(f'%{q}%'))
stmt = stmt.filter(File.filename.ilike(f'%{q}%'))
# Order by file changes
query = query.order_by(File.updated_at.desc(), File.id.asc())
stmt = stmt.order_by(File.updated_at.desc(), File.id.asc())
# Count before pagination
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit:
query = query.limit(limit)
stmt = stmt.limit(limit)
rows = query.all()
result = await db.execute(stmt)
rows = result.all()
items = []
for file, user, knowledge in rows:
@@ -332,7 +340,7 @@ class KnowledgeTable:
FileUserResponse(
**FileModel.model_validate(file).model_dump(),
user=(UserResponse(**UserModel.model_validate(user).model_dump()) if user else None),
collection=self._to_knowledge_model(knowledge, db=db).model_dump(),
collection=(await self._to_knowledge_model(knowledge, db=db)).model_dump(),
)
)
@@ -342,14 +350,15 @@ class KnowledgeTable:
print('search_knowledge_files error:', e)
return KnowledgeFileListResponse(items=[], total=0)
def check_access_by_user_id(self, id, user_id, permission='write', db: Optional[Session] = None) -> bool:
knowledge = self.get_knowledge_by_id(id, db=db)
async def check_access_by_user_id(self, id, user_id, permission='write', db: Optional[AsyncSession] = None) -> bool:
knowledge = await self.get_knowledge_by_id(id, db=db)
if not knowledge:
return False
if knowledge.user_id == user_id:
return True
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
return AccessGrants.has_access(
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
return await AccessGrants.has_access(
user_id=user_id,
resource_type='knowledge',
resource_id=knowledge.id,
@@ -358,45 +367,50 @@ class KnowledgeTable:
db=db,
)
def get_knowledge_bases_by_user_id(
self, user_id: str, permission: str = 'write', db: Optional[Session] = None
async def get_knowledge_bases_by_user_id(
self, user_id: str, permission: str = 'write', db: Optional[AsyncSession] = None
) -> list[KnowledgeUserModel]:
knowledge_bases = self.get_knowledge_bases(db=db)
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
return [
knowledge_base
for knowledge_base in knowledge_bases
if knowledge_base.user_id == user_id
or AccessGrants.has_access(
knowledge_bases = await self.get_knowledge_bases(db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
result = []
for knowledge_base in knowledge_bases:
if knowledge_base.user_id == user_id:
result.append(knowledge_base)
elif await AccessGrants.has_access(
user_id=user_id,
resource_type='knowledge',
resource_id=knowledge_base.id,
permission=permission,
user_group_ids=user_group_ids,
db=db,
)
]
):
result.append(knowledge_base)
return result
def get_knowledge_by_id(self, id: str, db: Optional[Session] = None) -> Optional[KnowledgeModel]:
async def get_knowledge_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[KnowledgeModel]:
try:
with get_db_context(db) as db:
knowledge = db.query(Knowledge).filter_by(id=id).first()
return self._to_knowledge_model(knowledge, db=db) if knowledge else None
async with get_async_db_context(db) as db:
result = await db.execute(select(Knowledge).filter_by(id=id))
knowledge = result.scalars().first()
return await self._to_knowledge_model(knowledge, db=db) if knowledge else None
except Exception:
return None
def get_knowledge_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[Session] = None
async def get_knowledge_by_id_and_user_id(
self, id: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[KnowledgeModel]:
knowledge = self.get_knowledge_by_id(id, db=db)
knowledge = await self.get_knowledge_by_id(id, db=db)
if not knowledge:
return None
if knowledge.user_id == user_id:
return knowledge
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
if AccessGrants.has_access(
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
if await AccessGrants.has_access(
user_id=user_id,
resource_type='knowledge',
resource_id=knowledge.id,
@@ -407,19 +421,19 @@ class KnowledgeTable:
return knowledge
return None
def get_knowledges_by_file_id(self, file_id: str, db: Optional[Session] = None) -> list[KnowledgeModel]:
async def get_knowledges_by_file_id(self, file_id: str, db: Optional[AsyncSession] = None) -> list[KnowledgeModel]:
try:
with get_db_context(db) as db:
knowledges = (
db.query(Knowledge)
async with get_async_db_context(db) as db:
result = await db.execute(
select(Knowledge)
.join(KnowledgeFile, Knowledge.id == KnowledgeFile.knowledge_id)
.filter(KnowledgeFile.file_id == file_id)
.all()
)
knowledges = result.scalars().all()
knowledge_ids = [k.id for k in knowledges]
grants_map = AccessGrants.get_grants_by_resources('knowledge', knowledge_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('knowledge', knowledge_ids, db=db)
return [
self._to_knowledge_model(
await self._to_knowledge_model(
knowledge,
access_grants=grants_map.get(knowledge.id, []),
db=db,
@@ -429,19 +443,19 @@ class KnowledgeTable:
except Exception:
return []
def search_files_by_id(
async def search_files_by_id(
self,
knowledge_id: str,
user_id: str,
filter: dict,
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> KnowledgeFileListResponse:
try:
with get_db_context(db) as db:
query = (
db.query(File, User)
async with get_async_db_context(db) as db:
stmt = (
select(File, User)
.join(KnowledgeFile, File.id == KnowledgeFile.file_id)
.outerjoin(User, User.id == KnowledgeFile.user_id)
.filter(KnowledgeFile.knowledge_id == knowledge_id)
@@ -453,13 +467,13 @@ class KnowledgeTable:
if filter:
query_key = filter.get('query')
if query_key:
query = query.filter(or_(File.filename.ilike(f'%{query_key}%')))
stmt = stmt.filter(or_(File.filename.ilike(f'%{query_key}%')))
view_option = filter.get('view_option')
if view_option == 'created':
query = query.filter(KnowledgeFile.user_id == user_id)
stmt = stmt.filter(KnowledgeFile.user_id == user_id)
elif view_option == 'shared':
query = query.filter(KnowledgeFile.user_id != user_id)
stmt = stmt.filter(KnowledgeFile.user_id != user_id)
order_by = filter.get('order_by')
direction = filter.get('direction')
@@ -473,17 +487,21 @@ class KnowledgeTable:
primary_sort = File.updated_at.asc() if is_asc else File.updated_at.desc()
# Apply sort with secondary key for deterministic pagination
query = query.order_by(primary_sort, File.id.asc())
stmt = stmt.order_by(primary_sort, File.id.asc())
# Count BEFORE pagination
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit:
query = query.limit(limit)
stmt = stmt.limit(limit)
items = query.all()
result = await db.execute(stmt)
items = result.all()
files = []
for file, user in items:
@@ -499,35 +517,34 @@ class KnowledgeTable:
print(e)
return KnowledgeFileListResponse(items=[], total=0)
def get_files_by_id(self, knowledge_id: str, db: Optional[Session] = None) -> list[FileModel]:
async def get_files_by_id(self, knowledge_id: str, db: Optional[AsyncSession] = None) -> list[FileModel]:
try:
with get_db_context(db) as db:
files = (
db.query(File)
async with get_async_db_context(db) as db:
result = await db.execute(
select(File)
.join(KnowledgeFile, File.id == KnowledgeFile.file_id)
.filter(KnowledgeFile.knowledge_id == knowledge_id)
.all()
)
files = result.scalars().all()
return [FileModel.model_validate(file) for file in files]
except Exception:
return []
def get_file_metadatas_by_id(self, knowledge_id: str, db: Optional[Session] = None) -> list[FileMetadataResponse]:
async def get_file_metadatas_by_id(self, knowledge_id: str, db: Optional[AsyncSession] = None) -> list[FileMetadataResponse]:
try:
with get_db_context(db) as db:
files = self.get_files_by_id(knowledge_id, db=db)
return [FileMetadataResponse(**file.model_dump()) for file in files]
files = await self.get_files_by_id(knowledge_id, db=db)
return [FileMetadataResponse(**file.model_dump()) for file in files]
except Exception:
return []
def add_file_to_knowledge_by_id(
async def add_file_to_knowledge_by_id(
self,
knowledge_id: str,
file_id: str,
user_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[KnowledgeFileModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
knowledge_file = KnowledgeFileModel(
**{
'id': str(uuid.uuid4()),
@@ -542,8 +559,8 @@ class KnowledgeTable:
try:
result = KnowledgeFile(**knowledge_file.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return KnowledgeFileModel.model_validate(result)
else:
@@ -551,103 +568,103 @@ class KnowledgeTable:
except Exception:
return None
def has_file(self, knowledge_id: str, file_id: str, db: Optional[Session] = None) -> bool:
async def has_file(self, knowledge_id: str, file_id: str, db: Optional[AsyncSession] = None) -> bool:
"""Check whether a file belongs to a knowledge base."""
try:
with get_db_context(db) as db:
return db.query(KnowledgeFile).filter_by(knowledge_id=knowledge_id, file_id=file_id).first() is not None
async with get_async_db_context(db) as db:
result = await db.execute(
select(KnowledgeFile).filter_by(knowledge_id=knowledge_id, file_id=file_id).limit(1)
)
return result.scalars().first() is not None
except Exception:
return False
def remove_file_from_knowledge_by_id(self, knowledge_id: str, file_id: str, db: Optional[Session] = None) -> bool:
async def remove_file_from_knowledge_by_id(self, knowledge_id: str, file_id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
db.query(KnowledgeFile).filter_by(knowledge_id=knowledge_id, file_id=file_id).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(KnowledgeFile).filter_by(knowledge_id=knowledge_id, file_id=file_id))
await db.commit()
return True
except Exception:
return False
def reset_knowledge_by_id(self, id: str, db: Optional[Session] = None) -> Optional[KnowledgeModel]:
async def reset_knowledge_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[KnowledgeModel]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Delete all knowledge_file entries for this knowledge_id
db.query(KnowledgeFile).filter_by(knowledge_id=id).delete()
db.commit()
await db.execute(delete(KnowledgeFile).filter_by(knowledge_id=id))
await db.commit()
# Update the knowledge entry's updated_at timestamp
db.query(Knowledge).filter_by(id=id).update(
{
'updated_at': int(time.time()),
}
await db.execute(
update(Knowledge).filter_by(id=id).values(updated_at=int(time.time()))
)
db.commit()
await db.commit()
return self.get_knowledge_by_id(id=id, db=db)
return await self.get_knowledge_by_id(id=id, db=db)
except Exception as e:
log.exception(e)
return None
def update_knowledge_by_id(
async def update_knowledge_by_id(
self,
id: str,
form_data: KnowledgeForm,
overwrite: bool = False,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[KnowledgeModel]:
try:
with get_db_context(db) as db:
knowledge = self.get_knowledge_by_id(id=id, db=db)
db.query(Knowledge).filter_by(id=id).update(
{
async with get_async_db_context(db) as db:
await db.execute(
update(Knowledge).filter_by(id=id).values(
**form_data.model_dump(exclude={'access_grants'}),
'updated_at': int(time.time()),
}
updated_at=int(time.time()),
)
)
db.commit()
await db.commit()
if form_data.access_grants is not None:
AccessGrants.set_access_grants('knowledge', id, form_data.access_grants, db=db)
return self.get_knowledge_by_id(id=id, db=db)
await AccessGrants.set_access_grants('knowledge', id, form_data.access_grants, db=db)
return await self.get_knowledge_by_id(id=id, db=db)
except Exception as e:
log.exception(e)
return None
def update_knowledge_data_by_id(
self, id: str, data: dict, db: Optional[Session] = None
async def update_knowledge_data_by_id(
self, id: str, data: dict, db: Optional[AsyncSession] = None
) -> Optional[KnowledgeModel]:
try:
with get_db_context(db) as db:
knowledge = self.get_knowledge_by_id(id=id, db=db)
db.query(Knowledge).filter_by(id=id).update(
{
'data': data,
'updated_at': int(time.time()),
}
async with get_async_db_context(db) as db:
await db.execute(
update(Knowledge).filter_by(id=id).values(
data=data,
updated_at=int(time.time()),
)
)
db.commit()
return self.get_knowledge_by_id(id=id, db=db)
await db.commit()
return await self.get_knowledge_by_id(id=id, db=db)
except Exception as e:
log.exception(e)
return None
def delete_knowledge_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_knowledge_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
AccessGrants.revoke_all_access('knowledge', id, db=db)
db.query(Knowledge).filter_by(id=id).delete()
db.commit()
async with get_async_db_context(db) as db:
await AccessGrants.revoke_all_access('knowledge', id, db=db)
await db.execute(delete(Knowledge).filter_by(id=id))
await db.commit()
return True
except Exception:
return False
def delete_all_knowledge(self, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def delete_all_knowledge(self, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
knowledge_ids = [row[0] for row in db.query(Knowledge.id).all()]
result = await db.execute(select(Knowledge.id))
knowledge_ids = [row[0] for row in result.all()]
for knowledge_id in knowledge_ids:
AccessGrants.revoke_all_access('knowledge', knowledge_id, db=db)
db.query(Knowledge).delete()
db.commit()
await AccessGrants.revoke_all_access('knowledge', knowledge_id, db=db)
await db.execute(delete(Knowledge))
await db.commit()
return True
except Exception:
backend/open_webui/models/memories.py
+39 -36
View File
@@ -2,8 +2,9 @@ import time
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db, get_db_context
from sqlalchemy import select, delete
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, String, Text
@@ -40,13 +41,13 @@ class MemoryModel(BaseModel):
class MemoriesTable:
def insert_new_memory(
async def insert_new_memory(
self,
user_id: str,
content: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[MemoryModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
id = str(uuid.uuid4())
memory = MemoryModel(
@@ -60,90 +61,92 @@ class MemoriesTable:
)
result = Memory(**memory.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return MemoryModel.model_validate(result)
else:
return None
def update_memory_by_id_and_user_id(
async def update_memory_by_id_and_user_id(
self,
id: str,
user_id: str,
content: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[MemoryModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
try:
memory = db.get(Memory, id)
memory = await db.get(Memory, id)
if not memory or memory.user_id != user_id:
return None
memory.content = content
memory.updated_at = int(time.time())
db.commit()
db.refresh(memory)
await db.commit()
await db.refresh(memory)
return MemoryModel.model_validate(memory)
except Exception:
return None
def get_memories(self, db: Optional[Session] = None) -> list[MemoryModel]:
with get_db_context(db) as db:
async def get_memories(self, db: Optional[AsyncSession] = None) -> list[MemoryModel]:
async with get_async_db_context(db) as db:
try:
memories = db.query(Memory).all()
result = await db.execute(select(Memory))
memories = result.scalars().all()
return [MemoryModel.model_validate(memory) for memory in memories]
except Exception:
return None
def get_memories_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[MemoryModel]:
with get_db_context(db) as db:
async def get_memories_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[MemoryModel]:
async with get_async_db_context(db) as db:
try:
memories = db.query(Memory).filter_by(user_id=user_id).all()
result = await db.execute(select(Memory).filter_by(user_id=user_id))
memories = result.scalars().all()
return [MemoryModel.model_validate(memory) for memory in memories]
except Exception:
return None
def get_memory_by_id(self, id: str, db: Optional[Session] = None) -> Optional[MemoryModel]:
with get_db_context(db) as db:
async def get_memory_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[MemoryModel]:
async with get_async_db_context(db) as db:
try:
memory = db.get(Memory, id)
return MemoryModel.model_validate(memory)
memory = await db.get(Memory, id)
return MemoryModel.model_validate(memory) if memory else None
except Exception:
return None
def delete_memory_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def delete_memory_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
db.query(Memory).filter_by(id=id).delete()
db.commit()
await db.execute(delete(Memory).filter_by(id=id))
await db.commit()
return True
except Exception:
return False
def delete_memories_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def delete_memories_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
db.query(Memory).filter_by(user_id=user_id).delete()
db.commit()
await db.execute(delete(Memory).filter_by(user_id=user_id))
await db.commit()
return True
except Exception:
return False
def delete_memory_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
async def delete_memory_by_id_and_user_id(self, id: str, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
try:
memory = db.get(Memory, id)
memory = await db.get(Memory, id)
if not memory or memory.user_id != user_id:
return None
# Delete the memory
db.delete(memory)
db.commit()
await db.delete(memory)
await db.commit()
return True
except Exception:
backend/open_webui/models/messages.py
+139 -151
View File
@@ -3,8 +3,9 @@ import time
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, func
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.tags import TagModel, Tag, Tags
from open_webui.models.users import Users, User, UserNameResponse
from open_webui.models.channels import Channels, ChannelMember
@@ -12,7 +13,7 @@ from open_webui.models.channels import Channels, ChannelMember
from pydantic import BaseModel, ConfigDict, field_validator
from sqlalchemy import BigInteger, Boolean, Column, String, Text, JSON
from sqlalchemy import or_, func, select, and_, text
from sqlalchemy import or_, func, and_, text
from sqlalchemy.sql import exists
####################
@@ -137,15 +138,15 @@ class MessageResponse(MessageReplyToResponse):
class MessageTable:
def insert_new_message(
async def insert_new_message(
self,
form_data: MessageForm,
channel_id: str,
user_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[MessageModel]:
with get_db_context(db) as db:
channel_member = Channels.join_channel(channel_id, user_id)
async with get_async_db_context(db) as db:
channel_member = await Channels.join_channel(channel_id, user_id)
id = str(uuid.uuid4())
ts = int(time.time_ns())
@@ -170,38 +171,38 @@ class MessageTable:
result = Message(**message.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
return MessageModel.model_validate(result) if result else None
def get_message_by_id(
async def get_message_by_id(
self,
id: str,
include_thread_replies: Optional[bool] = True,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[MessageResponse]:
with get_db_context(db) as db:
message = db.get(Message, id)
async with get_async_db_context(db) as db:
message = await db.get(Message, id)
if not message:
return None
reply_to_message = (
self.get_message_by_id(message.reply_to_id, include_thread_replies=False, db=db)
await self.get_message_by_id(message.reply_to_id, include_thread_replies=False, db=db)
if message.reply_to_id
else None
)
reactions = self.get_reactions_by_message_id(id, db=db)
reactions = await self.get_reactions_by_message_id(id, db=db)
thread_replies = []
if include_thread_replies:
thread_replies = self.get_thread_replies_by_message_id(id, db=db)
thread_replies = await self.get_thread_replies_by_message_id(id, db=db)
# Check if message was sent by webhook (webhook info in meta takes precedence)
webhook_info = message.meta.get('webhook') if message.meta else None
if webhook_info and webhook_info.get('id'):
# Look up webhook by ID to get current name
webhook = Channels.get_webhook_by_id(webhook_info.get('id'), db=db)
webhook = await Channels.get_webhook_by_id(webhook_info.get('id'), db=db)
if webhook:
user_info = {
'id': webhook.id,
@@ -216,7 +217,7 @@ class MessageTable:
'role': 'webhook',
}
else:
user = Users.get_user_by_id(message.user_id, db=db)
user = await Users.get_user_by_id(message.user_id, db=db)
user_info = user.model_dump() if user else None
return MessageResponse.model_validate(
@@ -230,34 +231,41 @@ class MessageTable:
}
)
def get_thread_replies_by_message_id(self, id: str, db: Optional[Session] = None) -> list[MessageReplyToResponse]:
with get_db_context(db) as db:
all_messages = db.query(Message).filter_by(parent_id=id).order_by(Message.created_at.desc()).all()
async def _resolve_user_info(self, message: Message, db: AsyncSession) -> Optional[dict]:
"""Resolve user info from message, handling webhook messages."""
webhook_info = message.meta.get('webhook') if message.meta else None
if webhook_info and webhook_info.get('id'):
webhook = await Channels.get_webhook_by_id(webhook_info.get('id'), db=db)
if webhook:
return {
'id': webhook.id,
'name': webhook.name,
'role': 'webhook',
}
else:
return {
'id': webhook_info.get('id'),
'name': 'Deleted Webhook',
'role': 'webhook',
}
return None
async def get_thread_replies_by_message_id(self, id: str, db: Optional[AsyncSession] = None) -> list[MessageReplyToResponse]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(Message).filter_by(parent_id=id).order_by(Message.created_at.desc())
)
all_messages = result.scalars().all()
messages = []
for message in all_messages:
reply_to_message = (
self.get_message_by_id(message.reply_to_id, include_thread_replies=False, db=db)
await self.get_message_by_id(message.reply_to_id, include_thread_replies=False, db=db)
if message.reply_to_id
else None
)
webhook_info = message.meta.get('webhook') if message.meta else None
user_info = None
if webhook_info and webhook_info.get('id'):
webhook = Channels.get_webhook_by_id(webhook_info.get('id'), db=db)
if webhook:
user_info = {
'id': webhook.id,
'name': webhook.name,
'role': 'webhook',
}
else:
user_info = {
'id': webhook_info.get('id'),
'name': 'Deleted Webhook',
'role': 'webhook',
}
user_info = await self._resolve_user_info(message, db)
messages.append(
MessageReplyToResponse.model_validate(
@@ -270,51 +278,37 @@ class MessageTable:
)
return messages
def get_reply_user_ids_by_message_id(self, id: str, db: Optional[Session] = None) -> list[str]:
with get_db_context(db) as db:
return [message.user_id for message in db.query(Message).filter_by(parent_id=id).all()]
async def get_reply_user_ids_by_message_id(self, id: str, db: Optional[AsyncSession] = None) -> list[str]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Message.user_id).filter_by(parent_id=id))
return [row[0] for row in result.all()]
def get_messages_by_channel_id(
async def get_messages_by_channel_id(
self,
channel_id: str,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[MessageReplyToResponse]:
with get_db_context(db) as db:
all_messages = (
db.query(Message)
async with get_async_db_context(db) as db:
result = await db.execute(
select(Message)
.filter_by(channel_id=channel_id, parent_id=None)
.order_by(Message.created_at.desc())
.offset(skip)
.limit(limit)
.all()
)
all_messages = result.scalars().all()
messages = []
for message in all_messages:
reply_to_message = (
self.get_message_by_id(message.reply_to_id, include_thread_replies=False, db=db)
await self.get_message_by_id(message.reply_to_id, include_thread_replies=False, db=db)
if message.reply_to_id
else None
)
webhook_info = message.meta.get('webhook') if message.meta else None
user_info = None
if webhook_info and webhook_info.get('id'):
webhook = Channels.get_webhook_by_id(webhook_info.get('id'), db=db)
if webhook:
user_info = {
'id': webhook.id,
'name': webhook.name,
'role': 'webhook',
}
else:
user_info = {
'id': webhook_info.get('id'),
'name': 'Deleted Webhook',
'role': 'webhook',
}
user_info = await self._resolve_user_info(message, db)
messages.append(
MessageReplyToResponse.model_validate(
@@ -327,28 +321,28 @@ class MessageTable:
)
return messages
def get_messages_by_parent_id(
async def get_messages_by_parent_id(
self,
channel_id: str,
parent_id: str,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[MessageReplyToResponse]:
with get_db_context(db) as db:
message = db.get(Message, parent_id)
async with get_async_db_context(db) as db:
message = await db.get(Message, parent_id)
if not message:
return []
all_messages = (
db.query(Message)
result = await db.execute(
select(Message)
.filter_by(channel_id=channel_id, parent_id=parent_id)
.order_by(Message.created_at.desc())
.offset(skip)
.limit(limit)
.all()
)
all_messages = list(result.scalars().all())
# If length of all_messages is less than limit, then add the parent message
if len(all_messages) < limit:
@@ -357,27 +351,12 @@ class MessageTable:
messages = []
for message in all_messages:
reply_to_message = (
self.get_message_by_id(message.reply_to_id, include_thread_replies=False, db=db)
await self.get_message_by_id(message.reply_to_id, include_thread_replies=False, db=db)
if message.reply_to_id
else None
)
webhook_info = message.meta.get('webhook') if message.meta else None
user_info = None
if webhook_info and webhook_info.get('id'):
webhook = Channels.get_webhook_by_id(webhook_info.get('id'), db=db)
if webhook:
user_info = {
'id': webhook.id,
'name': webhook.name,
'role': 'webhook',
}
else:
user_info = {
'id': webhook_info.get('id'),
'name': 'Deleted Webhook',
'role': 'webhook',
}
user_info = await self._resolve_user_info(message, db)
messages.append(
MessageReplyToResponse.model_validate(
@@ -390,34 +369,37 @@ class MessageTable:
)
return messages
def get_last_message_by_channel_id(self, channel_id: str, db: Optional[Session] = None) -> Optional[MessageModel]:
with get_db_context(db) as db:
message = db.query(Message).filter_by(channel_id=channel_id).order_by(Message.created_at.desc()).first()
async def get_last_message_by_channel_id(self, channel_id: str, db: Optional[AsyncSession] = None) -> Optional[MessageModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(Message).filter_by(channel_id=channel_id).order_by(Message.created_at.desc()).limit(1)
)
message = result.scalars().first()
return MessageModel.model_validate(message) if message else None
def get_pinned_messages_by_channel_id(
async def get_pinned_messages_by_channel_id(
self,
channel_id: str,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[MessageModel]:
with get_db_context(db) as db:
all_messages = (
db.query(Message)
async with get_async_db_context(db) as db:
result = await db.execute(
select(Message)
.filter_by(channel_id=channel_id, is_pinned=True)
.order_by(Message.pinned_at.desc())
.offset(skip)
.limit(limit)
.all()
)
all_messages = result.scalars().all()
return [MessageModel.model_validate(message) for message in all_messages]
def update_message_by_id(
self, id: str, form_data: MessageForm, db: Optional[Session] = None
async def update_message_by_id(
self, id: str, form_data: MessageForm, db: Optional[AsyncSession] = None
) -> Optional[MessageModel]:
with get_db_context(db) as db:
message = db.get(Message, id)
async with get_async_db_context(db) as db:
message = await db.get(Message, id)
message.content = form_data.content
message.data = {
**(message.data if message.data else {}),
@@ -428,49 +410,53 @@ class MessageTable:
**(form_data.meta if form_data.meta else {}),
}
message.updated_at = int(time.time_ns())
db.commit()
db.refresh(message)
await db.commit()
await db.refresh(message)
return MessageModel.model_validate(message) if message else None
def update_is_pinned_by_id(
async def update_is_pinned_by_id(
self,
id: str,
is_pinned: bool,
pinned_by: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[MessageModel]:
with get_db_context(db) as db:
message = db.get(Message, id)
async with get_async_db_context(db) as db:
message = await db.get(Message, id)
message.is_pinned = is_pinned
message.pinned_at = int(time.time_ns()) if is_pinned else None
message.pinned_by = pinned_by if is_pinned else None
db.commit()
db.refresh(message)
await db.commit()
await db.refresh(message)
return MessageModel.model_validate(message) if message else None
def get_unread_message_count(
async def get_unread_message_count(
self,
channel_id: str,
user_id: str,
last_read_at: Optional[int] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> int:
with get_db_context(db) as db:
query = db.query(Message).filter(
async with get_async_db_context(db) as db:
stmt = select(func.count(Message.id)).filter(
Message.channel_id == channel_id,
Message.parent_id == None, # only count top-level messages
Message.created_at > (last_read_at if last_read_at else 0),
)
if user_id:
query = query.filter(Message.user_id != user_id)
return query.count()
stmt = stmt.filter(Message.user_id != user_id)
result = await db.execute(stmt)
return result.scalar()
def add_reaction_to_message(
self, id: str, user_id: str, name: str, db: Optional[Session] = None
async def add_reaction_to_message(
self, id: str, user_id: str, name: str, db: Optional[AsyncSession] = None
) -> Optional[MessageReactionModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# check for existing reaction
existing_reaction = db.query(MessageReaction).filter_by(message_id=id, user_id=user_id, name=name).first()
result = await db.execute(
select(MessageReaction).filter_by(message_id=id, user_id=user_id, name=name)
)
existing_reaction = result.scalars().first()
if existing_reaction:
return MessageReactionModel.model_validate(existing_reaction)
@@ -484,19 +470,19 @@ class MessageTable:
)
result = MessageReaction(**reaction.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
return MessageReactionModel.model_validate(result) if result else None
def get_reactions_by_message_id(self, id: str, db: Optional[Session] = None) -> list[Reactions]:
with get_db_context(db) as db:
async def get_reactions_by_message_id(self, id: str, db: Optional[AsyncSession] = None) -> list[Reactions]:
async with get_async_db_context(db) as db:
# JOIN User so all user info is fetched in one query
results = (
db.query(MessageReaction, User)
result = await db.execute(
select(MessageReaction, User)
.join(User, MessageReaction.user_id == User.id)
.filter(MessageReaction.message_id == id)
.all()
)
results = result.all()
reactions = {}
@@ -518,58 +504,60 @@ class MessageTable:
return [Reactions(**reaction) for reaction in reactions.values()]
def remove_reaction_by_id_and_user_id_and_name(
self, id: str, user_id: str, name: str, db: Optional[Session] = None
async def remove_reaction_by_id_and_user_id_and_name(
self, id: str, user_id: str, name: str, db: Optional[AsyncSession] = None
) -> bool:
with get_db_context(db) as db:
db.query(MessageReaction).filter_by(message_id=id, user_id=user_id, name=name).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(MessageReaction).filter_by(message_id=id, user_id=user_id, name=name))
await db.commit()
return True
def delete_reactions_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
db.query(MessageReaction).filter_by(message_id=id).delete()
db.commit()
async def delete_reactions_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
await db.execute(delete(MessageReaction).filter_by(message_id=id))
await db.commit()
return True
def delete_replies_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
db.query(Message).filter_by(parent_id=id).delete()
db.commit()
async def delete_replies_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
await db.execute(delete(Message).filter_by(parent_id=id))
await db.commit()
return True
def delete_message_by_id(self, id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
db.query(Message).filter_by(id=id).delete()
async def delete_message_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
await db.execute(delete(Message).filter_by(id=id))
# Delete all reactions to this message
db.query(MessageReaction).filter_by(message_id=id).delete()
await db.execute(delete(MessageReaction).filter_by(message_id=id))
db.commit()
await db.commit()
return True
def search_messages_by_channel_ids(
async def search_messages_by_channel_ids(
self,
channel_ids: list[str],
query: str,
start_timestamp: Optional[int] = None,
end_timestamp: Optional[int] = None,
limit: int = 10,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[MessageModel]:
"""Search messages in specified channels by content."""
with get_db_context(db) as db:
query_builder = db.query(Message).filter(
async with get_async_db_context(db) as db:
stmt = select(Message).filter(
Message.channel_id.in_(channel_ids),
Message.content.ilike(f'%{query}%'),
)
if start_timestamp:
query_builder = query_builder.filter(Message.created_at >= start_timestamp)
stmt = stmt.filter(Message.created_at >= start_timestamp)
if end_timestamp:
query_builder = query_builder.filter(Message.created_at <= end_timestamp)
stmt = stmt.filter(Message.created_at <= end_timestamp)
messages = query_builder.order_by(Message.created_at.desc()).limit(limit).all()
stmt = stmt.order_by(Message.created_at.desc()).limit(limit)
result = await db.execute(stmt)
messages = result.scalars().all()
return [MessageModel.model_validate(msg) for msg in messages]
backend/open_webui/models/models.py
+140 -129
View File
@@ -2,8 +2,9 @@ import logging
import time
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, update, or_, func, String, cast
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.groups import Groups
from open_webui.models.users import User, UserModel, Users, UserResponse
@@ -12,9 +13,6 @@ from open_webui.models.access_grants import AccessGrantModel, AccessGrants
from pydantic import BaseModel, ConfigDict, Field, model_validator
from sqlalchemy import String, cast, or_, and_, func
from sqlalchemy.dialects import postgresql, sqlite
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy import BigInteger, Column, Text, Boolean
@@ -154,26 +152,26 @@ class ModelForm(BaseModel):
class ModelsTable:
def _get_access_grants(self, model_id: str, db: Optional[Session] = None) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource('model', model_id, db=db)
async def _get_access_grants(self, model_id: str, db: Optional[AsyncSession] = None) -> list[AccessGrantModel]:
return await AccessGrants.get_grants_by_resource('model', model_id, db=db)
def _to_model_model(
async def _to_model_model(
self,
model: Model,
access_grants: Optional[list[AccessGrantModel]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> ModelModel:
model_data = ModelModel.model_validate(model).model_dump(exclude={'access_grants'})
model_data['access_grants'] = (
access_grants if access_grants is not None else self._get_access_grants(model_data['id'], db=db)
access_grants if access_grants is not None else await self._get_access_grants(model_data['id'], db=db)
)
return ModelModel.model_validate(model_data)
def insert_new_model(
self, form_data: ModelForm, user_id: str, db: Optional[Session] = None
async def insert_new_model(
self, form_data: ModelForm, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[ModelModel]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
result = Model(
**{
**form_data.model_dump(exclude={'access_grants'}),
@@ -183,37 +181,39 @@ class ModelsTable:
}
)
db.add(result)
db.commit()
db.refresh(result)
AccessGrants.set_access_grants('model', result.id, form_data.access_grants, db=db)
await db.commit()
await db.refresh(result)
await AccessGrants.set_access_grants('model', result.id, form_data.access_grants, db=db)
if result:
return self._to_model_model(result, db=db)
return await self._to_model_model(result, db=db)
else:
return None
except Exception as e:
log.exception(f'Failed to insert a new model: {e}')
return None
def get_all_models(self, db: Optional[Session] = None) -> list[ModelModel]:
with get_db_context(db) as db:
all_models = db.query(Model).all()
async def get_all_models(self, db: Optional[AsyncSession] = None) -> list[ModelModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Model))
all_models = result.scalars().all()
model_ids = [model.id for model in all_models]
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('model', model_ids, db=db)
return [
self._to_model_model(model, access_grants=grants_map.get(model.id, []), db=db) for model in all_models
await self._to_model_model(model, access_grants=grants_map.get(model.id, []), db=db) for model in all_models
]
def get_models(self, db: Optional[Session] = None) -> list[ModelUserResponse]:
with get_db_context(db) as db:
all_models = db.query(Model).filter(Model.base_model_id != None).all()
async def get_models(self, db: Optional[AsyncSession] = None) -> list[ModelUserResponse]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Model).filter(Model.base_model_id != None))
all_models = result.scalars().all()
user_ids = list(set(model.user_id for model in all_models))
model_ids = [model.id for model in all_models]
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}
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('model', model_ids, db=db)
models = []
for model in all_models:
@@ -221,44 +221,48 @@ class ModelsTable:
models.append(
ModelUserResponse.model_validate(
{
**self._to_model_model(
**(await self._to_model_model(
model,
access_grants=grants_map.get(model.id, []),
db=db,
).model_dump(),
)).model_dump(),
'user': user.model_dump() if user else None,
}
)
)
return models
def get_base_models(self, db: Optional[Session] = None) -> list[ModelModel]:
with get_db_context(db) as db:
all_models = db.query(Model).filter(Model.base_model_id == None).all()
async def get_base_models(self, db: Optional[AsyncSession] = None) -> list[ModelModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Model).filter(Model.base_model_id == None))
all_models = result.scalars().all()
model_ids = [model.id for model in all_models]
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('model', model_ids, db=db)
return [
self._to_model_model(model, access_grants=grants_map.get(model.id, []), db=db) for model in all_models
await self._to_model_model(model, access_grants=grants_map.get(model.id, []), db=db) for model in all_models
]
def get_models_by_user_id(
self, user_id: str, permission: str = 'write', db: Optional[Session] = None
async def get_models_by_user_id(
self, user_id: str, permission: str = 'write', db: Optional[AsyncSession] = None
) -> list[ModelUserResponse]:
models = self.get_models(db=db)
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
return [
model
for model in models
if model.user_id == user_id
or AccessGrants.has_access(
models = await self.get_models(db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
result = []
for model in models:
if model.user_id == user_id:
result.append(model)
elif await AccessGrants.has_access(
user_id=user_id,
resource_type='model',
resource_id=model.id,
permission=permission,
user_group_ids=user_group_ids,
db=db,
)
]
):
result.append(model)
return result
def _has_permission(self, db, query, filter: dict, permission: str = 'read'):
return AccessGrants.has_permission_filter(
@@ -270,23 +274,22 @@ class ModelsTable:
permission=permission,
)
def search_models(
async def search_models(
self,
user_id: str,
filter: dict = {},
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> ModelListResponse:
with get_db_context(db) as db:
# Join GroupMember so we can order by group_id when requested
query = db.query(Model, User).outerjoin(User, User.id == Model.user_id)
query = query.filter(Model.base_model_id != None)
async with get_async_db_context(db) as db:
stmt = select(Model, User).outerjoin(User, User.id == Model.user_id)
stmt = stmt.filter(Model.base_model_id != None)
if filter:
query_key = filter.get('query')
if query_key:
query = query.filter(
stmt = stmt.filter(
or_(
Model.name.ilike(f'%{query_key}%'),
Model.base_model_id.ilike(f'%{query_key}%'),
@@ -298,92 +301,95 @@ class ModelsTable:
view_option = filter.get('view_option')
if view_option == 'created':
query = query.filter(Model.user_id == user_id)
stmt = stmt.filter(Model.user_id == user_id)
elif view_option == 'shared':
query = query.filter(Model.user_id != user_id)
stmt = stmt.filter(Model.user_id != user_id)
# Apply access control filtering
query = self._has_permission(
stmt = self._has_permission(
db,
query,
stmt,
filter,
permission='read',
)
tag = filter.get('tag')
if tag:
# TODO: This is a simple implementation and should be improved for performance
like_pattern = f'%"{tag.lower()}"%' # `"tag"` inside JSON array
like_pattern = f'%"{tag.lower()}"%'
meta_text = func.lower(cast(Model.meta, String))
query = query.filter(meta_text.like(like_pattern))
stmt = stmt.filter(meta_text.like(like_pattern))
order_by = filter.get('order_by')
direction = filter.get('direction')
if order_by == 'name':
if direction == 'asc':
query = query.order_by(Model.name.asc())
stmt = stmt.order_by(Model.name.asc())
else:
query = query.order_by(Model.name.desc())
stmt = stmt.order_by(Model.name.desc())
elif order_by == 'created_at':
if direction == 'asc':
query = query.order_by(Model.created_at.asc())
stmt = stmt.order_by(Model.created_at.asc())
else:
query = query.order_by(Model.created_at.desc())
stmt = stmt.order_by(Model.created_at.desc())
elif order_by == 'updated_at':
if direction == 'asc':
query = query.order_by(Model.updated_at.asc())
stmt = stmt.order_by(Model.updated_at.asc())
else:
query = query.order_by(Model.updated_at.desc())
stmt = stmt.order_by(Model.updated_at.desc())
else:
query = query.order_by(Model.created_at.desc())
stmt = stmt.order_by(Model.created_at.desc())
# Count BEFORE pagination
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit:
query = query.limit(limit)
stmt = stmt.limit(limit)
items = query.all()
result = await db.execute(stmt)
items = result.all()
model_ids = [model.id for model, _ in items]
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('model', model_ids, db=db)
models = []
for model, user in items:
models.append(
ModelUserResponse(
**self._to_model_model(
**(await self._to_model_model(
model,
access_grants=grants_map.get(model.id, []),
db=db,
).model_dump(),
)).model_dump(),
user=(UserResponse(**UserModel.model_validate(user).model_dump()) if user else None),
)
)
return ModelListResponse(items=models, total=total)
def get_model_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ModelModel]:
async def get_model_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ModelModel]:
try:
with get_db_context(db) as db:
model = db.get(Model, id)
return self._to_model_model(model, db=db) if model else None
async with get_async_db_context(db) as db:
model = await db.get(Model, id)
return await self._to_model_model(model, db=db) if model else None
except Exception:
return None
def get_models_by_ids(self, ids: list[str], db: Optional[Session] = None) -> list[ModelModel]:
async def get_models_by_ids(self, ids: list[str], db: Optional[AsyncSession] = None) -> list[ModelModel]:
try:
with get_db_context(db) as db:
models = db.query(Model).filter(Model.id.in_(ids)).all()
async with get_async_db_context(db) as db:
result = await db.execute(select(Model).filter(Model.id.in_(ids)))
models = result.scalars().all()
model_ids = [model.id for model in models]
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('model', model_ids, db=db)
return [
self._to_model_model(
await self._to_model_model(
model,
access_grants=grants_map.get(model.id, []),
db=db,
@@ -393,82 +399,86 @@ class ModelsTable:
except Exception:
return []
def toggle_model_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ModelModel]:
with get_db_context(db) as db:
async def toggle_model_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ModelModel]:
async with get_async_db_context(db) as db:
try:
model = db.query(Model).filter_by(id=id).first()
result = await db.execute(select(Model).filter_by(id=id))
model = result.scalars().first()
if not model:
return None
model.is_active = not model.is_active
model.updated_at = int(time.time())
db.commit()
db.refresh(model)
await db.commit()
await db.refresh(model)
return self._to_model_model(model, db=db)
return await self._to_model_model(model, db=db)
except Exception:
return None
def update_model_by_id(self, id: str, model: ModelForm, db: Optional[Session] = None) -> Optional[ModelModel]:
async def update_model_by_id(self, id: str, model: ModelForm, db: Optional[AsyncSession] = None) -> Optional[ModelModel]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# update only the fields that are present in the model
data = model.model_dump(exclude={'id', 'access_grants'})
data['updated_at'] = int(time.time())
result = db.query(Model).filter_by(id=id).update(data)
await db.execute(update(Model).filter_by(id=id).values(**data))
db.commit()
await db.commit()
if model.access_grants is not None:
AccessGrants.set_access_grants('model', id, model.access_grants, db=db)
await AccessGrants.set_access_grants('model', id, model.access_grants, db=db)
return self.get_model_by_id(id, db=db)
return await self.get_model_by_id(id, db=db)
except Exception as e:
log.exception(f'Failed to update the model by id {id}: {e}')
return None
def update_model_updated_at_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ModelModel]:
async def update_model_updated_at_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ModelModel]:
try:
with get_db_context(db) as db:
result = db.query(Model).filter_by(id=id).first()
if not result:
async with get_async_db_context(db) as db:
result = await db.execute(select(Model).filter_by(id=id))
model_obj = result.scalars().first()
if not model_obj:
return None
result.updated_at = int(time.time())
db.commit()
db.refresh(result)
return self._to_model_model(result, db=db)
model_obj.updated_at = int(time.time())
await db.commit()
await db.refresh(model_obj)
return await self._to_model_model(model_obj, db=db)
except Exception as e:
log.exception(f'Failed to update the model updated_at by id {id}: {e}')
return None
def delete_model_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_model_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
AccessGrants.revoke_all_access('model', id, db=db)
db.query(Model).filter_by(id=id).delete()
db.commit()
async with get_async_db_context(db) as db:
await AccessGrants.revoke_all_access('model', id, db=db)
await db.execute(delete(Model).filter_by(id=id))
await db.commit()
return True
except Exception:
return False
def delete_all_models(self, db: Optional[Session] = None) -> bool:
async def delete_all_models(self, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
model_ids = [row[0] for row in db.query(Model.id).all()]
async with get_async_db_context(db) as db:
result = await db.execute(select(Model.id))
model_ids = [row[0] for row in result.all()]
for model_id in model_ids:
AccessGrants.revoke_all_access('model', model_id, db=db)
db.query(Model).delete()
db.commit()
await AccessGrants.revoke_all_access('model', model_id, db=db)
await db.execute(delete(Model))
await db.commit()
return True
except Exception:
return False
def sync_models(self, user_id: str, models: list[ModelModel], db: Optional[Session] = None) -> list[ModelModel]:
async def sync_models(self, user_id: str, models: list[ModelModel], db: Optional[AsyncSession] = None) -> list[ModelModel]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Get existing models
existing_models = db.query(Model).all()
result = await db.execute(select(Model))
existing_models = result.scalars().all()
existing_ids = {model.id for model in existing_models}
# Prepare a set of new model IDs
@@ -477,12 +487,12 @@ class ModelsTable:
# Update or insert models
for model in models:
if model.id in existing_ids:
db.query(Model).filter_by(id=model.id).update(
{
await db.execute(
update(Model).filter_by(id=model.id).values(
**model.model_dump(exclude={'access_grants'}),
'user_id': user_id,
'updated_at': int(time.time()),
}
user_id=user_id,
updated_at=int(time.time()),
)
)
else:
new_model = Model(
@@ -493,21 +503,22 @@ class ModelsTable:
}
)
db.add(new_model)
AccessGrants.set_access_grants('model', model.id, model.access_grants, db=db)
await AccessGrants.set_access_grants('model', model.id, model.access_grants, db=db)
# Remove models that are no longer present
for model in existing_models:
if model.id not in new_model_ids:
AccessGrants.revoke_all_access('model', model.id, db=db)
db.delete(model)
await AccessGrants.revoke_all_access('model', model.id, db=db)
await db.delete(model)
db.commit()
await db.commit()
all_models = db.query(Model).all()
result = await db.execute(select(Model))
all_models = result.scalars().all()
model_ids = [model.id for model in all_models]
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('model', model_ids, db=db)
return [
self._to_model_model(
await self._to_model_model(
model,
access_grants=grants_map.get(model.id, []),
db=db,
backend/open_webui/models/notes.py
+81 -72
View File
@@ -4,8 +4,9 @@ import uuid
from typing import Optional
from functools import lru_cache
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db, get_db_context
from sqlalchemy import select, delete, update, or_, func, cast
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from open_webui.models.groups import Groups
from open_webui.models.users import User, UserModel, Users, UserResponse
from open_webui.models.access_grants import AccessGrantModel, AccessGrants
@@ -13,7 +14,6 @@ from open_webui.models.access_grants import AccessGrantModel, AccessGrants
from pydantic import BaseModel, ConfigDict, Field
from sqlalchemy import BigInteger, Column, Text, JSON
from sqlalchemy import or_, func, cast
####################
# Note DB Schema
@@ -88,18 +88,18 @@ class NoteListResponse(BaseModel):
class NoteTable:
def _get_access_grants(self, note_id: str, db: Optional[Session] = None) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource('note', note_id, db=db)
async def _get_access_grants(self, note_id: str, db: Optional[AsyncSession] = None) -> list[AccessGrantModel]:
return await AccessGrants.get_grants_by_resource('note', note_id, db=db)
def _to_note_model(
async def _to_note_model(
self,
note: Note,
access_grants: Optional[list[AccessGrantModel]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> NoteModel:
note_data = NoteModel.model_validate(note).model_dump(exclude={'access_grants'})
note_data['access_grants'] = (
access_grants if access_grants is not None else self._get_access_grants(note_data['id'], db=db)
access_grants if access_grants is not None else await self._get_access_grants(note_data['id'], db=db)
)
return NoteModel.model_validate(note_data)
@@ -113,8 +113,8 @@ class NoteTable:
permission=permission,
)
def insert_new_note(self, user_id: str, form_data: NoteForm, db: Optional[Session] = None) -> Optional[NoteModel]:
with get_db_context(db) as db:
async def insert_new_note(self, user_id: str, form_data: NoteForm, db: Optional[AsyncSession] = None) -> Optional[NoteModel]:
async with get_async_db_context(db) as db:
note = NoteModel(
**{
'id': str(uuid.uuid4()),
@@ -129,38 +129,39 @@ class NoteTable:
new_note = Note(**note.model_dump(exclude={'access_grants'}))
db.add(new_note)
db.commit()
AccessGrants.set_access_grants('note', note.id, form_data.access_grants, db=db)
return self._to_note_model(new_note, db=db)
await db.commit()
await AccessGrants.set_access_grants('note', note.id, form_data.access_grants, db=db)
return await self._to_note_model(new_note, db=db)
def get_notes(self, skip: int = 0, limit: int = 50, db: Optional[Session] = None) -> list[NoteModel]:
with get_db_context(db) as db:
query = db.query(Note).order_by(Note.updated_at.desc())
async def get_notes(self, skip: int = 0, limit: int = 50, db: Optional[AsyncSession] = None) -> list[NoteModel]:
async with get_async_db_context(db) as db:
stmt = select(Note).order_by(Note.updated_at.desc())
if skip is not None:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit is not None:
query = query.limit(limit)
notes = query.all()
stmt = stmt.limit(limit)
result = await db.execute(stmt)
notes = result.scalars().all()
note_ids = [note.id for note in notes]
grants_map = AccessGrants.get_grants_by_resources('note', note_ids, db=db)
return [self._to_note_model(note, access_grants=grants_map.get(note.id, []), db=db) for note in notes]
grants_map = await AccessGrants.get_grants_by_resources('note', note_ids, db=db)
return [await self._to_note_model(note, access_grants=grants_map.get(note.id, []), db=db) for note in notes]
def search_notes(
async def search_notes(
self,
user_id: str,
filter: dict = {},
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> NoteListResponse:
with get_db_context(db) as db:
query = db.query(Note, User).outerjoin(User, User.id == Note.user_id)
async with get_async_db_context(db) as db:
stmt = select(Note, User).outerjoin(User, User.id == Note.user_id)
if filter:
query_key = filter.get('query')
if query_key:
# Normalize search by removing hyphens and spaces (e.g., "todo" matches "to-do" and "to do")
normalized_query = query_key.replace('-', '').replace(' ', '')
query = query.filter(
stmt = stmt.filter(
or_(
func.replace(func.replace(Note.title, '-', ''), ' ', '').ilike(f'%{normalized_query}%'),
func.replace(
@@ -173,9 +174,9 @@ class NoteTable:
view_option = filter.get('view_option')
if view_option == 'created':
query = query.filter(Note.user_id == user_id)
stmt = stmt.filter(Note.user_id == user_id)
elif view_option == 'shared':
query = query.filter(Note.user_id != user_id)
stmt = stmt.filter(Note.user_id != user_id)
# Apply access control filtering
if 'permission' in filter:
@@ -183,9 +184,9 @@ class NoteTable:
else:
permission = 'write'
query = self._has_permission(
stmt = self._has_permission(
db,
query,
stmt,
filter,
permission=permission,
)
@@ -195,87 +196,95 @@ class NoteTable:
if order_by == 'name':
if direction == 'asc':
query = query.order_by(Note.title.asc())
stmt = stmt.order_by(Note.title.asc())
else:
query = query.order_by(Note.title.desc())
stmt = stmt.order_by(Note.title.desc())
elif order_by == 'created_at':
if direction == 'asc':
query = query.order_by(Note.created_at.asc())
stmt = stmt.order_by(Note.created_at.asc())
else:
query = query.order_by(Note.created_at.desc())
stmt = stmt.order_by(Note.created_at.desc())
elif order_by == 'updated_at':
if direction == 'asc':
query = query.order_by(Note.updated_at.asc())
stmt = stmt.order_by(Note.updated_at.asc())
else:
query = query.order_by(Note.updated_at.desc())
stmt = stmt.order_by(Note.updated_at.desc())
else:
query = query.order_by(Note.updated_at.desc())
stmt = stmt.order_by(Note.updated_at.desc())
else:
query = query.order_by(Note.updated_at.desc())
stmt = stmt.order_by(Note.updated_at.desc())
# Count BEFORE pagination
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit:
query = query.limit(limit)
stmt = stmt.limit(limit)
items = query.all()
result = await db.execute(stmt)
items = result.all()
note_ids = [note.id for note, _ in items]
grants_map = AccessGrants.get_grants_by_resources('note', note_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('note', note_ids, db=db)
notes = []
for note, user in items:
notes.append(
NoteUserResponse(
**self._to_note_model(
**(await self._to_note_model(
note,
access_grants=grants_map.get(note.id, []),
db=db,
).model_dump(),
)).model_dump(),
user=(UserResponse(**UserModel.model_validate(user).model_dump()) if user else None),
)
)
return NoteListResponse(items=notes, total=total)
def get_notes_by_user_id(
async def get_notes_by_user_id(
self,
user_id: str,
permission: str = 'read',
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[NoteModel]:
with get_db_context(db) as db:
user_group_ids = [group.id for group in Groups.get_groups_by_member_id(user_id, db=db)]
async with get_async_db_context(db) as db:
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = [group.id for group in user_groups]
query = db.query(Note).order_by(Note.updated_at.desc())
query = self._has_permission(db, query, {'user_id': user_id, 'group_ids': user_group_ids}, permission)
stmt = select(Note).order_by(Note.updated_at.desc())
stmt = self._has_permission(db, stmt, {'user_id': user_id, 'group_ids': user_group_ids}, permission)
if skip is not None:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit is not None:
query = query.limit(limit)
stmt = stmt.limit(limit)
notes = query.all()
result = await db.execute(stmt)
notes = result.scalars().all()
note_ids = [note.id for note in notes]
grants_map = AccessGrants.get_grants_by_resources('note', note_ids, db=db)
return [self._to_note_model(note, access_grants=grants_map.get(note.id, []), db=db) for note in notes]
grants_map = await AccessGrants.get_grants_by_resources('note', note_ids, db=db)
return [await self._to_note_model(note, access_grants=grants_map.get(note.id, []), db=db) for note in notes]
def get_note_by_id(self, id: str, db: Optional[Session] = None) -> Optional[NoteModel]:
with get_db_context(db) as db:
note = db.query(Note).filter(Note.id == id).first()
return self._to_note_model(note, db=db) if note else None
async def get_note_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[NoteModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Note).filter(Note.id == id))
note = result.scalars().first()
return await self._to_note_model(note, db=db) if note else None
def update_note_by_id(
self, id: str, form_data: NoteUpdateForm, db: Optional[Session] = None
async def update_note_by_id(
self, id: str, form_data: NoteUpdateForm, db: Optional[AsyncSession] = None
) -> Optional[NoteModel]:
with get_db_context(db) as db:
note = db.query(Note).filter(Note.id == id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Note).filter(Note.id == id))
note = result.scalars().first()
if not note:
return None
@@ -289,19 +298,19 @@ class NoteTable:
note.meta = {**note.meta, **form_data['meta']}
if 'access_grants' in form_data:
AccessGrants.set_access_grants('note', id, form_data['access_grants'], db=db)
await AccessGrants.set_access_grants('note', id, form_data['access_grants'], db=db)
note.updated_at = int(time.time_ns())
db.commit()
return self._to_note_model(note, db=db) if note else None
await db.commit()
return await self._to_note_model(note, db=db) if note else None
def delete_note_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_note_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
AccessGrants.revoke_all_access('note', id, db=db)
db.query(Note).filter(Note.id == id).delete()
db.commit()
async with get_async_db_context(db) as db:
await AccessGrants.revoke_all_access('note', id, db=db)
await db.execute(delete(Note).filter(Note.id == id))
await db.commit()
return True
except Exception:
return False
backend/open_webui/models/oauth_sessions.py
+110 -67
View File
@@ -8,8 +8,9 @@ import json
from cryptography.fernet import Fernet
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db, get_db_context
from sqlalchemy import select, delete, update
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from open_webui.env import OAUTH_SESSION_TOKEN_ENCRYPTION_KEY
from pydantic import BaseModel, ConfigDict
@@ -103,16 +104,16 @@ class OAuthSessionTable:
log.error(f'Error decrypting tokens: {type(e).__name__}: {e}')
raise
def create_session(
async def create_session(
self,
user_id: str,
provider: str,
token: dict,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[OAuthSessionModel]:
"""Create a new OAuth session"""
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
current_time = int(time.time())
id = str(uuid.uuid4())
@@ -129,91 +130,126 @@ class OAuthSessionTable:
)
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
db.expunge(result) # Detach so dict swap is never flushed
result.token = token # Return decrypted token
return OAuthSessionModel.model_validate(result)
# Make a copy of the model data before closing session
model = OAuthSessionModel(
id=result.id,
user_id=result.user_id,
provider=result.provider,
token=token, # Return decrypted token
expires_at=result.expires_at,
created_at=result.created_at,
updated_at=result.updated_at,
)
return model
else:
return None
except Exception as e:
log.error(f'Error creating OAuth session: {e}')
return None
def get_session_by_id(self, session_id: str, db: Optional[Session] = None) -> Optional[OAuthSessionModel]:
async def get_session_by_id(self, session_id: str, db: Optional[AsyncSession] = None) -> Optional[OAuthSessionModel]:
"""Get OAuth session by ID"""
try:
with get_db_context(db) as db:
session = db.query(OAuthSession).filter_by(id=session_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(OAuthSession).filter_by(id=session_id))
session = result.scalars().first()
if session:
db.expunge(session)
session.token = self._decrypt_token(session.token)
return OAuthSessionModel.model_validate(session)
return OAuthSessionModel(
id=session.id,
user_id=session.user_id,
provider=session.provider,
token=self._decrypt_token(session.token),
expires_at=session.expires_at,
created_at=session.created_at,
updated_at=session.updated_at,
)
return None
except Exception as e:
log.error(f'Error getting OAuth session by ID: {e}')
return None
def get_session_by_id_and_user_id(
self, session_id: str, user_id: str, db: Optional[Session] = None
async def get_session_by_id_and_user_id(
self, session_id: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[OAuthSessionModel]:
"""Get OAuth session by ID and user ID"""
try:
with get_db_context(db) as db:
session = db.query(OAuthSession).filter_by(id=session_id, user_id=user_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(OAuthSession).filter_by(id=session_id, user_id=user_id))
session = result.scalars().first()
if session:
db.expunge(session)
session.token = self._decrypt_token(session.token)
return OAuthSessionModel.model_validate(session)
return OAuthSessionModel(
id=session.id,
user_id=session.user_id,
provider=session.provider,
token=self._decrypt_token(session.token),
expires_at=session.expires_at,
created_at=session.created_at,
updated_at=session.updated_at,
)
return None
except Exception as e:
log.error(f'Error getting OAuth session by ID: {e}')
return None
def get_session_by_provider_and_user_id(
self, provider: str, user_id: str, db: Optional[Session] = None
async def get_session_by_provider_and_user_id(
self, provider: str, user_id: str, db: Optional[AsyncSession] = None
) -> Optional[OAuthSessionModel]:
"""Get OAuth session by provider and user ID"""
try:
with get_db_context(db) as db:
session = (
db.query(OAuthSession)
async with get_async_db_context(db) as db:
result = await db.execute(
select(OAuthSession)
.filter_by(provider=provider, user_id=user_id)
.order_by(OAuthSession.created_at.desc())
.first()
)
session = result.scalars().first()
if session:
db.expunge(session)
session.token = self._decrypt_token(session.token)
return OAuthSessionModel.model_validate(session)
return OAuthSessionModel(
id=session.id,
user_id=session.user_id,
provider=session.provider,
token=self._decrypt_token(session.token),
expires_at=session.expires_at,
created_at=session.created_at,
updated_at=session.updated_at,
)
return None
except Exception as e:
log.error(f'Error getting OAuth session by provider and user ID: {e}')
return None
def get_sessions_by_user_id(self, user_id: str, db: Optional[Session] = None) -> List[OAuthSessionModel]:
async def get_sessions_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> List[OAuthSessionModel]:
"""Get all OAuth sessions for a user"""
try:
with get_db_context(db) as db:
sessions = db.query(OAuthSession).filter_by(user_id=user_id).all()
async with get_async_db_context(db) as db:
result = await db.execute(select(OAuthSession).filter_by(user_id=user_id))
sessions = result.scalars().all()
results = []
for session in sessions:
try:
db.expunge(session)
session.token = self._decrypt_token(session.token)
results.append(OAuthSessionModel.model_validate(session))
results.append(OAuthSessionModel(
id=session.id,
user_id=session.user_id,
provider=session.provider,
token=self._decrypt_token(session.token),
expires_at=session.expires_at,
created_at=session.created_at,
updated_at=session.updated_at,
))
except Exception as e:
log.warning(
f'Skipping OAuth session {session.id} due to decryption failure, deleting corrupted session: {type(e).__name__}: {e}'
)
db.query(OAuthSession).filter_by(id=session.id).delete()
db.commit()
await db.execute(delete(OAuthSession).filter_by(id=session.id))
await db.commit()
return results
@@ -221,62 +257,69 @@ class OAuthSessionTable:
log.error(f'Error getting OAuth sessions by user ID: {e}')
return []
def update_session_by_id(
self, session_id: str, token: dict, db: Optional[Session] = None
async def update_session_by_id(
self, session_id: str, token: dict, db: Optional[AsyncSession] = None
) -> Optional[OAuthSessionModel]:
"""Update OAuth session tokens"""
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
current_time = int(time.time())
db.query(OAuthSession).filter_by(id=session_id).update(
{
'token': self._encrypt_token(token),
'expires_at': token.get('expires_at'),
'updated_at': current_time,
}
await db.execute(
update(OAuthSession).filter_by(id=session_id).values(
token=self._encrypt_token(token),
expires_at=token.get('expires_at'),
updated_at=current_time,
)
)
db.commit()
session = db.query(OAuthSession).filter_by(id=session_id).first()
await db.commit()
result = await db.execute(select(OAuthSession).filter_by(id=session_id))
session = result.scalars().first()
if session:
db.expunge(session)
session.token = self._decrypt_token(session.token)
return OAuthSessionModel.model_validate(session)
return OAuthSessionModel(
id=session.id,
user_id=session.user_id,
provider=session.provider,
token=self._decrypt_token(session.token),
expires_at=session.expires_at,
created_at=session.created_at,
updated_at=session.updated_at,
)
return None
except Exception as e:
log.error(f'Error updating OAuth session tokens: {e}')
return None
def delete_session_by_id(self, session_id: str, db: Optional[Session] = None) -> bool:
async def delete_session_by_id(self, session_id: str, db: Optional[AsyncSession] = None) -> bool:
"""Delete an OAuth session"""
try:
with get_db_context(db) as db:
result = db.query(OAuthSession).filter_by(id=session_id).delete()
db.commit()
return result > 0
async with get_async_db_context(db) as db:
result = await db.execute(delete(OAuthSession).filter_by(id=session_id))
await db.commit()
return result.rowcount > 0
except Exception as e:
log.error(f'Error deleting OAuth session: {e}')
return False
def delete_sessions_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
async def delete_sessions_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> bool:
"""Delete all OAuth sessions for a user"""
try:
with get_db_context(db) as db:
result = db.query(OAuthSession).filter_by(user_id=user_id).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(OAuthSession).filter_by(user_id=user_id))
await db.commit()
return True
except Exception as e:
log.error(f'Error deleting OAuth sessions by user ID: {e}')
return False
def delete_sessions_by_provider(self, provider: str, db: Optional[Session] = None) -> bool:
async def delete_sessions_by_provider(self, provider: str, db: Optional[AsyncSession] = None) -> bool:
"""Delete all OAuth sessions for a provider"""
try:
with get_db_context(db) as db:
db.query(OAuthSession).filter_by(provider=provider).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(OAuthSession).filter_by(provider=provider))
await db.commit()
return True
except Exception as e:
log.error(f'Error deleting OAuth sessions by provider {provider}: {e}')
backend/open_webui/models/prompt_history.py
+55 -45
View File
@@ -6,8 +6,9 @@ from typing import Optional
import json
import difflib
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db_context
from sqlalchemy import select, delete, func
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from open_webui.models.users import Users, UserResponse
from pydantic import BaseModel, ConfigDict
@@ -49,17 +50,17 @@ class PromptHistoryResponse(PromptHistoryModel):
class PromptHistoryTable:
def create_history_entry(
async def create_history_entry(
self,
prompt_id: str,
snapshot: dict,
user_id: str,
parent_id: Optional[str] = None,
commit_message: Optional[str] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[PromptHistoryModel]:
"""Create a new history entry (commit) for a prompt."""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
history = PromptHistory(
id=str(uuid.uuid4()),
prompt_id=prompt_id,
@@ -70,31 +71,31 @@ class PromptHistoryTable:
created_at=int(time.time()),
)
db.add(history)
db.commit()
db.refresh(history)
await db.commit()
await db.refresh(history)
return PromptHistoryModel.model_validate(history)
def get_history_by_prompt_id(
async def get_history_by_prompt_id(
self,
prompt_id: str,
limit: int = 50,
offset: int = 0,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[PromptHistoryResponse]:
"""Get all history entries for a prompt, ordered by created_at desc."""
with get_db_context(db) as db:
entries = (
db.query(PromptHistory)
async with get_async_db_context(db) as db:
result = await db.execute(
select(PromptHistory)
.filter(PromptHistory.prompt_id == prompt_id)
.order_by(PromptHistory.created_at.desc())
.offset(offset)
.limit(limit)
.all()
)
entries = result.scalars().all()
# Get user info for each entry
user_ids = list(set(e.user_id for e in entries))
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 [
@@ -105,54 +106,61 @@ class PromptHistoryTable:
for entry in entries
]
def get_history_entry_by_id(
async def get_history_entry_by_id(
self,
history_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[PromptHistoryModel]:
"""Get a specific history entry by ID."""
with get_db_context(db) as db:
entry = db.query(PromptHistory).filter(PromptHistory.id == history_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(PromptHistory).filter(PromptHistory.id == history_id))
entry = result.scalars().first()
if entry:
return PromptHistoryModel.model_validate(entry)
return None
def get_latest_history_entry(
async def get_latest_history_entry(
self,
prompt_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[PromptHistoryModel]:
"""Get the most recent history entry for a prompt."""
with get_db_context(db) as db:
entry = (
db.query(PromptHistory)
async with get_async_db_context(db) as db:
result = await db.execute(
select(PromptHistory)
.filter(PromptHistory.prompt_id == prompt_id)
.order_by(PromptHistory.created_at.desc())
.first()
.limit(1)
)
entry = result.scalars().first()
if entry:
return PromptHistoryModel.model_validate(entry)
return None
def get_history_count(
async def get_history_count(
self,
prompt_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> int:
"""Get the number of history entries for a prompt."""
with get_db_context(db) as db:
return db.query(PromptHistory).filter(PromptHistory.prompt_id == prompt_id).count()
async with get_async_db_context(db) as db:
result = await db.execute(
select(func.count()).select_from(PromptHistory).filter(PromptHistory.prompt_id == prompt_id)
)
return result.scalar()
def compute_diff(
async def compute_diff(
self,
from_id: str,
to_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[dict]:
"""Compute diff between two history entries."""
with get_db_context(db) as db:
from_entry = db.query(PromptHistory).filter(PromptHistory.id == from_id).first()
to_entry = db.query(PromptHistory).filter(PromptHistory.id == to_id).first()
async with get_async_db_context(db) as db:
result_from = await db.execute(select(PromptHistory).filter(PromptHistory.id == from_id))
from_entry = result_from.scalars().first()
result_to = await db.execute(select(PromptHistory).filter(PromptHistory.id == to_id))
to_entry = result_to.scalars().first()
if not from_entry or not to_entry:
return None
@@ -183,37 +191,39 @@ class PromptHistoryTable:
'name_changed': from_snapshot.get('name') != to_snapshot.get('name'),
}
def delete_history_by_prompt_id(
async def delete_history_by_prompt_id(
self,
prompt_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
"""Delete all history entries for a prompt."""
with get_db_context(db) as db:
db.query(PromptHistory).filter(PromptHistory.prompt_id == prompt_id).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(PromptHistory).filter(PromptHistory.prompt_id == prompt_id))
await db.commit()
return True
def delete_history_entry(
async def delete_history_entry(
self,
history_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
"""Delete a history entry and reparent its children to grandparent."""
with get_db_context(db) as db:
entry = db.query(PromptHistory).filter_by(id=history_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(PromptHistory).filter_by(id=history_id))
entry = result.scalars().first()
if not entry:
return False
# Find children that reference this entry as parent
children = db.query(PromptHistory).filter_by(parent_id=history_id).all()
children_result = await db.execute(select(PromptHistory).filter_by(parent_id=history_id))
children = children_result.scalars().all()
# Reparent children to grandparent
for child in children:
child.parent_id = entry.parent_id
db.delete(entry)
db.commit()
await db.delete(entry)
await db.commit()
return True
backend/open_webui/models/prompts.py
+151 -133
View File
@@ -2,16 +2,17 @@ import time
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete, update, or_, func, cast, String
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.groups import Groups
from open_webui.models.users import Users, UserResponse
from open_webui.models.users import Users, User, UserModel, UserResponse
from open_webui.models.prompt_history import PromptHistories
from open_webui.models.access_grants import AccessGrantModel, AccessGrants
from pydantic import BaseModel, ConfigDict, Field
from sqlalchemy import BigInteger, Boolean, Column, String, Text, JSON, or_, func, cast
from sqlalchemy import BigInteger, Boolean, Column, Text, JSON
####################
# Prompts DB Schema
@@ -92,23 +93,23 @@ class PromptForm(BaseModel):
class PromptsTable:
def _get_access_grants(self, prompt_id: str, db: Optional[Session] = None) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource('prompt', prompt_id, db=db)
async def _get_access_grants(self, prompt_id: str, db: Optional[AsyncSession] = None) -> list[AccessGrantModel]:
return await AccessGrants.get_grants_by_resource('prompt', prompt_id, db=db)
def _to_prompt_model(
async def _to_prompt_model(
self,
prompt: Prompt,
access_grants: Optional[list[AccessGrantModel]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> PromptModel:
prompt_data = PromptModel.model_validate(prompt).model_dump(exclude={'access_grants'})
prompt_data['access_grants'] = (
access_grants if access_grants is not None else self._get_access_grants(prompt_data['id'], db=db)
access_grants if access_grants is not None else await self._get_access_grants(prompt_data['id'], db=db)
)
return PromptModel.model_validate(prompt_data)
def insert_new_prompt(
self, user_id: str, form_data: PromptForm, db: Optional[Session] = None
async def insert_new_prompt(
self, user_id: str, form_data: PromptForm, db: Optional[AsyncSession] = None
) -> Optional[PromptModel]:
now = int(time.time())
prompt_id = str(uuid.uuid4())
@@ -129,15 +130,15 @@ class PromptsTable:
)
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
result = Prompt(**prompt.model_dump(exclude={'access_grants'}))
db.add(result)
db.commit()
db.refresh(result)
AccessGrants.set_access_grants('prompt', prompt_id, form_data.access_grants, db=db)
await db.commit()
await db.refresh(result)
await AccessGrants.set_access_grants('prompt', prompt_id, form_data.access_grants, db=db)
if result:
current_access_grants = self._get_access_grants(prompt_id, db=db)
current_access_grants = await self._get_access_grants(prompt_id, db=db)
snapshot = {
'name': form_data.name,
'content': form_data.content,
@@ -148,7 +149,7 @@ class PromptsTable:
'access_grants': [grant.model_dump() for grant in current_access_grants],
}
history_entry = PromptHistories.create_history_entry(
history_entry = await PromptHistories.create_history_entry(
prompt_id=prompt_id,
snapshot=snapshot,
user_id=user_id,
@@ -160,46 +161,51 @@ class PromptsTable:
# Set the initial version as the production version
if history_entry:
result.version_id = history_entry.id
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
return self._to_prompt_model(result, db=db)
return await self._to_prompt_model(result, db=db)
else:
return None
except Exception:
return None
def get_prompt_by_id(self, prompt_id: str, db: Optional[Session] = None) -> Optional[PromptModel]:
async def get_prompt_by_id(self, prompt_id: str, db: Optional[AsyncSession] = None) -> Optional[PromptModel]:
"""Get prompt by UUID."""
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(id=prompt_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(id=prompt_id))
prompt = result.scalars().first()
if prompt:
return self._to_prompt_model(prompt, db=db)
return await self._to_prompt_model(prompt, db=db)
return None
except Exception:
return None
def get_prompt_by_command(self, command: str, db: Optional[Session] = None) -> Optional[PromptModel]:
async def get_prompt_by_command(self, command: str, db: Optional[AsyncSession] = None) -> Optional[PromptModel]:
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(command=command).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(command=command))
prompt = result.scalars().first()
if prompt:
return self._to_prompt_model(prompt, db=db)
return await self._to_prompt_model(prompt, db=db)
return None
except Exception:
return None
def get_prompts(self, db: Optional[Session] = None) -> list[PromptUserResponse]:
with get_db_context(db) as db:
all_prompts = db.query(Prompt).filter(Prompt.is_active == True).order_by(Prompt.updated_at.desc()).all()
async def get_prompts(self, db: Optional[AsyncSession] = None) -> list[PromptUserResponse]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(Prompt).filter(Prompt.is_active == True).order_by(Prompt.updated_at.desc())
)
all_prompts = result.scalars().all()
user_ids = list(set(prompt.user_id for prompt in all_prompts))
prompt_ids = [prompt.id for prompt in all_prompts]
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}
grants_map = AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=db)
prompts = []
for prompt in all_prompts:
@@ -207,11 +213,11 @@ class PromptsTable:
prompts.append(
PromptUserResponse.model_validate(
{
**self._to_prompt_model(
**(await self._to_prompt_model(
prompt,
access_grants=grants_map.get(prompt.id, []),
db=db,
).model_dump(),
)).model_dump(),
'user': user.model_dump() if user else None,
}
)
@@ -219,44 +225,44 @@ class PromptsTable:
return prompts
def get_prompts_by_user_id(
self, user_id: str, permission: str = 'write', db: Optional[Session] = None
async def get_prompts_by_user_id(
self, user_id: str, permission: str = 'write', db: Optional[AsyncSession] = None
) -> list[PromptUserResponse]:
prompts = self.get_prompts(db=db)
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
prompts = await self.get_prompts(db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
return [
prompt
for prompt in prompts
if prompt.user_id == user_id
or AccessGrants.has_access(
result = []
for prompt in prompts:
if prompt.user_id == user_id:
result.append(prompt)
elif await AccessGrants.has_access(
user_id=user_id,
resource_type='prompt',
resource_id=prompt.id,
permission=permission,
user_group_ids=user_group_ids,
db=db,
)
]
):
result.append(prompt)
return result
def search_prompts(
async def search_prompts(
self,
user_id: str,
filter: dict = {},
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> PromptListResponse:
with get_db_context(db) as db:
from open_webui.models.users import User, UserModel
async with get_async_db_context(db) as db:
# Join with User table for user filtering and sorting
query = db.query(Prompt, User).outerjoin(User, User.id == Prompt.user_id)
stmt = select(Prompt, User).outerjoin(User, User.id == Prompt.user_id)
if filter:
query_key = filter.get('query')
if query_key:
query = query.filter(
stmt = stmt.filter(
or_(
Prompt.name.ilike(f'%{query_key}%'),
Prompt.command.ilike(f'%{query_key}%'),
@@ -268,14 +274,14 @@ class PromptsTable:
view_option = filter.get('view_option')
if view_option == 'created':
query = query.filter(Prompt.user_id == user_id)
stmt = stmt.filter(Prompt.user_id == user_id)
elif view_option == 'shared':
query = query.filter(Prompt.user_id != user_id)
stmt = stmt.filter(Prompt.user_id != user_id)
# Apply access grant filtering
query = AccessGrants.has_permission_filter(
stmt = AccessGrants.has_permission_filter(
db=db,
query=query,
query=stmt,
DocumentModel=Prompt,
filter=filter,
resource_type='prompt',
@@ -287,75 +293,80 @@ class PromptsTable:
# Search for tag in JSON array field
like_pattern = f'%"{tag.lower()}"%'
tags_text = func.lower(cast(Prompt.tags, String))
query = query.filter(tags_text.like(like_pattern))
stmt = stmt.filter(tags_text.like(like_pattern))
order_by = filter.get('order_by')
direction = filter.get('direction')
if order_by == 'name':
if direction == 'asc':
query = query.order_by(Prompt.name.asc())
stmt = stmt.order_by(Prompt.name.asc())
else:
query = query.order_by(Prompt.name.desc())
stmt = stmt.order_by(Prompt.name.desc())
elif order_by == 'created_at':
if direction == 'asc':
query = query.order_by(Prompt.created_at.asc())
stmt = stmt.order_by(Prompt.created_at.asc())
else:
query = query.order_by(Prompt.created_at.desc())
stmt = stmt.order_by(Prompt.created_at.desc())
elif order_by == 'updated_at':
if direction == 'asc':
query = query.order_by(Prompt.updated_at.asc())
stmt = stmt.order_by(Prompt.updated_at.asc())
else:
query = query.order_by(Prompt.updated_at.desc())
stmt = stmt.order_by(Prompt.updated_at.desc())
else:
query = query.order_by(Prompt.updated_at.desc())
stmt = stmt.order_by(Prompt.updated_at.desc())
else:
query = query.order_by(Prompt.updated_at.desc())
stmt = stmt.order_by(Prompt.updated_at.desc())
# Count BEFORE pagination
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit:
query = query.limit(limit)
stmt = stmt.limit(limit)
items = query.all()
result = await db.execute(stmt)
items = result.all()
prompt_ids = [prompt.id for prompt, _ in items]
grants_map = AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('prompt', prompt_ids, db=db)
prompts = []
for prompt, user in items:
prompts.append(
PromptUserResponse(
**self._to_prompt_model(
**(await self._to_prompt_model(
prompt,
access_grants=grants_map.get(prompt.id, []),
db=db,
).model_dump(),
)).model_dump(),
user=(UserResponse(**UserModel.model_validate(user).model_dump()) if user else None),
)
)
return PromptListResponse(items=prompts, total=total)
def update_prompt_by_command(
async def update_prompt_by_command(
self,
command: str,
form_data: PromptForm,
user_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[PromptModel]:
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(command=command).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(command=command))
prompt = result.scalars().first()
if not prompt:
return None
latest_history = PromptHistories.get_latest_history_entry(prompt.id, db=db)
latest_history = await PromptHistories.get_latest_history_entry(prompt.id, db=db)
parent_id = latest_history.id if latest_history else None
current_access_grants = self._get_access_grants(prompt.id, db=db)
current_access_grants = await self._get_access_grants(prompt.id, db=db)
# Check if content changed to decide on history creation
content_changed = (
@@ -371,10 +382,10 @@ class PromptsTable:
prompt.meta = form_data.meta or prompt.meta
prompt.updated_at = int(time.time())
if form_data.access_grants is not None:
AccessGrants.set_access_grants('prompt', prompt.id, form_data.access_grants, db=db)
current_access_grants = self._get_access_grants(prompt.id, db=db)
await AccessGrants.set_access_grants('prompt', prompt.id, form_data.access_grants, db=db)
current_access_grants = await self._get_access_grants(prompt.id, db=db)
db.commit()
await db.commit()
# Create history entry only if content changed
if content_changed:
@@ -387,7 +398,7 @@ class PromptsTable:
'access_grants': [grant.model_dump() for grant in current_access_grants],
}
history_entry = PromptHistories.create_history_entry(
history_entry = await PromptHistories.create_history_entry(
prompt_id=prompt.id,
snapshot=snapshot,
user_id=user_id,
@@ -399,28 +410,29 @@ class PromptsTable:
# Set as production if flag is True (default)
if form_data.is_production and history_entry:
prompt.version_id = history_entry.id
db.commit()
await db.commit()
return self._to_prompt_model(prompt, db=db)
return await self._to_prompt_model(prompt, db=db)
except Exception:
return None
def update_prompt_by_id(
async def update_prompt_by_id(
self,
prompt_id: str,
form_data: PromptForm,
user_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[PromptModel]:
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(id=prompt_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(id=prompt_id))
prompt = result.scalars().first()
if not prompt:
return None
latest_history = PromptHistories.get_latest_history_entry(prompt.id, db=db)
latest_history = await PromptHistories.get_latest_history_entry(prompt.id, db=db)
parent_id = latest_history.id if latest_history else None
current_access_grants = self._get_access_grants(prompt.id, db=db)
current_access_grants = await self._get_access_grants(prompt.id, db=db)
# Check if content changed to decide on history creation
content_changed = (
@@ -442,12 +454,12 @@ class PromptsTable:
prompt.tags = form_data.tags
if form_data.access_grants is not None:
AccessGrants.set_access_grants('prompt', prompt.id, form_data.access_grants, db=db)
current_access_grants = self._get_access_grants(prompt.id, db=db)
await AccessGrants.set_access_grants('prompt', prompt.id, form_data.access_grants, db=db)
current_access_grants = await self._get_access_grants(prompt.id, db=db)
prompt.updated_at = int(time.time())
db.commit()
await db.commit()
# Create history entry only if content changed
if content_changed:
@@ -461,7 +473,7 @@ class PromptsTable:
'access_grants': [grant.model_dump() for grant in current_access_grants],
}
history_entry = PromptHistories.create_history_entry(
history_entry = await PromptHistories.create_history_entry(
prompt_id=prompt.id,
snapshot=snapshot,
user_id=user_id,
@@ -473,24 +485,25 @@ class PromptsTable:
# Set as production if flag is True (default)
if form_data.is_production and history_entry:
prompt.version_id = history_entry.id
db.commit()
await db.commit()
return self._to_prompt_model(prompt, db=db)
return await self._to_prompt_model(prompt, db=db)
except Exception:
return None
def update_prompt_metadata(
async def update_prompt_metadata(
self,
prompt_id: str,
name: str,
command: str,
tags: Optional[list[str]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[PromptModel]:
"""Update only name, command, and tags (no history created)."""
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(id=prompt_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(id=prompt_id))
prompt = result.scalars().first()
if not prompt:
return None
@@ -501,26 +514,27 @@ class PromptsTable:
prompt.tags = tags
prompt.updated_at = int(time.time())
db.commit()
await db.commit()
return self._to_prompt_model(prompt, db=db)
return await self._to_prompt_model(prompt, db=db)
except Exception:
return None
def update_prompt_version(
async def update_prompt_version(
self,
prompt_id: str,
version_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[PromptModel]:
"""Set the active version of a prompt and restore content from that version's snapshot."""
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(id=prompt_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(id=prompt_id))
prompt = result.scalars().first()
if not prompt:
return None
history_entry = PromptHistories.get_history_entry_by_id(version_id, db=db)
history_entry = await PromptHistories.get_history_entry_by_id(version_id, db=db)
if not history_entry:
return None
@@ -537,63 +551,67 @@ class PromptsTable:
prompt.version_id = version_id
prompt.updated_at = int(time.time())
db.commit()
await db.commit()
return self._to_prompt_model(prompt, db=db)
return await self._to_prompt_model(prompt, db=db)
except Exception:
return None
def toggle_prompt_active(self, prompt_id: str, db: Optional[Session] = None) -> Optional[PromptModel]:
async def toggle_prompt_active(self, prompt_id: str, db: Optional[AsyncSession] = None) -> Optional[PromptModel]:
"""Toggle the is_active flag on a prompt."""
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(id=prompt_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(id=prompt_id))
prompt = result.scalars().first()
if prompt:
prompt.is_active = not prompt.is_active
prompt.updated_at = int(time.time())
db.commit()
db.refresh(prompt)
return self._to_prompt_model(prompt, db=db)
await db.commit()
await db.refresh(prompt)
return await self._to_prompt_model(prompt, db=db)
return None
except Exception:
return None
def delete_prompt_by_command(self, command: str, db: Optional[Session] = None) -> bool:
async def delete_prompt_by_command(self, command: str, db: Optional[AsyncSession] = None) -> bool:
"""Permanently delete a prompt and its history."""
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(command=command).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(command=command))
prompt = result.scalars().first()
if prompt:
PromptHistories.delete_history_by_prompt_id(prompt.id, db=db)
AccessGrants.revoke_all_access('prompt', prompt.id, db=db)
await PromptHistories.delete_history_by_prompt_id(prompt.id, db=db)
await AccessGrants.revoke_all_access('prompt', prompt.id, db=db)
db.delete(prompt)
db.commit()
await db.delete(prompt)
await db.commit()
return True
return False
except Exception:
return False
def delete_prompt_by_id(self, prompt_id: str, db: Optional[Session] = None) -> bool:
async def delete_prompt_by_id(self, prompt_id: str, db: Optional[AsyncSession] = None) -> bool:
"""Permanently delete a prompt and its history."""
try:
with get_db_context(db) as db:
prompt = db.query(Prompt).filter_by(id=prompt_id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(id=prompt_id))
prompt = result.scalars().first()
if prompt:
PromptHistories.delete_history_by_prompt_id(prompt.id, db=db)
AccessGrants.revoke_all_access('prompt', prompt.id, db=db)
await PromptHistories.delete_history_by_prompt_id(prompt.id, db=db)
await AccessGrants.revoke_all_access('prompt', prompt.id, db=db)
db.delete(prompt)
db.commit()
await db.delete(prompt)
await db.commit()
return True
return False
except Exception:
return False
def get_tags(self, db: Optional[Session] = None) -> list[str]:
async def get_tags(self, db: Optional[AsyncSession] = None) -> list[str]:
try:
with get_db_context(db) as db:
prompts = db.query(Prompt).filter_by(is_active=True).all()
async with get_async_db_context(db) as db:
result = await db.execute(select(Prompt).filter_by(is_active=True))
prompts = result.scalars().all()
tags = set()
for prompt in prompts:
if prompt.tags:
backend/open_webui/models/skills.py
+88 -80
View File
@@ -2,14 +2,15 @@ import logging
import time
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db, get_db_context
from open_webui.models.users import Users, UserResponse
from sqlalchemy import select, delete, update, or_
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from open_webui.models.users import Users, User, UserModel, UserResponse
from open_webui.models.groups import Groups
from open_webui.models.access_grants import AccessGrantModel, AccessGrants
from pydantic import BaseModel, ConfigDict, Field
from sqlalchemy import JSON, BigInteger, Boolean, Column, String, Text, or_
from sqlalchemy import JSON, BigInteger, Boolean, Column, String, Text, func
log = logging.getLogger(__name__)
@@ -105,28 +106,28 @@ class SkillAccessListResponse(BaseModel):
class SkillsTable:
def _get_access_grants(self, skill_id: str, db: Optional[Session] = None) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource('skill', skill_id, db=db)
async def _get_access_grants(self, skill_id: str, db: Optional[AsyncSession] = None) -> list[AccessGrantModel]:
return await AccessGrants.get_grants_by_resource('skill', skill_id, db=db)
def _to_skill_model(
async def _to_skill_model(
self,
skill: Skill,
access_grants: Optional[list[AccessGrantModel]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> SkillModel:
skill_data = SkillModel.model_validate(skill).model_dump(exclude={'access_grants'})
skill_data['access_grants'] = (
access_grants if access_grants is not None else self._get_access_grants(skill_data['id'], db=db)
access_grants if access_grants is not None else await self._get_access_grants(skill_data['id'], db=db)
)
return SkillModel.model_validate(skill_data)
def insert_new_skill(
async def insert_new_skill(
self,
user_id: str,
form_data: SkillForm,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[SkillModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
try:
result = Skill(
**{
@@ -137,43 +138,45 @@ class SkillsTable:
}
)
db.add(result)
db.commit()
db.refresh(result)
AccessGrants.set_access_grants('skill', result.id, form_data.access_grants, db=db)
await db.commit()
await db.refresh(result)
await AccessGrants.set_access_grants('skill', result.id, form_data.access_grants, db=db)
if result:
return self._to_skill_model(result, db=db)
return await self._to_skill_model(result, db=db)
else:
return None
except Exception as e:
log.exception(f'Error creating a new skill: {e}')
return None
def get_skill_by_id(self, id: str, db: Optional[Session] = None) -> Optional[SkillModel]:
async def get_skill_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[SkillModel]:
try:
with get_db_context(db) as db:
skill = db.get(Skill, id)
return self._to_skill_model(skill, db=db) if skill else None
async with get_async_db_context(db) as db:
skill = await db.get(Skill, id)
return await self._to_skill_model(skill, db=db) if skill else None
except Exception:
return None
def get_skill_by_name(self, name: str, db: Optional[Session] = None) -> Optional[SkillModel]:
async def get_skill_by_name(self, name: str, db: Optional[AsyncSession] = None) -> Optional[SkillModel]:
try:
with get_db_context(db) as db:
skill = db.query(Skill).filter_by(name=name).first()
return self._to_skill_model(skill, db=db) if skill else None
async with get_async_db_context(db) as db:
result = await db.execute(select(Skill).filter_by(name=name))
skill = result.scalars().first()
return await self._to_skill_model(skill, db=db) if skill else None
except Exception:
return None
def get_skills(self, db: Optional[Session] = None) -> list[SkillUserModel]:
with get_db_context(db) as db:
all_skills = db.query(Skill).order_by(Skill.updated_at.desc()).all()
async def get_skills(self, db: Optional[AsyncSession] = None) -> list[SkillUserModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Skill).order_by(Skill.updated_at.desc()))
all_skills = result.scalars().all()
user_ids = list(set(skill.user_id for skill in all_skills))
skill_ids = [skill.id for skill in all_skills]
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}
grants_map = AccessGrants.get_grants_by_resources('skill', skill_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('skill', skill_ids, db=db)
skills = []
for skill in all_skills:
@@ -181,56 +184,56 @@ class SkillsTable:
skills.append(
SkillUserModel.model_validate(
{
**self._to_skill_model(
**(await self._to_skill_model(
skill,
access_grants=grants_map.get(skill.id, []),
db=db,
).model_dump(),
)).model_dump(),
'user': user.model_dump() if user else None,
}
)
)
return skills
def get_skills_by_user_id(
self, user_id: str, permission: str = 'write', db: Optional[Session] = None
async def get_skills_by_user_id(
self, user_id: str, permission: str = 'write', db: Optional[AsyncSession] = None
) -> list[SkillUserModel]:
skills = self.get_skills(db=db)
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
skills = await self.get_skills(db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
return [
skill
for skill in skills
if skill.user_id == user_id
or AccessGrants.has_access(
result = []
for skill in skills:
if skill.user_id == user_id:
result.append(skill)
elif await AccessGrants.has_access(
user_id=user_id,
resource_type='skill',
resource_id=skill.id,
permission=permission,
user_group_ids=user_group_ids,
db=db,
)
]
):
result.append(skill)
return result
def search_skills(
async def search_skills(
self,
user_id: str,
filter: dict = {},
skip: int = 0,
limit: int = 30,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> SkillListResponse:
try:
with get_db_context(db) as db:
from open_webui.models.users import User, UserModel
async with get_async_db_context(db) as db:
# Join with User table for user filtering
query = db.query(Skill, User).outerjoin(User, User.id == Skill.user_id)
stmt = select(Skill, User).outerjoin(User, User.id == Skill.user_id)
if filter:
query_key = filter.get('query')
if query_key:
query = query.filter(
stmt = stmt.filter(
or_(
Skill.name.ilike(f'%{query_key}%'),
Skill.description.ilike(f'%{query_key}%'),
@@ -242,44 +245,48 @@ class SkillsTable:
view_option = filter.get('view_option')
if view_option == 'created':
query = query.filter(Skill.user_id == user_id)
stmt = stmt.filter(Skill.user_id == user_id)
elif view_option == 'shared':
query = query.filter(Skill.user_id != user_id)
stmt = stmt.filter(Skill.user_id != user_id)
# Apply access grant filtering
query = AccessGrants.has_permission_filter(
stmt = AccessGrants.has_permission_filter(
db=db,
query=query,
query=stmt,
DocumentModel=Skill,
filter=filter,
resource_type='skill',
permission='read',
)
query = query.order_by(Skill.updated_at.desc())
stmt = stmt.order_by(Skill.updated_at.desc())
# Count BEFORE pagination
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
if skip:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit:
query = query.limit(limit)
stmt = stmt.limit(limit)
items = query.all()
result = await db.execute(stmt)
items = result.all()
skill_ids = [skill.id for skill, _ in items]
grants_map = AccessGrants.get_grants_by_resources('skill', skill_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('skill', skill_ids, db=db)
skills = []
for skill, user in items:
skills.append(
SkillUserResponse(
**self._to_skill_model(
**(await self._to_skill_model(
skill,
access_grants=grants_map.get(skill.id, []),
db=db,
).model_dump(),
)).model_dump(),
user=(UserResponse(**UserModel.model_validate(user).model_dump()) if user else None),
)
)
@@ -289,43 +296,44 @@ class SkillsTable:
log.exception(f'Error searching skills: {e}')
return SkillListResponse(items=[], total=0)
def update_skill_by_id(self, id: str, updated: dict, db: Optional[Session] = None) -> Optional[SkillModel]:
async def update_skill_by_id(self, id: str, updated: dict, db: Optional[AsyncSession] = None) -> Optional[SkillModel]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
access_grants = updated.pop('access_grants', None)
db.query(Skill).filter_by(id=id).update({**updated, 'updated_at': int(time.time())})
db.commit()
await db.execute(update(Skill).filter_by(id=id).values(**updated, updated_at=int(time.time())))
await db.commit()
if access_grants is not None:
AccessGrants.set_access_grants('skill', id, access_grants, db=db)
await AccessGrants.set_access_grants('skill', id, access_grants, db=db)
skill = db.query(Skill).get(id)
db.refresh(skill)
return self._to_skill_model(skill, db=db)
skill = await db.get(Skill, id)
await db.refresh(skill)
return await self._to_skill_model(skill, db=db)
except Exception:
return None
def toggle_skill_by_id(self, id: str, db: Optional[Session] = None) -> Optional[SkillModel]:
with get_db_context(db) as db:
async def toggle_skill_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[SkillModel]:
async with get_async_db_context(db) as db:
try:
skill = db.query(Skill).filter_by(id=id).first()
result = await db.execute(select(Skill).filter_by(id=id))
skill = result.scalars().first()
if not skill:
return None
skill.is_active = not skill.is_active
skill.updated_at = int(time.time())
db.commit()
db.refresh(skill)
await db.commit()
await db.refresh(skill)
return self._to_skill_model(skill, db=db)
return await self._to_skill_model(skill, db=db)
except Exception:
return None
def delete_skill_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_skill_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
AccessGrants.revoke_all_access('skill', id, db=db)
db.query(Skill).filter_by(id=id).delete()
db.commit()
async with get_async_db_context(db) as db:
await AccessGrants.revoke_all_access('skill', id, db=db)
await db.execute(delete(Skill).filter_by(id=id))
await db.commit()
return True
except Exception:
backend/open_webui/models/tags.py
+34 -32
View File
@@ -3,8 +3,9 @@ import time
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from sqlalchemy import select, delete
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from pydantic import BaseModel, ConfigDict
@@ -53,15 +54,15 @@ class TagChatIdForm(BaseModel):
class TagTable:
def insert_new_tag(self, name: str, user_id: str, db: Optional[Session] = None) -> Optional[TagModel]:
with get_db_context(db) as db:
async def insert_new_tag(self, name: str, user_id: str, db: Optional[AsyncSession] = None) -> Optional[TagModel]:
async with get_async_db_context(db) as db:
id = name.replace(' ', '_').lower()
tag = TagModel(**{'id': id, 'user_id': user_id, 'name': name})
try:
result = Tag(**tag.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return TagModel.model_validate(result)
else:
@@ -70,64 +71,65 @@ class TagTable:
log.exception(f'Error inserting a new tag: {e}')
return None
def get_tag_by_name_and_user_id(self, name: str, user_id: str, db: Optional[Session] = None) -> Optional[TagModel]:
async def get_tag_by_name_and_user_id(self, name: str, user_id: str, db: Optional[AsyncSession] = None) -> Optional[TagModel]:
try:
id = name.replace(' ', '_').lower()
with get_db_context(db) as db:
tag = db.query(Tag).filter_by(id=id, user_id=user_id).first()
return TagModel.model_validate(tag)
async with get_async_db_context(db) as db:
result = await db.execute(select(Tag).filter_by(id=id, user_id=user_id))
tag = result.scalars().first()
return TagModel.model_validate(tag) if tag else None
except Exception:
return None
def get_tags_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[TagModel]:
with get_db_context(db) as db:
return [TagModel.model_validate(tag) for tag in (db.query(Tag).filter_by(user_id=user_id).all())]
async def get_tags_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[TagModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Tag).filter_by(user_id=user_id))
return [TagModel.model_validate(tag) for tag in result.scalars().all()]
def get_tags_by_ids_and_user_id(self, ids: list[str], user_id: str, db: Optional[Session] = None) -> list[TagModel]:
with get_db_context(db) as db:
return [
TagModel.model_validate(tag)
for tag in (db.query(Tag).filter(Tag.id.in_(ids), Tag.user_id == user_id).all())
]
async def get_tags_by_ids_and_user_id(self, ids: list[str], user_id: str, db: Optional[AsyncSession] = None) -> list[TagModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Tag).filter(Tag.id.in_(ids), Tag.user_id == user_id))
return [TagModel.model_validate(tag) for tag in result.scalars().all()]
def delete_tag_by_name_and_user_id(self, name: str, user_id: str, db: Optional[Session] = None) -> bool:
async def delete_tag_by_name_and_user_id(self, name: str, user_id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
id = name.replace(' ', '_').lower()
res = db.query(Tag).filter_by(id=id, user_id=user_id).delete()
log.debug(f'res: {res}')
db.commit()
result = await db.execute(delete(Tag).filter_by(id=id, user_id=user_id))
log.debug(f'res: {result.rowcount}')
await db.commit()
return True
except Exception as e:
log.error(f'delete_tag: {e}')
return False
def delete_tags_by_ids_and_user_id(self, ids: list[str], user_id: str, db: Optional[Session] = None) -> bool:
async def delete_tags_by_ids_and_user_id(self, ids: list[str], user_id: str, db: Optional[AsyncSession] = None) -> bool:
"""Delete all tags whose id is in *ids* for the given user, in one query."""
if not ids:
return True
try:
with get_db_context(db) as db:
db.query(Tag).filter(Tag.id.in_(ids), Tag.user_id == user_id).delete(synchronize_session=False)
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(Tag).filter(Tag.id.in_(ids), Tag.user_id == user_id))
await db.commit()
return True
except Exception as e:
log.error(f'delete_tags_by_ids: {e}')
return False
def ensure_tags_exist(self, names: list[str], user_id: str, db: Optional[Session] = None) -> None:
async def ensure_tags_exist(self, names: list[str], user_id: str, db: Optional[AsyncSession] = None) -> None:
"""Create tag rows for any *names* that don't already exist for *user_id*."""
if not names:
return
ids = [n.replace(' ', '_').lower() for n in names]
with get_db_context(db) as db:
existing = {t.id for t in db.query(Tag.id).filter(Tag.id.in_(ids), Tag.user_id == user_id).all()}
async with get_async_db_context(db) as db:
result = await db.execute(select(Tag.id).filter(Tag.id.in_(ids), Tag.user_id == user_id))
existing = {row[0] for row in result.all()}
new_tags = [
Tag(id=tag_id, name=name, user_id=user_id) for tag_id, name in zip(ids, names) if tag_id not in existing
]
if new_tags:
db.add_all(new_tags)
db.commit()
await db.commit()
Tags = TagTable()
backend/open_webui/models/tools.py
+73 -65
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, UserResponse
from open_webui.models.groups import Groups
from open_webui.models.access_grants import AccessGrantModel, AccessGrants
@@ -97,29 +98,29 @@ class ToolValves(BaseModel):
class ToolsTable:
def _get_access_grants(self, tool_id: str, db: Optional[Session] = None) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource('tool', tool_id, db=db)
async def _get_access_grants(self, tool_id: str, db: Optional[AsyncSession] = None) -> list[AccessGrantModel]:
return await AccessGrants.get_grants_by_resource('tool', tool_id, db=db)
def _to_tool_model(
async def _to_tool_model(
self,
tool: Tool,
access_grants: Optional[list[AccessGrantModel]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> ToolModel:
tool_data = ToolModel.model_validate(tool).model_dump(exclude={'access_grants'})
tool_data['access_grants'] = (
access_grants if access_grants is not None else self._get_access_grants(tool_data['id'], db=db)
access_grants if access_grants is not None else await self._get_access_grants(tool_data['id'], db=db)
)
return ToolModel.model_validate(tool_data)
def insert_new_tool(
async def insert_new_tool(
self,
user_id: str,
form_data: ToolForm,
specs: list[dict],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[ToolModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
try:
result = Tool(
**{
@@ -131,38 +132,39 @@ class ToolsTable:
}
)
db.add(result)
db.commit()
db.refresh(result)
AccessGrants.set_access_grants('tool', result.id, form_data.access_grants, db=db)
await db.commit()
await db.refresh(result)
await AccessGrants.set_access_grants('tool', result.id, form_data.access_grants, db=db)
if result:
return self._to_tool_model(result, db=db)
return await self._to_tool_model(result, db=db)
else:
return None
except Exception as e:
log.exception(f'Error creating a new tool: {e}')
return None
def get_tool_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ToolModel]:
async def get_tool_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ToolModel]:
try:
with get_db_context(db) as db:
tool = db.get(Tool, id)
return self._to_tool_model(tool, db=db) if tool else None
async with get_async_db_context(db) as db:
tool = await db.get(Tool, id)
return await self._to_tool_model(tool, db=db) if tool else None
except Exception:
return None
def get_tools(self, defer_content: bool = False, db: Optional[Session] = None) -> list[ToolUserModel]:
with get_db_context(db) as db:
query = db.query(Tool).order_by(Tool.updated_at.desc())
async def get_tools(self, defer_content: bool = False, db: Optional[AsyncSession] = None) -> list[ToolUserModel]:
async with get_async_db_context(db) as db:
stmt = select(Tool).order_by(Tool.updated_at.desc())
if defer_content:
query = query.options(defer(Tool.content), defer(Tool.specs))
all_tools = query.all()
stmt = stmt
result = await db.execute(stmt)
all_tools = result.scalars().all()
user_ids = list(set(tool.user_id for tool in all_tools))
tool_ids = [tool.id for tool in all_tools]
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}
grants_map = AccessGrants.get_grants_by_resources('tool', tool_ids, db=db)
grants_map = await AccessGrants.get_grants_by_resources('tool', tool_ids, db=db)
tools = []
for tool in all_tools:
@@ -170,62 +172,66 @@ class ToolsTable:
tools.append(
ToolUserModel.model_validate(
{
**self._to_tool_model(
**(await self._to_tool_model(
tool,
access_grants=grants_map.get(tool.id, []),
db=db,
).model_dump(),
)).model_dump(),
'user': user.model_dump() if user else None,
}
)
)
return tools
def get_tools_by_user_id(
async def get_tools_by_user_id(
self,
user_id: str,
permission: str = 'write',
defer_content: bool = False,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[ToolUserModel]:
tools = self.get_tools(defer_content=defer_content, db=db)
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
tools = await self.get_tools(defer_content=defer_content, db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
return [
tool
for tool in tools
if tool.user_id == user_id
or AccessGrants.has_access(
result = []
for tool in tools:
if tool.user_id == user_id:
result.append(tool)
elif await AccessGrants.has_access(
user_id=user_id,
resource_type='tool',
resource_id=tool.id,
permission=permission,
user_group_ids=user_group_ids,
db=db,
)
]
):
result.append(tool)
return result
def get_tool_valves_by_id(self, id: str, db: Optional[Session] = None) -> Optional[dict]:
async def get_tool_valves_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[dict]:
try:
with get_db_context(db) as db:
tool = db.get(Tool, id)
async with get_async_db_context(db) as db:
tool = await db.get(Tool, id)
return tool.valves if tool.valves else {}
except Exception as e:
log.exception(f'Error getting tool valves by id {id}')
return None
def update_tool_valves_by_id(self, id: str, valves: dict, db: Optional[Session] = None) -> Optional[ToolValves]:
async def update_tool_valves_by_id(self, id: str, valves: dict, db: Optional[AsyncSession] = None) -> Optional[ToolValves]:
try:
with get_db_context(db) as db:
db.query(Tool).filter_by(id=id).update({'valves': valves, 'updated_at': int(time.time())})
db.commit()
return self.get_tool_by_id(id, db=db)
async with get_async_db_context(db) as db:
await db.execute(
update(Tool).filter_by(id=id).values(valves=valves, updated_at=int(time.time()))
)
await db.commit()
return await self.get_tool_by_id(id, db=db)
except Exception:
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 "tools" and "valves" settings
@@ -239,11 +245,11 @@ class ToolsTable:
log.exception(f'Error getting user values by id {id} and user_id {user_id}: {e}')
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 "tools" and "valves" settings
@@ -255,34 +261,36 @@ class ToolsTable:
user_settings['tools']['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['tools']['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_tool_by_id(self, id: str, updated: dict, db: Optional[Session] = None) -> Optional[ToolModel]:
async def update_tool_by_id(self, id: str, updated: dict, db: Optional[AsyncSession] = None) -> Optional[ToolModel]:
try:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
access_grants = updated.pop('access_grants', None)
db.query(Tool).filter_by(id=id).update({**updated, 'updated_at': int(time.time())})
db.commit()
await db.execute(
update(Tool).filter_by(id=id).values(**updated, updated_at=int(time.time()))
)
await db.commit()
if access_grants is not None:
AccessGrants.set_access_grants('tool', id, access_grants, db=db)
await AccessGrants.set_access_grants('tool', id, access_grants, db=db)
tool = db.query(Tool).get(id)
db.refresh(tool)
return self._to_tool_model(tool, db=db)
tool = await db.get(Tool, id)
await db.refresh(tool)
return await self._to_tool_model(tool, db=db)
except Exception:
return None
def delete_tool_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_tool_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
AccessGrants.revoke_all_access('tool', id, db=db)
db.query(Tool).filter_by(id=id).delete()
db.commit()
async with get_async_db_context(db) as db:
await AccessGrants.revoke_all_access('tool', id, db=db)
await db.execute(delete(Tool).filter_by(id=id))
await db.commit()
return True
except Exception:
backend/open_webui/models/users.py
+207 -198
View File
@@ -1,20 +1,15 @@
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, func, or_, case, exists
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.env import DATABASE_USER_ACTIVE_STATUS_UPDATE_INTERVAL
from open_webui.models.chats import Chats
from open_webui.models.groups import Groups, GroupMember
from open_webui.models.channels import ChannelMember
from open_webui.utils.misc import throttle
from open_webui.utils.validate import validate_profile_image_url
from pydantic import BaseModel, ConfigDict, field_validator, model_validator
from sqlalchemy import (
BigInteger,
@@ -24,11 +19,8 @@ from sqlalchemy import (
Boolean,
Text,
Date,
exists,
select,
cast,
)
from sqlalchemy import or_, case, func
from sqlalchemy.dialects.postgresql import JSONB
import datetime
@@ -39,13 +31,11 @@ import datetime
# daily bread of every session. Let none go hungry.
####################
class UserSettings(BaseModel):
ui: Optional[dict] = {}
model_config = ConfigDict(extra='allow')
pass
class User(Base):
__tablename__ = 'user'
@@ -79,7 +69,6 @@ class User(Base):
updated_at = Column(BigInteger)
created_at = Column(BigInteger)
class UserModel(BaseModel):
id: str
@@ -120,13 +109,11 @@ class UserModel(BaseModel):
self.profile_image_url = f'/api/v1/users/{self.id}/profile/image'
return self
class UserStatusModel(UserModel):
is_active: bool = False
model_config = ConfigDict(from_attributes=True)
class ApiKey(Base):
__tablename__ = 'api_key'
@@ -139,7 +126,6 @@ class ApiKey(Base):
created_at = Column(BigInteger, nullable=False)
updated_at = Column(BigInteger, nullable=False)
class ApiKeyModel(BaseModel):
id: str
user_id: str
@@ -152,12 +138,10 @@ class ApiKeyModel(BaseModel):
model_config = ConfigDict(from_attributes=True)
####################
# Forms
####################
class UpdateProfileForm(BaseModel):
profile_image_url: str
name: str
@@ -170,31 +154,25 @@ class UpdateProfileForm(BaseModel):
def check_profile_image_url(cls, v: str) -> str:
return validate_profile_image_url(v)
class UserGroupIdsModel(UserModel):
group_ids: list[str] = []
class UserModelResponse(UserModel):
model_config = ConfigDict(extra='allow')
class UserListResponse(BaseModel):
users: list[UserModelResponse]
total: int
class UserGroupIdsListResponse(BaseModel):
users: list[UserGroupIdsModel]
total: int
class UserStatus(BaseModel):
status_emoji: Optional[str] = None
status_message: Optional[str] = None
status_expires_at: Optional[int] = None
class UserInfoResponse(UserStatus):
id: str
name: str
@@ -204,48 +182,39 @@ class UserInfoResponse(UserStatus):
groups: Optional[list] = []
is_active: bool = False
class UserIdNameResponse(BaseModel):
id: str
name: str
class UserIdNameStatusResponse(UserStatus):
id: str
name: str
is_active: Optional[bool] = None
class UserInfoListResponse(BaseModel):
users: list[UserInfoResponse]
total: int
class UserIdNameListResponse(BaseModel):
users: list[UserIdNameResponse]
total: int
class UserNameResponse(BaseModel):
id: str
name: str
role: str
class UserResponse(UserNameResponse):
email: str
class UserProfileImageResponse(UserNameResponse):
email: str
profile_image_url: str
class UserRoleUpdateForm(BaseModel):
id: str
role: str
class UserUpdateForm(BaseModel):
role: str
name: str
@@ -258,9 +227,8 @@ class UserUpdateForm(BaseModel):
def check_profile_image_url(cls, v: str) -> str:
return validate_profile_image_url(v)
class UsersTable:
def insert_new_user(
async def insert_new_user(
self,
id: str,
name: str,
@@ -269,9 +237,9 @@ class UsersTable:
role: str = 'pending',
username: Optional[str] = None,
oauth: Optional[dict] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[UserModel]:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
user = UserModel(
**{
'id': id,
@@ -288,87 +256,98 @@ class UsersTable:
)
result = User(**user.model_dump())
db.add(result)
db.commit()
db.refresh(result)
await db.commit()
await db.refresh(result)
if result:
return user
else:
return None
def get_user_by_id(self, id: str, db: Optional[Session] = None) -> Optional[UserModel]:
async def get_user_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
except Exception:
return None
def get_user_by_api_key(self, api_key: str, db: Optional[Session] = None) -> Optional[UserModel]:
try:
with get_db_context(db) as db:
user = db.query(User).join(ApiKey, User.id == ApiKey.user_id).filter(ApiKey.key == api_key).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
return UserModel.model_validate(user) if user else None
except Exception:
return None
def get_user_by_email(self, email: str, db: Optional[Session] = None) -> Optional[UserModel]:
async def get_user_by_api_key(self, api_key: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
try:
with get_db_context(db) as db:
user = db.query(User).filter(func.lower(User.email) == email.lower()).first()
async with get_async_db_context(db) as db:
result = await db.execute(
select(User).join(ApiKey, User.id == ApiKey.user_id).filter(ApiKey.key == api_key)
)
user = result.scalars().first()
return UserModel.model_validate(user) if user else None
except Exception:
return None
def get_user_by_oauth_sub(self, provider: str, sub: str, db: Optional[Session] = None) -> Optional[UserModel]:
async def get_user_by_email(self, email: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
try:
with get_db_context(db) as db: # type: Session
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter(func.lower(User.email) == email.lower()))
user = result.scalars().first()
return UserModel.model_validate(user) if user else None
except Exception:
return None
async def get_user_by_oauth_sub(self, provider: str, sub: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
try:
async with get_async_db_context(db) as db:
dialect_name = db.bind.dialect.name
query = db.query(User)
stmt = select(User)
if dialect_name == 'sqlite':
query = query.filter(User.oauth.contains({provider: {'sub': sub}}))
stmt = stmt.filter(User.oauth.contains({provider: {'sub': sub}}))
elif dialect_name == 'postgresql':
query = query.filter(User.oauth[provider].cast(JSONB)['sub'].astext == sub)
stmt = stmt.filter(User.oauth[provider].cast(JSONB)['sub'].astext == sub)
user = query.first()
result = await db.execute(stmt)
user = result.scalars().first()
return UserModel.model_validate(user) if user else None
except Exception as e:
# You may want to log the exception here
return None
def get_user_by_scim_external_id(
self, provider: str, external_id: str, db: Optional[Session] = None
async def get_user_by_scim_external_id(
self, provider: str, external_id: str, db: Optional[AsyncSession] = None
) -> Optional[UserModel]:
try:
with get_db_context(db) as db: # type: Session
async with get_async_db_context(db) as db:
dialect_name = db.bind.dialect.name
query = db.query(User)
stmt = select(User)
if dialect_name == 'sqlite':
query = query.filter(User.scim.contains({provider: {'external_id': external_id}}))
stmt = stmt.filter(User.scim.contains({provider: {'external_id': external_id}}))
elif dialect_name == 'postgresql':
query = query.filter(User.scim[provider].cast(JSONB)['external_id'].astext == external_id)
stmt = stmt.filter(User.scim[provider].cast(JSONB)['external_id'].astext == external_id)
user = query.first()
result = await db.execute(stmt)
user = result.scalars().first()
return UserModel.model_validate(user) if user else None
except Exception:
return None
def get_users(
async def get_users(
self,
filter: Optional[dict] = None,
skip: Optional[int] = None,
limit: Optional[int] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Import here to avoid circular imports
from open_webui.models.groups import GroupMember
from open_webui.models.channels import ChannelMember
# Join GroupMember so we can order by group_id when requested
query = db.query(User).options(defer(User.profile_image_url))
stmt = select(User)
if filter:
query_key = filter.get('query')
if query_key:
query = query.filter(
stmt = stmt.filter(
or_(
User.name.ilike(f'%{query_key}%'),
User.email.ilike(f'%{query_key}%'),
@@ -377,7 +356,7 @@ class UsersTable:
channel_id = filter.get('channel_id')
if channel_id:
query = query.filter(
stmt = stmt.filter(
exists(
select(ChannelMember.id).where(
ChannelMember.user_id == User.id,
@@ -395,10 +374,10 @@ class UsersTable:
return {'users': [], 'total': 0}
if user_ids:
query = query.filter(User.id.in_(user_ids))
stmt = stmt.filter(User.id.in_(user_ids))
if group_ids:
query = query.filter(
stmt = stmt.filter(
exists(
select(GroupMember.id).where(
GroupMember.user_id == User.id,
@@ -413,9 +392,9 @@ class UsersTable:
exclude_roles = [role[1:] for role in roles if role.startswith('!')]
if include_roles:
query = query.filter(User.role.in_(include_roles))
stmt = stmt.filter(User.role.in_(include_roles))
if exclude_roles:
query = query.filter(~User.role.in_(exclude_roles))
stmt = stmt.filter(~User.role.in_(exclude_roles))
order_by = filter.get('order_by')
direction = filter.get('direction')
@@ -435,99 +414,111 @@ class UsersTable:
group_sort = case((membership_exists, 1), else_=0)
if direction == 'asc':
query = query.order_by(group_sort.asc(), User.name.asc())
stmt = stmt.order_by(group_sort.asc(), User.name.asc())
else:
query = query.order_by(group_sort.desc(), User.name.asc())
stmt = stmt.order_by(group_sort.desc(), User.name.asc())
elif order_by == 'name':
if direction == 'asc':
query = query.order_by(User.name.asc())
stmt = stmt.order_by(User.name.asc())
else:
query = query.order_by(User.name.desc())
stmt = stmt.order_by(User.name.desc())
elif order_by == 'email':
if direction == 'asc':
query = query.order_by(User.email.asc())
stmt = stmt.order_by(User.email.asc())
else:
query = query.order_by(User.email.desc())
stmt = stmt.order_by(User.email.desc())
elif order_by == 'created_at':
if direction == 'asc':
query = query.order_by(User.created_at.asc())
stmt = stmt.order_by(User.created_at.asc())
else:
query = query.order_by(User.created_at.desc())
stmt = stmt.order_by(User.created_at.desc())
elif order_by == 'last_active_at':
if direction == 'asc':
query = query.order_by(User.last_active_at.asc())
stmt = stmt.order_by(User.last_active_at.asc())
else:
query = query.order_by(User.last_active_at.desc())
stmt = stmt.order_by(User.last_active_at.desc())
elif order_by == 'updated_at':
if direction == 'asc':
query = query.order_by(User.updated_at.asc())
stmt = stmt.order_by(User.updated_at.asc())
else:
query = query.order_by(User.updated_at.desc())
stmt = stmt.order_by(User.updated_at.desc())
elif order_by == 'role':
if direction == 'asc':
query = query.order_by(User.role.asc())
stmt = stmt.order_by(User.role.asc())
else:
query = query.order_by(User.role.desc())
stmt = stmt.order_by(User.role.desc())
else:
query = query.order_by(User.created_at.desc())
stmt = stmt.order_by(User.created_at.desc())
# Count BEFORE pagination
total = query.count()
count_result = await db.execute(
select(func.count()).select_from(stmt.subquery())
)
total = count_result.scalar()
# correct pagination logic
if skip is not None:
query = query.offset(skip)
stmt = stmt.offset(skip)
if limit is not None:
query = query.limit(limit)
stmt = stmt.limit(limit)
users = query.all()
result = await db.execute(stmt)
users = result.scalars().all()
return {
'users': [UserModel.model_validate(user) for user in users],
'total': total,
}
def get_users_by_group_id(self, group_id: str, db: Optional[Session] = None) -> list[UserModel]:
with get_db_context(db) as db:
users = (
db.query(User)
.options(defer(User.profile_image_url))
async def get_users_by_group_id(self, group_id: str, db: Optional[AsyncSession] = None) -> list[UserModel]:
async with get_async_db_context(db) as db:
from open_webui.models.groups import GroupMember
result = await db.execute(
select(User)
.join(GroupMember, User.id == GroupMember.user_id)
.filter(GroupMember.group_id == group_id)
.all()
)
users = result.scalars().all()
return [UserModel.model_validate(user) for user in users]
def get_users_by_user_ids(self, user_ids: list[str], db: Optional[Session] = None) -> list[UserStatusModel]:
with get_db_context(db) as db:
users = db.query(User).options(defer(User.profile_image_url)).filter(User.id.in_(user_ids)).all()
async def get_users_by_user_ids(self, user_ids: list[str], db: Optional[AsyncSession] = None) -> list[UserStatusModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
select(User).filter(User.id.in_(user_ids))
)
users = result.scalars().all()
return [UserModel.model_validate(user) for user in users]
def get_num_users(self, db: Optional[Session] = None) -> Optional[int]:
with get_db_context(db) as db:
return db.query(User).count()
async def get_num_users(self, db: Optional[AsyncSession] = None) -> Optional[int]:
async with get_async_db_context(db) as db:
result = await db.execute(select(func.count()).select_from(User))
return result.scalar()
def has_users(self, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
return db.query(db.query(User).exists()).scalar()
async def has_users(self, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(select(exists(select(User))))
return result.scalar()
def get_first_user(self, db: Optional[Session] = None) -> UserModel:
async def get_first_user(self, db: Optional[AsyncSession] = None) -> UserModel:
try:
with get_db_context(db) as db:
user = db.query(User).order_by(User.created_at).first()
return UserModel.model_validate(user)
async with get_async_db_context(db) as db:
result = await db.execute(select(User).order_by(User.created_at).limit(1))
user = result.scalars().first()
return UserModel.model_validate(user) if user else None
except Exception:
return None
def get_user_webhook_url_by_id(self, id: str, db: Optional[Session] = None) -> Optional[str]:
async def get_user_webhook_url_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[str]:
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
if user.settings is None:
return None
@@ -536,68 +527,73 @@ class UsersTable:
except Exception:
return None
def get_num_users_active_today(self, db: Optional[Session] = None) -> Optional[int]:
with get_db_context(db) as db:
async def get_num_users_active_today(self, db: Optional[AsyncSession] = None) -> Optional[int]:
async with get_async_db_context(db) as db:
current_timestamp = int(datetime.datetime.now().timestamp())
today_midnight_timestamp = current_timestamp - (current_timestamp % 86400)
query = db.query(User).filter(User.last_active_at > today_midnight_timestamp)
return query.count()
result = await db.execute(
select(func.count()).select_from(User).filter(User.last_active_at > today_midnight_timestamp)
)
return result.scalar()
def update_user_role_by_id(self, id: str, role: str, db: Optional[Session] = None) -> Optional[UserModel]:
async def update_user_role_by_id(self, id: str, role: str, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
if not user:
return None
user.role = role
db.commit()
db.refresh(user)
await db.commit()
await db.refresh(user)
return UserModel.model_validate(user)
except Exception:
return None
def update_user_status_by_id(
self, id: str, form_data: UserStatus, db: Optional[Session] = None
async def update_user_status_by_id(
self, id: str, form_data: UserStatus, db: Optional[AsyncSession] = None
) -> Optional[UserModel]:
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
if not user:
return None
for key, value in form_data.model_dump(exclude_none=True).items():
setattr(user, key, value)
db.commit()
db.refresh(user)
await db.commit()
await db.refresh(user)
return UserModel.model_validate(user)
except Exception:
return None
def update_user_profile_image_url_by_id(
self, id: str, profile_image_url: str, db: Optional[Session] = None
async def update_user_profile_image_url_by_id(
self, id: str, profile_image_url: str, db: Optional[AsyncSession] = None
) -> Optional[UserModel]:
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
if not user:
return None
user.profile_image_url = profile_image_url
db.commit()
db.refresh(user)
await db.commit()
await db.refresh(user)
return UserModel.model_validate(user)
except Exception:
return None
@throttle(DATABASE_USER_ACTIVE_STATUS_UPDATE_INTERVAL)
def update_last_active_by_id(self, id: str, db: Optional[Session] = None) -> None:
async def update_last_active_by_id(self, id: str, db: Optional[AsyncSession] = None) -> None:
try:
with get_db_context(db) as db:
db.query(User).filter_by(id=id).update({'last_active_at': int(time.time())})
db.commit()
async with get_async_db_context(db) as db:
await db.execute(update(User).filter_by(id=id).values(last_active_at=int(time.time())))
await db.commit()
except Exception:
pass
def update_user_oauth_by_id(
self, id: str, provider: str, sub: str, db: Optional[Session] = None
async def update_user_oauth_by_id(
self, id: str, provider: str, sub: str, db: Optional[AsyncSession] = None
) -> Optional[UserModel]:
"""
Update or insert an OAuth provider/sub pair into the user's oauth JSON field.
@@ -608,8 +604,9 @@ class UsersTable:
}
"""
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
if not user:
return None
@@ -620,20 +617,20 @@ class UsersTable:
oauth[provider] = {'sub': sub}
# Persist updated JSON
db.query(User).filter_by(id=id).update({'oauth': oauth})
db.commit()
await db.execute(update(User).filter_by(id=id).values(oauth=oauth))
await db.commit()
return UserModel.model_validate(user)
except Exception:
return None
def update_user_scim_by_id(
async def update_user_scim_by_id(
self,
id: str,
provider: str,
external_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[UserModel]:
"""
Update or insert a SCIM provider/external_id pair into the user's scim JSON field.
@@ -644,41 +641,44 @@ class UsersTable:
}
"""
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
if not user:
return None
scim = user.scim or {}
scim[provider] = {'external_id': external_id}
db.query(User).filter_by(id=id).update({'scim': scim})
db.commit()
await db.execute(update(User).filter_by(id=id).values(scim=scim))
await db.commit()
return UserModel.model_validate(user)
except Exception:
return None
def update_user_by_id(self, id: str, updated: dict, db: Optional[Session] = None) -> Optional[UserModel]:
async def update_user_by_id(self, id: str, updated: dict, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
if not user:
return None
for key, value in updated.items():
setattr(user, key, value)
db.commit()
db.refresh(user)
await db.commit()
await db.refresh(user)
return UserModel.model_validate(user)
except Exception as e:
print(e)
return None
def update_user_settings_by_id(self, id: str, updated: dict, db: Optional[Session] = None) -> Optional[UserModel]:
async def update_user_settings_by_id(self, id: str, updated: dict, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
try:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
if not user:
return None
@@ -689,26 +689,30 @@ class UsersTable:
user_settings.update(updated)
db.query(User).filter_by(id=id).update({'settings': user_settings})
db.commit()
await db.execute(update(User).filter_by(id=id).values(settings=user_settings))
await db.commit()
user = db.query(User).filter_by(id=id).first()
result = await db.execute(select(User).filter_by(id=id))
user = result.scalars().first()
return UserModel.model_validate(user)
except Exception:
return None
def delete_user_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_user_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
from open_webui.models.groups import Groups
from open_webui.models.chats import Chats
# Remove User from Groups
Groups.remove_user_from_all_groups(id)
await Groups.remove_user_from_all_groups(id)
# Delete User Chats
result = Chats.delete_chats_by_user_id(id, db=db)
result = await Chats.delete_chats_by_user_id(id, db=db)
if result:
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Delete User
db.query(User).filter_by(id=id).delete()
db.commit()
await db.execute(delete(User).filter_by(id=id))
await db.commit()
return True
else:
@@ -716,19 +720,20 @@ class UsersTable:
except Exception:
return False
def get_user_api_key_by_id(self, id: str, db: Optional[Session] = None) -> Optional[str]:
async def get_user_api_key_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[str]:
try:
with get_db_context(db) as db:
api_key = db.query(ApiKey).filter_by(user_id=id).first()
async with get_async_db_context(db) as db:
result = await db.execute(select(ApiKey).filter_by(user_id=id))
api_key = result.scalars().first()
return api_key.key if api_key else None
except Exception:
return None
def update_user_api_key_by_id(self, id: str, api_key: str, db: Optional[Session] = None) -> bool:
async def update_user_api_key_by_id(self, id: str, api_key: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
db.query(ApiKey).filter_by(user_id=id).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(ApiKey).filter_by(user_id=id))
await db.commit()
now = int(time.time())
new_api_key = ApiKey(
@@ -739,41 +744,45 @@ class UsersTable:
updated_at=now,
)
db.add(new_api_key)
db.commit()
await db.commit()
return True
except Exception:
return False
def delete_user_api_key_by_id(self, id: str, db: Optional[Session] = None) -> bool:
async def delete_user_api_key_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool:
try:
with get_db_context(db) as db:
db.query(ApiKey).filter_by(user_id=id).delete()
db.commit()
async with get_async_db_context(db) as db:
await db.execute(delete(ApiKey).filter_by(user_id=id))
await db.commit()
return True
except Exception:
return False
def get_valid_user_ids(self, user_ids: list[str], db: Optional[Session] = None) -> list[str]:
with get_db_context(db) as db:
users = db.query(User).filter(User.id.in_(user_ids)).all()
async def get_valid_user_ids(self, user_ids: list[str], db: Optional[AsyncSession] = None) -> list[str]:
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter(User.id.in_(user_ids)))
users = result.scalars().all()
return [user.id for user in users]
def get_super_admin_user(self, db: Optional[Session] = None) -> Optional[UserModel]:
with get_db_context(db) as db:
user = db.query(User).filter_by(role='admin').first()
async def get_super_admin_user(self, db: Optional[AsyncSession] = None) -> Optional[UserModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(role='admin').limit(1))
user = result.scalars().first()
if user:
return UserModel.model_validate(user)
else:
return None
def get_active_user_count(self, db: Optional[Session] = None) -> int:
with get_db_context(db) as db:
async def get_active_user_count(self, db: Optional[AsyncSession] = None) -> int:
async with get_async_db_context(db) as db:
# Consider user active if last_active_at within the last 3 minutes
three_minutes_ago = int(time.time()) - 180
count = db.query(User).filter(User.last_active_at >= three_minutes_ago).count()
return count
result = await db.execute(
select(func.count()).select_from(User).filter(User.last_active_at >= three_minutes_ago)
)
return result.scalar()
@staticmethod
def is_active(user: UserModel) -> bool:
@@ -783,14 +792,14 @@ class UsersTable:
return user.last_active_at >= three_minutes_ago
return False
def is_user_active(self, user_id: str, db: Optional[Session] = None) -> bool:
with get_db_context(db) as db:
user = db.query(User).filter_by(id=user_id).first()
async def is_user_active(self, user_id: str, db: Optional[AsyncSession] = None) -> bool:
async with get_async_db_context(db) as db:
result = await db.execute(select(User).filter_by(id=user_id))
user = result.scalars().first()
if user and user.last_active_at:
# Consider user active if last_active_at within the last 3 minutes
three_minutes_ago = int(time.time()) - 180
return user.last_active_at >= three_minutes_ago
return False
Users = UsersTable()