refac

backend/open_webui/models/memories.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import time
 import uuid
-from typing import Optional
+from typing import Literal, Optional
 
 from open_webui.internal.db import Base, get_async_db_context
 from pydantic import BaseModel, ConfigDict
@@ -19,6 +19,7 @@ class Memory(Base):  # user memory store
 
     id = Column(String, primary_key=True, unique=True)
     user_id = Column(String, index=True)
+    type = Column(String, default='context', server_default='context', index=True)
     content = Column(Text)  # free-form text learned from conversation
     updated_at = Column(BigInteger)  # epoch seconds
     created_at = Column(BigInteger)  # epoch seconds
@@ -29,6 +30,7 @@ class MemoryModel(BaseModel):
 
     id: str
     user_id: str
+    type: Literal['user', 'context'] = 'context'
     content: str
     updated_at: int  # timestamp in epoch
     created_at: int  # timestamp in epoch
@@ -36,10 +38,15 @@ class MemoryModel(BaseModel):
 
 
 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,
         db: AsyncSession | None = None,
     ) -> MemoryModel | None:
         """Persist a new memory entry and return the created model."""
@@ -48,6 +55,7 @@ class MemoriesTable:
             record = Memory(
                 id=str(uuid.uuid4()),
                 user_id=user_id,
+                type=self.normalize_memory_type(memory_type),
                 content=content,
                 created_at=now,
                 updated_at=now,
@@ -61,7 +69,8 @@ class MemoriesTable:
         self,
         id: str,
         user_id: str,
-        content: str,
+        content: str | None,
+        memory_type: str | None = None,
         db: AsyncSession | None = None,
     ) -> MemoryModel | None:
         async with get_async_db_context(db) as db:
@@ -70,7 +79,10 @@ class MemoriesTable:
                 if not memory or memory.user_id != user_id:
                     return None
 
-                memory.content = content
+                if content is not None:
+                    memory.content = content
+                if memory_type is not None:
+                    memory.type = self.normalize_memory_type(memory_type)
                 memory.updated_at = int(time.time())
 
                 await db.commit()
@@ -139,5 +151,78 @@ class MemoriesTable:
             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'))
+                    result = await db.execute(
+                        select(Memory).filter_by(user_id=user_id, content=content, type=memory_type)
+                    )
+                    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,
+                        content=content,
+                        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'))
+                    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