From 26246926dfb5890636e006460e55faf78a4bf39d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E8=89=AF?= <841369634@qq.com> Date: Sun, 14 Jun 2026 23:24:59 +0800 Subject: [PATCH] perf: `ai-chat-node`: Improve the performance of reading data --- .../chat_pipeline/I_base_chat_pipeline.py | 2 +- .../ai_chat_step_node/impl/base_chat_node.py | 26 +++++++++++++------ apps/knowledge/vector/base_vector.py | 9 ++++--- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/apps/application/chat_pipeline/I_base_chat_pipeline.py b/apps/application/chat_pipeline/I_base_chat_pipeline.py index adeb3cef2c4..f231c2c4514 100644 --- a/apps/application/chat_pipeline/I_base_chat_pipeline.py +++ b/apps/application/chat_pipeline/I_base_chat_pipeline.py @@ -26,7 +26,7 @@ def __init__(self, _id: str, document_id: str, knowledge_id: str, content: str, self.knowledge_id = knowledge_id self.content = content self.title = title - self.status = status, + self.status = status self.is_active = is_active self.comprehensive_score = comprehensive_score self.similarity = similarity diff --git a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py index 0019cb312c6..b784f3fdfba 100644 --- a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py +++ b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py @@ -123,6 +123,8 @@ def write_context(node_variable: Dict, workflow_variable: Dict, node: INode, wor def get_default_model_params_setting(model_id): model = QuerySet(Model).filter(id=model_id).first() + if model is None: + raise Exception(f"{_('Model does not exist')}，model_id = {model_id}") credential = get_model_credential(model.provider, model.model_type, model.model_name) model_params_setting = credential.get_model_params_setting_form( model.model_name).get_default_form_data() @@ -273,8 +275,10 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids workspace_id) if tool_ids and len(tool_ids) > 0: # 如果有工具ID，则将其转换为MCP self.context['tool_ids'] = tool_ids + custom_tools_map = {str(t.id): t for t in + QuerySet(Tool).filter(id__in=tool_ids, tool_type=ToolType.CUSTOM)} for tool_id in tool_ids: - tool = QuerySet(Tool).filter(id=tool_id, tool_type=ToolType.CUSTOM).first() + tool = custom_tools_map.get(str(tool_id)) if tool is None or not tool.is_active: continue executor = ToolExecutor() @@ -288,16 +292,21 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids if application_ids and len(application_ids) > 0: self.context['application_ids'] = application_ids + apps_map = {str(a.id): a for a in + QuerySet(Application).filter(id__in=application_ids, is_publish=True)} + app_keys_map = {str(ak.application_id): ak for ak in + QuerySet(ApplicationApiKey).filter(application_id__in=application_ids, is_active=True)} + app_access_tokens_map = {str(at.application_id): at for at in + QuerySet(ApplicationAccessToken).filter( + application_id__in=application_ids)} for application_id in application_ids: - app = QuerySet(Application).filter(id=application_id, is_publish=True).first() + app = apps_map.get(str(application_id)) if app is None: continue - app_key = QuerySet(ApplicationApiKey).filter(application_id=application_id, is_active=True).first() + app_key = app_keys_map.get(str(application_id)) if app_key is not None: api_key = app_key.secret_key - application_access_token = QuerySet(ApplicationAccessToken).filter( - application_id=app_key.application_id - ).first() + application_access_token = app_access_tokens_map.get(str(app_key.application_id)) if application_access_token is not None and application_access_token.authentication: raise AppApiException( 500, @@ -316,9 +325,10 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids if skill_tool_ids and len(skill_tool_ids) > 0: self.context['skill_tool_ids'] = skill_tool_ids skill_file_items = [] - + skill_tools_map = {str(t.id): t for t in + QuerySet(Tool).filter(id__in=skill_tool_ids, is_active=True)} for tool_id in skill_tool_ids: - tool = QuerySet(Tool).filter(id=tool_id, is_active=True).first() + tool = skill_tools_map.get(str(tool_id)) if tool is None or tool.is_active is False: continue init_params_default_value = {i["field"]: i.get('default_value') for i in tool.init_field_list} diff --git a/apps/knowledge/vector/base_vector.py b/apps/knowledge/vector/base_vector.py index 75718775c42..32478f2c37a 100644 --- a/apps/knowledge/vector/base_vector.py +++ b/apps/knowledge/vector/base_vector.py @@ -170,11 +170,14 @@ def search( if knowledge_id_list is None or len(knowledge_id_list) == 0: return [] query_text = normalize_for_embedding(query_text) - embedding_query = embedding.embed_query(query_text) + query_embedding = embedding.embed_query(query_text) result = self.query( - embedding_query, knowledge_id_list, exclude_document_id_list, exclude_paragraph_list, is_active, 1, 3, 0.65 + query_text, query_embedding, + knowledge_id_list, None, + exclude_document_id_list, exclude_paragraph_list, + is_active, 3, 0.65, SearchMode.embedding ) - return result[0] + return result[0] if result else None @abstractmethod def query(