feat: Knowledge base search based on authorized users (#3438)

Author: shaohuzhang1

apps/application/chat_pipeline/step/search_dataset_step/impl/base_search_dataset_step.py

@@ -16,6 +16,8 @@
 from application.chat_pipeline.I_base_chat_pipeline import ParagraphPipelineModel
 from application.chat_pipeline.step.search_dataset_step.i_search_dataset_step import ISearchDatasetStep
 from common.config.embedding_config import VectorStore, ModelManage
+from common.constants.permission_constants import RoleConstants
+from common.database_model_manage.database_model_manage import DatabaseModelManage
 from common.db.search import native_search
 from common.utils.common import get_file_content
 from knowledge.models import Paragraph, Knowledge
@@ -52,6 +54,11 @@ def execute(self, problem_text: str, knowledge_id_list: list[str], exclude_docum
                 search_mode: str = None,
                 workspace_id=None,
                 **kwargs) -> List[ParagraphPipelineModel]:
+        get_knowledge_list_of_authorized = DatabaseModelManage.get_model('get_knowledge_list_of_authorized')
+        chat_user_type = self.context.get('chat_user_type')
+        if get_knowledge_list_of_authorized is not None and RoleConstants.CHAT_USER.value.name == chat_user_type:
+            knowledge_id_list = get_knowledge_list_of_authorized(self.context.get('chat_user_id'),
+                                                                 knowledge_id_list)
         if len(knowledge_id_list) == 0:
             return []
         exec_problem_text = padding_problem_text if padding_problem_text is not None else problem_text

apps/application/flow/step_node/search_knowledge_node/impl/base_search_knowledge_node.py

@@ -15,6 +15,8 @@
 from application.flow.i_step_node import NodeResult
 from application.flow.step_node.search_knowledge_node.i_search_knowledge_node import ISearchKnowledgeStepNode
 from common.config.embedding_config import VectorStore
+from common.constants.permission_constants import RoleConstants
+from common.database_model_manage.database_model_manage import DatabaseModelManage
 from common.db.search import native_search
 from common.utils.common import get_file_content
 from knowledge.models import Document, Paragraph, Knowledge, SearchMode
@@ -64,6 +66,11 @@ def execute(self, knowledge_id_list, knowledge_setting, question,
                 exclude_paragraph_id_list=None,
                 **kwargs) -> NodeResult:
         self.context['question'] = question
+        get_knowledge_list_of_authorized = DatabaseModelManage.get_model('get_knowledge_list_of_authorized')
+        chat_user_type = self.workflow_manage.get_body().get('chat_user_type')
+        if get_knowledge_list_of_authorized is not None and RoleConstants.CHAT_USER.value.name == chat_user_type:
+            knowledge_id_list = get_knowledge_list_of_authorized(self.workflow_manage.get_body().get('chat_user_id'),
+                                                                 knowledge_id_list)
         if len(knowledge_id_list) == 0:
             return get_none_result(question)
         model_id = get_embedding_id(knowledge_id_list)