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

feat: template transform

parent 71ff2a83
from os import environ
from typing import Literal
from httpx import post
from pydantic import BaseModel
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_EXECUTION_ENDPOINT = environ.get('CODE_EXECUTION_ENDPOINT', '')
......@@ -24,7 +25,7 @@ class CodeExecutionResponse(BaseModel):
class CodeExecutor:
@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
:param language: code language
......@@ -32,7 +33,13 @@ class CodeExecutor:
:param inputs: inputs
: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'
headers = {
......@@ -67,4 +74,4 @@ class CodeExecutor:
if response.data.stderr:
raise CodeExecutionException(response.data.stderr)
return PythonTemplateTransformer.transform_response(response.data.stdout)
\ No newline at end of file
return template_transformer.transform_response(response.data.stdout)
\ No newline at end of file
import json
import re
from core.helper.code_executor.template_transformer import TemplateTransformer
PYTHON_RUNNER = """# declare main function here
{{code}}
......@@ -19,7 +21,7 @@ print(result)
"""
class PythonTemplateTransformer:
class PythonTemplateTransformer(TemplateTransformer):
@classmethod
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 core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor
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.code.code_executor import CodeExecutionException, CodeExecutor
from core.workflow.nodes.code.entities import CodeNodeData
from models.workflow import WorkflowNodeExecutionStatus
......
......@@ -16,6 +16,6 @@ class CodeNodeData(BaseNodeData):
variables: list[VariableSelector]
answer: str
code_language: str
code_language: Literal['python3', 'javascript']
code: str
outputs: dict[str, Output]
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.variable_pool import VariablePool
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.template_transform.entities import TemplateTransformNodeData
from models.workflow import WorkflowNodeExecutionStatus
class TemplateTransformNode(BaseNode):
_node_data_cls = TemplateTransformNodeData
_node_type = NodeType.TEMPLATE_TRANSFORM
@classmethod
def get_default_config(cls, filters: Optional[dict] = None) -> dict:
"""
......@@ -23,3 +32,51 @@ class TemplateTransformNode(BaseNode):
"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