mirror of
https://github.com/open-webui/open-webui.git
synced 2025-04-02 17:10:02 +02:00
refac: endpoints regarding db operations
This commit is contained in:
parent
bbeed7cd85
commit
062649e483
@ -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):
|
||||
|
@ -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
|
||||
############################
|
||||
|
@ -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',
|
||||
|
@ -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}`
|
||||
|
Loading…
x
Reference in New Issue
Block a user