improve: test CodeExecutor with code templates and extract CodeLangua…

…ge enum (langgenius#4098)
rennokki · May 9, 2024 · ec31632 · ec31632
1 parent 9f0760e
commit ec31632
Show file tree

Hide file tree

Showing 7 changed files with 79 additions and 33 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: 'nodejs',
+        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/core/workflow/nodes/code/code_node.py b/api/core/workflow/nodes/code/code_node.py
@@ -1,7 +1,7 @@
 import os
 from typing import Optional, Union, cast
 
-from core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor
+from core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor, CodeLanguage
 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
@@ -39,7 +39,7 @@ def get_default_config(cls, filters: Optional[dict] = None) -> dict:
         :param filters: filter by node config parameters.
         :return:
         """
-        if filters and filters.get("code_language") == "javascript":
+        if filters and filters.get("code_language") == CodeLanguage.JAVASCRIPT:
             return {
                 "type": "code",
                 "config": {
@@ -53,7 +53,7 @@ def get_default_config(cls, filters: Optional[dict] = None) -> dict:
                             "value_selector": []
                         }
                     ],
-                    "code_language": "javascript",
+                    "code_language": CodeLanguage.JAVASCRIPT,
                     "code": JAVASCRIPT_DEFAULT_CODE,
                     "outputs": {
                         "result": {
@@ -77,7 +77,7 @@ def get_default_config(cls, filters: Optional[dict] = None) -> dict:
                         "value_selector": []
                     }
                 ],
-                "code_language": "python3",
+                "code_language": CodeLanguage.PYTHON3,
                 "code": PYTHON_DEFAULT_CODE,
                 "outputs": {
                     "result": {

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,9 @@
-from core.helper.code_executor.code_executor import CodeExecutor
+from textwrap import dedent
 
-CODE_LANGUAGE = 'javascript'
+from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
+from core.workflow.nodes.code.code_node import JAVASCRIPT_DEFAULT_CODE
+
+CODE_LANGUAGE = CodeLanguage.JAVASCRIPT
 
 
 def test_javascript_plain():
@@ -10,9 +13,15 @@ def test_javascript_plain():
 
 
 def test_javascript_json():
-    code = """
-obj = {'Hello': 'World'}
-console.log(JSON.stringify(obj))
-    """
+    code = dedent("""
+    obj = {'Hello': 'World'}
+    console.log(JSON.stringify(obj))
+    """)
     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,9 @@
-from core.helper.code_executor.code_executor import CodeExecutor
+from textwrap import dedent
 
-CODE_LANGUAGE = 'python3'
+from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
+from core.workflow.nodes.code.code_node import PYTHON_DEFAULT_CODE
+
+CODE_LANGUAGE = CodeLanguage.PYTHON3
 
 
 def test_python3_plain():
@@ -10,9 +13,15 @@ def test_python3_plain():
 
 
 def test_python3_json():
-    code = """
-import json
-print(json.dumps({'Hello': 'World'}))
-    """
+    code = dedent("""
+    import json
+    print(json.dumps({'Hello': 'World'}))
+    """)
     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'}