diff --git a/api/core/helper/code_executor/code_executor.py b/api/core/helper/code_executor/code_executor.py index fb0ad9642a..a62cf4de95 100644 --- a/api/core/helper/code_executor/code_executor.py +++ b/api/core/helper/code_executor/code_executor.py @@ -4,6 +4,7 @@ from typing import Literal, Optional from httpx import post from pydantic import BaseModel from yarl import URL +from core.helper.code_executor.jina2_transformer import Jinja2TemplateTransformer from core.helper.code_executor.python_transformer import PythonTemplateTransformer @@ -25,7 +26,7 @@ class CodeExecutionResponse(BaseModel): class CodeExecutor: @classmethod - def execute_code(cls, language: Literal['python3', 'javascript', 'jina2'], code: str, inputs: dict) -> dict: + def execute_code(cls, language: Literal['python3', 'javascript', 'jinja2'], code: str, inputs: dict) -> dict: """ Execute code :param language: code language @@ -36,6 +37,8 @@ class CodeExecutor: template_transformer = None if language == 'python3': template_transformer = PythonTemplateTransformer + elif language == 'jinja2': + template_transformer = Jinja2TemplateTransformer else: raise CodeExecutionException('Unsupported language') @@ -46,7 +49,7 @@ class CodeExecutor: 'X-Api-Key': CODE_EXECUTION_API_KEY } data = { - 'language': language, + 'language': language if language != 'jinja2' else 'python3', 'code': runner, } diff --git a/api/core/helper/code_executor/jina2_transformer.py b/api/core/helper/code_executor/jina2_transformer.py index f87f5c14cb..87e8ce130f 100644 --- a/api/core/helper/code_executor/jina2_transformer.py +++ b/api/core/helper/code_executor/jina2_transformer.py @@ -1 +1,54 @@ -# TODO \ No newline at end of file +import json +import re + +from core.helper.code_executor.template_transformer import TemplateTransformer + +PYTHON_RUNNER = """ +import jinja2 + +template = jinja2.Template('''{{code}}''') + +def main(**inputs): + return template.render(**inputs) + +# execute main function, and return the result +output = main(**{{inputs}}) + +result = f'''<>{output}<>''' + +print(result) + +""" + +class Jinja2TemplateTransformer(TemplateTransformer): + @classmethod + def transform_caller(cls, code: str, inputs: dict) -> str: + """ + Transform code to python runner + :param code: code + :param inputs: inputs + :return: + """ + + # transform jinja2 template to python code + runner = PYTHON_RUNNER.replace('{{code}}', code) + runner = runner.replace('{{inputs}}', json.dumps(inputs, indent=4)) + + return runner + + @classmethod + def transform_response(cls, response: str) -> dict: + """ + Transform response to dict + :param response: response + :return: + """ + # extract result + result = re.search(r'<>(.*)<>', response, re.DOTALL) + if not result: + raise ValueError('Failed to parse result') + result = result.group(1) + + return { + 'result': result + } \ No newline at end of file diff --git a/api/core/workflow/nodes/template_transform/template_transform_node.py b/api/core/workflow/nodes/template_transform/template_transform_node.py index 724b84495c..a037332f4b 100644 --- a/api/core/workflow/nodes/template_transform/template_transform_node.py +++ b/api/core/workflow/nodes/template_transform/template_transform_node.py @@ -52,7 +52,7 @@ class TemplateTransformNode(BaseNode): # Run code try: result = CodeExecutor.execute_code( - language='jina2', + language='jinja2', code=node_data.template, inputs=variables ) @@ -66,7 +66,9 @@ class TemplateTransformNode(BaseNode): return NodeRunResult( status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs=variables, - outputs=result['result'] + outputs={ + 'output': result['result'] + } ) @classmethod diff --git a/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py b/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py index b95c76b133..a1c8eb71dc 100644 --- a/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py +++ b/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py @@ -9,7 +9,7 @@ MOCK = os.getenv('MOCK_SWITCH', 'false') == 'true' class MockedCodeExecutor: @classmethod - def invoke(cls, language: Literal['python3', 'javascript', 'jina2'], code: str, inputs: dict) -> dict: + def invoke(cls, language: Literal['python3', 'javascript', 'jinja2'], code: str, inputs: dict) -> dict: # invoke directly if language == 'python3': return {