refac: async db
This commit is contained in:
Timothy Jaeryang Baek
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user