improve: code upgrade (#4231)

api/core/helper/code_executor/code_executor.py

@@ -1,14 +1,20 @@
+import logging
+import time
 from enum import Enum
+from threading import Lock
 from typing import Literal, Optional
 
-from httpx import post
+from httpx import get, post
 from pydantic import BaseModel
 from yarl import URL
 
 from config import get_env
+from core.helper.code_executor.entities import CodeDependency
 from core.helper.code_executor.javascript_transformer import NodeJsTemplateTransformer
 from core.helper.code_executor.jinja2_transformer import Jinja2TemplateTransformer
-from core.helper.code_executor.python_transformer import PythonTemplateTransformer
+from core.helper.code_executor.python_transformer import PYTHON_STANDARD_PACKAGES, PythonTemplateTransformer
+
+logger = logging.getLogger(__name__)
 
 # Code Executor
 CODE_EXECUTION_ENDPOINT = get_env('CODE_EXECUTION_ENDPOINT')
@@ -28,7 +34,6 @@ class CodeExecutionResponse(BaseModel):
     message: str
     data: Data
 
-
 class CodeLanguage(str, Enum):
     PYTHON3 = 'python3'
     JINJA2 = 'jinja2'
@@ -36,6 +41,9 @@ class CodeLanguage(str, Enum):
 
 
 class CodeExecutor:
+    dependencies_cache = {}
+    dependencies_cache_lock = Lock()
+
     code_template_transformers = {
         CodeLanguage.PYTHON3: PythonTemplateTransformer,
         CodeLanguage.JINJA2: Jinja2TemplateTransformer,
@@ -49,7 +57,11 @@ class CodeExecutor:
     }
 
     @classmethod
-    def execute_code(cls, language: Literal['python3', 'javascript', 'jinja2'], preload: str, code: str) -> str:
+    def execute_code(cls, 
+                     language: Literal['python3', 'javascript', 'jinja2'], 
+                     preload: str, 
+                     code: str, 
+                     dependencies: Optional[list[CodeDependency]] = None) -> str:
         """
         Execute code
         :param language: code language
@@ -65,9 +77,13 @@ class CodeExecutor:
         data = {
             'language': cls.code_language_to_running_language.get(language),
             'code': code,
-            'preload': preload
+            'preload': preload,
+            'enable_network': True
         }
 
+        if dependencies:
+            data['dependencies'] = [dependency.dict() for dependency in dependencies]
+
         try:
             response = post(str(url), json=data, headers=headers, timeout=CODE_EXECUTION_TIMEOUT)
             if response.status_code == 503:
@@ -95,7 +111,7 @@ class CodeExecutor:
         return response.data.stdout
 
     @classmethod
-    def execute_workflow_code_template(cls, language: Literal['python3', 'javascript', 'jinja2'], code: str, inputs: dict) -> dict:
+    def execute_workflow_code_template(cls, language: Literal['python3', 'javascript', 'jinja2'], code: str, inputs: dict, dependencies: Optional[list[CodeDependency]] = None) -> dict:
         """
         Execute code
         :param language: code language
@@ -107,11 +123,63 @@ class CodeExecutor:
         if not template_transformer:
             raise CodeExecutionException(f'Unsupported language {language}')
 
-        runner, preload = template_transformer.transform_caller(code, inputs)
+        runner, preload, dependencies = template_transformer.transform_caller(code, inputs, dependencies)
 
         try:
-            response = cls.execute_code(language, preload, runner)
+            response = cls.execute_code(language, preload, runner, dependencies)
         except CodeExecutionException as e:
             raise e
 
-        return template_transformer.transform_response(response)
+        return template_transformer.transform_response(response)
+    
+    @classmethod
+    def list_dependencies(cls, language: Literal['python3']) -> list[CodeDependency]:
+        with cls.dependencies_cache_lock:
+            if language in cls.dependencies_cache:
+                # check expiration
+                dependencies = cls.dependencies_cache[language]
+                if dependencies['expiration'] > time.time():
+                    return dependencies['data']
+                # remove expired cache
+                del cls.dependencies_cache[language]
+        
+        dependencies = cls._get_dependencies(language)
+        with cls.dependencies_cache_lock:
+            cls.dependencies_cache[language] = {
+                'data': dependencies,
+                'expiration': time.time() + 60
+            }
+        
+        return dependencies
+        
+    @classmethod
+    def _get_dependencies(cls, language: Literal['python3']) -> list[CodeDependency]:
+        """
+        List dependencies
+        """
+        url = URL(CODE_EXECUTION_ENDPOINT) / 'v1' / 'sandbox' / 'dependencies'
+
+        headers = {
+            'X-Api-Key': CODE_EXECUTION_API_KEY
+        }
+
+        running_language = cls.code_language_to_running_language.get(language)
+        if isinstance(running_language, Enum):
+            running_language = running_language.value
+
+        data = {
+            'language': running_language,
+        }
+
+        try:
+            response = get(str(url), params=data, headers=headers, timeout=CODE_EXECUTION_TIMEOUT)
+            if response.status_code != 200:
+                raise Exception(f'Failed to list dependencies, got status code {response.status_code}, please check if the sandbox service is running')
+            response = response.json()
+            dependencies = response.get('data', {}).get('dependencies', [])
+            return [
+                CodeDependency(**dependency) for dependency in dependencies if dependency.get('name') not in PYTHON_STANDARD_PACKAGES
+            ]
+        except Exception as e:
+            logger.exception(f'Failed to list dependencies: {e}')
+            return []

api/core/helper/code_executor/entities.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class CodeDependency(BaseModel):
+    name: str
+    version: str

api/core/helper/code_executor/javascript_transformer.py

@@ -1,6 +1,8 @@
 import json
 import re
+from typing import Optional
 
+from core.helper.code_executor.entities import CodeDependency
 from core.helper.code_executor.template_transformer import TemplateTransformer
 
 NODEJS_RUNNER = """// declare main function here
@@ -22,7 +24,8 @@ NODEJS_PRELOAD = """"""
 
 class NodeJsTemplateTransformer(TemplateTransformer):
     @classmethod
-    def transform_caller(cls, code: str, inputs: dict) -> tuple[str, str]:
+    def transform_caller(cls, code: str, inputs: dict, 
+                         dependencies: Optional[list[CodeDependency]] = None) -> tuple[str, str, list[CodeDependency]]:
         """
         Transform code to python runner
         :param code: code
@@ -37,7 +40,7 @@ class NodeJsTemplateTransformer(TemplateTransformer):
         runner = NODEJS_RUNNER.replace('{{code}}', code)
         runner = runner.replace('{{inputs}}', inputs_str)
 
-        return runner, NODEJS_PRELOAD
+        return runner, NODEJS_PRELOAD, []
 
     @classmethod
     def transform_response(cls, response: str) -> dict:

api/core/helper/code_executor/jinja2_transformer.py

@@ -1,7 +1,10 @@
 import json
 import re
 from base64 import b64encode
+from typing import Optional
 
+from core.helper.code_executor.entities import CodeDependency
+from core.helper.code_executor.python_transformer import PYTHON_STANDARD_PACKAGES
 from core.helper.code_executor.template_transformer import TemplateTransformer
 
 PYTHON_RUNNER = """
@@ -58,7 +61,8 @@ if __name__ == '__main__':
 
 class Jinja2TemplateTransformer(TemplateTransformer):
     @classmethod
-    def transform_caller(cls, code: str, inputs: dict) -> tuple[str, str]:
+    def transform_caller(cls, code: str, inputs: dict, 
+                         dependencies: Optional[list[CodeDependency]] = None) -> tuple[str, str, list[CodeDependency]]:
         """
         Transform code to python runner
         :param code: code
@@ -72,7 +76,19 @@ class Jinja2TemplateTransformer(TemplateTransformer):
         runner = PYTHON_RUNNER.replace('{{code}}', code)
         runner = runner.replace('{{inputs}}', inputs_str)
 
-        return runner, JINJA2_PRELOAD
+        if not dependencies:
+            dependencies = []
+
+        # add native packages and jinja2
+        for package in PYTHON_STANDARD_PACKAGES.union(['jinja2']):
+            dependencies.append(CodeDependency(name=package, version=''))
+
+        # deduplicate
+        dependencies = list({
+            dep.name: dep for dep in dependencies if dep.name
+        }.values())
+
+        return runner, JINJA2_PRELOAD, dependencies
 
     @classmethod
     def transform_response(cls, response: str) -> dict:

api/core/helper/code_executor/python_transformer.py

@@ -1,7 +1,9 @@
 import json
 import re
 from base64 import b64encode
+from typing import Optional
 
+from core.helper.code_executor.entities import CodeDependency
 from core.helper.code_executor.template_transformer import TemplateTransformer
 
 PYTHON_RUNNER = """# declare main function here
@@ -25,32 +27,17 @@ result = f'''<<RESULT>>
 print(result)
 """
 
-PYTHON_PRELOAD = """
-# prepare general imports
-import json
-import datetime
-import math
-import random
-import re
-import string
-import sys
-import time
-import traceback
-import uuid
-import os
-import base64
-import hashlib
-import hmac
-import binascii
-import collections
-import functools
-import operator
-import itertools
-"""
+PYTHON_PRELOAD = """"""
+
+PYTHON_STANDARD_PACKAGES = set([
+    'json', 'datetime', 'math', 'random', 're', 'string', 'sys', 'time', 'traceback', 'uuid', 'os', 'base64',
+    'hashlib', 'hmac', 'binascii', 'collections', 'functools', 'operator', 'itertools', 'uuid', 
+])
 
 class PythonTemplateTransformer(TemplateTransformer):
     @classmethod
-    def transform_caller(cls, code: str, inputs: dict) -> tuple[str, str]:
+    def transform_caller(cls, code: str, inputs: dict, 
+                         dependencies: Optional[list[CodeDependency]] = None) -> tuple[str, str, list[CodeDependency]]:
         """
         Transform code to python runner
         :param code: code
@@ -65,7 +52,18 @@ class PythonTemplateTransformer(TemplateTransformer):
         runner = PYTHON_RUNNER.replace('{{code}}', code)
         runner = runner.replace('{{inputs}}', inputs_str)
 
-        return runner, PYTHON_PRELOAD
+        # add standard packages
+        if dependencies is None:
+            dependencies = []
+
+        for package in PYTHON_STANDARD_PACKAGES:
+            if package not in dependencies:
+                dependencies.append(CodeDependency(name=package, version=''))
+
+        # deduplicate
+        dependencies = list({dep.name: dep for dep in dependencies if dep.name}.values())
+
+        return runner, PYTHON_PRELOAD, dependencies
     
     @classmethod
     def transform_response(cls, response: str) -> dict:

api/core/helper/code_executor/template_transformer.py

@@ -1,10 +1,14 @@
 from abc import ABC, abstractmethod
+from typing import Optional
+
+from core.helper.code_executor.entities import CodeDependency
 
 
 class TemplateTransformer(ABC):
     @classmethod
     @abstractmethod
-    def transform_caller(cls, code: str, inputs: dict) -> tuple[str, str]:
+    def transform_caller(cls, code: str, inputs: dict, 
+                         dependencies: Optional[list[CodeDependency]] = None) -> tuple[str, str, list[CodeDependency]]:
         """
         Transform code to python runner
         :param code: code