Skip to content

Commit b77735a

Browse files
perf: ai-chat-node: Improve the performance of reading data (#6222)
* perf: `ai-chat-node`: Improve the performance of reading data * chore: simplify tool existence checks in base_chat_node.py Refactored tool validation logic to simplify checks. * chore: update tool filtering to include only active tools Filter custom tools to include only active ones. --------- Co-authored-by: 刘瑞斌 <bin@fit2cloud.com>
1 parent 092e85d commit b77735a

2 files changed

Lines changed: 24 additions & 13 deletions

File tree

apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -273,9 +273,11 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids
273273
workspace_id)
274274
if tool_ids and len(tool_ids) > 0: # 如果有工具ID,则将其转换为MCP
275275
self.context['tool_ids'] = tool_ids
276+
custom_tools_map = {str(t.id): t for t in
277+
QuerySet(Tool).filter(id__in=tool_ids, tool_type=ToolType.CUSTOM, is_active=True)}
276278
for tool_id in tool_ids:
277-
tool = QuerySet(Tool).filter(id=tool_id, tool_type=ToolType.CUSTOM).first()
278-
if tool is None or not tool.is_active:
279+
tool = custom_tools_map.get(str(tool_id))
280+
if tool is None:
279281
continue
280282
executor = ToolExecutor()
281283
if tool.init_params is not None:
@@ -288,16 +290,21 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids
288290

289291
if application_ids and len(application_ids) > 0:
290292
self.context['application_ids'] = application_ids
293+
apps_map = {str(a.id): a for a in
294+
QuerySet(Application).filter(id__in=application_ids, is_publish=True)}
295+
app_keys_map = {str(ak.application_id): ak for ak in
296+
QuerySet(ApplicationApiKey).filter(application_id__in=application_ids, is_active=True)}
297+
app_access_tokens_map = {str(at.application_id): at for at in
298+
QuerySet(ApplicationAccessToken).filter(
299+
application_id__in=application_ids)}
291300
for application_id in application_ids:
292-
app = QuerySet(Application).filter(id=application_id, is_publish=True).first()
301+
app = apps_map.get(str(application_id))
293302
if app is None:
294303
continue
295-
app_key = QuerySet(ApplicationApiKey).filter(application_id=application_id, is_active=True).first()
304+
app_key = app_keys_map.get(str(application_id))
296305
if app_key is not None:
297306
api_key = app_key.secret_key
298-
application_access_token = QuerySet(ApplicationAccessToken).filter(
299-
application_id=app_key.application_id
300-
).first()
307+
application_access_token = app_access_tokens_map.get(str(app_key.application_id))
301308
if application_access_token is not None and application_access_token.authentication:
302309
raise AppApiException(
303310
500,
@@ -316,10 +323,11 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids
316323
if skill_tool_ids and len(skill_tool_ids) > 0:
317324
self.context['skill_tool_ids'] = skill_tool_ids
318325
skill_file_items = []
319-
326+
skill_tools_map = {str(t.id): t for t in
327+
QuerySet(Tool).filter(id__in=skill_tool_ids, is_active=True)}
320328
for tool_id in skill_tool_ids:
321-
tool = QuerySet(Tool).filter(id=tool_id, is_active=True).first()
322-
if tool is None or tool.is_active is False:
329+
tool = skill_tools_map.get(str(tool_id))
330+
if tool is None:
323331
continue
324332
init_params_default_value = {i["field"]: i.get('default_value') for i in tool.init_field_list}
325333
if tool.init_params is not None:

apps/knowledge/vector/base_vector.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -170,11 +170,14 @@ def search(
170170
if knowledge_id_list is None or len(knowledge_id_list) == 0:
171171
return []
172172
query_text = normalize_for_embedding(query_text)
173-
embedding_query = embedding.embed_query(query_text)
173+
query_embedding = embedding.embed_query(query_text)
174174
result = self.query(
175-
embedding_query, knowledge_id_list, exclude_document_id_list, exclude_paragraph_list, is_active, 1, 3, 0.65
175+
query_text, query_embedding,
176+
knowledge_id_list, None,
177+
exclude_document_id_list, exclude_paragraph_list,
178+
is_active, 3, 0.65, SearchMode.embedding
176179
)
177-
return result[0]
180+
return result[0] if result else None
178181

179182
@abstractmethod
180183
def query(

0 commit comments

Comments
 (0)