diff --git a/api/core/helper/code_executor/code_executor.py b/api/core/helper/code_executor/code_executor.py index 063a21b192..ec685ae814 100644 --- 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 CodeExecutionResponse(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 @@ class CodeExecutor: } 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 @@ class CodeExecutor: :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 index ae9b1cb612..37645bf0d0 100644 --- 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 @@ class SimpleCode(BuiltinTool): 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 index 2c1529f492..12e7ae940f 100644 --- 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 @@ class CodeNode(BaseNode): :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 @@ class CodeNode(BaseNode): "value_selector": [] } ], - "code_language": "javascript", + "code_language": CodeLanguage.JAVASCRIPT, "code": JAVASCRIPT_DEFAULT_CODE, "outputs": { "result": { @@ -77,7 +77,7 @@ class CodeNode(BaseNode): "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 new file mode 100644 index 0000000000..ae6e7ceaa7 --- /dev/null +++ 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 index c794ae8e4b..19c9d18307 100644 --- 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/api/tests/integration_tests/workflow/nodes/code_executor/test_code_jina2.py b/api/tests/integration_tests/workflow/nodes/code_executor/test_code_jinja2.py similarity index 65% rename from api/tests/integration_tests/workflow/nodes/code_executor/test_code_jina2.py rename to api/tests/integration_tests/workflow/nodes/code_executor/test_code_jinja2.py index aae3c7acec..6793cd3cc2 100644 --- a/api/tests/integration_tests/workflow/nodes/code_executor/test_code_jina2.py +++ b/api/tests/integration_tests/workflow/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 == '<>Hello World<>\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 index 1983bc5e6b..b5c59c93fc 100644 --- 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'}