enh/refac: kb pagination

backend/open_webui/routers/knowledge.py

@@ -4,6 +4,7 @@ from fastapi import APIRouter, Depends, HTTPException, status, Request, Query
 from fastapi.concurrency import run_in_threadpool
 import logging
 
+from open_webui.models.groups import Groups
 from open_webui.models.knowledge import (
     KnowledgeFileListResponse,
     Knowledges,
@@ -40,53 +41,115 @@ router = APIRouter()
 # getKnowledgeBases
 ############################
 
+PAGE_ITEM_COUNT = 30
+
 
 class KnowledgeAccessResponse(KnowledgeUserResponse):
     write_access: Optional[bool] = False
 
 
-@router.get("/", response_model=list[KnowledgeAccessResponse])
-async def get_knowledge(user=Depends(get_verified_user)):
-    # Return knowledge bases with read access
-    knowledge_bases = []
-    if user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL:
-        knowledge_bases = Knowledges.get_knowledge_bases()
-    else:
-        knowledge_bases = Knowledges.get_knowledge_bases_by_user_id(user.id, "read")
-
-    return [
-        KnowledgeAccessResponse(
-            **knowledge_base.model_dump(),
-            files=Knowledges.get_file_metadatas_by_id(knowledge_base.id),
-            write_access=(
-                user.id == knowledge_base.user_id
-                or has_access(user.id, "write", knowledge_base.access_control)
-            ),
-        )
-        for knowledge_base in knowledge_bases
-    ]
+class KnowledgeAccessListResponse(BaseModel):
+    items: list[KnowledgeAccessResponse]
+    total: int
 
 
-@router.get("/list", response_model=list[KnowledgeAccessResponse])
-async def get_knowledge_list(user=Depends(get_verified_user)):
-    # Return knowledge bases with write access
-    knowledge_bases = []
-    if user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL:
-        knowledge_bases = Knowledges.get_knowledge_bases()
-    else:
-        knowledge_bases = Knowledges.get_knowledge_bases_by_user_id(user.id, "read")
+@router.get("/", response_model=KnowledgeAccessListResponse)
+async def get_knowledge_bases(page: Optional[int] = 1, user=Depends(get_verified_user)):
+    page = max(page, 1)
+    limit = PAGE_ITEM_COUNT
+    skip = (page - 1) * limit
 
-    return [
-        KnowledgeAccessResponse(
-            **knowledge_base.model_dump(),
-            files=Knowledges.get_file_metadatas_by_id(knowledge_base.id),
-            write_access=(
-                user.id == knowledge_base.user_id
-                or has_access(user.id, "write", knowledge_base.access_control)
-            ),
-        )
-        for knowledge_base in knowledge_bases
-    ]
+    filter = {}
+    if not user.role == "admin" or not BYPASS_ADMIN_ACCESS_CONTROL:
+        groups = Groups.get_groups_by_member_id(user.id)
+        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
+    )
+
+    return KnowledgeAccessListResponse(
+        items=[
+            KnowledgeAccessResponse(
+                **knowledge_base.model_dump(),
+                write_access=(
+                    user.id == knowledge_base.user_id
+                    or has_access(user.id, "write", knowledge_base.access_control)
+                ),
+            )
+            for knowledge_base in result.items
+        ],
+        total=result.total,
+    )
+
+
+@router.get("/search", response_model=KnowledgeAccessListResponse)
+async def search_knowledge_bases(
+    query: Optional[str] = None,
+    view_option: Optional[str] = None,
+    page: Optional[int] = 1,
+    user=Depends(get_verified_user),
+):
+    page = max(page, 1)
+    limit = PAGE_ITEM_COUNT
+    skip = (page - 1) * limit
+
+    filter = {}
+    if query:
+        filter["query"] = query
+    if view_option:
+        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)
+        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
+    )
+
+    return KnowledgeAccessListResponse(
+        items=[
+            KnowledgeAccessResponse(
+                **knowledge_base.model_dump(),
+                write_access=(
+                    user.id == knowledge_base.user_id
+                    or has_access(user.id, "write", knowledge_base.access_control)
+                ),
+            )
+            for knowledge_base in result.items
+        ],
+        total=result.total,
+    )
+
+
+@router.get("/search/files", response_model=KnowledgeFileListResponse)
+async def search_knowledge_files(
+    query: Optional[str] = None,
+    page: Optional[int] = 1,
+    user=Depends(get_verified_user),
+):
+    page = max(page, 1)
+    limit = PAGE_ITEM_COUNT
+    skip = (page - 1) * limit
+
+    filter = {}
+    if query:
+        filter["query"] = query
+
+    groups = Groups.get_groups_by_member_id(user.id)
+    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)
 
 
 ############################
@@ -198,7 +261,7 @@ async def reindex_knowledge_files(request: Request, user=Depends(get_verified_us
 
 
 class KnowledgeFilesResponse(KnowledgeResponse):
-    files: list[FileMetadataResponse]
+    files: Optional[list[FileMetadataResponse]] = None
     write_access: Optional[bool] = False
 
 
@@ -215,7 +278,6 @@ async def get_knowledge_by_id(id: str, user=Depends(get_verified_user)):
 
             return KnowledgeFilesResponse(
                 **knowledge.model_dump(),
-                files=Knowledges.get_file_metadatas_by_id(knowledge.id),
                 write_access=(
                     user.id == knowledge.user_id
                     or has_access(user.id, "write", knowledge.access_control)