update

api/core/helper/code_executor/code_executor.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Literal, Optional
+from typing import Optional
 
 from httpx import post
 from pydantic import BaseModel
@@ -9,6 +9,7 @@
 from core.helper.code_executor.javascript_transformer import NodeJsTemplateTransformer
 from core.helper.code_executor.jinja2_transformer import Jinja2TemplateTransformer
 from core.helper.code_executor.python_transformer import PythonTemplateTransformer
+from core.helper.code_executor.template_transformer import TemplateTransformer
 
 # Code Executor
 CODE_EXECUTION_ENDPOINT = get_env('CODE_EXECUTION_ENDPOINT')
@@ -36,20 +37,8 @@ class CodeLanguage(str, Enum):
 
 
 class CodeExecutor:
-    code_template_transformers = {
-        CodeLanguage.PYTHON3: PythonTemplateTransformer,
-        CodeLanguage.JINJA2: Jinja2TemplateTransformer,
-        CodeLanguage.JAVASCRIPT: NodeJsTemplateTransformer,
-    }
-
-    code_language_to_running_language = {
-        CodeLanguage.JAVASCRIPT: 'nodejs',
-        CodeLanguage.JINJA2: CodeLanguage.PYTHON3,
-        CodeLanguage.PYTHON3: CodeLanguage.PYTHON3,
-    }
-
     @classmethod
-    def execute_code(cls, language: Literal['python3', 'javascript', 'jinja2'], preload: str, code: str) -> str:
+    def execute_code(cls, language: CodeLanguage | str, preload: str, code: str) -> str:
         """
         Execute code
         :param language: code language
@@ -63,7 +52,7 @@ def execute_code(cls, language: Literal['python3', 'javascript', 'jinja2'], prel
         }
 
         data = {
-            'language': cls.code_language_to_running_language.get(language),
+            'language': cls.get_execution_language(language),
             'code': code,
             'preload': preload
         }
@@ -95,15 +84,15 @@ def execute_code(cls, language: Literal['python3', 'javascript', 'jinja2'], prel
         return response.data.stdout
 
     @classmethod
-    def execute_workflow_code_template(cls, language: Literal['python3', 'javascript', 'jinja2'], code: str, inputs: dict) -> dict:
+    def execute_workflow_code_template(cls, language: CodeLanguage | str, code: str, inputs: dict) -> dict:
         """
         Execute code
         :param language: code language
         :param code: code
         :param inputs: inputs
         :return:
         """
-        template_transformer = cls.code_template_transformers.get(language)
+        template_transformer = cls.get_template_transformer(language)
         if not template_transformer:
             raise CodeExecutionException(f'Unsupported language {language}')
 
@@ -114,4 +103,27 @@ def execute_workflow_code_template(cls, language: Literal['python3', 'javascript
         except CodeExecutionException as e:
             raise e
 
-        return template_transformer.transform_response(response)
+        return template_transformer.transform_response(response)
+
+    @classmethod
+    def get_template_transformer(cls, language: CodeLanguage | str) -> type[TemplateTransformer]:
+        """
+        get the template transformer
+        :param language:
+        :return:
+        """
+        code_template_transformers = {
+            CodeLanguage.PYTHON3: PythonTemplateTransformer,
+            CodeLanguage.JINJA2: Jinja2TemplateTransformer,
+            CodeLanguage.JAVASCRIPT: NodeJsTemplateTransformer,
+        }
+        return code_template_transformers.get(language)
+
+    @classmethod
+    def get_execution_language(cls, language: CodeLanguage | str) -> str:
+        code_language_to_running_language = {
+            CodeLanguage.PYTHON3: CodeLanguage.PYTHON3,
+            CodeLanguage.JINJA2: CodeLanguage.PYTHON3,
+            CodeLanguage.JAVASCRIPT: 'nodejs',
+        }
+        return code_language_to_running_language.get(language)