enh/refac: kb pagination

backend/open_webui/models/knowledge.py

@@ -5,6 +5,7 @@ from typing import Optional
 import uuid
 
 from open_webui.internal.db import Base, get_db
+
 from open_webui.env import SRC_LOG_LEVELS
 
 from open_webui.models.files import (
@@ -30,6 +31,8 @@ from sqlalchemy import (
 )
 
 from open_webui.utils.access_control import has_access
+from open_webui.utils.db.access_control import has_permission
+
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["MODELS"])
@@ -145,6 +148,11 @@ class FileUserResponse(FileModelResponse):
     user: Optional[UserResponse] = None
 
 
+class KnowledgeListResponse(BaseModel):
+    items: list[KnowledgeUserModel]
+    total: int
+
+
 class KnowledgeFileListResponse(BaseModel):
     items: list[FileUserResponse]
     total: int
@@ -177,12 +185,13 @@ class KnowledgeTable:
             except Exception:
                 return None
 
-    def get_knowledge_bases(self) -> list[KnowledgeUserModel]:
+    def get_knowledge_bases(
+        self, skip: int = 0, limit: int = 30
+    ) -> list[KnowledgeUserModel]:
         with get_db() as db:
             all_knowledge = (
                 db.query(Knowledge).order_by(Knowledge.updated_at.desc()).all()
             )
-
             user_ids = list(set(knowledge.user_id for knowledge in all_knowledge))
 
             users = Users.get_users_by_user_ids(user_ids) if user_ids else []
@@ -201,6 +210,126 @@ class KnowledgeTable:
                 )
             return knowledge_bases
 
+    def search_knowledge_bases(
+        self, user_id: str, filter: dict, skip: int = 0, limit: int = 30
+    ) -> KnowledgeListResponse:
+        try:
+            with get_db() as db:
+                query = db.query(Knowledge, User).outerjoin(
+                    User, User.id == Knowledge.user_id
+                )
+
+                if filter:
+                    query_key = filter.get("query")
+                    if query_key:
+                        query = query.filter(
+                            or_(
+                                Knowledge.name.ilike(f"%{query_key}%"),
+                                Knowledge.description.ilike(f"%{query_key}%"),
+                            )
+                        )
+
+                    view_option = filter.get("view_option")
+                    if view_option == "created":
+                        query = query.filter(Knowledge.user_id == user_id)
+                    elif view_option == "shared":
+                        query = query.filter(Knowledge.user_id != user_id)
+
+                    query = has_permission(db, Knowledge, query, filter)
+
+                query = query.order_by(Knowledge.updated_at.desc())
+
+                total = query.count()
+                if skip:
+                    query = query.offset(skip)
+                if limit:
+                    query = query.limit(limit)
+
+                items = query.all()
+
+                knowledge_bases = []
+                for knowledge_base, user in items:
+                    knowledge_bases.append(
+                        KnowledgeUserModel.model_validate(
+                            {
+                                **KnowledgeModel.model_validate(
+                                    knowledge_base
+                                ).model_dump(),
+                                "user": (
+                                    UserModel.model_validate(user).model_dump()
+                                    if user
+                                    else None
+                                ),
+                            }
+                        )
+                    )
+
+                return KnowledgeListResponse(items=knowledge_bases, total=total)
+        except Exception as e:
+            print(e)
+            return KnowledgeListResponse(items=[], total=0)
+
+    def search_knowledge_files(
+        self, filter: dict, skip: int = 0, limit: int = 30
+    ) -> KnowledgeFileListResponse:
+        """
+        Scalable version: search files across all knowledge bases the user has
+        READ access to, without loading all KBs or using large IN() lists.
+        """
+        try:
+            with get_db() as db:
+                # Base query: join Knowledge → KnowledgeFile → File
+                query = (
+                    db.query(File, User)
+                    .join(KnowledgeFile, File.id == KnowledgeFile.file_id)
+                    .join(Knowledge, KnowledgeFile.knowledge_id == Knowledge.id)
+                    .outerjoin(User, User.id == KnowledgeFile.user_id)
+                )
+
+                # Apply access-control directly to the joined query
+                # This makes the database handle filtering, even with 10k+ KBs
+                query = has_permission(db, Knowledge, query, filter)
+
+                # Apply filename search
+                if filter:
+                    q = filter.get("query")
+                    if q:
+                        query = query.filter(File.filename.ilike(f"%{q}%"))
+
+                # Order by file changes
+                query = query.order_by(File.updated_at.desc())
+
+                # Count before pagination
+                total = query.count()
+
+                if skip:
+                    query = query.offset(skip)
+                if limit:
+                    query = query.limit(limit)
+
+                rows = query.all()
+
+                items = []
+                for file, user in rows:
+                    items.append(
+                        FileUserResponse(
+                            **FileModel.model_validate(file).model_dump(),
+                            user=(
+                                UserResponse(
+                                    **UserModel.model_validate(user).model_dump()
+                                )
+                                if user
+                                else None
+                            ),
+                        )
+                    )
+
+                return KnowledgeFileListResponse(items=items, total=total)
+
+        except Exception as e:
+            print("search_knowledge_files error:", e)
+            return KnowledgeFileListResponse(items=[], total=0)
+
     def check_access_by_user_id(self, id, user_id, permission="write") -> bool:
         knowledge = self.get_knowledge_by_id(id)
         if not knowledge: