Unverified Commit 460c0da1 authored by Yeuoly's avatar Yeuoly

feat: jinja2

parent 295a2485
......@@ -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,
}
......
# 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'''<<RESULT>>{output}<<RESULT>>'''
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'<<RESULT>>(.*)<<RESULT>>', 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
......@@ -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
......
......@@ -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 {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment