refac: endpoints regarding db operations

This commit is contained in:
Timothy J. Baek 2024-08-27 14:01:00 +02:00
parent bbeed7cd85
commit 062649e483
4 changed files with 76 additions and 71 deletions

View File

@ -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):

View File

@ -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
############################

View File

@ -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',

View File

@ -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}`