feat: template transform

api/core/helper/code_executor/code_executor.py

@@ -0,0 +1,77 @@
+from os import environ
+from typing import Literal
+
+from httpx import post
+from pydantic import BaseModel
+from yarl import URL
+
+from core.helper.code_executor.python_transformer import PythonTemplateTransformer
+
+# Code Executor
+CODE_EXECUTION_ENDPOINT = environ.get('CODE_EXECUTION_ENDPOINT', '')
+CODE_EXECUTION_API_KEY = environ.get('CODE_EXECUTION_API_KEY', '')
+
+class CodeExecutionException(Exception):
+    pass
+
+class CodeExecutionResponse(BaseModel):
+    class Data(BaseModel):
+        stdout: str
+        stderr: str
+
+    code: int
+    message: str
+    data: Data
+
+class CodeExecutor:
+    @classmethod
+    def execute_code(cls, language: Literal['python3', 'javascript', 'jina2'], code: str, inputs: dict) -> dict:
+        """
+        Execute code
+        :param language: code language
+        :param code: code
+        :param inputs: inputs
+        :return:
+        """
+        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 = {
+            'X-Api-Key': CODE_EXECUTION_API_KEY
+        }
+        data = {
+            'language': language,
+            'code': runner,
+        }
+
+        try:
+            response = post(str(url), json=data, headers=headers)
+            if response.status_code == 503:
+                raise CodeExecutionException('Code execution service is unavailable')
+            elif response.status_code != 200:
+                raise Exception('Failed to execute code')
+        except CodeExecutionException as e:
+            raise e
+        except Exception:
+            raise CodeExecutionException('Failed to execute code')
+
+        try:
+            response = response.json()
+        except:
+            raise CodeExecutionException('Failed to parse response')
+
+        response = CodeExecutionResponse(**response)
+
+        if response.code != 0:
+            raise CodeExecutionException(response.message)
+        
+        if response.data.stderr:
+            raise CodeExecutionException(response.data.stderr)
+        
+        return template_transformer.transform_response(response.data.stdout)

api/core/helper/code_executor/javascript_transformer.py

@@ -0,0 +1 @@
+# TODO

api/core/helper/code_executor/jina2_transformer.py

@@ -0,0 +1 @@
+# TODO

api/core/helper/code_executor/python_transformer.py

@@ -0,0 +1,57 @@
+import json
+import re
+
+from core.helper.code_executor.template_transformer import TemplateTransformer
+
+PYTHON_RUNNER = """# declare main function here
+{{code}}
+
+# execute main function, and return the result
+# inputs is a dict, and it
+output = main(**{{inputs}})
+
+# convert output to json and print
+result = '''
+<<RESULT>>
+{output}
+<<RESULT>>
+'''
+
+print(result)
+"""
+
+
+class PythonTemplateTransformer(TemplateTransformer):
+    @classmethod
+    def transform_caller(cls, code: str, inputs: dict) -> str:
+        """
+        Transform code to python runner
+        :param code: code
+        :param inputs: inputs
+        :return:
+        """
+        
+        # transform inputs to json string
+        inputs_str = json.dumps(inputs, indent=4)
+
+        # replace code and inputs
+        runner = PYTHON_RUNNER.replace('{{code}}', code)
+        runner = runner.replace('{{inputs}}', inputs_str)
+
+        return runner
+    
+    @classmethod
+    def transform_response(cls, response: str) -> dict:
+        """
+        Transform response to dict
+        :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)

api/core/helper/code_executor/template_transformer.py

@@ -0,0 +1,24 @@
+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