mirror of
https://github.com/open-webui/open-webui.git
synced 2026-08-28 05:27:35 +00:00
perf: write a chat's messages in one transaction instead of one per message (#28806)
Saving a chat rewrote its message rows one at a time. Each message took its own session out of the pool and committed on its own, and the save endpoint hands over the entire merged history rather than only what changed, so a two hundred message chat cost two hundred sessions and two hundred commits on every save. The messages now go through a single select and a single commit. The field mapping for the insert and the update branch moved into two small helpers, so the batch and the single-message path cannot drift apart. Measured on a two hundred message chat with one message edited: 201 queries and 200 transactions before, 2 queries and 1 transaction after, ~149 ms against ~6 ms. Re-saving an unchanged history now costs one select and no writes at all. One behaviour change worth stating: a message the database cannot store used to be skipped on its own, and now costs the rest of that same save. This table is a rebuildable fast path, so the reader falls back to the history on the chat row and re-triggers the backfill, and the next save reconciles everything still present. A per-message retry was tried and dropped, because a commit that lands but still raises would re-apply the usage merge and double the recorded token counts.
This commit is contained in:
Classic298
GitHub
parent
bcb50fe7b0
commit
81fe43f210
2 changed files with 104 additions and 77 deletions
|
|
@ -206,6 +206,66 @@ class ChatMessageModel(BaseModel):
|
|||
|
||||
|
||||
class ChatMessageTable:
|
||||
@staticmethod
|
||||
def _apply_message_data(message: ChatMessage, data: dict, now: int) -> None:
|
||||
"""Overwrite only the fields the payload carries."""
|
||||
if 'role' in data:
|
||||
message.role = data['role']
|
||||
if 'parent_id' in data or 'parentId' in data:
|
||||
message.parent_id = data.get('parent_id') or data.get('parentId')
|
||||
if 'content' in data:
|
||||
message.content = data.get('content')
|
||||
if 'output' in data:
|
||||
message.output = data.get('output')
|
||||
if 'model_id' in data or 'model' in data:
|
||||
message.model_id = data.get('model_id') or data.get('model')
|
||||
if 'files' in data:
|
||||
message.files = data.get('files')
|
||||
if 'sources' in data:
|
||||
message.sources = data.get('sources')
|
||||
if 'embeds' in data:
|
||||
message.embeds = data.get('embeds')
|
||||
if 'meta' in data:
|
||||
message.meta = data.get('meta')
|
||||
if 'done' in data:
|
||||
message.done = data['done']
|
||||
if 'status_history' in data or 'statusHistory' in data:
|
||||
message.status_history = data.get('status_history') or data.get('statusHistory')
|
||||
if 'error' in data:
|
||||
message.error = data.get('error')
|
||||
if 'context_summary' in data or 'contextSummary' in data:
|
||||
message.context_summary = data.get('context_summary') or data.get('contextSummary')
|
||||
|
||||
usage = get_usage(data)
|
||||
if usage:
|
||||
existing_usage = normalize_usage(message.usage)
|
||||
message.usage = existing_usage if usage == existing_usage else merge_usage(existing_usage, usage)
|
||||
message.updated_at = now
|
||||
|
||||
@staticmethod
|
||||
def _build_message(composite_id: str, chat_id: str, user_id: str, data: dict, now: int) -> ChatMessage:
|
||||
return 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'),
|
||||
meta=data.get('meta'),
|
||||
done=data.get('done', True),
|
||||
status_history=data.get('status_history') or data.get('statusHistory'),
|
||||
error=data.get('error'),
|
||||
usage=get_usage(data),
|
||||
context_summary=data.get('context_summary') or data.get('contextSummary'),
|
||||
created_at=data.get('timestamp', now),
|
||||
updated_at=now,
|
||||
)
|
||||
|
||||
async def upsert_message(
|
||||
self,
|
||||
message_id: str,
|
||||
|
|
@ -217,76 +277,46 @@ class ChatMessageTable:
|
|||
"""Insert or update a chat message."""
|
||||
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 = await db.get(ChatMessage, composite_id)
|
||||
if existing:
|
||||
# Update existing
|
||||
if 'role' in data:
|
||||
existing.role = data['role']
|
||||
if 'parent_id' in data or 'parentId' 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 'meta' in data:
|
||||
existing.meta = data.get('meta')
|
||||
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 'context_summary' in data or 'contextSummary' in data:
|
||||
existing.context_summary = data.get('context_summary') or data.get('contextSummary')
|
||||
# Extract and normalize usage
|
||||
usage = get_usage(data)
|
||||
if usage:
|
||||
existing_usage = normalize_usage(existing.usage or {}) if existing.usage else {}
|
||||
existing.usage = existing_usage if usage == existing_usage else merge_usage(existing_usage, usage)
|
||||
existing.updated_at = now
|
||||
await db.commit()
|
||||
return ChatMessageModel.model_validate(existing)
|
||||
message = await db.get(ChatMessage, composite_id)
|
||||
if message:
|
||||
self._apply_message_data(message, data, now)
|
||||
else:
|
||||
# Insert new
|
||||
# Extract and normalize usage
|
||||
usage = get_usage(data)
|
||||
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'),
|
||||
meta=data.get('meta'),
|
||||
done=data.get('done', True),
|
||||
status_history=data.get('status_history') or data.get('statusHistory'),
|
||||
error=data.get('error'),
|
||||
usage=usage,
|
||||
context_summary=data.get('context_summary') or data.get('contextSummary'),
|
||||
created_at=timestamp,
|
||||
updated_at=now,
|
||||
)
|
||||
message = self._build_message(composite_id, chat_id, user_id, data, now)
|
||||
db.add(message)
|
||||
await db.commit()
|
||||
return ChatMessageModel.model_validate(message)
|
||||
|
||||
await db.commit()
|
||||
return ChatMessageModel.model_validate(message)
|
||||
|
||||
async def upsert_messages(
|
||||
self,
|
||||
chat_id: str,
|
||||
user_id: str,
|
||||
messages: dict[str, dict],
|
||||
db: AsyncSession | None = None,
|
||||
) -> None:
|
||||
"""Insert or update the given messages of one chat."""
|
||||
if not messages:
|
||||
return
|
||||
|
||||
async with get_async_db_context(db) as db:
|
||||
now = int(time.time())
|
||||
result = await db.execute(
|
||||
select(ChatMessage).filter(ChatMessage.id.in_([f'{chat_id}-{message_id}' for message_id in messages]))
|
||||
)
|
||||
existing_by_id = {row.id: row for row in result.scalars().all()}
|
||||
|
||||
for message_id, data in messages.items():
|
||||
composite_id = f'{chat_id}-{message_id}'
|
||||
message = existing_by_id.get(composite_id)
|
||||
if message:
|
||||
self._apply_message_data(message, data, now)
|
||||
else:
|
||||
db.add(self._build_message(composite_id, chat_id, user_id, data, now))
|
||||
|
||||
await db.commit()
|
||||
|
||||
async def get_message_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ChatMessageModel]:
|
||||
async with get_async_db_context(db) as db:
|
||||
|
|
|
|||
|
|
@ -959,18 +959,15 @@ class ChatTable:
|
|||
"""Write messages to the ``chat_message`` table so future lookups
|
||||
use the fast path. Errors are logged but never raised.
|
||||
"""
|
||||
for message_id, message in messages.items():
|
||||
if not isinstance(message, dict) or not message.get('role'):
|
||||
continue
|
||||
try:
|
||||
await ChatMessages.upsert_message(
|
||||
message_id=message_id,
|
||||
chat_id=chat_id,
|
||||
user_id=user_id,
|
||||
data=message,
|
||||
)
|
||||
except Exception as e:
|
||||
log.warning('Backfill failed for message %s in chat %s: %s', message_id, chat_id, e)
|
||||
writable = {
|
||||
message_id: message
|
||||
for message_id, message in messages.items()
|
||||
if isinstance(message, dict) and message.get('role')
|
||||
}
|
||||
try:
|
||||
await ChatMessages.upsert_messages(chat_id, user_id, writable)
|
||||
except Exception as e:
|
||||
log.warning('Backfill failed for chat %s: %s', chat_id, e)
|
||||
|
||||
async def reconcile_messages_by_chat_id(self, chat_id: str, user_id: str, messages: dict[str, dict]) -> None:
|
||||
"""Sync ``chat_message`` rows with the committed JSON blob.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue