掘金 人工智能 04月30日 17:08
【Dify(v1.x) 核心源码深入解析】Code Executor 模块
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

本文深入解析了 Dify 的 Code Executor 模块,这是一个用于在沙盒环境中安全执行代码的强大工具。它支持 Python3、Jinja2 和 JavaScript 等多种语言,通过模板转换器、代码提供者和代码执行器等核心组件,实现了代码的安全执行与结果处理。文章详细介绍了模块的工作机制,帮助读者全面理解其功能和实现原理。

💡模板转换器是核心组件,负责将用户代码和输入参数转换为可执行脚本。针对 Python3、JavaScript 和 Jinja2,分别实现了不同的转换逻辑,包括脚本的生成、输入参数的解析以及输出结果的格式化。

🔑代码提供者为不同编程语言提供默认代码模板,方便用户快速上手。Python3 和 JavaScript 代码提供者分别定义了各自语言的默认代码模板,简化了代码的编写过程。

⚙️Code Executor 模块还包括抽象基类和辅助类,定义了模板转换器和代码提供者的接口,以及辅助功能,确保了模块的灵活性和可扩展性。

重磅推荐专栏:《大模型AIGC》《课程大纲》《知识星球》

本专栏致力于探索和讨论当今最前沿的技术趋势和应用领域,包括但不限于ChatGPT和Stable Diffusion等。我们将深入研究大型模型的开发和应用,以及与之相关的人工智能生成内容(AIGC)技术。通过深入的技术解析和实践经验分享,旨在帮助读者更好地理解和应用这些领域的最新进展

一、引言

Dify 的 Code Executor 模块是一个功能强大的工具,用于在沙盒环境中安全地执行代码。它支持多种编程语言,包括 Python3、Jinja2 和 JavaScript,并提供了丰富的功能来处理代码执行的各个方面。本文将深入剖析 Code Executor 模块的内部工作机制,帮助读者全面理解其功能和实现原理。

二、模块概览

Code Executor 模块主要包括以下几个核心部分:

    模板转换器(Template Transformers):将代码和输入参数转换为可执行的脚本。代码提供者(Code Providers):为不同编程语言提供默认代码模板。代码执行器(Code Executor):负责在沙盒环境中执行代码,并处理执行结果。抽象基类和辅助类:定义了代码提供者的接口和辅助功能。

三、详细解读

1. 模板转换器(Template Transformers)

模板转换器是 Code Executor 模块的核心组件之一,负责将用户提供的代码和输入参数转换为可执行的脚本。每个模板转换器都针对特定的编程语言进行实现。

Python3 模板转换器

from textwrap import dedentfrom core.helper.code_executor.template_transformer import TemplateTransformerclass Python3TemplateTransformer(TemplateTransformer):    @classmethod    def get_runner_script(cls) -> str:        runner_script = dedent(f"""            # declare main function            {cls._code_placeholder}                        import json            from base64 import b64decode                        # decode and prepare input dict            inputs_obj = json.loads(b64decode('{cls._inputs_placeholder}').decode('utf-8'))                        # execute main function            output_obj = main(**inputs_obj)                        # convert output to json and print            output_json = json.dumps(output_obj, indent=4)            result = f'''<<RESULT>>{{output_json}}<<RESULT>>'''            print(result)            """)        return runner_script

JavaScript 模板转换器

from textwrap import dedentfrom core.helper.code_executor.template_transformer import TemplateTransformerclass NodeJsTemplateTransformer(TemplateTransformer):    @classmethod    def get_runner_script(cls) -> str:        runner_script = dedent(            f"""            // declare main function            {cls._code_placeholder}                        // decode and prepare input object            var inputs_obj = JSON.parse(Buffer.from('{cls._inputs_placeholder}', 'base64').toString('utf-8'))                        // execute main function            var output_obj = main(inputs_obj)                        // convert output to json and print            var output_json = JSON.stringify(output_obj)            var result = `<<RESULT>>${{output_json}}<<RESULT>>`            console.log(result)            """        )        return runner_script

Jinja2 模板转换器

from textwrap import dedentfrom core.helper.code_executor.template_transformer import TemplateTransformerclass Jinja2TemplateTransformer(TemplateTransformer):    @classmethod    def transform_response(cls, response: str) -> dict:        """        Transform response to dict        :param response: response        :return:        """        return {"result": cls.extract_result_str_from_response(response)}    @classmethod    def get_runner_script(cls) -> str:        runner_script = dedent(f"""            # declare main function            def main(**inputs):                import jinja2                template = jinja2.Template('''{cls._code_placeholder}''')                return template.render(**inputs)                            import json            from base64 import b64decode                        # decode and prepare input dict            inputs_obj = json.loads(b64decode('{cls._inputs_placeholder}').decode('utf-8'))                        # execute main function            output = main(**inputs_obj)                        # convert output and print            result = f'''<<RESULT>>{{output}}<<RESULT>>'''            print(result)                        """)        return runner_script    @classmethod    def get_preload_script(cls) -> str:        preload_script = dedent("""            import jinja2            from base64 import b64decode                        def _jinja2_preload_():                # prepare jinja2 environment, load template and render before to avoid sandbox issue                template = jinja2.Template('{{s}}')                template.render(s='a')                        if __name__ == '__main__':                _jinja2_preload_()                        """)        return preload_script

模板转换器的抽象基类

import jsonimport refrom abc import ABC, abstractmethodfrom base64 import b64encodefrom collections.abc import Mappingfrom typing import Anyclass TemplateTransformer(ABC):    _code_placeholder: str = "{{code}}"    _inputs_placeholder: str = "{{inputs}}"    _result_tag: str = "<<RESULT>>"    @classmethod    def transform_caller(cls, code: str, inputs: Mapping[str, Any]) -> tuple[str, str]:        """        Transform code to python runner        :param code: code        :param inputs: inputs        :return: runner, preload        """        runner_script = cls.assemble_runner_script(code, inputs)        preload_script = cls.get_preload_script()        return runner_script, preload_script    @classmethod    def extract_result_str_from_response(cls, response: str):        result = re.search(rf"{cls._result_tag}(.*){cls._result_tag}", response, re.DOTALL)        if not result:            raise ValueError("Failed to parse result")        return result.group(1)    @classmethod    def transform_response(cls, response: str) -> Mapping[str, Any]:        """        Transform response to dict        :param response: response        :return:        """        try:            result = json.loads(cls.extract_result_str_from_response(response))        except json.JSONDecodeError:            raise ValueError("failed to parse response")        if not isinstance(result, dict):            raise ValueError("result must be a dict")        if not all(isinstance(k, str) for k in result):            raise ValueError("result keys must be strings")        return result    @classmethod    @abstractmethod    def get_runner_script(cls) -> str:        """        Get runner script        """        pass    @classmethod    def serialize_inputs(cls, inputs: Mapping[str, Any]) -> str:        inputs_json_str = json.dumps(inputs, ensure_ascii=False).encode()        input_base64_encoded = b64encode(inputs_json_str).decode("utf-8")        return input_base64_encoded    @classmethod    def assemble_runner_script(cls, code: str, inputs: Mapping[str, Any]) -> str:        # assemble runner script        script = cls.get_runner_script()        script = script.replace(cls._code_placeholder, code)        inputs_str = cls.serialize_inputs(inputs)        script = script.replace(cls._inputs_placeholder, inputs_str)        return script    @classmethod    def get_preload_script(cls) -> str:        """        Get preload script        """        return ""

2. 代码提供者(Code Providers)

代码提供者为不同编程语言提供默认代码模板,帮助用户快速上手。

Python3 代码提供者

from textwrap import dedentfrom core.helper.code_executor.code_executor import CodeLanguagefrom core.helper.code_executor.code_node_provider import CodeNodeProviderclass Python3CodeProvider(CodeNodeProvider):    @staticmethod    def get_language() -> str:        return CodeLanguage.PYTHON3    @classmethod    def get_default_code(cls) -> str:        return dedent(            """            def main(arg1: str, arg2: str) -> dict:                return {                    "result": arg1 + arg2,                }            """        )

JavaScript 代码提供者

from textwrap import dedentfrom core.helper.code_executor.code_executor import CodeLanguagefrom core.helper.code_executor.code_node_provider import CodeNodeProviderclass JavascriptCodeProvider(CodeNodeProvider):    @staticmethod    def get_language() -> str:        return CodeLanguage.JAVASCRIPT    @classmethod    def get_default_code(cls) -> str:        return dedent(            """            function main({arg1, arg2}) {                return {                    result: arg1 + arg2                }            }            """        )

代码提供者的抽象基类

from abc import abstractmethodfrom pydantic import BaseModelclass CodeNodeProvider(BaseModel):    @staticmethod    @abstractmethod    def get_language() -> str:        pass    @classmethod    def is_accept_language(cls, language: str) -> bool:        return language == cls.get_language()    @classmethod    @abstractmethod    def get_default_code(cls) -> str:        """        get default code in specific programming language for the code node        """        pass    @classmethod    def get_default_config(cls) -> dict:        return {            "type": "code",            "config": {                "variables": [{"variable": "arg1", "value_selector": []}, {"variable": "arg2", "value_selector": []}],                "code_language": cls.get_language(),                "code": cls.get_default_code(),                "outputs": {"result": {"type": "string", "children": None}},            },        }

3. 代码执行器(Code Executor)

代码执行器负责在沙盒环境中执行代码,并处理执行结果。

import loggingfrom collections.abc import Mappingfrom enum import StrEnumfrom threading import Lockfrom typing import Any, Optionalfrom httpx import Timeout, postfrom pydantic import BaseModelfrom yarl import URLfrom configs import dify_configfrom core.helper.code_executor.javascript.javascript_transformer import NodeJsTemplateTransformerfrom core.helper.code_executor.jinja2.jinja2_transformer import Jinja2TemplateTransformerfrom core.helper.code_executor.python3.python3_transformer import Python3TemplateTransformerfrom core.helper.code_executor.template_transformer import TemplateTransformerlogger = logging.getLogger(__name__)class CodeExecutionError(Exception):    passclass CodeExecutionResponse(BaseModel):    class Data(BaseModel):        stdout: Optional[str] = None        error: Optional[str] = None    code: int    message: str    data: Dataclass CodeLanguage(StrEnum):    PYTHON3 = "python3"    JINJA2 = "jinja2"    JAVASCRIPT = "javascript"class CodeExecutor:    dependencies_cache: dict[str, str] = {}    dependencies_cache_lock = Lock()    code_template_transformers: dict[CodeLanguage, type[TemplateTransformer]] = {        CodeLanguage.PYTHON3: Python3TemplateTransformer,        CodeLanguage.JINJA2: Jinja2TemplateTransformer,        CodeLanguage.JAVASCRIPT: NodeJsTemplateTransformer,    }    code_language_to_running_language = {        CodeLanguage.JAVASCRIPT: "nodejs",        CodeLanguage.JINJA2: CodeLanguage.PYTHON3,        CodeLanguage.PYTHON3: CodeLanguage.PYTHON3,    }    supported_dependencies_languages: set[CodeLanguage] = {CodeLanguage.PYTHON3}    @classmethod    def execute_code(cls, language: CodeLanguage, preload: str, code: str) -> str:        """        Execute code        :param language: code language        :param preload: the preload script        :param code: code        :return:        """        url = URL(str(dify_config.CODE_EXECUTION_ENDPOINT)) / "v1" / "sandbox" / "run"        headers = {"X-Api-Key": dify_config.CODE_EXECUTION_API_KEY}        data = {            "language": cls.code_language_to_running_language.get(language),            "code": code,            "preload": preload,            "enable_network": True,        }        try:            response = post(                str(url),                json=data,                headers=headers,                timeout=Timeout(                    connect=dify_config.CODE_EXECUTION_CONNECT_TIMEOUT,                    read=dify_config.CODE_EXECUTION_READ_TIMEOUT,                    write=dify_config.CODE_EXECUTION_WRITE_TIMEOUT,                    pool=None,                ),            )            if response.status_code == 503:                raise CodeExecutionError("Code execution service is unavailable")            elif response.status_code != 200:                raise Exception(                    f"Failed to execute code, got status code {response.status_code},"                    f" please check if the sandbox service is running"                )        except CodeExecutionError as e:            raise e        except Exception as e:            raise CodeExecutionError(                "Failed to execute code, which is likely a network issue,"                " please check if the sandbox service is running."                f" ( Error: {str(e)} )"            )        try:            response_data = response.json()        except:            raise CodeExecutionError("Failed to parse response")        if (code := response_data.get("code")) != 0:            raise CodeExecutionError(f"Got error code: {code}. Got error msg: {response_data.get('message')}")        response_code = CodeExecutionResponse(**response_data)        if response_code.data.error:            raise CodeExecutionError(response_code.data.error)        return response_code.data.stdout or ""    @classmethod    def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]):        """        Execute code        :param language: code language        :param code: code        :param inputs: inputs        :return:        """        template_transformer = cls.code_template_transformers.get(language)        if not template_transformer:            raise CodeExecutionError(f"Unsupported language {language}")        runner, preload = template_transformer.transform_caller(code, inputs)        try:            response = cls.execute_code(language, preload, runner)        except CodeExecutionError as e:            raise e        return template_transformer.transform_response(response)

4. 模板转换器和代码提供者的协作

模板转换器和代码提供者协同工作,为用户提供了一个完整的代码执行流程。用户可以选择不同的编程语言,使用默认代码模板或自定义代码,然后通过模板转换器生成可执行的脚本,最后由代码执行器在沙盒环境中执行。

5. 代码执行流程

代码执行流程可以分为以下几个步骤:

    选择编程语言:用户选择要使用的编程语言,如 Python3、Jinja2 或 JavaScript。获取默认代码(可选):用户可以使用代码提供者提供的默认代码模板。编写代码:用户编写要执行的代码。准备输入参数:用户准备代码执行所需的输入参数。转换为运行脚本:模板转换器将用户代码和输入参数转换为可执行的脚本和预加载脚本。执行代码:代码执行器在沙盒环境中执行运行脚本。处理执行结果:代码执行器处理执行结果,并将其转换为用户友好的格式。

时序图

sequenceDiagram    participant User    participant CodeProvider    participant TemplateTransformer    participant CodeExecutor    User->>CodeProvider: 选择编程语言    CodeProvider->>User: 返回默认代码模板    User->>TemplateTransformer: 提供代码和输入参数    TemplateTransformer->>TemplateTransformer: 生成运行脚本和预加载脚本    TemplateTransformer->>CodeExecutor: 提交运行脚本和预加载脚本    CodeExecutor->>CodeExecutor: 在沙盒环境中执行代码    CodeExecutor->>TemplateTransformer: 获取执行结果    TemplateTransformer->>TemplateTransformer: 转换执行结果    TemplateTransformer->>User: 返回处理后的结果

四、总结

Dify 的 Code Executor 模块是一个功能强大且灵活的工具,它支持多种编程语言,提供了安全的代码执行环境,并通过模板转换器和代码提供者简化了代码执行的流程。本文详细介绍了 Code Executor 模块的各个核心组件及其协作方式,希望能帮助读者更好地理解和使用这一模块。

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

Dify Code Executor 代码执行 沙盒环境
相关文章