From 40ebf8cd62746cb4b9bd01d9b0b31a9013c21384 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 20 Jul 2025 15:00:24 +0400 Subject: [PATCH] refac: memory handling --- backend/open_webui/models/memories.py | 17 +++++++++++++---- backend/open_webui/routers/memories.py | 4 ++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/backend/open_webui/models/memories.py b/backend/open_webui/models/memories.py index 8b10a77cf9..253371c680 100644 --- a/backend/open_webui/models/memories.py +++ b/backend/open_webui/models/memories.py @@ -71,9 +71,13 @@ class MemoriesTable: ) -> Optional[MemoryModel]: with get_db() as db: try: - db.query(Memory).filter_by(id=id, user_id=user_id).update( - {"content": content, "updated_at": int(time.time())} - ) + memory = db.get(Memory, id) + if not memory or memory.user_id != user_id: + return None + + memory.content = content + memory.updated_at = int(time.time()) + db.commit() return self.get_memory_by_id(id) except Exception: @@ -127,7 +131,12 @@ class MemoriesTable: def delete_memory_by_id_and_user_id(self, id: str, user_id: str) -> bool: with get_db() as db: try: - db.query(Memory).filter_by(id=id, user_id=user_id).delete() + memory = db.get(Memory, id) + if not memory or memory.user_id != user_id: + return None + + # Delete the memory + db.delete(memory) db.commit() return True diff --git a/backend/open_webui/routers/memories.py b/backend/open_webui/routers/memories.py index 333e9ecc6a..11b3d0c96c 100644 --- a/backend/open_webui/routers/memories.py +++ b/backend/open_webui/routers/memories.py @@ -82,6 +82,10 @@ class QueryMemoryForm(BaseModel): async def query_memory( request: Request, form_data: QueryMemoryForm, user=Depends(get_verified_user) ): + memories = Memories.get_memories_by_user_id(user.id) + if not memories: + raise HTTPException(status_code=404, detail="No memories found for user") + results = VECTOR_DB_CLIENT.search( collection_name=f"user-memory-{user.id}", vectors=[request.app.state.EMBEDDING_FUNCTION(form_data.content, user=user)],