improve: test CodeExecutor with code templates and extract CodeLanguage enum (#4098)

This commit is contained in:
Bowen Liang 2024-05-07 12:37:18 +08:00 committed by GitHub
parent 45d21677a0
commit 049abd698f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 79 additions and 33 deletions

View File

@ -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)

View File

@ -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)

View File

@ -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": {

View File

@ -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}'

View File

@ -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'}

View File

@ -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'}

View File

@ -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'}