浏览代码

fix: incorrect type parser (#3682)

Yeuoly 1 年之前
父节点
当前提交
2a213c6af7

+ 9 - 2
api/core/helper/code_executor/jina2_transformer.py

@@ -1,10 +1,13 @@
 import json
 import re
+from base64 import b64encode
 
 from core.helper.code_executor.template_transformer import TemplateTransformer
 
 PYTHON_RUNNER = """
 import jinja2
+from json import loads
+from base64 import b64decode
 
 template = jinja2.Template('''{{code}}''')
 
@@ -12,7 +15,8 @@ def main(**inputs):
     return template.render(**inputs)
 
 # execute main function, and return the result
-output = main(**{{inputs}})
+inputs = b64decode('{{inputs}}').decode('utf-8')
+output = main(**loads(inputs))
 
 result = f'''<<RESULT>>{output}<<RESULT>>'''
 
@@ -39,6 +43,7 @@ JINJA2_PRELOAD_TEMPLATE = """{% set fruits = ['Apple'] %}
 
 JINJA2_PRELOAD = f"""
 import jinja2
+from base64 import b64decode
 
 def _jinja2_preload_():
     # prepare jinja2 environment, load template and render before to avoid sandbox issue
@@ -60,9 +65,11 @@ class Jinja2TemplateTransformer(TemplateTransformer):
         :return:
         """
 
+        inputs_str = b64encode(json.dumps(inputs, ensure_ascii=False).encode()).decode('utf-8')
+
         # transform jinja2 template to python code
         runner = PYTHON_RUNNER.replace('{{code}}', code)
-        runner = runner.replace('{{inputs}}', json.dumps(inputs, indent=4, ensure_ascii=False))
+        runner = runner.replace('{{inputs}}', inputs_str)
 
         return runner, JINJA2_PRELOAD
 

+ 8 - 3
api/core/helper/code_executor/python_transformer.py

@@ -1,17 +1,22 @@
 import json
 import re
+from base64 import b64encode
 
 from core.helper.code_executor.template_transformer import TemplateTransformer
 
 PYTHON_RUNNER = """# declare main function here
 {{code}}
 
+from json import loads, dumps
+from base64 import b64decode
+
 # execute main function, and return the result
 # inputs is a dict, and it
-output = main(**{{inputs}})
+inputs = b64decode('{{inputs}}').decode('utf-8')
+output = main(**json.loads(inputs))
 
 # convert output to json and print
-output = json.dumps(output, indent=4)
+output = dumps(output, indent=4)
 
 result = f'''<<RESULT>>
 {output}
@@ -54,7 +59,7 @@ class PythonTemplateTransformer(TemplateTransformer):
         """
         
         # transform inputs to json string
-        inputs_str = json.dumps(inputs, indent=4, ensure_ascii=False)
+        inputs_str = b64encode(json.dumps(inputs, ensure_ascii=False).encode()).decode('utf-8')
 
         # replace code and inputs
         runner = PYTHON_RUNNER.replace('{{code}}', code)