refac/enh: db session sharing

backend/open_webui/routers/knowledge.py

@@ -4,6 +4,8 @@ from fastapi import APIRouter, Depends, HTTPException, status, Request, Query
 from fastapi.concurrency import run_in_threadpool
 import logging
 
+from sqlalchemy.orm import Session
+from open_webui.internal.db import get_session
 from open_webui.models.groups import Groups
 from open_webui.models.knowledge import (
     KnowledgeFileListResponse,
@@ -52,21 +54,25 @@ class KnowledgeAccessListResponse(BaseModel):
 
 
 @router.get("/", response_model=KnowledgeAccessListResponse)
-async def get_knowledge_bases(page: Optional[int] = 1, user=Depends(get_verified_user)):
+async def get_knowledge_bases(
+    page: Optional[int] = 1,
+    user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
+):
     page = max(page, 1)
     limit = PAGE_ITEM_COUNT
     skip = (page - 1) * limit
 
     filter = {}
     if not user.role == "admin" or not BYPASS_ADMIN_ACCESS_CONTROL:
-        groups = Groups.get_groups_by_member_id(user.id)
+        groups = Groups.get_groups_by_member_id(user.id, db=db)
         if groups:
             filter["group_ids"] = [group.id for group in groups]
 
         filter["user_id"] = user.id
 
     result = Knowledges.search_knowledge_bases(
-        user.id, filter=filter, skip=skip, limit=limit
+        user.id, filter=filter, skip=skip, limit=limit, db=db
     )
 
     return KnowledgeAccessListResponse(
@@ -75,7 +81,9 @@ async def get_knowledge_bases(page: Optional[int] = 1, user=Depends(get_verified
                 **knowledge_base.model_dump(),
                 write_access=(
                     user.id == knowledge_base.user_id
-                    or has_access(user.id, "write", knowledge_base.access_control)
+                    or has_access(
+                        user.id, "write", knowledge_base.access_control, db=db
+                    )
                 ),
             )
             for knowledge_base in result.items
@@ -90,6 +98,7 @@ async def search_knowledge_bases(
     view_option: Optional[str] = None,
     page: Optional[int] = 1,
     user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
     page = max(page, 1)
     limit = PAGE_ITEM_COUNT
@@ -102,14 +111,14 @@ async def search_knowledge_bases(
         filter["view_option"] = view_option
 
     if not user.role == "admin" or not BYPASS_ADMIN_ACCESS_CONTROL:
-        groups = Groups.get_groups_by_member_id(user.id)
+        groups = Groups.get_groups_by_member_id(user.id, db=db)
         if groups:
             filter["group_ids"] = [group.id for group in groups]
 
         filter["user_id"] = user.id
 
     result = Knowledges.search_knowledge_bases(
-        user.id, filter=filter, skip=skip, limit=limit
+        user.id, filter=filter, skip=skip, limit=limit, db=db
     )
 
     return KnowledgeAccessListResponse(
@@ -118,7 +127,9 @@ async def search_knowledge_bases(
                 **knowledge_base.model_dump(),
                 write_access=(
                     user.id == knowledge_base.user_id
-                    or has_access(user.id, "write", knowledge_base.access_control)
+                    or has_access(
+                        user.id, "write", knowledge_base.access_control, db=db
+                    )
                 ),
             )
             for knowledge_base in result.items
@@ -132,6 +143,7 @@ async def search_knowledge_files(
     query: Optional[str] = None,
     page: Optional[int] = 1,
     user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
     page = max(page, 1)
     limit = PAGE_ITEM_COUNT
@@ -141,13 +153,15 @@ async def search_knowledge_files(
     if query:
         filter["query"] = query
 
-    groups = Groups.get_groups_by_member_id(user.id)
+    groups = Groups.get_groups_by_member_id(user.id, db=db)
     if groups:
         filter["group_ids"] = [group.id for group in groups]
 
     filter["user_id"] = user.id
 
-    return Knowledges.search_knowledge_files(filter=filter, skip=skip, limit=limit)
+    return Knowledges.search_knowledge_files(
+        filter=filter, skip=skip, limit=limit, db=db
+    )
 
 
 ############################
@@ -157,10 +171,13 @@ async def search_knowledge_files(
 
 @router.post("/create", response_model=Optional[KnowledgeResponse])
 async def create_new_knowledge(
-    request: Request, form_data: KnowledgeForm, user=Depends(get_verified_user)
+    request: Request,
+    form_data: KnowledgeForm,
+    user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
     if user.role != "admin" and not has_permission(
-        user.id, "workspace.knowledge", request.app.state.config.USER_PERMISSIONS
+        user.id, "workspace.knowledge", request.app.state.config.USER_PERMISSIONS, db=db
     ):
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
@@ -175,11 +192,12 @@ async def create_new_knowledge(
             user.id,
             "sharing.public_knowledge",
             request.app.state.config.USER_PERMISSIONS,
+            db=db,
         )
     ):
         form_data.access_control = {}
 
-    knowledge = Knowledges.insert_new_knowledge(user.id, form_data)
+    knowledge = Knowledges.insert_new_knowledge(user.id, form_data, db=db)
 
     if knowledge:
         return knowledge
@@ -196,20 +214,24 @@ async def create_new_knowledge(
 
 
 @router.post("/reindex", response_model=bool)
-async def reindex_knowledge_files(request: Request, user=Depends(get_verified_user)):
+async def reindex_knowledge_files(
+    request: Request,
+    user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
+):
     if user.role != "admin":
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail=ERROR_MESSAGES.UNAUTHORIZED,
         )
 
-    knowledge_bases = Knowledges.get_knowledge_bases()
+    knowledge_bases = Knowledges.get_knowledge_bases(db=db)
 
     log.info(f"Starting reindexing for {len(knowledge_bases)} knowledge bases")
 
     for knowledge_base in knowledge_bases:
         try:
-            files = Knowledges.get_files_by_id(knowledge_base.id)
+            files = Knowledges.get_files_by_id(knowledge_base.id, db=db)
             try:
                 if VECTOR_DB_CLIENT.has_collection(collection_name=knowledge_base.id):
                     VECTOR_DB_CLIENT.delete_collection(
@@ -229,6 +251,7 @@ async def reindex_knowledge_files(request: Request, user=Depends(get_verified_us
                             file_id=file.id, collection_name=knowledge_base.id
                         ),
                         user=user,
+                        db=db,
                     )
                 except Exception as e:
                     log.error(
@@ -264,21 +287,23 @@ class KnowledgeFilesResponse(KnowledgeResponse):
 
 
 @router.get("/{id}", response_model=Optional[KnowledgeFilesResponse])
-async def get_knowledge_by_id(id: str, user=Depends(get_verified_user)):
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+async def get_knowledge_by_id(
+    id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)
+):
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
 
     if knowledge:
         if (
             user.role == "admin"
             or knowledge.user_id == user.id
-            or has_access(user.id, "read", knowledge.access_control)
+            or has_access(user.id, "read", knowledge.access_control, db=db)
         ):
 
             return KnowledgeFilesResponse(
                 **knowledge.model_dump(),
                 write_access=(
                     user.id == knowledge.user_id
-                    or has_access(user.id, "write", knowledge.access_control)
+                    or has_access(user.id, "write", knowledge.access_control, db=db)
                 ),
             )
     else:
@@ -299,8 +324,9 @@ async def update_knowledge_by_id(
     id: str,
     form_data: KnowledgeForm,
     user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
     if not knowledge:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -309,7 +335,7 @@ async def update_knowledge_by_id(
     # Is the user the original creator, in a group with write access, or an admin
     if (
         knowledge.user_id != user.id
-        and not has_access(user.id, "write", knowledge.access_control)
+        and not has_access(user.id, "write", knowledge.access_control, db=db)
         and user.role != "admin"
     ):
         raise HTTPException(
@@ -325,15 +351,16 @@ async def update_knowledge_by_id(
             user.id,
             "sharing.public_knowledge",
             request.app.state.config.USER_PERMISSIONS,
+            db=db,
         )
     ):
         form_data.access_control = {}
 
-    knowledge = Knowledges.update_knowledge_by_id(id=id, form_data=form_data)
+    knowledge = Knowledges.update_knowledge_by_id(id=id, form_data=form_data, db=db)
     if knowledge:
         return KnowledgeFilesResponse(
             **knowledge.model_dump(),
-            files=Knowledges.get_file_metadatas_by_id(knowledge.id),
+            files=Knowledges.get_file_metadatas_by_id(knowledge.id, db=db),
         )
     else:
         raise HTTPException(
@@ -356,9 +383,10 @@ async def get_knowledge_files_by_id(
     direction: Optional[str] = None,
     page: Optional[int] = 1,
     user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
 
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
     if not knowledge:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -368,7 +396,7 @@ async def get_knowledge_files_by_id(
     if not (
         user.role == "admin"
         or knowledge.user_id == user.id
-        or has_access(user.id, "read", knowledge.access_control)
+        or has_access(user.id, "read", knowledge.access_control, db=db)
     ):
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -391,7 +419,7 @@ async def get_knowledge_files_by_id(
         filter["direction"] = direction
 
     return Knowledges.search_files_by_id(
-        id, user.id, filter=filter, skip=skip, limit=limit
+        id, user.id, filter=filter, skip=skip, limit=limit, db=db
     )
 
 
@@ -410,8 +438,9 @@ def add_file_to_knowledge_by_id(
     id: str,
     form_data: KnowledgeFileIdForm,
     user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
     if not knowledge:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -420,7 +449,7 @@ def add_file_to_knowledge_by_id(
 
     if (
         knowledge.user_id != user.id
-        and not has_access(user.id, "write", knowledge.access_control)
+        and not has_access(user.id, "write", knowledge.access_control, db=db)
         and user.role != "admin"
     ):
         raise HTTPException(
@@ -428,7 +457,7 @@ def add_file_to_knowledge_by_id(
             detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
 
-    file = Files.get_file_by_id(form_data.file_id)
+    file = Files.get_file_by_id(form_data.file_id, db=db)
     if not file:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -446,11 +475,12 @@ def add_file_to_knowledge_by_id(
             request,
             ProcessFileForm(file_id=form_data.file_id, collection_name=id),
             user=user,
+            db=db,
         )
 
         # Add file to knowledge base
         Knowledges.add_file_to_knowledge_by_id(
-            knowledge_id=id, file_id=form_data.file_id, user_id=user.id
+            knowledge_id=id, file_id=form_data.file_id, user_id=user.id, db=db
         )
     except Exception as e:
         log.debug(e)
@@ -462,7 +492,7 @@ def add_file_to_knowledge_by_id(
     if knowledge:
         return KnowledgeFilesResponse(
             **knowledge.model_dump(),
-            files=Knowledges.get_file_metadatas_by_id(knowledge.id),
+            files=Knowledges.get_file_metadatas_by_id(knowledge.id, db=db),
         )
     else:
         raise HTTPException(
@@ -477,8 +507,9 @@ def update_file_from_knowledge_by_id(
     id: str,
     form_data: KnowledgeFileIdForm,
     user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
     if not knowledge:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -487,7 +518,7 @@ def update_file_from_knowledge_by_id(
 
     if (
         knowledge.user_id != user.id
-        and not has_access(user.id, "write", knowledge.access_control)
+        and not has_access(user.id, "write", knowledge.access_control, db=db)
         and user.role != "admin"
     ):
 
@@ -496,7 +527,7 @@ def update_file_from_knowledge_by_id(
             detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
 
-    file = Files.get_file_by_id(form_data.file_id)
+    file = Files.get_file_by_id(form_data.file_id, db=db)
     if not file:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -514,6 +545,7 @@ def update_file_from_knowledge_by_id(
             request,
             ProcessFileForm(file_id=form_data.file_id, collection_name=id),
             user=user,
+            db=db,
         )
     except Exception as e:
         raise HTTPException(
@@ -524,7 +556,7 @@ def update_file_from_knowledge_by_id(
     if knowledge:
         return KnowledgeFilesResponse(
             **knowledge.model_dump(),
-            files=Knowledges.get_file_metadatas_by_id(knowledge.id),
+            files=Knowledges.get_file_metadatas_by_id(knowledge.id, db=db),
         )
     else:
         raise HTTPException(
@@ -544,8 +576,9 @@ def remove_file_from_knowledge_by_id(
     form_data: KnowledgeFileIdForm,
     delete_file: bool = Query(True),
     user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
     if not knowledge:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -554,7 +587,7 @@ def remove_file_from_knowledge_by_id(
 
     if (
         knowledge.user_id != user.id
-        and not has_access(user.id, "write", knowledge.access_control)
+        and not has_access(user.id, "write", knowledge.access_control, db=db)
         and user.role != "admin"
     ):
         raise HTTPException(
@@ -562,7 +595,7 @@ def remove_file_from_knowledge_by_id(
             detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
 
-    file = Files.get_file_by_id(form_data.file_id)
+    file = Files.get_file_by_id(form_data.file_id, db=db)
     if not file:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -570,7 +603,7 @@ def remove_file_from_knowledge_by_id(
         )
 
     Knowledges.remove_file_from_knowledge_by_id(
-        knowledge_id=id, file_id=form_data.file_id
+        knowledge_id=id, file_id=form_data.file_id, db=db
     )
 
     # Remove content from the vector database
@@ -599,12 +632,12 @@ def remove_file_from_knowledge_by_id(
             pass
 
         # Delete file from database
-        Files.delete_file_by_id(form_data.file_id)
+        Files.delete_file_by_id(form_data.file_id, db=db)
 
     if knowledge:
         return KnowledgeFilesResponse(
             **knowledge.model_dump(),
-            files=Knowledges.get_file_metadatas_by_id(knowledge.id),
+            files=Knowledges.get_file_metadatas_by_id(knowledge.id, db=db),
         )
     else:
         raise HTTPException(
@@ -619,8 +652,10 @@ def remove_file_from_knowledge_by_id(
 
 
 @router.delete("/{id}/delete", response_model=bool)
-async def delete_knowledge_by_id(id: str, user=Depends(get_verified_user)):
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+async def delete_knowledge_by_id(
+    id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)
+):
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
     if not knowledge:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -629,7 +664,7 @@ async def delete_knowledge_by_id(id: str, user=Depends(get_verified_user)):
 
     if (
         knowledge.user_id != user.id
-        and not has_access(user.id, "write", knowledge.access_control)
+        and not has_access(user.id, "write", knowledge.access_control, db=db)
         and user.role != "admin"
     ):
         raise HTTPException(
@@ -640,7 +675,7 @@ async def delete_knowledge_by_id(id: str, user=Depends(get_verified_user)):
     log.info(f"Deleting knowledge base: {id} (name: {knowledge.name})")
 
     # Get all models
-    models = Models.get_all_models()
+    models = Models.get_all_models(db=db)
     log.info(f"Found {len(models)} models to check for knowledge base {id}")
 
     # Update models that reference this knowledge base
@@ -664,7 +699,7 @@ async def delete_knowledge_by_id(id: str, user=Depends(get_verified_user)):
                     access_control=model.access_control,
                     is_active=model.is_active,
                 )
-                Models.update_model_by_id(model.id, model_form)
+                Models.update_model_by_id(model.id, model_form, db=db)
 
     # Clean up vector DB
     try:
@@ -672,7 +707,7 @@ async def delete_knowledge_by_id(id: str, user=Depends(get_verified_user)):
     except Exception as e:
         log.debug(e)
         pass
-    result = Knowledges.delete_knowledge_by_id(id=id)
+    result = Knowledges.delete_knowledge_by_id(id=id, db=db)
     return result
 
 
@@ -682,8 +717,10 @@ async def delete_knowledge_by_id(id: str, user=Depends(get_verified_user)):
 
 
 @router.post("/{id}/reset", response_model=Optional[KnowledgeResponse])
-async def reset_knowledge_by_id(id: str, user=Depends(get_verified_user)):
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+async def reset_knowledge_by_id(
+    id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)
+):
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
     if not knowledge:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -692,7 +729,7 @@ async def reset_knowledge_by_id(id: str, user=Depends(get_verified_user)):
 
     if (
         knowledge.user_id != user.id
-        and not has_access(user.id, "write", knowledge.access_control)
+        and not has_access(user.id, "write", knowledge.access_control, db=db)
         and user.role != "admin"
     ):
         raise HTTPException(
@@ -706,7 +743,7 @@ async def reset_knowledge_by_id(id: str, user=Depends(get_verified_user)):
         log.debug(e)
         pass
 
-    knowledge = Knowledges.reset_knowledge_by_id(id=id)
+    knowledge = Knowledges.reset_knowledge_by_id(id=id, db=db)
     return knowledge
 
 
@@ -721,11 +758,12 @@ async def add_files_to_knowledge_batch(
     id: str,
     form_data: list[KnowledgeFileIdForm],
     user=Depends(get_verified_user),
+    db: Session = Depends(get_session),
 ):
     """
     Add multiple files to a knowledge base
     """
-    knowledge = Knowledges.get_knowledge_by_id(id=id)
+    knowledge = Knowledges.get_knowledge_by_id(id=id, db=db)
     if not knowledge:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -734,7 +772,7 @@ async def add_files_to_knowledge_batch(
 
     if (
         knowledge.user_id != user.id
-        and not has_access(user.id, "write", knowledge.access_control)
+        and not has_access(user.id, "write", knowledge.access_control, db=db)
         and user.role != "admin"
     ):
         raise HTTPException(
@@ -746,7 +784,7 @@ async def add_files_to_knowledge_batch(
     log.info(f"files/batch/add - {len(form_data)} files")
     files: List[FileModel] = []
     for form in form_data:
-        file = Files.get_file_by_id(form.file_id)
+        file = Files.get_file_by_id(form.file_id, db=db)
         if not file:
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
@@ -760,6 +798,7 @@ async def add_files_to_knowledge_batch(
             request=request,
             form_data=BatchProcessFilesForm(files=files, collection_name=id),
             user=user,
+            db=db,
         )
     except Exception as e:
         log.error(
@@ -771,7 +810,7 @@ async def add_files_to_knowledge_batch(
     successful_file_ids = [r.file_id for r in result.results if r.status == "completed"]
     for file_id in successful_file_ids:
         Knowledges.add_file_to_knowledge_by_id(
-            knowledge_id=id, file_id=file_id, user_id=user.id
+            knowledge_id=id, file_id=file_id, user_id=user.id, db=db
         )
 
     # If there were any errors, include them in the response
@@ -779,7 +818,7 @@ async def add_files_to_knowledge_batch(
         error_details = [f"{err.file_id}: {err.error}" for err in result.errors]
         return KnowledgeFilesResponse(
             **knowledge.model_dump(),
-            files=Knowledges.get_file_metadatas_by_id(knowledge.id),
+            files=Knowledges.get_file_metadatas_by_id(knowledge.id, db=db),
             warnings={
                 "message": "Some files failed to process",
                 "errors": error_details,
@@ -788,5 +827,5 @@ async def add_files_to_knowledge_batch(
 
     return KnowledgeFilesResponse(
         **knowledge.model_dump(),
-        files=Knowledges.get_file_metadatas_by_id(knowledge.id),
+        files=Knowledges.get_file_metadatas_by_id(knowledge.id, db=db),
     )