From 2895c3bc8c997efcaef70f6008917e38c4366d22 Mon Sep 17 00:00:00 2001 From: Yeuoly Date: Sat, 9 Mar 2024 22:49:53 +0800 Subject: [PATCH] feat: template transform --- .../code_executor}/code_executor.py | 15 +++-- .../code_executor/javascript_transformer.py | 1 + .../helper/code_executor/jina2_transformer.py | 1 + .../code_executor/python_transformer.py} | 4 +- .../code_executor/template_transformer.py | 24 ++++++++ api/core/workflow/nodes/code/code_node.py | 2 +- api/core/workflow/nodes/code/entities.py | 2 +- .../nodes/http_request/http_request_node.py | 1 - .../nodes/template_transform/entities.py | 14 +++++ .../template_transform_node.py | 59 ++++++++++++++++++- 10 files changed, 114 insertions(+), 9 deletions(-) rename api/core/{workflow/nodes/code => helper/code_executor}/code_executor.py (75%) create mode 100644 api/core/helper/code_executor/javascript_transformer.py create mode 100644 api/core/helper/code_executor/jina2_transformer.py rename api/core/{workflow/nodes/code/python_template.py => helper/code_executor/python_transformer.py} (90%) create mode 100644 api/core/helper/code_executor/template_transformer.py create mode 100644 api/core/workflow/nodes/template_transform/entities.py diff --git a/api/core/workflow/nodes/code/code_executor.py b/api/core/helper/code_executor/code_executor.py similarity index 75% rename from api/core/workflow/nodes/code/code_executor.py rename to api/core/helper/code_executor/code_executor.py index 058ee83d46..f1bc4fbdaf 100644 --- a/api/core/workflow/nodes/code/code_executor.py +++ b/api/core/helper/code_executor/code_executor.py @@ -1,10 +1,11 @@ from os import environ +from typing import Literal from httpx import post from pydantic import BaseModel from yarl import URL -from core.workflow.nodes.code.python_template import PythonTemplateTransformer +from core.helper.code_executor.python_transformer import PythonTemplateTransformer # Code Executor CODE_EXECUTION_ENDPOINT = environ.get('CODE_EXECUTION_ENDPOINT', '') @@ -24,7 +25,7 @@ class CodeExecutionResponse(BaseModel): class CodeExecutor: @classmethod - def execute_code(cls, language: str, code: str, inputs: dict) -> dict: + def execute_code(cls, language: Literal['python3', 'javascript', 'jina2'], code: str, inputs: dict) -> dict: """ Execute code :param language: code language @@ -32,7 +33,13 @@ class CodeExecutor: :param inputs: inputs :return: """ - runner = PythonTemplateTransformer.transform_caller(code, inputs) + template_transformer = None + if language == 'python3': + template_transformer = PythonTemplateTransformer + else: + raise CodeExecutionException('Unsupported language') + + runner = template_transformer.transform_caller(code, inputs) url = URL(CODE_EXECUTION_ENDPOINT) / 'v1' / 'sandbox' / 'run' headers = { @@ -67,4 +74,4 @@ class CodeExecutor: if response.data.stderr: raise CodeExecutionException(response.data.stderr) - return PythonTemplateTransformer.transform_response(response.data.stdout) \ No newline at end of file + return template_transformer.transform_response(response.data.stdout) \ No newline at end of file diff --git a/api/core/helper/code_executor/javascript_transformer.py b/api/core/helper/code_executor/javascript_transformer.py new file mode 100644 index 0000000000..f87f5c14cb --- /dev/null +++ b/api/core/helper/code_executor/javascript_transformer.py @@ -0,0 +1 @@ +# TODO \ No newline at end of file diff --git a/api/core/helper/code_executor/jina2_transformer.py b/api/core/helper/code_executor/jina2_transformer.py new file mode 100644 index 0000000000..f87f5c14cb --- /dev/null +++ b/api/core/helper/code_executor/jina2_transformer.py @@ -0,0 +1 @@ +# TODO \ No newline at end of file diff --git a/api/core/workflow/nodes/code/python_template.py b/api/core/helper/code_executor/python_transformer.py similarity index 90% rename from api/core/workflow/nodes/code/python_template.py rename to api/core/helper/code_executor/python_transformer.py index 03dfee36f3..7b862649d8 100644 --- a/api/core/workflow/nodes/code/python_template.py +++ b/api/core/helper/code_executor/python_transformer.py @@ -1,6 +1,8 @@ import json import re +from core.helper.code_executor.template_transformer import TemplateTransformer + PYTHON_RUNNER = """# declare main function here {{code}} @@ -19,7 +21,7 @@ print(result) """ -class PythonTemplateTransformer: +class PythonTemplateTransformer(TemplateTransformer): @classmethod def transform_caller(cls, code: str, inputs: dict) -> str: """ diff --git a/api/core/helper/code_executor/template_transformer.py b/api/core/helper/code_executor/template_transformer.py new file mode 100644 index 0000000000..5505df8749 --- /dev/null +++ b/api/core/helper/code_executor/template_transformer.py @@ -0,0 +1,24 @@ +from abc import ABC, abstractmethod + + +class TemplateTransformer(ABC): + @classmethod + @abstractmethod + def transform_caller(cls, code: str, inputs: dict) -> str: + """ + Transform code to python runner + :param code: code + :param inputs: inputs + :return: + """ + pass + + @classmethod + @abstractmethod + def transform_response(cls, response: str) -> dict: + """ + Transform response to dict + :param response: response + :return: + """ + pass \ No newline at end of file diff --git a/api/core/workflow/nodes/code/code_node.py b/api/core/workflow/nodes/code/code_node.py index 3d3c475d06..7d3162d983 100644 --- a/api/core/workflow/nodes/code/code_node.py +++ b/api/core/workflow/nodes/code/code_node.py @@ -1,9 +1,9 @@ from typing import Optional, Union, cast +from core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor from core.workflow.entities.node_entities import NodeRunResult, NodeType from core.workflow.entities.variable_pool import VariablePool from core.workflow.nodes.base_node import BaseNode -from core.workflow.nodes.code.code_executor import CodeExecutionException, CodeExecutor from core.workflow.nodes.code.entities import CodeNodeData from models.workflow import WorkflowNodeExecutionStatus diff --git a/api/core/workflow/nodes/code/entities.py b/api/core/workflow/nodes/code/entities.py index 2212d77e2d..6a18d181cb 100644 --- a/api/core/workflow/nodes/code/entities.py +++ b/api/core/workflow/nodes/code/entities.py @@ -16,6 +16,6 @@ class CodeNodeData(BaseNodeData): variables: list[VariableSelector] answer: str - code_language: str + code_language: Literal['python3', 'javascript'] code: str outputs: dict[str, Output] diff --git a/api/core/workflow/nodes/http_request/http_request_node.py b/api/core/workflow/nodes/http_request/http_request_node.py index e3e864b6b0..4ee76deb83 100644 --- a/api/core/workflow/nodes/http_request/http_request_node.py +++ b/api/core/workflow/nodes/http_request/http_request_node.py @@ -1,6 +1,5 @@ from typing import cast -from core.workflow.entities.base_node_data_entities import BaseNodeData from core.workflow.entities.node_entities import NodeRunResult, NodeType from core.workflow.entities.variable_pool import VariablePool from core.workflow.nodes.base_node import BaseNode diff --git a/api/core/workflow/nodes/template_transform/entities.py b/api/core/workflow/nodes/template_transform/entities.py new file mode 100644 index 0000000000..2d3d35b84c --- /dev/null +++ b/api/core/workflow/nodes/template_transform/entities.py @@ -0,0 +1,14 @@ +from typing import Literal, Union + +from pydantic import BaseModel + +from core.workflow.entities.base_node_data_entities import BaseNodeData +from core.workflow.entities.variable_entities import VariableSelector + + +class TemplateTransformNodeData(BaseNodeData): + """ + Code Node Data. + """ + variables: list[VariableSelector] + template: str \ No newline at end of file diff --git a/api/core/workflow/nodes/template_transform/template_transform_node.py b/api/core/workflow/nodes/template_transform/template_transform_node.py index 2bf26e307e..3fb880d926 100644 --- a/api/core/workflow/nodes/template_transform/template_transform_node.py +++ b/api/core/workflow/nodes/template_transform/template_transform_node.py @@ -1,9 +1,18 @@ -from typing import Optional +from typing import Optional, cast +from core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor +from core.workflow.entities.base_node_data_entities import BaseNodeData +from core.workflow.entities.node_entities import NodeRunResult, NodeType +from core.workflow.entities.variable_pool import VariablePool from core.workflow.nodes.base_node import BaseNode +from core.workflow.nodes.template_transform.entities import TemplateTransformNodeData +from models.workflow import WorkflowNodeExecutionStatus class TemplateTransformNode(BaseNode): + _node_data_cls = TemplateTransformNodeData + _node_type = NodeType.TEMPLATE_TRANSFORM + @classmethod def get_default_config(cls, filters: Optional[dict] = None) -> dict: """ @@ -23,3 +32,51 @@ class TemplateTransformNode(BaseNode): "template": "{{ arg1 }}" } } + + def _run(self, variable_pool: VariablePool) -> NodeRunResult: + """ + Run node + """ + node_data = self.node_data + node_data: TemplateTransformNodeData = cast(self._node_data_cls, node_data) + + # Get variables + variables = {} + for variable_selector in node_data.variables: + variable = variable_selector.variable + value = variable_pool.get_variable_value( + variable_selector=variable_selector.value_selector + ) + + variables[variable] = value + + # Run code + try: + result = CodeExecutor.execute_code( + language='jina2', + code=node_data.template, + inputs=variables + ) + except CodeExecutionException as e: + return NodeRunResult( + inputs=variables, + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e) + ) + + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=variables, + outputs=result['result'] + ) + + @classmethod + def _extract_variable_selector_to_variable_mapping(cls, node_data: TemplateTransformNodeData) -> dict[list[str], str]: + """ + Extract variable selector to variable mapping + :param node_data: node data + :return: + """ + return { + variable_selector.value_selector: variable_selector.variable for variable_selector in node_data.variables + } \ No newline at end of file