refac

backend/open_webui/models/knowledge.py

@@ -144,13 +144,18 @@ class KnowledgeTable:
         return AccessGrants.get_grants_by_resource("knowledge", knowledge_id, db=db)
 
     def _to_knowledge_model(
-        self, knowledge: Knowledge, db: Optional[Session] = None
+        self,
+        knowledge: Knowledge,
+        access_grants: Optional[list[AccessGrantModel]] = None,
+        db: Optional[Session] = None,
     ) -> KnowledgeModel:
         knowledge_data = KnowledgeModel.model_validate(knowledge).model_dump(
             exclude={"access_grants"}
         )
-        knowledge_data["access_grants"] = self._get_access_grants(
-            knowledge_data["id"], db=db
+        knowledge_data["access_grants"] = (
+            access_grants
+            if access_grants is not None
+            else self._get_access_grants(knowledge_data["id"], db=db)
         )
         return KnowledgeModel.model_validate(knowledge_data)
 
@@ -192,9 +197,13 @@ class KnowledgeTable:
                 db.query(Knowledge).order_by(Knowledge.updated_at.desc()).all()
             )
             user_ids = list(set(knowledge.user_id for knowledge in all_knowledge))
+            knowledge_ids = [knowledge.id for knowledge in all_knowledge]
 
             users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else []
             users_dict = {user.id: user for user in users}
+            grants_map = AccessGrants.get_grants_by_resources(
+                "knowledge", knowledge_ids, db=db
+            )
 
             knowledge_bases = []
             for knowledge in all_knowledge:
@@ -202,7 +211,11 @@ class KnowledgeTable:
                 knowledge_bases.append(
                     KnowledgeUserModel.model_validate(
                         {
-                            **self._to_knowledge_model(knowledge, db=db).model_dump(),
+                            **self._to_knowledge_model(
+                                knowledge,
+                                access_grants=grants_map.get(knowledge.id, []),
+                                db=db,
+                            ).model_dump(),
                             "user": user.model_dump() if user else None,
                         }
                     )
@@ -261,13 +274,22 @@ class KnowledgeTable:
 
                 items = query.all()
 
+                knowledge_ids = [kb.id for kb, _ in items]
+                grants_map = AccessGrants.get_grants_by_resources(
+                    "knowledge", knowledge_ids, db=db
+                )
+
                 knowledge_bases = []
                 for knowledge_base, user in items:
                     knowledge_bases.append(
                         KnowledgeUserModel.model_validate(
                             {
                                 **self._to_knowledge_model(
-                                    knowledge_base, db=db
+                                    knowledge_base,
+                                    access_grants=grants_map.get(
+                                        knowledge_base.id, []
+                                    ),
+                                    db=db,
                                 ).model_dump(),
                                 "user": (
                                     UserModel.model_validate(user).model_dump()
@@ -440,8 +462,16 @@ class KnowledgeTable:
                     .filter(KnowledgeFile.file_id == file_id)
                     .all()
                 )
+                knowledge_ids = [k.id for k in knowledges]
+                grants_map = AccessGrants.get_grants_by_resources(
+                    "knowledge", knowledge_ids, db=db
+                )
                 return [
-                    self._to_knowledge_model(knowledge, db=db)
+                    self._to_knowledge_model(
+                        knowledge,
+                        access_grants=grants_map.get(knowledge.id, []),
+                        db=db,
+                    )
                     for knowledge in knowledges
                 ]
         except Exception: