improve: code upgrade (#4231)

.github/workflows/api-tests.yml

@@ -46,6 +46,7 @@ jobs:
             docker/docker-compose.middleware.yaml
           services: |
             sandbox
+            ssrf_proxy
 
       - name: Run Workflow
         run: dev/pytest/pytest_workflow.sh

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 = """
+PYTHON_PRELOAD = """"""
-# prepare general imports
+
-import json
+PYTHON_STANDARD_PACKAGES = set([
-import datetime
+    'json', 'datetime', 'math', 'random', 're', 'string', 'sys', 'time', 'traceback', 'uuid', 'os', 'base64',
-import math
+    'hashlib', 'hmac', 'binascii', 'collections', 'functools', 'operator', 'itertools', 'uuid', 
-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
-"""
 
 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

api/core/workflow/nodes/code/code_node.py

@@ -2,6 +2,7 @@ import os
 from typing import Optional, Union, cast
 
 from core.helper.code_executor.code_executor import CodeExecutionException, CodeExecutor, CodeLanguage
+from core.model_runtime.utils.encoders import jsonable_encoder
 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
@@ -61,7 +62,8 @@ class CodeNode(BaseNode):
                             "children": None
                         }
                     }
-                }
+                },
+                "available_dependencies": []
             }
 
         return {
@@ -84,8 +86,11 @@ class CodeNode(BaseNode):
                         "type": "string",
                         "children": None
                     }
-                }
+                },
-            }
+                "dependencies": [
+                ]
+            },
+            "available_dependencies": jsonable_encoder(CodeExecutor.list_dependencies('python3'))
         }
 
     def _run(self, variable_pool: VariablePool) -> NodeRunResult:
@@ -115,7 +120,8 @@ class CodeNode(BaseNode):
             result = CodeExecutor.execute_workflow_code_template(
                 language=code_language,
                 code=code,
-                inputs=variables
+                inputs=variables,
+                dependencies=node_data.dependencies
             )
 
             # Transform result

api/core/workflow/nodes/code/entities.py

@@ -2,6 +2,7 @@ from typing import Literal, Optional
 
 from pydantic import BaseModel
 
+from core.helper.code_executor.entities import CodeDependency
 from core.workflow.entities.base_node_data_entities import BaseNodeData
 from core.workflow.entities.variable_entities import VariableSelector
 
@@ -17,4 +18,5 @@ class CodeNodeData(BaseNodeData):
     variables: list[VariableSelector]
     code_language: Literal['python3', 'javascript']
     code: str
-    outputs: dict[str, Output]
+    outputs: dict[str, Output]
+    dependencies: Optional[list[CodeDependency]] = None

api/tests/integration_tests/workflow/nodes/__mock/code_executor.py

@@ -1,17 +1,19 @@
 import os
-from typing import Literal
+from typing import Literal, Optional
 
 import pytest
 from _pytest.monkeypatch import MonkeyPatch
 from jinja2 import Template
 
 from core.helper.code_executor.code_executor import CodeExecutor
+from core.helper.code_executor.entities import CodeDependency
 
 MOCK = os.getenv('MOCK_SWITCH', 'false') == 'true'
 
 class MockedCodeExecutor:
     @classmethod
-    def invoke(cls, language: Literal['python3', 'javascript', 'jinja2'], code: str, inputs: dict) -> dict:
+    def invoke(cls, language: Literal['python3', 'javascript', 'jinja2'], 
+               code: str, inputs: dict, dependencies: Optional[list[CodeDependency]] = None) -> dict:
         # invoke directly
         if language == 'python3':
             return {

docker/docker-compose.middleware.yaml

@@ -53,20 +53,38 @@ services:
 
   # The DifySandbox
   sandbox:
-    image: langgenius/dify-sandbox:0.1.0
+    image: langgenius/dify-sandbox:0.2.0
     restart: always
-    cap_add:
-      # Why is sys_admin permission needed?
-      # https://docs.dify.ai/getting-started/install-self-hosted/install-faq#id-16.-why-is-sys_admin-permission-needed
-      - SYS_ADMIN
     environment:
       # The DifySandbox configurations
+      # Make sure you are changing this key for your deployment with a strong key.
+      # You can generate a strong key using `openssl rand -base64 42`.
       API_KEY: dify-sandbox
       GIN_MODE: 'release'
       WORKER_TIMEOUT: 15
+      ENABLE_NETWORK: 'true'
+      HTTP_PROXY: 'http://ssrf_proxy:3128'
+      HTTPS_PROXY: 'http://ssrf_proxy:3128'
+    volumes:
+      - ./volumes/sandbox/dependencies:/dependencies
+    networks:
+      - ssrf_proxy_network
+
+  # ssrf_proxy server
+  # for more information, please refer to
+  # https://docs.dify.ai/getting-started/install-self-hosted/install-faq#id-16.-why-is-ssrf_proxy-needed
+  ssrf_proxy:
+    image: ubuntu/squid:latest
+    restart: always
     ports:
+      - "3128:3128"
       - "8194:8194"
-
+    volumes:
+      # pls clearly modify the squid.conf file to fit your network environment.
+      - ./volumes/ssrf_proxy/squid.conf:/etc/squid/squid.conf
+    networks:
+      - ssrf_proxy_network
+      - default
   # Qdrant vector store.
   # uncomment to use qdrant as vector store.
   # (if uncommented, you need to comment out the weaviate service above,
@@ -81,3 +99,10 @@ services:
   #   ports:
   #     - "6333:6333"
   #     - "6334:6334"
+
+
+networks:
+  # create a network between sandbox, api and ssrf_proxy, and can not access outside.
+  ssrf_proxy_network:
+    driver: bridge
+    internal: true

docker/docker-compose.yaml

@@ -161,6 +161,9 @@ services:
       CODE_MAX_STRING_ARRAY_LENGTH: 30
       CODE_MAX_OBJECT_ARRAY_LENGTH: 30
       CODE_MAX_NUMBER_ARRAY_LENGTH: 1000
+      # SSRF Proxy server
+      SSRF_PROXY_HTTP_URL: 'http://ssrf_proxy:3128'
+      SSRF_PROXY_HTTPS_URL: 'http://ssrf_proxy:3128'
     depends_on:
       - db
       - redis
@@ -170,6 +173,9 @@ services:
     # uncomment to expose dify-api port to host
     # ports:
     #   - "5001:5001"
+    networks:
+      - ssrf_proxy_network
+      - default
 
   # worker service
   # The Celery worker for processing the queue.
@@ -283,6 +289,9 @@ services:
     volumes:
       # Mount the storage directory to the container, for storing user files.
       - ./volumes/app/storage:/app/api/storage
+    networks:
+      - ssrf_proxy_network
+      - default
 
   # Frontend web application.
   web:
@@ -367,18 +376,35 @@ services:
 
   # The DifySandbox
   sandbox:
-    image: langgenius/dify-sandbox:0.1.0
+    image: langgenius/dify-sandbox:0.2.0
     restart: always
-    cap_add:
-    # Why is sys_admin permission needed?
-    # https://docs.dify.ai/getting-started/install-self-hosted/install-faq#id-16.-why-is-sys_admin-permission-needed
-      - SYS_ADMIN
     environment:
       # The DifySandbox configurations
+      # Make sure you are changing this key for your deployment with a strong key.
+      # You can generate a strong key using `openssl rand -base64 42`.
       API_KEY: dify-sandbox
-      GIN_MODE: release
+      GIN_MODE: 'release'
       WORKER_TIMEOUT: 15
+      ENABLE_NETWORK: 'true'
+      HTTP_PROXY: 'http://ssrf_proxy:3128'
+      HTTPS_PROXY: 'http://ssrf_proxy:3128'
+    volumes:
+      - ./volumes/sandbox/dependencies:/dependencies
+    networks:
+      - ssrf_proxy_network
 
+  # ssrf_proxy server
+  # for more information, please refer to
+  # https://docs.dify.ai/getting-started/install-self-hosted/install-faq#id-16.-why-is-ssrf_proxy-needed
+  ssrf_proxy:
+    image: ubuntu/squid:latest
+    restart: always
+    volumes:
+      # pls clearly modify the squid.conf file to fit your network environment.
+      - ./volumes/ssrf_proxy/squid.conf:/etc/squid/squid.conf
+    networks:
+      - ssrf_proxy_network
+      - default
   # Qdrant vector store.
   # uncomment to use qdrant as vector store.
   # (if uncommented, you need to comment out the weaviate service above,
@@ -436,3 +462,8 @@ services:
     ports:
       - "80:80"
       #- "443:443"
+networks:
+  # create a network between sandbox, api and ssrf_proxy, and can not access outside.
+  ssrf_proxy_network:
+    driver: bridge
+    internal: true

docker/volumes/ssrf_proxy/squid.conf

@@ -0,0 +1,50 @@
+acl localnet src 0.0.0.1-0.255.255.255	# RFC 1122 "this" network (LAN)
+acl localnet src 10.0.0.0/8		# RFC 1918 local private network (LAN)
+acl localnet src 100.64.0.0/10		# RFC 6598 shared address space (CGN)
+acl localnet src 169.254.0.0/16 	# RFC 3927 link-local (directly plugged) machines
+acl localnet src 172.16.0.0/12		# RFC 1918 local private network (LAN)
+acl localnet src 192.168.0.0/16		# RFC 1918 local private network (LAN)
+acl localnet src fc00::/7       	# RFC 4193 local private network range
+acl localnet src fe80::/10      	# RFC 4291 link-local (directly plugged) machines
+acl SSL_ports port 443
+acl Safe_ports port 80		# http
+acl Safe_ports port 21		# ftp
+acl Safe_ports port 443		# https
+acl Safe_ports port 70		# gopher
+acl Safe_ports port 210		# wais
+acl Safe_ports port 1025-65535	# unregistered ports
+acl Safe_ports port 280		# http-mgmt
+acl Safe_ports port 488		# gss-http
+acl Safe_ports port 591		# filemaker
+acl Safe_ports port 777		# multiling http
+acl CONNECT method CONNECT
+http_access deny !Safe_ports
+http_access deny CONNECT !SSL_ports
+http_access allow localhost manager
+http_access deny manager
+http_access allow localhost
+http_access allow localnet
+http_access deny all
+
+################################## Proxy Server ################################
+http_port 3128
+coredump_dir /var/spool/squid
+refresh_pattern ^ftp:		1440	20%	10080
+refresh_pattern ^gopher:	1440	0%	1440
+refresh_pattern -i (/cgi-bin/|\?) 0	0%	0
+refresh_pattern \/(Packages|Sources)(|\.bz2|\.gz|\.xz)$ 0 0% 0 refresh-ims
+refresh_pattern \/Release(|\.gpg)$ 0 0% 0 refresh-ims
+refresh_pattern \/InRelease$ 0 0% 0 refresh-ims
+refresh_pattern \/(Translation-.*)(|\.bz2|\.gz|\.xz)$ 0 0% 0 refresh-ims
+refresh_pattern .		0	20%	4320
+logfile_rotate 0
+
+# upstream proxy, set to your own upstream proxy IP to avoid SSRF attacks
+# cache_peer 172.1.1.1 parent 3128 0 no-query no-digest no-netdb-exchange default 
+
+
+################################## Reverse Proxy To Sandbox ################################
+http_port 8194 accel vhost
+cache_peer sandbox parent 8194 0 no-query originserver
+acl all src all
+http_access allow all

web/app/components/workflow/nodes/code/dependency-picker.tsx

@@ -0,0 +1,94 @@
+import type { FC } from 'react'
+import React, { useCallback, useState } from 'react'
+import { t } from 'i18next'
+import type { CodeDependency } from './types'
+import { ChevronDown } from '@/app/components/base/icons/src/vender/line/arrows'
+import { PortalToFollowElem, PortalToFollowElemContent, PortalToFollowElemTrigger } from '@/app/components/base/portal-to-follow-elem'
+import { Check, SearchLg } from '@/app/components/base/icons/src/vender/line/general'
+import { XCircle } from '@/app/components/base/icons/src/vender/solid/general'
+
+type Props = {
+  value: CodeDependency
+  available_dependencies: CodeDependency[]
+  onChange: (dependency: CodeDependency) => void
+}
+
+const DependencyPicker: FC<Props> = ({
+  available_dependencies,
+  value,
+  onChange,
+}) => {
+  const [open, setOpen] = useState(false)
+  const [searchText, setSearchText] = useState('')
+
+  const handleChange = useCallback((dependency: CodeDependency) => {
+    return () => {
+      setOpen(false)
+      onChange(dependency)
+    }
+  }, [onChange])
+
+  return (
+    <PortalToFollowElem
+      open={open}
+      onOpenChange={setOpen}
+      placement='bottom-start'
+      offset={4}
+    >
+      <PortalToFollowElemTrigger onClick={() => setOpen(!open)} className='flex-grow cursor-pointer'>
+        <div className='flex items-center h-8 justify-between px-2.5 rounded-lg border-0 bg-gray-100 text-gray-900 text-[13px]'>
+          <div className='grow w-0 truncate' title={value.name}>{value.name}</div>
+          <ChevronDown className='shrink-0 w-3.5 h-3.5 text-gray-700' />
+        </div>
+      </PortalToFollowElemTrigger>
+      <PortalToFollowElemContent style={{
+        zIndex: 100,
+      }}>
+        <div className='p-1 bg-white rounded-lg shadow-sm' style={{
+          width: 350,
+        }}>
+          <div
+            className='shadow-sm bg-white mb-2 mx-1 flex items-center px-2 rounded-lg bg-gray-100'
+          >
+            <SearchLg className='shrink-0 ml-[1px] mr-[5px] w-3.5 h-3.5 text-gray-400' />
+            <input
+              value={searchText}
+              className='grow px-0.5 py-[7px] text-[13px] text-gray-700 bg-transparent appearance-none outline-none caret-primary-600 placeholder:text-gray-400'
+              placeholder={t('workflow.nodes.code.searchDependencies') || ''}
+              onChange={e => setSearchText(e.target.value)}
+              autoFocus
+            />
+            {
+              searchText && (
+                <div
+                  className='flex items-center justify-center ml-[5px] w-[18px] h-[18px] cursor-pointer'
+                  onClick={() => setSearchText('')}
+                >
+                  <XCircle className='w-[14px] h-[14px] text-gray-400' />
+                </div>
+              )
+            }
+          </div>
+          <div className='max-h-[30vh] overflow-y-auto'>
+            {available_dependencies.filter((v) => {
+              if (!searchText)
+                return true
+              return v.name.toLowerCase().includes(searchText.toLowerCase())
+            }).map(dependency => (
+              <div
+                key={dependency.name}
+                className='flex items-center h-[30px] justify-between pl-3 pr-2 rounded-lg hover:bg-gray-100 text-gray-900 text-[13px] cursor-pointer'
+                onClick={handleChange(dependency)}
+              >
+                <div className='w-0 grow truncate'>{dependency.name}</div>
+                {dependency.name === value.name && <Check className='shrink-0 w-4 h-4 text-primary-600' />}
+              </div>
+            ))}
+          </div>
+        </div>
+      </PortalToFollowElemContent>
+    </PortalToFollowElem>
+  )
+}
+
+export default React.memo(DependencyPicker)

web/app/components/workflow/nodes/code/dependency.tsx

@@ -0,0 +1,36 @@
+import type { FC } from 'react'
+import React from 'react'
+import RemoveButton from '../_base/components/remove-button'
+import type { CodeDependency } from './types'
+import DependencyPicker from './dependency-picker'
+
+type Props = {
+  available_dependencies: CodeDependency[]
+  dependencies: CodeDependency[]
+  handleRemove: (index: number) => void
+  handleChange: (index: number, dependency: CodeDependency) => void
+}
+
+const Dependencies: FC<Props> = ({
+  available_dependencies, dependencies, handleRemove, handleChange,
+}) => {
+  return (
+    <div className='space-y-2'>
+      {dependencies.map((dependency, index) => (
+        <div className='flex items-center space-x-1' key={index}>
+          <DependencyPicker
+            value={dependency}
+            available_dependencies={available_dependencies}
+            onChange={dependency => handleChange(index, dependency)}
+          />
+          <RemoveButton
+            className='!p-2 !bg-gray-100 hover:!bg-gray-200'
+            onClick={() => handleRemove(index)}
+          />
+        </div>
+      ))}
+    </div>
+  )
+}
+
+export default React.memo(Dependencies)

web/app/components/workflow/nodes/code/panel.tsx

@@ -5,6 +5,7 @@ import RemoveEffectVarConfirm from '../_base/components/remove-effect-var-confir
 import useConfig from './use-config'
 import type { CodeNodeType } from './types'
 import { CodeLanguage } from './types'
+import Dependencies from './dependency'
 import VarList from '@/app/components/workflow/nodes/_base/components/variable/var-list'
 import OutputVarList from '@/app/components/workflow/nodes/_base/components/variable/output-var-list'
 import AddButton from '@/app/components/base/button/add-button'
@@ -59,6 +60,11 @@ const Panel: FC<NodePanelProps<CodeNodeType>> = ({
     varInputs,
     inputVarValues,
     setInputVarValues,
+    allowDependencies,
+    availableDependencies,
+    handleAddDependency,
+    handleRemoveDependency,
+    handleChangeDependency,
   } = useConfig(id, data)
 
   return (
@@ -78,6 +84,31 @@ const Panel: FC<NodePanelProps<CodeNodeType>> = ({
             filterVar={filterVar}
           />
         </Field>
+        {
+          allowDependencies
+            ? (
+              <div>
+                <Split />
+                <div className='pt-4'>
+                  <Field
+                    title={t(`${i18nPrefix}.advancedDependencies`)}
+                    operations={
+                      <AddButton onClick={() => handleAddDependency({ name: '', version: '' })} />
+                    }
+                    tooltip={t(`${i18nPrefix}.advancedDependenciesTip`)!}
+                  >
+                    <Dependencies
+                      available_dependencies={availableDependencies}
+                      dependencies={inputs.dependencies || []}
+                      handleRemove={index => handleRemoveDependency(index)}
+                      handleChange={(index, dependency) => handleChangeDependency(index, dependency)}
+                    />
+                  </Field>
+                </div>
+              </div>
+            )
+            : null
+        }
         <Split />
         <CodeEditor
           isInNode

web/app/components/workflow/nodes/code/types.ts

@@ -16,4 +16,10 @@ export type CodeNodeType = CommonNodeType & {
   code_language: CodeLanguage
   code: string
   outputs: OutputVar
+  dependencies?: CodeDependency[]
+}
+
+export type CodeDependency = {
+  name: string
+  version: string
 }

web/app/components/workflow/nodes/code/use-config.ts

@@ -5,7 +5,7 @@ import useOutputVarList from '../_base/hooks/use-output-var-list'
 import { BlockEnum, VarType } from '../../types'
 import type { Var } from '../../types'
 import { useStore } from '../../store'
-import type { CodeNodeType, OutputVar } from './types'
+import type { CodeDependency, CodeNodeType, OutputVar } from './types'
 import { CodeLanguage } from './types'
 import useNodeCrud from '@/app/components/workflow/nodes/_base/hooks/use-node-crud'
 import useOneStepRun from '@/app/components/workflow/nodes/_base/hooks/use-one-step-run'
@@ -21,15 +21,19 @@ const useConfig = (id: string, payload: CodeNodeType) => {
   const appId = useAppStore.getState().appDetail?.id
 
   const [allLanguageDefault, setAllLanguageDefault] = useState<Record<CodeLanguage, CodeNodeType> | null>(null)
+  const [allLanguageDependencies, setAllLanguageDependencies] = useState<Record<CodeLanguage, CodeDependency[]> | null>(null)
   useEffect(() => {
     if (appId) {
       (async () => {
         const { config: javaScriptConfig } = await fetchNodeDefault(appId, BlockEnum.Code, { code_language: CodeLanguage.javascript }) as any
-        const { config: pythonConfig } = await fetchNodeDefault(appId, BlockEnum.Code, { code_language: CodeLanguage.python3 }) as any
+        const { config: pythonConfig, available_dependencies: pythonDependencies } = await fetchNodeDefault(appId, BlockEnum.Code, { code_language: CodeLanguage.python3 }) as any
         setAllLanguageDefault({
           [CodeLanguage.javascript]: javaScriptConfig as CodeNodeType,
           [CodeLanguage.python3]: pythonConfig as CodeNodeType,
         } as any)
+        setAllLanguageDependencies({
+          [CodeLanguage.python3]: pythonDependencies as CodeDependency[],
+        } as any)
       })()
     }
   }, [appId])
@@ -41,6 +45,62 @@ const useConfig = (id: string, payload: CodeNodeType) => {
     setInputs,
   })
 
+  const handleAddDependency = useCallback((dependency: CodeDependency) => {
+    const newInputs = produce(inputs, (draft) => {
+      if (!draft.dependencies)
+        draft.dependencies = []
+      draft.dependencies.push(dependency)
+    })
+    setInputs(newInputs)
+  }, [inputs, setInputs])
+
+  const handleRemoveDependency = useCallback((index: number) => {
+    const newInputs = produce(inputs, (draft) => {
+      if (!draft.dependencies)
+        draft.dependencies = []
+      draft.dependencies.splice(index, 1)
+    })
+    setInputs(newInputs)
+  }, [inputs, setInputs])
+
+  const handleChangeDependency = useCallback((index: number, dependency: CodeDependency) => {
+    const newInputs = produce(inputs, (draft) => {
+      if (!draft.dependencies)
+        draft.dependencies = []
+      draft.dependencies[index] = dependency
+    })
+    setInputs(newInputs)
+  }, [inputs, setInputs])
+
+  const [allowDependencies, setAllowDependencies] = useState<boolean>(false)
+  useEffect(() => {
+    if (!inputs.code_language)
+      return
+    if (!allLanguageDependencies)
+      return
+
+    const newAllowDependencies = !!allLanguageDependencies[inputs.code_language]
+    setAllowDependencies(newAllowDependencies)
+  }, [allLanguageDependencies, inputs.code_language])
+
+  const [availableDependencies, setAvailableDependencies] = useState<CodeDependency[]>([])
+  useEffect(() => {
+    if (!inputs.code_language)
+      return
+    if (!allLanguageDependencies)
+      return
+
+    const newAvailableDependencies = produce(allLanguageDependencies[inputs.code_language], (draft) => {
+      const currentLanguage = inputs.code_language
+      if (!currentLanguage || !draft || !inputs.dependencies)
+        return []
+      return draft.filter((dependency) => {
+        return !inputs.dependencies?.find(d => d.name === dependency.name)
+      })
+    })
+    setAvailableDependencies(newAvailableDependencies || [])
+  }, [allLanguageDependencies, inputs.code_language, inputs.dependencies])
+
   const [outputKeyOrders, setOutputKeyOrders] = useState<string[]>([])
   const syncOutputKeyOrders = useCallback((outputs: OutputVar) => {
     setOutputKeyOrders(Object.keys(outputs))
@@ -163,6 +223,11 @@ const useConfig = (id: string, payload: CodeNodeType) => {
     inputVarValues,
     setInputVarValues,
     runResult,
+    availableDependencies,
+    allowDependencies,
+    handleAddDependency,
+    handleRemoveDependency,
+    handleChangeDependency,
   }
 }
 

web/i18n/en-US/workflow.ts

@@ -273,6 +273,9 @@ const translation = {
     code: {
       inputVars: 'Input Variables',
       outputVars: 'Output Variables',
+      advancedDependencies: 'Advanced Dependencies',
+      advancedDependenciesTip: 'Add some preloaded dependencies that take more time to consume or are not default built-in here',
+      searchDependencies: 'Search Dependencies',
     },
     templateTransform: {
       inputVars: 'Input Variables',

web/i18n/zh-Hans/workflow.ts

@@ -273,6 +273,9 @@ const translation = {
     code: {
       inputVars: '输入变量',
       outputVars: '输出变量',
+      advancedDependencies: '高级依赖',
+      advancedDependenciesTip: '在这里添加一些预加载需要消耗较多时间或非默认内置的依赖包',
+      searchDependencies: '搜索依赖',
     },
     templateTransform: {
       inputVars: '输入变量',