update

langgenius · May 5, 2024 · 6729425 · 6729425
1 parent 51a9e67
commit 6729425
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 21 deletions.
diff --git a/api/core/helper/code_executor/code_executor.py b/api/core/helper/code_executor/code_executor.py
@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import Literal, Optional
 
 from httpx import post
@@ -28,7 +29,25 @@ class Data(BaseModel):
     data: Data
 
 
+class CodeLanguage(str, Enum):
+    PYTHON3 = 'python3'
+    JINJA2 = 'jinja2'
+    JAVASCRIPT = 'javascript'
+
+
 class CodeExecutor:
+    code_template_transformers = {
+        CodeLanguage.PYTHON3: PythonTemplateTransformer,
+        CodeLanguage.JINJA2: Jinja2TemplateTransformer,
+        CodeLanguage.JAVASCRIPT: NodeJsTemplateTransformer,
+    }
+
+    code_language_to_running_language = {
+        CodeLanguage.JAVASCRIPT: CodeLanguage.JAVASCRIPT,
+        CodeLanguage.JINJA2: CodeLanguage.PYTHON3,
+        CodeLanguage.PYTHON3: CodeLanguage.PYTHON3,
+    }
+
     @classmethod
     def execute_code(cls, language: Literal['python3', 'javascript', 'jinja2'], preload: str, code: str) -> str:
         """
@@ -44,9 +63,7 @@ def execute_code(cls, language: Literal['python3', 'javascript', 'jinja2'], prel
         }
 
         data = {
-            'language': 'python3' if language == 'jinja2' else
-                        'nodejs' if language == 'javascript' else
-                        'python3' if language == 'python3' else None,
+            'language': cls.code_language_to_running_language.get(language),
             'code': code,
             'preload': preload
         }
@@ -86,15 +103,9 @@ def execute_workflow_code_template(cls, language: Literal['python3', 'javascript
         :param inputs: inputs
         :return:
         """
-        template_transformer = None
-        if language == 'python3':
-            template_transformer = PythonTemplateTransformer
-        elif language == 'jinja2':
-            template_transformer = Jinja2TemplateTransformer
-        elif language == 'javascript':
-            template_transformer = NodeJsTemplateTransformer
-        else:
-            raise CodeExecutionException('Unsupported language')
+        template_transformer = cls.code_template_transformers.get(language)
+        if not template_transformer:
+            raise CodeExecutionException(f'Unsupported language {language}')
 
         runner, preload = template_transformer.transform_caller(code, inputs)
 

diff --git a/api/core/tools/provider/builtin/code/tools/simple_code.py b/api/core/tools/provider/builtin/code/tools/simple_code.py
@@ -1,6 +1,6 @@
 from typing import Any
 
-from core.helper.code_executor.code_executor import CodeExecutor
+from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
 from core.tools.entities.tool_entities import ToolInvokeMessage
 from core.tools.tool.builtin_tool import BuiltinTool
 
@@ -11,10 +11,10 @@ def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMe
             invoke simple code
         """
 
-        language = tool_parameters.get('language', 'python3')
+        language = tool_parameters.get('language', CodeLanguage.PYTHON3)
         code = tool_parameters.get('code', '')
 
-        if language not in ['python3', 'javascript']:
+        if language not in [CodeLanguage.PYTHON3, CodeLanguage.JAVASCRIPT]:
             raise ValueError(f'Only python3 and javascript are supported, not {language}')
 
         result = CodeExecutor.execute_code(language, '', code)

diff --git a/api/tests/integration_tests/workflow/nodes/code_executor/test_code_executor.py b/api/tests/integration_tests/workflow/nodes/code_executor/test_code_executor.py
@@ -0,0 +1,11 @@
+import pytest
+
+from core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor
+
+CODE_LANGUAGE = 'unsupported_language'
+
+
+def test_unsupported_with_code_template():
+    with pytest.raises(CodeExecutionException) as e:
+        CodeExecutor.execute_workflow_code_template(language=CODE_LANGUAGE, code='', inputs={})
+    assert str(e.value) == f'Unsupported language {CODE_LANGUAGE}'
diff --git a/api/tests/integration_tests/workflow/nodes/code_executor/test_code_javascript.py b/api/tests/integration_tests/workflow/nodes/code_executor/test_code_javascript.py
@@ -1,6 +1,7 @@
-from core.helper.code_executor.code_executor import CodeExecutor
+from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
+from core.workflow.nodes.code.code_node import JAVASCRIPT_DEFAULT_CODE
 
-CODE_LANGUAGE = 'javascript'
+CODE_LANGUAGE = CodeLanguage.JAVASCRIPT
 
 
 def test_javascript_plain():
@@ -16,3 +17,9 @@ def test_javascript_json():
     """
     result = CodeExecutor.execute_code(language=CODE_LANGUAGE, preload='', code=code)
     assert result == '{"Hello":"World"}\n'
+
+
+def test_javascript_with_code_template():
+    result = CodeExecutor.execute_workflow_code_template(
+        language=CODE_LANGUAGE, code=JAVASCRIPT_DEFAULT_CODE, inputs={'arg1': 'Hello', 'arg2': 'World'})
+    assert result == {'result': 'HelloWorld'}
diff --git a/...ow/nodes/code_executor/test_code_jina2.py → ...w/nodes/code_executor/test_code_jinja2.py b/...ow/nodes/code_executor/test_code_jina2.py → ...w/nodes/code_executor/test_code_jinja2.py
@@ -1,9 +1,9 @@
 import base64
 
-from core.helper.code_executor.code_executor import CodeExecutor
+from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
 from core.helper.code_executor.jinja2_transformer import JINJA2_PRELOAD, PYTHON_RUNNER
 
-CODE_LANGUAGE = 'jinja2'
+CODE_LANGUAGE = CodeLanguage.JINJA2
 
 
 def test_jinja2():
@@ -12,3 +12,9 @@ def test_jinja2():
     code = PYTHON_RUNNER.replace('{{code}}', template).replace('{{inputs}}', inputs)
     result = CodeExecutor.execute_code(language=CODE_LANGUAGE, preload=JINJA2_PRELOAD, code=code)
     assert result == '<<RESULT>>Hello World<<RESULT>>\n'
+
+
+def test_jinja2_with_code_template():
+    result = CodeExecutor.execute_workflow_code_template(
+        language=CODE_LANGUAGE, code='Hello {{template}}', inputs={'template': 'World'})
+    assert result == {'result': 'Hello World'}
diff --git a/api/tests/integration_tests/workflow/nodes/code_executor/test_code_python3.py b/api/tests/integration_tests/workflow/nodes/code_executor/test_code_python3.py
@@ -1,6 +1,7 @@
-from core.helper.code_executor.code_executor import CodeExecutor
+from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
+from core.workflow.nodes.code.code_node import PYTHON_DEFAULT_CODE
 
-CODE_LANGUAGE = 'python3'
+CODE_LANGUAGE = CodeLanguage.PYTHON3
 
 
 def test_python3_plain():
@@ -16,3 +17,9 @@ def test_python3_json():
     """
     result = CodeExecutor.execute_code(language=CODE_LANGUAGE, preload='', code=code)
     assert result == '{"Hello": "World"}\n'
+
+
+def test_python3_with_code_template():
+    result = CodeExecutor.execute_workflow_code_template(
+        language=CODE_LANGUAGE, code=PYTHON_DEFAULT_CODE, inputs={'arg1': 'Hello', 'arg2': 'World'})
+    assert result == {'result': 'HelloWorld'}