fix: code node dose not work as expected

This commit is contained in:
Yeuoly 2024-03-10 17:55:24 +08:00 committed by takatost
parent 2d8497f79b
commit 1e6feadc7e
3 changed files with 16 additions and 18 deletions

View File

@ -1,5 +1,5 @@
from os import environ from os import environ
from typing import Literal from typing import Literal, Optional
from httpx import post from httpx import post
from pydantic import BaseModel from pydantic import BaseModel
@ -16,8 +16,8 @@ class CodeExecutionException(Exception):
class CodeExecutionResponse(BaseModel): class CodeExecutionResponse(BaseModel):
class Data(BaseModel): class Data(BaseModel):
stdout: str stdout: Optional[str]
stderr: str error: Optional[str]
code: int code: int
message: str message: str
@ -58,9 +58,9 @@ class CodeExecutor:
raise Exception('Failed to execute code') raise Exception('Failed to execute code')
except CodeExecutionException as e: except CodeExecutionException as e:
raise e raise e
except Exception: except Exception as e:
raise CodeExecutionException('Failed to execute code') raise CodeExecutionException('Failed to execute code')
try: try:
response = response.json() response = response.json()
except: except:
@ -71,7 +71,7 @@ class CodeExecutor:
if response.code != 0: if response.code != 0:
raise CodeExecutionException(response.message) raise CodeExecutionException(response.message)
if response.data.stderr: if response.data.error:
raise CodeExecutionException(response.data.stderr) raise CodeExecutionException(response.data.error)
return template_transformer.transform_response(response.data.stdout) return template_transformer.transform_response(response.data.stdout)

View File

@ -11,11 +11,11 @@ PYTHON_RUNNER = """# declare main function here
output = main(**{{inputs}}) output = main(**{{inputs}})
# convert output to json and print # convert output to json and print
result = ''' output = json.dumps(output, indent=4)
<<RESULT>>
result = f'''<<RESULT>>
{output} {output}
<<RESULT>> <<RESULT>>'''
'''
print(result) print(result)
""" """
@ -47,11 +47,9 @@ class PythonTemplateTransformer(TemplateTransformer):
:param response: response :param response: response
:return: :return:
""" """
# extract result # extract result
result = re.search(r'<<RESULT>>(.*)<<RESULT>>', response, re.DOTALL) result = re.search(r'<<RESULT>>(.*)<<RESULT>>', response, re.DOTALL)
if not result: if not result:
raise ValueError('Failed to parse result') raise ValueError('Failed to parse result')
result = result.group(1) result = result.group(1)
return json.loads(result) return json.loads(result)

View File

@ -101,7 +101,6 @@ class CodeNode(BaseNode):
) )
variables[variable] = value variables[variable] = value
# Run code # Run code
try: try:
result = CodeExecutor.execute_code( result = CodeExecutor.execute_code(
@ -109,15 +108,16 @@ class CodeNode(BaseNode):
code=code, code=code,
inputs=variables inputs=variables
) )
except CodeExecutionException as e:
# Transform result
result = self._transform_result(result, node_data.outputs)
except (CodeExecutionException, ValueError) as e:
return NodeRunResult( return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED, status=WorkflowNodeExecutionStatus.FAILED,
inputs=variables,
error=str(e) error=str(e)
) )
# Transform result
result = self._transform_result(result, node_data.outputs)
return NodeRunResult( return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED, status=WorkflowNodeExecutionStatus.SUCCEEDED,
inputs=variables, inputs=variables,