refac: endpoints regarding db operations

2024-08-27 14:01:00 +02:00
parent bbeed7cd85
commit 062649e483
4 changed files with 76 additions and 71 deletions
@@ -68,6 +68,76 @@ async def add_memory(
    return memory


+############################
+# QueryMemory
+############################
+
+
+class QueryMemoryForm(BaseModel):
+    content: str
+    k: Optional[int] = 1
+
+
+@router.post("/query")
+async def query_memory(
+    request: Request, form_data: QueryMemoryForm, user=Depends(get_verified_user)
+):
+    query_embedding = request.app.state.EMBEDDING_FUNCTION(form_data.content)
+    collection = CHROMA_CLIENT.get_or_create_collection(name=f"user-memory-{user.id}")
+
+    results = collection.query(
+        query_embeddings=[query_embedding],
+        n_results=form_data.k,  # how many results to return
+    )
+
+    return results
+
+
+############################
+# ResetMemoryFromVectorDB
+############################
+@router.post("/reset", response_model=bool)
+async def reset_memory_from_vector_db(
+    request: Request, user=Depends(get_verified_user)
+):
+    CHROMA_CLIENT.delete_collection(f"user-memory-{user.id}")
+    collection = CHROMA_CLIENT.get_or_create_collection(name=f"user-memory-{user.id}")
+
+    memories = Memories.get_memories_by_user_id(user.id)
+    for memory in memories:
+        memory_embedding = request.app.state.EMBEDDING_FUNCTION(memory.content)
+        collection.upsert(
+            documents=[memory.content],
+            ids=[memory.id],
+            embeddings=[memory_embedding],
+        )
+    return True
+
+
+############################
+# DeleteMemoriesByUserId
+############################
+
+
+@router.delete("/delete/user", response_model=bool)
+async def delete_memory_by_user_id(user=Depends(get_verified_user)):
+    result = Memories.delete_memories_by_user_id(user.id)
+
+    if result:
+        try:
+            CHROMA_CLIENT.delete_collection(f"user-memory-{user.id}")
+        except Exception as e:
+            log.error(e)
+        return True
+
+    return False
+
+
+############################
+# UpdateMemoryById
+############################
+
+
@router.post("/{memory_id}/update", response_model=Optional[MemoryModel])
 async def update_memory_by_id(
    memory_id: str,
@@ -96,71 +166,6 @@ async def update_memory_by_id(
    return memory


-############################
-# QueryMemory
-############################
-
-
-class QueryMemoryForm(BaseModel):
-    content: str
-    k: Optional[int] = 1
-
-
-@router.post("/query")
-async def query_memory(
-    request: Request, form_data: QueryMemoryForm, user=Depends(get_verified_user)
-):
-    query_embedding = request.app.state.EMBEDDING_FUNCTION(form_data.content)
-    collection = CHROMA_CLIENT.get_or_create_collection(name=f"user-memory-{user.id}")
-
-    results = collection.query(
-        query_embeddings=[query_embedding],
-        n_results=form_data.k,  # how many results to return
-    )
-
-    return results
-
-
-############################
-# ResetMemoryFromVectorDB
-############################
-@router.get("/reset", response_model=bool)
-async def reset_memory_from_vector_db(
-    request: Request, user=Depends(get_verified_user)
-):
-    CHROMA_CLIENT.delete_collection(f"user-memory-{user.id}")
-    collection = CHROMA_CLIENT.get_or_create_collection(name=f"user-memory-{user.id}")
-
-    memories = Memories.get_memories_by_user_id(user.id)
-    for memory in memories:
-        memory_embedding = request.app.state.EMBEDDING_FUNCTION(memory.content)
-        collection.upsert(
-            documents=[memory.content],
-            ids=[memory.id],
-            embeddings=[memory_embedding],
-        )
-    return True
-
-
-############################
-# DeleteMemoriesByUserId
-############################
-
-
-@router.delete("/user", response_model=bool)
-async def delete_memory_by_user_id(user=Depends(get_verified_user)):
-    result = Memories.delete_memories_by_user_id(user.id)
-
-    if result:
-        try:
-            CHROMA_CLIENT.delete_collection(f"user-memory-{user.id}")
-        except Exception as e:
-            log.error(e)
-        return True
-
-    return False
-
-
 ############################
 # DeleteMemoryById
 ############################