feat: template transform

api/core/workflow/nodes/code/code_executor.py

@@ -1,70 +0,0 @@
-from os import environ
-
-from httpx import post
-from pydantic import BaseModel
-from yarl import URL
-
-from core.workflow.nodes.code.python_template import PythonTemplateTransformer
-
-# Code Executor
-CODE_EXECUTION_ENDPOINT = environ.get('CODE_EXECUTION_ENDPOINT', '')
-CODE_EXECUTION_API_KEY = environ.get('CODE_EXECUTION_API_KEY', '')
-
-class CodeExecutionException(Exception):
-    pass
-
-class CodeExecutionResponse(BaseModel):
-    class Data(BaseModel):
-        stdout: str
-        stderr: str
-
-    code: int
-    message: str
-    data: Data
-
-class CodeExecutor:
-    @classmethod
-    def execute_code(cls, language: str, code: str, inputs: dict) -> dict:
-        """
-        Execute code
-        :param language: code language
-        :param code: code
-        :param inputs: inputs
-        :return:
-        """
-        runner = PythonTemplateTransformer.transform_caller(code, inputs)
-
-        url = URL(CODE_EXECUTION_ENDPOINT) / 'v1' / 'sandbox' / 'run'
-        headers = {
-            'X-Api-Key': CODE_EXECUTION_API_KEY
-        }
-        data = {
-            'language': language,
-            'code': runner,
-        }
-
-        try:
-            response = post(str(url), json=data, headers=headers)
-            if response.status_code == 503:
-                raise CodeExecutionException('Code execution service is unavailable')
-            elif response.status_code != 200:
-                raise Exception('Failed to execute code')
-        except CodeExecutionException as e:
-            raise e
-        except Exception:
-            raise CodeExecutionException('Failed to execute code')
-
-        try:
-            response = response.json()
-        except:
-            raise CodeExecutionException('Failed to parse response')
-
-        response = CodeExecutionResponse(**response)
-
-        if response.code != 0:
-            raise CodeExecutionException(response.message)
-        
-        if response.data.stderr:
-            raise CodeExecutionException(response.data.stderr)
-        
-        return PythonTemplateTransformer.transform_response(response.data.stdout)

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
 

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]

api/core/workflow/nodes/code/python_template.py

@@ -1,55 +0,0 @@
-import json
-import re
-
-PYTHON_RUNNER = """# declare main function here
-{{code}}
-
-# execute main function, and return the result
-# inputs is a dict, and it
-output = main(**{{inputs}})
-
-# convert output to json and print
-result = '''
-<<RESULT>>
-{output}
-<<RESULT>>
-'''
-
-print(result)
-"""
-
-
-class PythonTemplateTransformer:
-    @classmethod
-    def transform_caller(cls, code: str, inputs: dict) -> str:
-        """
-        Transform code to python runner
-        :param code: code
-        :param inputs: inputs
-        :return:
-        """
-        
-        # transform inputs to json string
-        inputs_str = json.dumps(inputs, indent=4)
-
-        # replace code and inputs
-        runner = PYTHON_RUNNER.replace('{{code}}', code)
-        runner = runner.replace('{{inputs}}', inputs_str)
-
-        return runner
-    
-    @classmethod
-    def transform_response(cls, response: str) -> dict:
-        """
-        Transform response to dict
-        :param response: response
-        :return:
-        """
-        
-        # extract result
-        result = re.search(r'<<RESULT>>(.*)<<RESULT>>', response, re.DOTALL)
-        if not result:
-            raise ValueError('Failed to parse result')
-        
-        result = result.group(1)
-        return json.loads(result)

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

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

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
+        }