From 86a32517e560cda354570c24c4fd8ab96a430dc8 Mon Sep 17 00:00:00 2001 From: Yeuoly Date: Mon, 1 Apr 2024 15:42:00 +0800 Subject: [PATCH] fix: tool variable selectors --- api/core/workflow/nodes/tool/tool_node.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/api/core/workflow/nodes/tool/tool_node.py b/api/core/workflow/nodes/tool/tool_node.py index bbad1f1c36..83e51a04c0 100644 --- a/api/core/workflow/nodes/tool/tool_node.py +++ b/api/core/workflow/nodes/tool/tool_node.py @@ -7,7 +7,6 @@ from core.tools.entities.tool_entities import ToolInvokeMessage from core.tools.tool_engine import ToolEngine from core.tools.tool_manager import ToolManager from core.tools.utils.message_transformer import ToolFileMessageTransformer -from core.workflow.entities.base_node_data_entities import BaseNodeData from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult, NodeType from core.workflow.entities.variable_pool import VariablePool from core.workflow.nodes.base_node import BaseNode @@ -178,10 +177,21 @@ class ToolNode(BaseNode): ]) @classmethod - def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[str, list[str]]: + def _extract_variable_selector_to_variable_mapping(cls, node_data: ToolNodeData) -> dict[str, list[str]]: """ Extract variable selector to variable mapping :param node_data: node data :return: """ - return {} \ No newline at end of file + result = {} + for parameter in node_data.tool_parameters: + if parameter.value_type == 'variable': + if isinstance(parameter.variable_value, str): + parser = VariableTemplateParser(parameter.variable_value) + variable_selectors = parser.extract_variable_selectors() + for selector in variable_selectors: + result[selector.variable] = selector.value_selector + elif isinstance(parameter.variable_value, list): + result[parameter.parameter_name] = parameter.variable_value + + return result \ No newline at end of file