From 59ba7917c4001123e3ca41fcf5ce4672b67fa65f Mon Sep 17 00:00:00 2001
From: Yeuoly <admin@srmxy.cn>
Date: Sun, 10 Mar 2024 17:55:24 +0800
Subject: [PATCH] fix: code node dose not work as expected

---
 api/core/helper/code_executor/code_executor.py     | 14 +++++++-------
 .../helper/code_executor/python_transformer.py     | 10 ++++------
 api/core/workflow/nodes/code/code_node.py          | 10 +++++-----
 3 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/api/core/helper/code_executor/code_executor.py b/api/core/helper/code_executor/code_executor.py
index f1bc4fbdaf..fb0ad9642a 100644
--- a/api/core/helper/code_executor/code_executor.py
+++ b/api/core/helper/code_executor/code_executor.py
@@ -1,5 +1,5 @@
 from os import environ
-from typing import Literal
+from typing import Literal, Optional
 
 from httpx import post
 from pydantic import BaseModel
@@ -16,8 +16,8 @@ class CodeExecutionException(Exception):
 
 class CodeExecutionResponse(BaseModel):
     class Data(BaseModel):
-        stdout: str
-        stderr: str
+        stdout: Optional[str]
+        error: Optional[str]
 
     code: int
     message: str
@@ -58,9 +58,9 @@ class CodeExecutor:
                 raise Exception('Failed to execute code')
         except CodeExecutionException as e:
             raise e
-        except Exception:
+        except Exception as e:
             raise CodeExecutionException('Failed to execute code')
-
+        
         try:
             response = response.json()
         except:
@@ -71,7 +71,7 @@ class CodeExecutor:
         if response.code != 0:
             raise CodeExecutionException(response.message)
         
-        if response.data.stderr:
-            raise CodeExecutionException(response.data.stderr)
+        if response.data.error:
+            raise CodeExecutionException(response.data.error)
         
         return template_transformer.transform_response(response.data.stdout)
\ No newline at end of file
diff --git a/api/core/helper/code_executor/python_transformer.py b/api/core/helper/code_executor/python_transformer.py
index 7b862649d8..27863ee443 100644
--- a/api/core/helper/code_executor/python_transformer.py
+++ b/api/core/helper/code_executor/python_transformer.py
@@ -11,11 +11,11 @@ PYTHON_RUNNER = """# declare main function here
 output = main(**{{inputs}})
 
 # convert output to json and print
-result = '''
-<<RESULT>>
+output = json.dumps(output, indent=4)
+
+result = f'''<<RESULT>>
 {output}
-<<RESULT>>
-'''
+<<RESULT>>'''
 
 print(result)
 """
@@ -47,11 +47,9 @@ class PythonTemplateTransformer(TemplateTransformer):
         :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 json.loads(result)
diff --git a/api/core/workflow/nodes/code/code_node.py b/api/core/workflow/nodes/code/code_node.py
index 7d3162d983..9cc5865133 100644
--- a/api/core/workflow/nodes/code/code_node.py
+++ b/api/core/workflow/nodes/code/code_node.py
@@ -101,7 +101,6 @@ class CodeNode(BaseNode):
             )
 
             variables[variable] = value
-
         # Run code
         try:
             result = CodeExecutor.execute_code(
@@ -109,15 +108,16 @@ class CodeNode(BaseNode):
                 code=code,
                 inputs=variables
             )
-        except CodeExecutionException as e:
+
+            # Transform result
+            result = self._transform_result(result, node_data.outputs)
+        except (CodeExecutionException, ValueError) as e:
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.FAILED,
+                inputs=variables,
                 error=str(e)
             )
 
-        # Transform result
-        result = self._transform_result(result, node_data.outputs)
-
         return NodeRunResult(
             status=WorkflowNodeExecutionStatus.SUCCEEDED,
             inputs=variables,