Commit 3407b4d8 authored by Yeuoly's avatar Yeuoly Committed by takatost

feat: template transform

parent 71ff2a83
from os import environ from os import environ
from typing import Literal
from httpx import post from httpx import post
from pydantic import BaseModel from pydantic import BaseModel
from yarl import URL from yarl import URL
from core.workflow.nodes.code.python_template import PythonTemplateTransformer from core.helper.code_executor.python_transformer import PythonTemplateTransformer
# Code Executor # Code Executor
CODE_EXECUTION_ENDPOINT = environ.get('CODE_EXECUTION_ENDPOINT', '') CODE_EXECUTION_ENDPOINT = environ.get('CODE_EXECUTION_ENDPOINT', '')
...@@ -24,7 +25,7 @@ class CodeExecutionResponse(BaseModel): ...@@ -24,7 +25,7 @@ class CodeExecutionResponse(BaseModel):
class CodeExecutor: class CodeExecutor:
@classmethod @classmethod
def execute_code(cls, language: str, code: str, inputs: dict) -> dict: def execute_code(cls, language: Literal['python3', 'javascript', 'jina2'], code: str, inputs: dict) -> dict:
""" """
Execute code Execute code
:param language: code language :param language: code language
...@@ -32,7 +33,13 @@ class CodeExecutor: ...@@ -32,7 +33,13 @@ class CodeExecutor:
:param inputs: inputs :param inputs: inputs
:return: :return:
""" """
runner = PythonTemplateTransformer.transform_caller(code, inputs) template_transformer = None
if language == 'python3':
template_transformer = PythonTemplateTransformer
else:
raise CodeExecutionException('Unsupported language')
runner = template_transformer.transform_caller(code, inputs)
url = URL(CODE_EXECUTION_ENDPOINT) / 'v1' / 'sandbox' / 'run' url = URL(CODE_EXECUTION_ENDPOINT) / 'v1' / 'sandbox' / 'run'
headers = { headers = {
...@@ -67,4 +74,4 @@ class CodeExecutor: ...@@ -67,4 +74,4 @@ class CodeExecutor:
if response.data.stderr: if response.data.stderr:
raise CodeExecutionException(response.data.stderr) raise CodeExecutionException(response.data.stderr)
return PythonTemplateTransformer.transform_response(response.data.stdout) return template_transformer.transform_response(response.data.stdout)
\ No newline at end of file \ No newline at end of file
import json import json
import re import re
from core.helper.code_executor.template_transformer import TemplateTransformer
PYTHON_RUNNER = """# declare main function here PYTHON_RUNNER = """# declare main function here
{{code}} {{code}}
...@@ -19,7 +21,7 @@ print(result) ...@@ -19,7 +21,7 @@ print(result)
""" """
class PythonTemplateTransformer: class PythonTemplateTransformer(TemplateTransformer):
@classmethod @classmethod
def transform_caller(cls, code: str, inputs: dict) -> str: def transform_caller(cls, code: str, inputs: dict) -> str:
""" """
......
from abc import ABC, abstractmethod
class TemplateTransformer(ABC):
@classmethod
@abstractmethod
def transform_caller(cls, code: str, inputs: dict) -> str:
"""
Transform code to python runner
:param code: code
:param inputs: inputs
:return:
"""
pass
@classmethod
@abstractmethod
def transform_response(cls, response: str) -> dict:
"""
Transform response to dict
:param response: response
:return:
"""
pass
\ No newline at end of file
from typing import Optional, Union, cast from typing import Optional, Union, cast
from core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor
from core.workflow.entities.node_entities import NodeRunResult, NodeType from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.entities.variable_pool import VariablePool from core.workflow.entities.variable_pool import VariablePool
from core.workflow.nodes.base_node import BaseNode from core.workflow.nodes.base_node import BaseNode
from core.workflow.nodes.code.code_executor import CodeExecutionException, CodeExecutor
from core.workflow.nodes.code.entities import CodeNodeData from core.workflow.nodes.code.entities import CodeNodeData
from models.workflow import WorkflowNodeExecutionStatus from models.workflow import WorkflowNodeExecutionStatus
......
...@@ -16,6 +16,6 @@ class CodeNodeData(BaseNodeData): ...@@ -16,6 +16,6 @@ class CodeNodeData(BaseNodeData):
variables: list[VariableSelector] variables: list[VariableSelector]
answer: str answer: str
code_language: str code_language: Literal['python3', 'javascript']
code: str code: str
outputs: dict[str, Output] outputs: dict[str, Output]
from typing import cast from typing import cast
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.entities.variable_pool import VariablePool from core.workflow.entities.variable_pool import VariablePool
from core.workflow.nodes.base_node import BaseNode from core.workflow.nodes.base_node import BaseNode
......
from typing import Literal, Union
from pydantic import BaseModel
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.variable_entities import VariableSelector
class TemplateTransformNodeData(BaseNodeData):
"""
Code Node Data.
"""
variables: list[VariableSelector]
template: str
\ No newline at end of file
from typing import Optional from typing import Optional, cast
from core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.nodes.base_node import BaseNode from core.workflow.nodes.base_node import BaseNode
from core.workflow.nodes.template_transform.entities import TemplateTransformNodeData
from models.workflow import WorkflowNodeExecutionStatus
class TemplateTransformNode(BaseNode): class TemplateTransformNode(BaseNode):
_node_data_cls = TemplateTransformNodeData
_node_type = NodeType.TEMPLATE_TRANSFORM
@classmethod @classmethod
def get_default_config(cls, filters: Optional[dict] = None) -> dict: def get_default_config(cls, filters: Optional[dict] = None) -> dict:
""" """
...@@ -23,3 +32,51 @@ class TemplateTransformNode(BaseNode): ...@@ -23,3 +32,51 @@ class TemplateTransformNode(BaseNode):
"template": "{{ arg1 }}" "template": "{{ arg1 }}"
} }
} }
def _run(self, variable_pool: VariablePool) -> NodeRunResult:
"""
Run node
"""
node_data = self.node_data
node_data: TemplateTransformNodeData = cast(self._node_data_cls, node_data)
# Get variables
variables = {}
for variable_selector in node_data.variables:
variable = variable_selector.variable
value = variable_pool.get_variable_value(
variable_selector=variable_selector.value_selector
)
variables[variable] = value
# Run code
try:
result = CodeExecutor.execute_code(
language='jina2',
code=node_data.template,
inputs=variables
)
except CodeExecutionException as e:
return NodeRunResult(
inputs=variables,
status=WorkflowNodeExecutionStatus.FAILED,
error=str(e)
)
return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED,
inputs=variables,
outputs=result['result']
)
@classmethod
def _extract_variable_selector_to_variable_mapping(cls, node_data: TemplateTransformNodeData) -> dict[list[str], str]:
"""
Extract variable selector to variable mapping
:param node_data: node data
:return:
"""
return {
variable_selector.value_selector: variable_selector.variable for variable_selector in node_data.variables
}
\ No newline at end of file
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