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

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-03-17 17:58:01 -05:00
parent fcf7208352
commit de3317e26b
220 changed files with 17200 additions and 22836 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/models/chat_messages.py
+112 -175
View File
@@ -47,13 +47,11 @@ def _normalize_timestamp(timestamp: int) -> float:
class ChatMessage(Base):
__tablename__ = "chat_message"
__tablename__ = 'chat_message'
# Identity
id = Column(Text, primary_key=True)
chat_id = Column(
Text, ForeignKey("chat.id", ondelete="CASCADE"), nullable=False, index=True
)
chat_id = Column(Text, ForeignKey('chat.id', ondelete='CASCADE'), nullable=False, index=True)
user_id = Column(Text, index=True)
# Structure
@@ -85,9 +83,9 @@ class ChatMessage(Base):
updated_at = Column(BigInteger)
__table_args__ = (
Index("chat_message_chat_parent_idx", "chat_id", "parent_id"),
Index("chat_message_model_created_idx", "model_id", "created_at"),
Index("chat_message_user_created_idx", "user_id", "created_at"),
Index('chat_message_chat_parent_idx', 'chat_id', 'parent_id'),
Index('chat_message_model_created_idx', 'model_id', 'created_at'),
Index('chat_message_user_created_idx', 'user_id', 'created_at'),
)
@@ -135,43 +133,41 @@ class ChatMessageTable:
"""Insert or update a chat message."""
with get_db_context(db) as db:
now = int(time.time())
timestamp = data.get("timestamp", now)
timestamp = data.get('timestamp', now)
# Use composite ID: {chat_id}-{message_id}
composite_id = f"{chat_id}-{message_id}"
composite_id = f'{chat_id}-{message_id}'
existing = db.get(ChatMessage, composite_id)
if existing:
# Update existing
if "role" in data:
existing.role = data["role"]
if "parent_id" in data:
existing.parent_id = data.get("parent_id") or data.get("parentId")
if "content" in data:
existing.content = data.get("content")
if "output" in data:
existing.output = data.get("output")
if "model_id" in data or "model" in data:
existing.model_id = data.get("model_id") or data.get("model")
if "files" in data:
existing.files = data.get("files")
if "sources" in data:
existing.sources = data.get("sources")
if "embeds" in data:
existing.embeds = data.get("embeds")
if "done" in data:
existing.done = data.get("done", True)
if "status_history" in data or "statusHistory" in data:
existing.status_history = data.get("status_history") or data.get(
"statusHistory"
)
if "error" in data:
existing.error = data.get("error")
if 'role' in data:
existing.role = data['role']
if 'parent_id' in data:
existing.parent_id = data.get('parent_id') or data.get('parentId')
if 'content' in data:
existing.content = data.get('content')
if 'output' in data:
existing.output = data.get('output')
if 'model_id' in data or 'model' in data:
existing.model_id = data.get('model_id') or data.get('model')
if 'files' in data:
existing.files = data.get('files')
if 'sources' in data:
existing.sources = data.get('sources')
if 'embeds' in data:
existing.embeds = data.get('embeds')
if 'done' in data:
existing.done = data.get('done', True)
if 'status_history' in data or 'statusHistory' in data:
existing.status_history = data.get('status_history') or data.get('statusHistory')
if 'error' in data:
existing.error = data.get('error')
# Extract usage - check direct field first, then info.usage
usage = data.get("usage")
usage = data.get('usage')
if not usage:
info = data.get("info", {})
usage = info.get("usage") if info else None
info = data.get('info', {})
usage = info.get('usage') if info else None
if usage:
existing.usage = usage
existing.updated_at = now
@@ -181,26 +177,25 @@ class ChatMessageTable:
else:
# Insert new
# Extract usage - check direct field first, then info.usage
usage = data.get("usage")
usage = data.get('usage')
if not usage:
info = data.get("info", {})
usage = info.get("usage") if info else None
info = data.get('info', {})
usage = info.get('usage') if info else None
message = ChatMessage(
id=composite_id,
chat_id=chat_id,
user_id=user_id,
role=data.get("role", "user"),
parent_id=data.get("parent_id") or data.get("parentId"),
content=data.get("content"),
output=data.get("output"),
model_id=data.get("model_id") or data.get("model"),
files=data.get("files"),
sources=data.get("sources"),
embeds=data.get("embeds"),
done=data.get("done", True),
status_history=data.get("status_history")
or data.get("statusHistory"),
error=data.get("error"),
role=data.get('role', 'user'),
parent_id=data.get('parent_id') or data.get('parentId'),
content=data.get('content'),
output=data.get('output'),
model_id=data.get('model_id') or data.get('model'),
files=data.get('files'),
sources=data.get('sources'),
embeds=data.get('embeds'),
done=data.get('done', True),
status_history=data.get('status_history') or data.get('statusHistory'),
error=data.get('error'),
usage=usage,
created_at=timestamp,
updated_at=now,
@@ -210,23 +205,14 @@ class ChatMessageTable:
db.refresh(message)
return ChatMessageModel.model_validate(message)
def get_message_by_id(
self, id: str, db: Optional[Session] = None
) -> Optional[ChatMessageModel]:
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)
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]:
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()
)
messages = db.query(ChatMessage).filter_by(chat_id=chat_id).order_by(ChatMessage.created_at.asc()).all()
return [ChatMessageModel.model_validate(message) for message in messages]
def get_messages_by_user_id(
@@ -262,12 +248,7 @@ class ChatMessageTable:
query = query.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()
)
messages = query.order_by(ChatMessage.created_at.desc()).offset(skip).limit(limit).all()
return [ChatMessageModel.model_validate(message) for message in messages]
def get_chat_ids_by_model_id(
@@ -284,7 +265,7 @@ class ChatMessageTable:
with get_db_context(db) as db:
query = db.query(
ChatMessage.chat_id,
func.max(ChatMessage.created_at).label("last_message_at"),
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)
@@ -303,9 +284,7 @@ class ChatMessageTable:
)
return [chat_id for chat_id, _ in chat_ids]
def delete_messages_by_chat_id(
self, chat_id: str, db: Optional[Session] = None
) -> bool:
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()
@@ -323,12 +302,10 @@ class ChatMessageTable:
from sqlalchemy import func
from open_webui.models.groups import GroupMember
query = db.query(
ChatMessage.model_id, func.count(ChatMessage.id).label("count")
).filter(
ChatMessage.role == "assistant",
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-%"),
~ChatMessage.user_id.like('shared-%'),
)
if start_date:
@@ -336,11 +313,7 @@ class ChatMessageTable:
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = (
db.query(GroupMember.user_id)
.filter(GroupMember.group_id == group_id)
.subquery()
)
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.model_id).all()
@@ -360,36 +333,32 @@ class ChatMessageTable:
dialect = db.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":
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"),
func.json_extract_path_text(ChatMessage.usage, 'input_tokens'),
Integer,
)
output_tokens = cast(
func.json_extract_path_text(ChatMessage.usage, "output_tokens"),
func.json_extract_path_text(ChatMessage.usage, 'output_tokens'),
Integer,
)
else:
raise NotImplementedError(f"Unsupported dialect: {dialect}")
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"),
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.role == 'assistant',
ChatMessage.model_id.isnot(None),
ChatMessage.usage.isnot(None),
~ChatMessage.user_id.like("shared-%"),
~ChatMessage.user_id.like('shared-%'),
)
if start_date:
@@ -397,21 +366,17 @@ class ChatMessageTable:
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = (
db.query(GroupMember.user_id)
.filter(GroupMember.group_id == group_id)
.subquery()
)
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.model_id).all()
return {
row.model_id: {
"input_tokens": row.input_tokens,
"output_tokens": row.output_tokens,
"total_tokens": row.input_tokens + row.output_tokens,
"message_count": row.message_count,
'input_tokens': row.input_tokens,
'output_tokens': row.output_tokens,
'total_tokens': row.input_tokens + row.output_tokens,
'message_count': row.message_count,
}
for row in results
}
@@ -430,36 +395,32 @@ class ChatMessageTable:
dialect = db.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":
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"),
func.json_extract_path_text(ChatMessage.usage, 'input_tokens'),
Integer,
)
output_tokens = cast(
func.json_extract_path_text(ChatMessage.usage, "output_tokens"),
func.json_extract_path_text(ChatMessage.usage, 'output_tokens'),
Integer,
)
else:
raise NotImplementedError(f"Unsupported dialect: {dialect}")
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"),
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.role == 'assistant',
ChatMessage.user_id.isnot(None),
ChatMessage.usage.isnot(None),
~ChatMessage.user_id.like("shared-%"),
~ChatMessage.user_id.like('shared-%'),
)
if start_date:
@@ -467,21 +428,17 @@ class ChatMessageTable:
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = (
db.query(GroupMember.user_id)
.filter(GroupMember.group_id == group_id)
.subquery()
)
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.user_id).all()
return {
row.user_id: {
"input_tokens": row.input_tokens,
"output_tokens": row.output_tokens,
"total_tokens": row.input_tokens + row.output_tokens,
"message_count": row.message_count,
'input_tokens': row.input_tokens,
'output_tokens': row.output_tokens,
'total_tokens': row.input_tokens + row.output_tokens,
'message_count': row.message_count,
}
for row in results
}
@@ -497,20 +454,16 @@ class ChatMessageTable:
from sqlalchemy import func
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-%"))
query = db.query(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)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = (
db.query(GroupMember.user_id)
.filter(GroupMember.group_id == group_id)
.subquery()
)
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.user_id).all()
@@ -527,20 +480,16 @@ class ChatMessageTable:
from sqlalchemy import func
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-%"))
query = db.query(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)
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = (
db.query(GroupMember.user_id)
.filter(GroupMember.group_id == group_id)
.subquery()
)
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
results = query.group_by(ChatMessage.chat_id).all()
@@ -559,9 +508,9 @@ class ChatMessageTable:
from open_webui.models.groups import GroupMember
query = db.query(ChatMessage.created_at, ChatMessage.model_id).filter(
ChatMessage.role == "assistant",
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
~ChatMessage.user_id.like("shared-%"),
~ChatMessage.user_id.like('shared-%'),
)
if start_date:
@@ -569,11 +518,7 @@ class ChatMessageTable:
if end_date:
query = query.filter(ChatMessage.created_at <= end_date)
if group_id:
group_users = (
db.query(GroupMember.user_id)
.filter(GroupMember.group_id == group_id)
.subquery()
)
group_users = db.query(GroupMember.user_id).filter(GroupMember.group_id == group_id).subquery()
query = query.filter(ChatMessage.user_id.in_(group_users))
results = query.all()
@@ -581,21 +526,17 @@ class ChatMessageTable:
# Group by date -> model -> count
daily_counts: dict[str, dict[str, int]] = {}
for timestamp, model_id in results:
date_str = datetime.fromtimestamp(
_normalize_timestamp(timestamp)
).strftime("%Y-%m-%d")
date_str = datetime.fromtimestamp(_normalize_timestamp(timestamp)).strftime('%Y-%m-%d')
if date_str not in daily_counts:
daily_counts[date_str] = {}
daily_counts[date_str][model_id] = (
daily_counts[date_str].get(model_id, 0) + 1
)
daily_counts[date_str][model_id] = daily_counts[date_str].get(model_id, 0) + 1
# Fill in missing days
if start_date and end_date:
current = datetime.fromtimestamp(_normalize_timestamp(start_date))
end_dt = datetime.fromtimestamp(_normalize_timestamp(end_date))
while current <= end_dt:
date_str = current.strftime("%Y-%m-%d")
date_str = current.strftime('%Y-%m-%d')
if date_str not in daily_counts:
daily_counts[date_str] = {}
current += timedelta(days=1)
@@ -613,9 +554,9 @@ class ChatMessageTable:
from datetime import datetime, timedelta
query = db.query(ChatMessage.created_at, ChatMessage.model_id).filter(
ChatMessage.role == "assistant",
ChatMessage.role == 'assistant',
ChatMessage.model_id.isnot(None),
~ChatMessage.user_id.like("shared-%"),
~ChatMessage.user_id.like('shared-%'),
)
if start_date:
@@ -628,23 +569,19 @@ class ChatMessageTable:
# Group by hour -> model -> count
hourly_counts: dict[str, dict[str, int]] = {}
for timestamp, model_id in results:
hour_str = datetime.fromtimestamp(
_normalize_timestamp(timestamp)
).strftime("%Y-%m-%d %H:00")
hour_str = datetime.fromtimestamp(_normalize_timestamp(timestamp)).strftime('%Y-%m-%d %H:00')
if hour_str not in hourly_counts:
hourly_counts[hour_str] = {}
hourly_counts[hour_str][model_id] = (
hourly_counts[hour_str].get(model_id, 0) + 1
)
hourly_counts[hour_str][model_id] = hourly_counts[hour_str].get(model_id, 0) + 1
# Fill in missing hours
if start_date and end_date:
current = datetime.fromtimestamp(
_normalize_timestamp(start_date)
).replace(minute=0, second=0, microsecond=0)
current = datetime.fromtimestamp(_normalize_timestamp(start_date)).replace(
minute=0, second=0, microsecond=0
)
end_dt = datetime.fromtimestamp(_normalize_timestamp(end_date))
while current <= end_dt:
hour_str = current.strftime("%Y-%m-%d %H:00")
hour_str = current.strftime('%Y-%m-%d %H:00')
if hour_str not in hourly_counts:
hourly_counts[hour_str] = {}
current += timedelta(hours=1)