Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac/enh: db session sharing

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2025-12-29 00:21:18 +04:00
parent 6dd0f99b90
commit b1d0f00d8c
23 changed files with 1173 additions and 663 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/routers/retrieval.py
+17 -9
View File
@@ -39,6 +39,8 @@ from langchain_core.documents import Document
from open_webui.models.files import FileModel, FileUpdateForm, Files
from open_webui.models.knowledge import Knowledges
from open_webui.storage.provider import Storage
from open_webui.internal.db import get_session
from sqlalchemy.orm import Session
from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT
@@ -1484,14 +1486,15 @@ def process_file(
request: Request,
form_data: ProcessFileForm,
user=Depends(get_verified_user),
db: Session = Depends(get_session),
):
"""
Process a file and save its content to the vector database.
"""
if user.role == "admin":
file = Files.get_file_by_id(form_data.file_id)
file = Files.get_file_by_id(form_data.file_id, db=db)
else:
file = Files.get_file_by_id_and_user_id(form_data.file_id, user.id)
file = Files.get_file_by_id_and_user_id(form_data.file_id, user.id, db=db)
if file:
try:
@@ -1633,12 +1636,13 @@ def process_file(
Files.update_file_data_by_id(
file.id,
{"content": text_content},
db=db,
)
hash = calculate_sha256_string(text_content)
Files.update_file_hash_by_id(file.id, hash)
Files.update_file_hash_by_id(file.id, hash, db=db)
if request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL:
Files.update_file_data_by_id(file.id, {"status": "completed"})
Files.update_file_data_by_id(file.id, {"status": "completed"}, db=db)
return {
"status": True,
"collection_name": None,
@@ -1667,11 +1671,13 @@ def process_file(
{
"collection_name": collection_name,
},
db=db,
)
Files.update_file_data_by_id(
file.id,
{"status": "completed"},
db=db,
)
return {
@@ -1690,6 +1696,7 @@ def process_file(
Files.update_file_data_by_id(
file.id,
{"status": "failed"},
db=db,
)
if "No pandoc was found" in str(e):
@@ -2417,10 +2424,10 @@ class DeleteForm(BaseModel):
@router.post("/delete")
def delete_entries_from_collection(form_data: DeleteForm, user=Depends(get_admin_user)):
def delete_entries_from_collection(form_data: DeleteForm, user=Depends(get_admin_user), db: Session = Depends(get_session)):
try:
if VECTOR_DB_CLIENT.has_collection(collection_name=form_data.collection_name):
file = Files.get_file_by_id(form_data.file_id)
file = Files.get_file_by_id(form_data.file_id, db=db)
hash = file.hash
VECTOR_DB_CLIENT.delete(
@@ -2436,9 +2443,9 @@ def delete_entries_from_collection(form_data: DeleteForm, user=Depends(get_admin
@router.post("/reset/db")
def reset_vector_db(user=Depends(get_admin_user)):
def reset_vector_db(user=Depends(get_admin_user), db: Session = Depends(get_session)):
VECTOR_DB_CLIENT.reset()
Knowledges.delete_all_knowledge()
Knowledges.delete_all_knowledge(db=db)
@router.post("/reset/uploads")
@@ -2496,6 +2503,7 @@ async def process_files_batch(
request: Request,
form_data: BatchProcessFilesForm,
user=Depends(get_verified_user),
db: Session = Depends(get_session),
) -> BatchProcessFilesResponse:
"""
Process a batch of files and save them to the vector database.
@@ -2558,7 +2566,7 @@ async def process_files_batch(
# Update all files with collection name
for file_update, file_result in zip(file_updates, file_results):
Files.update_file_by_id(id=file_result.file_id, form_data=file_update)
Files.update_file_by_id(id=file_result.file_id, form_data=file_update, db=db)
file_result.status = "completed"
except Exception as e: