diff --git a/.github/workflows/api-tests.yml b/.github/workflows/api-tests.yml index 7d24b15bdf..bfb1054639 100644 --- a/.github/workflows/api-tests.yml +++ b/.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 diff --git a/api/core/helper/code_executor/code_executor.py b/api/core/helper/code_executor/code_executor.py index ec685ae814..f078563658 100644 --- a/api/core/helper/code_executor/code_executor.py +++ b/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) \ No newline at end of file + 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 [] \ No newline at end of file diff --git a/api/core/helper/code_executor/entities.py b/api/core/helper/code_executor/entities.py new file mode 100644 index 0000000000..55464d2ff7 --- /dev/null +++ b/api/core/helper/code_executor/entities.py @@ -0,0 +1,6 @@ +from pydantic import BaseModel + + +class CodeDependency(BaseModel): + name: str + version: str \ No newline at end of file diff --git a/api/core/helper/code_executor/javascript_transformer.py b/api/core/helper/code_executor/javascript_transformer.py index 29b8e06e86..8da16b568f 100644 --- a/api/core/helper/code_executor/javascript_transformer.py +++ b/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: diff --git a/api/core/helper/code_executor/jinja2_transformer.py b/api/core/helper/code_executor/jinja2_transformer.py index 27a3579493..3d557372f1 100644 --- a/api/core/helper/code_executor/jinja2_transformer.py +++ b/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: diff --git a/api/core/helper/code_executor/python_transformer.py b/api/core/helper/code_executor/python_transformer.py index f44acbb9bf..fd28b06187 100644 --- a/api/core/helper/code_executor/python_transformer.py +++ b/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'''<> 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: diff --git a/api/core/helper/code_executor/template_transformer.py b/api/core/helper/code_executor/template_transformer.py index c3564afd04..b83d3df30a 100644 --- a/api/core/helper/code_executor/template_transformer.py +++ b/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 diff --git a/api/core/workflow/nodes/code/code_node.py b/api/core/workflow/nodes/code/code_node.py index 12e7ae940f..3e00e501ac 100644 --- a/api/core/workflow/nodes/code/code_node.py +++ b/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 diff --git a/api/core/workflow/nodes/code/entities.py b/api/core/workflow/nodes/code/entities.py index 555bb3918e..4f957e5afb 100644 --- a/api/core/workflow/nodes/code/entities.py +++ b/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] \ No newline at end of file + outputs: dict[str, Output] + dependencies: Optional[list[CodeDependency]] = None \ No newline at end of file diff --git a/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py b/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py index ef84c92625..a7252d3036 100644 --- a/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py +++ b/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 { diff --git a/docker/docker-compose.middleware.yaml b/docker/docker-compose.middleware.yaml index 6bf45da9e0..d786e7d4c1 100644 --- a/docker/docker-compose.middleware.yaml +++ b/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 - ports: - - "8194:8194" + 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 diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 0c3a0c202f..e232ef436a 100644 --- a/docker/docker-compose.yaml +++ b/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 diff --git a/docker/volumes/sandbox/dependencies/python-requirements.txt b/docker/volumes/sandbox/dependencies/python-requirements.txt new file mode 100644 index 0000000000..e69de29bb2 diff --git a/docker/volumes/ssrf_proxy/squid.conf b/docker/volumes/ssrf_proxy/squid.conf new file mode 100644 index 0000000000..3028bf35c6 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/web/app/components/workflow/nodes/code/dependency-picker.tsx b/web/app/components/workflow/nodes/code/dependency-picker.tsx new file mode 100644 index 0000000000..8fe9fa01bc --- /dev/null +++ b/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 = ({ + 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 ( + + setOpen(!open)} className='flex-grow cursor-pointer'> +
+
{value.name}
+ +
+
+ +
+
+ + setSearchText(e.target.value)} + autoFocus + /> + { + searchText && ( +
setSearchText('')} + > + +
+ ) + } +
+
+ {available_dependencies.filter((v) => { + if (!searchText) + return true + return v.name.toLowerCase().includes(searchText.toLowerCase()) + }).map(dependency => ( +
+
{dependency.name}
+ {dependency.name === value.name && } +
+ ))} +
+
+
+
+ ) +} + +export default React.memo(DependencyPicker) diff --git a/web/app/components/workflow/nodes/code/dependency.tsx b/web/app/components/workflow/nodes/code/dependency.tsx new file mode 100644 index 0000000000..5e868efe31 --- /dev/null +++ b/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 = ({ + available_dependencies, dependencies, handleRemove, handleChange, +}) => { + return ( +
+ {dependencies.map((dependency, index) => ( +
+ handleChange(index, dependency)} + /> + handleRemove(index)} + /> +
+ ))} +
+ ) +} + +export default React.memo(Dependencies) diff --git a/web/app/components/workflow/nodes/code/panel.tsx b/web/app/components/workflow/nodes/code/panel.tsx index 838e7190d3..8ab9b3d0e5 100644 --- a/web/app/components/workflow/nodes/code/panel.tsx +++ b/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> = ({ varInputs, inputVarValues, setInputVarValues, + allowDependencies, + availableDependencies, + handleAddDependency, + handleRemoveDependency, + handleChangeDependency, } = useConfig(id, data) return ( @@ -78,6 +84,31 @@ const Panel: FC> = ({ filterVar={filterVar} /> + { + allowDependencies + ? ( +
+ +
+ handleAddDependency({ name: '', version: '' })} /> + } + tooltip={t(`${i18nPrefix}.advancedDependenciesTip`)!} + > + handleRemoveDependency(index)} + handleChange={(index, dependency) => handleChangeDependency(index, dependency)} + /> + +
+
+ ) + : null + } { const appId = useAppStore.getState().appDetail?.id const [allLanguageDefault, setAllLanguageDefault] = useState | null>(null) + const [allLanguageDependencies, setAllLanguageDependencies] = useState | 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(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([]) + 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([]) 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, } } diff --git a/web/i18n/en-US/workflow.ts b/web/i18n/en-US/workflow.ts index 206bae5400..3ffad0198d 100644 --- a/web/i18n/en-US/workflow.ts +++ b/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', diff --git a/web/i18n/zh-Hans/workflow.ts b/web/i18n/zh-Hans/workflow.ts index 781ff3b49d..f9ae082f6f 100644 --- a/web/i18n/zh-Hans/workflow.ts +++ b/web/i18n/zh-Hans/workflow.ts @@ -273,6 +273,9 @@ const translation = { code: { inputVars: '输入变量', outputVars: '输出变量', + advancedDependencies: '高级依赖', + advancedDependenciesTip: '在这里添加一些预加载需要消耗较多时间或非默认内置的依赖包', + searchDependencies: '搜索依赖', }, templateTransform: { inputVars: '输入变量',