From 062649e4838f003c13c5cba076531d60de896068 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Tue, 27 Aug 2024 14:01:00 +0200 Subject: [PATCH] refac: endpoints regarding db operations --- backend/apps/rag/main.py | 6 +- backend/apps/webui/routers/memories.py | 135 +++++++++++++------------ src/lib/apis/memories/index.ts | 2 +- src/lib/apis/rag/index.ts | 4 +- 4 files changed, 76 insertions(+), 71 deletions(-) diff --git a/backend/apps/rag/main.py b/backend/apps/rag/main.py index 7b2fbc679..31445f791 100644 --- a/backend/apps/rag/main.py +++ b/backend/apps/rag/main.py @@ -1373,12 +1373,12 @@ def scan_docs_dir(user=Depends(get_admin_user)): return True -@app.get("/reset/db") +@app.post("/reset/db") def reset_vector_db(user=Depends(get_admin_user)): CHROMA_CLIENT.reset() -@app.get("/reset/uploads") +@app.post("/reset/uploads") def reset_upload_dir(user=Depends(get_admin_user)) -> bool: folder = f"{UPLOAD_DIR}" try: @@ -1402,7 +1402,7 @@ def reset_upload_dir(user=Depends(get_admin_user)) -> bool: return True -@app.get("/reset") +@app.post("/reset") def reset(user=Depends(get_admin_user)) -> bool: folder = f"{UPLOAD_DIR}" for filename in os.listdir(folder): diff --git a/backend/apps/webui/routers/memories.py b/backend/apps/webui/routers/memories.py index a7b5474f0..ae0a9efcb 100644 --- a/backend/apps/webui/routers/memories.py +++ b/backend/apps/webui/routers/memories.py @@ -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 ############################ diff --git a/src/lib/apis/memories/index.ts b/src/lib/apis/memories/index.ts index c3c122adf..3fd83ca9e 100644 --- a/src/lib/apis/memories/index.ts +++ b/src/lib/apis/memories/index.ts @@ -156,7 +156,7 @@ export const deleteMemoryById = async (token: string, id: string) => { export const deleteMemoriesByUserId = async (token: string) => { let error = null; - const res = await fetch(`${WEBUI_API_BASE_URL}/memories/user`, { + const res = await fetch(`${WEBUI_API_BASE_URL}/memories/delete/user`, { method: 'DELETE', headers: { Accept: 'application/json', diff --git a/src/lib/apis/rag/index.ts b/src/lib/apis/rag/index.ts index 5c0a47b35..3c0dba4b5 100644 --- a/src/lib/apis/rag/index.ts +++ b/src/lib/apis/rag/index.ts @@ -400,7 +400,7 @@ export const resetUploadDir = async (token: string) => { let error = null; const res = await fetch(`${RAG_API_BASE_URL}/reset/uploads`, { - method: 'GET', + method: 'POST', headers: { Accept: 'application/json', authorization: `Bearer ${token}` @@ -426,7 +426,7 @@ export const resetVectorDB = async (token: string) => { let error = null; const res = await fetch(`${RAG_API_BASE_URL}/reset/db`, { - method: 'GET', + method: 'POST', headers: { Accept: 'application/json', authorization: `Bearer ${token}`