This commit is contained in:
takatost
2024-07-22 19:57:32 +08:00
372 changed files with 9779 additions and 1678 deletions

View File

@ -1,7 +1,5 @@
import json
from typing import cast
from core.file.file_obj import FileVar
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.nodes.answer.answer_stream_generate_router import AnswerStreamGeneratorRouter
@ -18,7 +16,7 @@ from models.workflow import WorkflowNodeExecutionStatus
class AnswerNode(BaseNode):
_node_data_cls = AnswerNodeData
node_type = NodeType.ANSWER
_node_type: NodeType = NodeType.ANSWER
def _run(self) -> NodeRunResult:
"""
@ -36,31 +34,12 @@ class AnswerNode(BaseNode):
if part.type == GenerateRouteChunk.ChunkType.VAR:
part = cast(VarGenerateRouteChunk, part)
value_selector = part.value_selector
value = self.graph_runtime_state.variable_pool.get_variable_value(
value = self.graph_runtime_state.variable_pool.get(
variable_selector=value_selector
)
text = ''
if isinstance(value, str | int | float):
text = str(value)
elif isinstance(value, dict):
# other types
text = json.dumps(value, ensure_ascii=False)
elif isinstance(value, FileVar):
# convert file to markdown
text = value.to_markdown()
elif isinstance(value, list):
for item in value:
if isinstance(item, FileVar):
text += item.to_markdown() + ' '
text = text.strip()
if not text and value:
# other types
text = json.dumps(value, ensure_ascii=False)
answer += text
if value:
answer += value.markdown
else:
part = cast(TextGenerateRouteChunk, part)
answer += part.text

View File

@ -1,6 +1,6 @@
from abc import ABC, abstractmethod
from collections.abc import Generator
from typing import Optional
from collections.abc import Generator, Mapping
from typing import Any, Optional
from core.workflow.entities.base_node_data_entities import BaseIterationState, BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType
@ -17,7 +17,7 @@ class BaseNode(ABC):
_node_type: NodeType
def __init__(self,
config: dict,
config: Mapping[str, Any],
graph_init_params: GraphInitParams,
graph: Graph,
graph_runtime_state: GraphRuntimeState,

View File

@ -57,11 +57,8 @@ class CodeNode(BaseNode):
variables = {}
for variable_selector in node_data.variables:
variable = variable_selector.variable
value = self.graph_runtime_state.variable_pool.get_variable_value(
variable_selector=variable_selector.value_selector
)
variables[variable] = value
value = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector)
variables[variable] = value.value if value else None
# Run code
try:
result = CodeExecutor.execute_workflow_code_template(

View File

@ -9,7 +9,7 @@ from models.workflow import WorkflowNodeExecutionStatus
class EndNode(BaseNode):
_node_data_cls = EndNodeData
node_type = NodeType.END
_node_type = NodeType.END
def _run(self) -> NodeRunResult:
"""
@ -22,11 +22,8 @@ class EndNode(BaseNode):
outputs = {}
for variable_selector in output_variables:
value = self.graph_runtime_state.variable_pool.get_variable_value(
variable_selector=variable_selector.value_selector
)
outputs[variable_selector.variable] = value
value = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector)
outputs[variable_selector.variable] = value.value if value else None
return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED,
@ -43,7 +40,7 @@ class EndNode(BaseNode):
:return:
"""
node_data = cls._node_data_cls(**config.get("data", {}))
node_data = cast(cls._node_data_cls, node_data)
node_data = cast(EndNodeData, node_data)
return cls.extract_generate_nodes_from_node_data(graph, node_data)
@ -55,7 +52,7 @@ class EndNode(BaseNode):
:param node_data: node data object
:return:
"""
nodes = graph.get('nodes')
nodes = graph.get('nodes', [])
node_mapping = {node.get('id'): node for node in nodes}
variable_selectors = node_data.outputs

View File

@ -58,4 +58,3 @@ class HttpRequestNodeData(BaseNodeData):
params: str
body: Optional[HttpRequestNodeBody] = None
timeout: Optional[HttpRequestNodeTimeout] = None
mask_authorization_header: Optional[bool] = True

View File

@ -9,7 +9,7 @@ import httpx
import core.helper.ssrf_proxy as ssrf_proxy
from configs import dify_config
from core.workflow.entities.variable_entities import VariableSelector
from core.workflow.entities.variable_pool import ValueType, VariablePool
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.nodes.http_request.entities import (
HttpRequestNodeAuthorization,
HttpRequestNodeBody,
@ -212,13 +212,11 @@ class HttpExecutor:
raise ValueError('self.authorization config is required')
if authorization.config is None:
raise ValueError('authorization config is required')
if authorization.config.type != 'bearer' and authorization.config.header is None:
raise ValueError('authorization config header is required')
if self.authorization.config.api_key is None:
raise ValueError('api_key is required')
if not self.authorization.config.header:
if not authorization.config.header:
authorization.config.header = 'Authorization'
if self.authorization.config.type == 'bearer':
@ -283,7 +281,7 @@ class HttpExecutor:
# validate response
return self._validate_and_parse_response(response)
def to_raw_request(self, mask_authorization_header: Optional[bool] = True) -> str:
def to_raw_request(self) -> str:
"""
convert to raw request
"""
@ -295,16 +293,15 @@ class HttpExecutor:
headers = self._assembling_headers()
for k, v in headers.items():
if mask_authorization_header:
# get authorization header
if self.authorization.type == 'api-key':
authorization_header = 'Authorization'
if self.authorization.config and self.authorization.config.header:
authorization_header = self.authorization.config.header
# get authorization header
if self.authorization.type == 'api-key':
authorization_header = 'Authorization'
if self.authorization.config and self.authorization.config.header:
authorization_header = self.authorization.config.header
if k.lower() == authorization_header.lower():
raw_request += f'{k}: {"*" * len(v)}\n'
continue
if k.lower() == authorization_header.lower():
raw_request += f'{k}: {"*" * len(v)}\n'
continue
raw_request += f'{k}: {v}\n'
@ -336,16 +333,13 @@ class HttpExecutor:
if variable_pool:
variable_value_mapping = {}
for variable_selector in variable_selectors:
value = variable_pool.get_variable_value(
variable_selector=variable_selector.value_selector, target_value_type=ValueType.STRING
)
if value is None:
variable = variable_pool.get(variable_selector.value_selector)
if variable is None:
raise ValueError(f'Variable {variable_selector.variable} not found')
if escape_quotes and isinstance(value, str):
value = value.replace('"', '\\"')
if escape_quotes and isinstance(variable.value, str):
value = variable.value.replace('"', '\\"')
else:
value = variable.value
variable_value_mapping[variable_selector.variable] = value
return variable_template_parser.format(variable_value_mapping), variable_selectors

View File

@ -3,6 +3,7 @@ from mimetypes import guess_extension
from os import path
from typing import cast
from core.app.segments import parser
from core.file.file_obj import FileTransferMethod, FileType, FileVar
from core.tools.tool_file_manager import ToolFileManager
from core.workflow.entities.base_node_data_entities import BaseNodeData
@ -50,6 +51,9 @@ class HttpRequestNode(BaseNode):
def _run(self) -> NodeRunResult:
node_data: HttpRequestNodeData = cast(HttpRequestNodeData, self.node_data)
# TODO: Switch to use segment directly
if node_data.authorization.config and node_data.authorization.config.api_key:
node_data.authorization.config.api_key = parser.convert_template(template=node_data.authorization.config.api_key, variable_pool=variable_pool).text
# init http executor
http_executor = None
@ -66,9 +70,7 @@ class HttpRequestNode(BaseNode):
process_data = {}
if http_executor:
process_data = {
'request': http_executor.to_raw_request(
mask_authorization_header=node_data.mask_authorization_header
),
'request': http_executor.to_raw_request(),
}
return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
@ -87,9 +89,7 @@ class HttpRequestNode(BaseNode):
'files': files,
},
process_data={
'request': http_executor.to_raw_request(
mask_authorization_header=node_data.mask_authorization_header,
),
'request': http_executor.to_raw_request(),
},
)

View File

@ -10,7 +10,7 @@ from models.workflow import WorkflowNodeExecutionStatus
class IfElseNode(BaseNode):
_node_data_cls = IfElseNodeData
node_type = NodeType.IF_ELSE
_node_type = NodeType.IF_ELSE
def _run(self) -> NodeRunResult:
"""

View File

@ -21,7 +21,8 @@ class IterationNode(BaseIterationNode):
"""
Run the node.
"""
iterator = variable_pool.get_variable_value(cast(IterationNodeData, self.node_data).iterator_selector)
self.node_data = cast(IterationNodeData, self.node_data)
iterator = variable_pool.get_any(self.node_data.iterator_selector)
if not isinstance(iterator, list):
raise ValueError(f"Invalid iterator value: {iterator}, please provide a list.")
@ -64,15 +65,15 @@ class IterationNode(BaseIterationNode):
"""
node_data = cast(IterationNodeData, self.node_data)
variable_pool.append_variable(self.node_id, ['index'], state.index)
variable_pool.add((self.node_id, 'index'), state.index)
# get the iterator value
iterator = variable_pool.get_variable_value(node_data.iterator_selector)
iterator = variable_pool.get_any(node_data.iterator_selector)
if iterator is None or not isinstance(iterator, list):
return
if state.index < len(iterator):
variable_pool.append_variable(self.node_id, ['item'], iterator[state.index])
variable_pool.add((self.node_id, 'item'), iterator[state.index])
def _next_iteration(self, variable_pool: VariablePool, state: IterationState):
"""
@ -88,7 +89,7 @@ class IterationNode(BaseIterationNode):
:return: True if iteration limit is reached, False otherwise
"""
node_data = cast(IterationNodeData, self.node_data)
iterator = variable_pool.get_variable_value(node_data.iterator_selector)
iterator = variable_pool.get_any(node_data.iterator_selector)
if iterator is None or not isinstance(iterator, list):
return True
@ -101,9 +102,9 @@ class IterationNode(BaseIterationNode):
:param variable_pool: variable pool
"""
output_selector = cast(IterationNodeData, self.node_data).output_selector
output = variable_pool.get_variable_value(output_selector)
output = variable_pool.get_any(output_selector)
# clear the output for this iteration
variable_pool.append_variable(self.node_id, output_selector[1:], None)
variable_pool.remove([self.node_id] + output_selector[1:])
state.current_output = output
if output is not None:
state.outputs.append(output)

View File

@ -21,7 +21,7 @@ from models.dataset import Dataset, Document, DocumentSegment
from models.workflow import WorkflowNodeExecutionStatus
default_retrieval_model = {
'search_method': RetrievalMethod.SEMANTIC_SEARCH,
'search_method': RetrievalMethod.SEMANTIC_SEARCH.value,
'reranking_enable': False,
'reranking_model': {
'reranking_provider_name': '',
@ -40,7 +40,8 @@ class KnowledgeRetrievalNode(BaseNode):
node_data = cast(KnowledgeRetrievalNodeData, self.node_data)
# extract variables
query = self.graph_runtime_state.variable_pool.get_variable_value(variable_selector=node_data.query_variable_selector)
variable = self.graph_runtime_state.variable_pool.get(node_data.query_variable_selector)
query = variable.value if variable else None
variables = {
'query': query
}

View File

@ -44,7 +44,7 @@ from models.workflow import WorkflowNodeExecutionStatus
class LLMNode(BaseNode):
_node_data_cls = LLMNodeData
node_type = NodeType.LLM
_node_type = NodeType.LLM
def _run(self) -> Generator[RunEvent, None, None]:
"""
@ -98,7 +98,7 @@ class LLMNode(BaseNode):
# fetch prompt messages
prompt_messages, stop = self._fetch_prompt_messages(
node_data=node_data,
query=variable_pool.get_variable_value(['sys', SystemVariable.QUERY.value]) # type: ignore
query=variable_pool.get_any(['sys', SystemVariable.QUERY.value])
if node_data.memory else None,
query_prompt_template=node_data.memory.query_prompt_template if node_data.memory else None,
inputs=inputs,
@ -276,8 +276,8 @@ class LLMNode(BaseNode):
for variable_selector in node_data.prompt_config.jinja2_variables or []:
variable = variable_selector.variable
value = variable_pool.get_variable_value(
variable_selector=variable_selector.value_selector
value = variable_pool.get_any(
variable_selector.value_selector
)
def parse_dict(d: dict) -> str:
@ -340,7 +340,7 @@ class LLMNode(BaseNode):
variable_selectors = variable_template_parser.extract_variable_selectors()
for variable_selector in variable_selectors:
variable_value = variable_pool.get_variable_value(variable_selector.value_selector)
variable_value = variable_pool.get_any(variable_selector.value_selector)
if variable_value is None:
raise ValueError(f'Variable {variable_selector.variable} not found')
@ -351,7 +351,7 @@ class LLMNode(BaseNode):
query_variable_selectors = (VariableTemplateParser(template=memory.query_prompt_template)
.extract_variable_selectors())
for variable_selector in query_variable_selectors:
variable_value = variable_pool.get_variable_value(variable_selector.value_selector)
variable_value = variable_pool.get_any(variable_selector.value_selector)
if variable_value is None:
raise ValueError(f'Variable {variable_selector.variable} not found')
@ -369,7 +369,7 @@ class LLMNode(BaseNode):
if not node_data.vision.enabled:
return []
files = variable_pool.get_variable_value(['sys', SystemVariable.FILES.value])
files = variable_pool.get_any(['sys', SystemVariable.FILES.value])
if not files:
return []
@ -388,7 +388,7 @@ class LLMNode(BaseNode):
if not node_data.context.variable_selector:
return
context_value = variable_pool.get_variable_value(node_data.context.variable_selector)
context_value = variable_pool.get_any(node_data.context.variable_selector)
if context_value:
if isinstance(context_value, str):
yield RunRetrieverResourceEvent(
@ -530,7 +530,7 @@ class LLMNode(BaseNode):
return None
# get conversation id
conversation_id = variable_pool.get_variable_value(['sys', SystemVariable.CONVERSATION_ID.value])
conversation_id = variable_pool.get_any(['sys', SystemVariable.CONVERSATION_ID.value])
if conversation_id is None:
return None

View File

@ -71,9 +71,10 @@ class ParameterExtractorNode(LLMNode):
Run the node.
"""
node_data = cast(ParameterExtractorNodeData, self.node_data)
query = self.graph_runtime_state.variable_pool.get_variable_value(node_data.query)
if not query:
variable = self.graph_runtime_state.variable_pool.get(node_data.query)
if not variable:
raise ValueError("Input variable content not found or is empty")
query = variable.value
inputs = {
'query': query,
@ -567,7 +568,8 @@ class ParameterExtractorNode(LLMNode):
variable_template_parser = VariableTemplateParser(instruction)
inputs = {}
for selector in variable_template_parser.extract_variable_selectors():
inputs[selector.variable] = variable_pool.get_variable_value(selector.value_selector)
variable = variable_pool.get(selector.value_selector)
inputs[selector.variable] = variable.value if variable else None
return variable_template_parser.format(inputs)

View File

@ -43,7 +43,8 @@ class QuestionClassifierNode(LLMNode):
variable_pool = self.graph_runtime_state.variable_pool
# extract variables
query = variable_pool.get_variable_value(variable_selector=node_data.query_variable_selector)
variable = variable_pool.get(node_data.query_variable_selector)
query = variable.value if variable else None
variables = {
'query': query
}
@ -305,7 +306,8 @@ class QuestionClassifierNode(LLMNode):
variable_template_parser = VariableTemplateParser(template=instruction)
variable_selectors.extend(variable_template_parser.extract_variable_selectors())
for variable_selector in variable_selectors:
variable_value = variable_pool.get_variable_value(variable_selector.value_selector)
variable = variable_pool.get(variable_selector.value_selector)
variable_value = variable.value if variable else None
if variable_value is None:
raise ValueError(f'Variable {variable_selector.variable} not found')

View File

@ -8,7 +8,7 @@ from models.workflow import WorkflowNodeExecutionStatus
class StartNode(BaseNode):
_node_data_cls = StartNodeData
node_type = NodeType.START
_node_type = NodeType.START
def _run(self) -> NodeRunResult:
"""
@ -16,7 +16,7 @@ class StartNode(BaseNode):
:return:
"""
# Get cleaned inputs
cleaned_inputs = self.graph_runtime_state.variable_pool.user_inputs
cleaned_inputs = dict(self.graph_runtime_state.variable_pool.user_inputs)
for var in self.graph_runtime_state.variable_pool.system_variables:
cleaned_inputs['sys.' + var.value] = self.graph_runtime_state.variable_pool.system_variables[var]

View File

@ -44,12 +44,9 @@ class TemplateTransformNode(BaseNode):
# Get variables
variables = {}
for variable_selector in node_data.variables:
variable = variable_selector.variable
value = self.graph_runtime_state.variable_pool.get_variable_value(
variable_selector=variable_selector.value_selector
)
variables[variable] = value
variable_name = variable_selector.variable
value = self.graph_runtime_state.variable_pool.get_any(variable_selector.value_selector)
variables[variable_name] = value
# Run code
try:
result = CodeExecutor.execute_workflow_code_template(

View File

@ -29,6 +29,7 @@ class ToolEntity(BaseModel):
class ToolNodeData(BaseNodeData, ToolEntity):
class ToolInput(BaseModel):
# TODO: check this type
value: Union[Any, list[str]]
type: Literal['mixed', 'variable', 'constant']

View File

@ -1,10 +1,11 @@
from collections.abc import Mapping, Sequence
from os import path
from typing import Optional, cast
from typing import Any, cast
from core.app.segments import parser
from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler
from core.file.file_obj import FileTransferMethod, FileType, FileVar
from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
from core.tools.tool.tool import Tool
from core.tools.tool_engine import ToolEngine
from core.tools.tool_manager import ToolManager
from core.tools.utils.message_transformer import ToolFileMessageTransformer
@ -20,6 +21,7 @@ class ToolNode(BaseNode):
"""
Tool Node
"""
_node_data_cls = ToolNodeData
_node_type = NodeType.TOOL
@ -50,23 +52,24 @@ class ToolNode(BaseNode):
},
error=f'Failed to get tool runtime: {str(e)}'
)
# get parameters
parameters = self._generate_parameters(self.graph_runtime_state.variable_pool, node_data, tool_runtime)
tool_parameters = tool_runtime.get_runtime_parameters() or []
parameters = self._generate_parameters(tool_parameters=tool_parameters, variable_pool=self.graph_runtime_state.variable_pool, node_data=node_data)
parameters_for_log = self._generate_parameters(tool_parameters=tool_parameters, variable_pool=self.graph_runtime_state.variable_pool, node_data=node_data, for_log=True)
try:
messages = ToolEngine.workflow_invoke(
tool=tool_runtime,
tool_parameters=parameters,
user_id=self.user_id,
workflow_id=self.workflow_id,
workflow_tool_callback=DifyWorkflowCallbackHandler(),
workflow_call_depth=self.workflow_call_depth,
)
except Exception as e:
return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
inputs=parameters,
inputs=parameters_for_log,
metadata={
NodeRunMetadataKey.TOOL_INFO: tool_info
},
@ -86,21 +89,34 @@ class ToolNode(BaseNode):
metadata={
NodeRunMetadataKey.TOOL_INFO: tool_info
},
inputs=parameters
inputs=parameters_for_log
)
def _generate_parameters(self, variable_pool: VariablePool, node_data: ToolNodeData, tool_runtime: Tool) -> dict:
def _generate_parameters(
self,
*,
tool_parameters: Sequence[ToolParameter],
variable_pool: VariablePool,
node_data: ToolNodeData,
for_log: bool = False,
) -> Mapping[str, Any]:
"""
Generate parameters
"""
tool_parameters = tool_runtime.get_all_runtime_parameters()
Generate parameters based on the given tool parameters, variable pool, and node data.
def fetch_parameter(name: str) -> Optional[ToolParameter]:
return next((parameter for parameter in tool_parameters if parameter.name == name), None)
Args:
tool_parameters (Sequence[ToolParameter]): The list of tool parameters.
variable_pool (VariablePool): The variable pool containing the variables.
node_data (ToolNodeData): The data associated with the tool node.
Returns:
Mapping[str, Any]: A dictionary containing the generated parameters.
"""
tool_parameters_dictionary = {parameter.name: parameter for parameter in tool_parameters}
result = {}
for parameter_name in node_data.tool_parameters:
parameter = fetch_parameter(parameter_name)
parameter = tool_parameters_dictionary.get(parameter_name)
if not parameter:
continue
if parameter.type == ToolParameter.ToolParameterType.FILE:
@ -108,35 +124,21 @@ class ToolNode(BaseNode):
v.to_dict() for v in self._fetch_files(variable_pool)
]
else:
input = node_data.tool_parameters[parameter_name]
if input.type == 'mixed':
result[parameter_name] = self._format_variable_template(input.value, variable_pool)
elif input.type == 'variable':
result[parameter_name] = variable_pool.get_variable_value(input.value)
elif input.type == 'constant':
result[parameter_name] = input.value
tool_input = node_data.tool_parameters[parameter_name]
segment_group = parser.convert_template(
template=str(tool_input.value),
variable_pool=variable_pool,
)
result[parameter_name] = segment_group.log if for_log else segment_group.text
return result
def _format_variable_template(self, template: str, variable_pool: VariablePool) -> str:
"""
Format variable template
"""
inputs = {}
template_parser = VariableTemplateParser(template)
for selector in template_parser.extract_variable_selectors():
inputs[selector.variable] = variable_pool.get_variable_value(selector.value_selector)
return template_parser.format(inputs)
def _fetch_files(self, variable_pool: VariablePool) -> list[FileVar]:
files = variable_pool.get_variable_value(['sys', SystemVariable.FILES.value])
if not files:
return []
return files
def _convert_tool_messages(self, messages: list[ToolInvokeMessage]) \
def _fetch_files(self, variable_pool: VariablePool) -> list[FileVar]:
# FIXME: ensure this is a ArrayVariable contains FileVariable.
variable = variable_pool.get(['sys', SystemVariable.FILES.value])
return [file_var.value for file_var in variable.value] if variable else []
def _convert_tool_messages(self, messages: list[ToolInvokeMessage])\
-> tuple[str, list[FileVar], list[dict]]:
"""
Convert ToolInvokeMessages into tuple[plain_text, files]

View File

@ -18,28 +18,27 @@ class VariableAggregatorNode(BaseNode):
inputs = {}
if not node_data.advanced_settings or not node_data.advanced_settings.group_enabled:
for variable in node_data.variables:
value = self.graph_runtime_state.variable_pool.get_variable_value(variable)
if value is not None:
for selector in node_data.variables:
variable = self.graph_runtime_state.variable_pool.get(selector)
if variable is not None:
outputs = {
"output": value
"output": variable.value
}
inputs = {
'.'.join(variable[1:]): value
'.'.join(selector[1:]): variable.value
}
break
else:
for group in node_data.advanced_settings.groups:
for variable in group.variables:
value = self.graph_runtime_state.variable_pool.get_variable_value(variable)
for selector in group.variables:
variable = self.graph_runtime_state.variable_pool.get(selector)
if value is not None:
if variable is not None:
outputs[group.group_name] = {
'output': value
'output': variable.value
}
inputs['.'.join(variable[1:])] = value
inputs['.'.join(selector[1:])] = variable.value
break
return NodeRunResult(