From 2a213c6af7ef87c5b6bccbc9f380bf773d788a6b Mon Sep 17 00:00:00 2001
From: Yeuoly <45712896+Yeuoly@users.noreply.github.com>
Date: Mon, 22 Apr 2024 19:32:41 +0800
Subject: [PATCH] fix: incorrect type parser (#3682)

---
 api/core/helper/code_executor/jina2_transformer.py  | 11 +++++++++--
 api/core/helper/code_executor/python_transformer.py | 11 ++++++++---
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/api/core/helper/code_executor/jina2_transformer.py b/api/core/helper/code_executor/jina2_transformer.py
index d7b46b0e25..8d97a28e85 100644
--- a/api/core/helper/code_executor/jina2_transformer.py
+++ b/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
 
diff --git a/api/core/helper/code_executor/python_transformer.py b/api/core/helper/code_executor/python_transformer.py
index d6a13e9923..f44acbb9bf 100644
--- a/api/core/helper/code_executor/python_transformer.py
+++ b/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)