Classic298
699d512e2f
Both session factories run with expire_on_commit=False, so ORM objects keep their attribute values after commit. Every session.refresh issued right after a commit therefore re-SELECTed a row whose values the session already held, including full chat JSON blobs and user settings, purely to overwrite identical data. Fifty such calls existed across the model layer, covering nearly every write path in the app (chat inserts, title updates, pin/archive toggles, user role and settings updates, tool, prompt, function, model, file, tag, feedback, memory, automation and grant writes). All fifty are removed. The only refreshes with an actual job were the two update-then-reload paths in tools and skills, where a Core UPDATE statement bypasses the identity map; those now use session.get(..., populate_existing=True), which guarantees a fresh row in one SELECT whether or not the row was already present in the session (the previous code issued get plus refresh, two SELECTs, on the default configuration). Benchmark (real SQLite DB, per write): | write path | before | after | | --- | --- | --- | | chat title update, ~600 KB chat blob | 2.08 ms | 1.24 ms | | user role update, small row | 1.21 ms | 0.68 ms | On Postgres each removed refresh is additionally a network round trip. The chat-blob case also skips re-parsing the entire JSON document per write. Functionally verified against a fresh database: user insert, role and settings updates, chat insert (including the server-default meta column, which is always provided client-side), title update and pin toggle, tool insert and the Core-update reload path, tag insert and the prompt insert flow that pins version_id after history creation all return correct values and persist correctly.
268 lines
10 KiB
Python
268 lines
10 KiB
Python
"""Long-term memory storage for per-user context recall."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
import uuid
|
|
from typing import Literal
|
|
|
|
from open_webui.internal.db import Base, get_async_db_context
|
|
from pydantic import BaseModel, ConfigDict
|
|
from sqlalchemy import JSON, BigInteger, Column, String, Text, delete, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
|
|
class Memory(Base): # user memory store
|
|
"""Stores user-created memory entries linked to a vector collection."""
|
|
|
|
__tablename__ = 'memory'
|
|
|
|
id = Column(String, primary_key=True, unique=True)
|
|
user_id = Column(String, index=True)
|
|
type = Column(String, default='context', server_default='context', index=True)
|
|
path = Column(Text, nullable=True)
|
|
content = Column(Text) # free-form text learned from conversation
|
|
meta = Column(JSON, nullable=True)
|
|
updated_at = Column(BigInteger) # epoch seconds
|
|
created_at = Column(BigInteger) # epoch seconds
|
|
|
|
|
|
class MemoryModel(BaseModel):
|
|
"""Pydantic mirror of the Memory table row."""
|
|
|
|
id: str
|
|
user_id: str
|
|
type: Literal['user', 'context'] = 'context'
|
|
path: str | None = None
|
|
content: str
|
|
meta: dict | None = None
|
|
updated_at: int # timestamp in epoch
|
|
created_at: int # timestamp in epoch
|
|
model_config = ConfigDict(from_attributes=True) # allows ORM mapping
|
|
|
|
|
|
class MemoriesTable:
|
|
@staticmethod
|
|
def normalize_memory_type(memory_type: str | None = None) -> str:
|
|
return 'user' if memory_type == 'user' else 'context'
|
|
|
|
async def insert_new_memory(
|
|
self,
|
|
user_id: str,
|
|
content: str,
|
|
memory_type: str | None = None,
|
|
path: str | None = None,
|
|
meta: dict | None = None,
|
|
db: AsyncSession | None = None,
|
|
) -> MemoryModel | None:
|
|
"""Persist a new memory entry and return the created model."""
|
|
async with get_async_db_context(db) as db:
|
|
now = int(time.time())
|
|
record = Memory(
|
|
id=str(uuid.uuid4()),
|
|
user_id=user_id,
|
|
type=self.normalize_memory_type(memory_type),
|
|
path=path,
|
|
content=content,
|
|
meta=meta,
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
db.add(record)
|
|
await db.commit()
|
|
return MemoryModel.model_validate(record) if record else None
|
|
|
|
async def update_memory_by_id_and_user_id(
|
|
self,
|
|
id: str,
|
|
user_id: str,
|
|
content: str | None,
|
|
memory_type: str | None = None,
|
|
path: str | None = None,
|
|
update_path: bool = False,
|
|
meta: dict | None = None,
|
|
db: AsyncSession | None = None,
|
|
) -> MemoryModel | None:
|
|
async with get_async_db_context(db) as db:
|
|
try:
|
|
memory = await db.get(Memory, id)
|
|
if not memory or memory.user_id != user_id:
|
|
return None
|
|
|
|
if content is not None:
|
|
memory.content = content
|
|
if memory_type is not None:
|
|
memory.type = self.normalize_memory_type(memory_type)
|
|
if update_path:
|
|
memory.path = path
|
|
if meta is not None:
|
|
memory.meta = {**(memory.meta or {}), **meta}
|
|
memory.updated_at = int(time.time())
|
|
|
|
await db.commit()
|
|
return MemoryModel.model_validate(memory)
|
|
except Exception:
|
|
return None
|
|
|
|
async def get_memories(self, db: AsyncSession | None = None) -> list[MemoryModel]:
|
|
async with get_async_db_context(db) as db:
|
|
try:
|
|
result = await db.execute(select(Memory))
|
|
memories = result.scalars().all()
|
|
return [MemoryModel.model_validate(memory) for memory in memories]
|
|
except Exception:
|
|
return None
|
|
|
|
async def get_memories_by_user_id(self, user_id: str, db: AsyncSession | None = None) -> list[MemoryModel]:
|
|
async with get_async_db_context(db) as db:
|
|
try:
|
|
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
|
|
|
|
async def get_memory_by_id(self, id: str, db: AsyncSession | None = None) -> MemoryModel | None:
|
|
async with get_async_db_context(db) as db:
|
|
try:
|
|
memory = await db.get(Memory, id)
|
|
return MemoryModel.model_validate(memory) if memory else None
|
|
except Exception:
|
|
return None
|
|
|
|
async def delete_memory_by_id(self, id: str, db: AsyncSession | None = None) -> bool:
|
|
async with get_async_db_context(db) as db:
|
|
try:
|
|
await db.execute(delete(Memory).filter_by(id=id))
|
|
await db.commit()
|
|
|
|
return True
|
|
|
|
except Exception:
|
|
return False
|
|
|
|
async def delete_memories_by_user_id(self, user_id: str, db: AsyncSession | None = None) -> bool:
|
|
async with get_async_db_context(db) as db:
|
|
try:
|
|
await db.execute(delete(Memory).filter_by(user_id=user_id))
|
|
await db.commit()
|
|
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
async def delete_memory_by_id_and_user_id(self, id: str, user_id: str, db: AsyncSession | None = None) -> bool:
|
|
async with get_async_db_context(db) as db:
|
|
try:
|
|
memory = await db.get(Memory, id)
|
|
if not memory or memory.user_id != user_id:
|
|
return False
|
|
|
|
await db.delete(memory)
|
|
await db.commit()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
async def apply_memory_operations(
|
|
self,
|
|
user_id: str,
|
|
operations: list[dict],
|
|
db: AsyncSession | None = None,
|
|
) -> list[dict]:
|
|
now = int(time.time())
|
|
results: list[dict] = []
|
|
|
|
async with get_async_db_context(db) as db:
|
|
for operation in operations:
|
|
action = operation.get('action')
|
|
|
|
if action == 'add':
|
|
content = operation.get('content', '').strip()
|
|
memory_type = self.normalize_memory_type(operation.get('type'))
|
|
path = operation.get('path')
|
|
result = await db.execute(
|
|
select(Memory).filter_by(user_id=user_id, content=content, type=memory_type, path=path)
|
|
)
|
|
existing = result.scalars().first()
|
|
if existing:
|
|
results.append(
|
|
{
|
|
'action': action,
|
|
'status': 'skipped',
|
|
'memory': MemoryModel.model_validate(existing),
|
|
'reason': 'duplicate',
|
|
}
|
|
)
|
|
continue
|
|
|
|
memory = Memory(
|
|
id=str(uuid.uuid4()),
|
|
user_id=user_id,
|
|
type=memory_type,
|
|
path=path,
|
|
content=content,
|
|
meta=operation.get('meta'),
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
db.add(memory)
|
|
await db.flush()
|
|
results.append(
|
|
{'action': action, 'status': 'created', 'memory': MemoryModel.model_validate(memory)}
|
|
)
|
|
|
|
elif action == 'replace':
|
|
memory_id = operation.get('id')
|
|
content = operation.get('content', '').strip()
|
|
memory = await db.get(Memory, memory_id)
|
|
if not memory or memory.user_id != user_id:
|
|
raise ValueError(f'Memory not found: {memory_id}')
|
|
|
|
memory.content = content
|
|
if operation.get('type') is not None:
|
|
memory.type = self.normalize_memory_type(operation.get('type'))
|
|
if 'path' in operation:
|
|
memory.path = operation.get('path')
|
|
if operation.get('meta') is not None:
|
|
memory.meta = {**(memory.meta or {}), **operation.get('meta')}
|
|
memory.updated_at = now
|
|
await db.flush()
|
|
results.append(
|
|
{'action': action, 'status': 'updated', 'memory': MemoryModel.model_validate(memory)}
|
|
)
|
|
|
|
elif action == 'move':
|
|
memory_id = operation.get('id')
|
|
memory = await db.get(Memory, memory_id)
|
|
if not memory or memory.user_id != user_id:
|
|
raise ValueError(f'Memory not found: {memory_id}')
|
|
|
|
memory.path = operation.get('path')
|
|
if operation.get('meta') is not None:
|
|
memory.meta = {**(memory.meta or {}), **operation.get('meta')}
|
|
memory.updated_at = now
|
|
await db.flush()
|
|
results.append(
|
|
{'action': action, 'status': 'updated', 'memory': MemoryModel.model_validate(memory)}
|
|
)
|
|
|
|
elif action == 'remove':
|
|
memory_id = operation.get('id')
|
|
memory = await db.get(Memory, memory_id)
|
|
if not memory or memory.user_id != user_id:
|
|
raise ValueError(f'Memory not found: {memory_id}')
|
|
|
|
await db.delete(memory)
|
|
results.append({'action': action, 'status': 'deleted', 'id': memory_id})
|
|
|
|
else:
|
|
raise ValueError(f'Unsupported memory operation: {action}')
|
|
|
|
await db.commit()
|
|
|
|
return results
|
|
|
|
|
|
Memories = MemoriesTable() # user memory registry
|