当前位置：首页 > article >正文

Autogen_core： test_code_executor.py

article 2025/2/1 5:05:47

代码

import textwrap

import pytest
from autogen_core.code_executor import (
    Alias,
    FunctionWithRequirements,
    FunctionWithRequirementsStr,
    ImportFromModule,
)
from autogen_core.code_executor._func_with_reqs import build_python_functions_file
from pandas import DataFrame, concat


def template_function() -> DataFrame:  # type: ignore
    data1 = {
        "name": ["John", "Anna"],
        "location": ["New York", "Paris"],
        "age": [24, 13],
    }
    data2 = {
        "name": ["Peter", "Linda"],
        "location": ["Berlin", "London"],
        "age": [53, 33],
    }
    df1 = DataFrame.from_dict(data1)  # type: ignore
    df2 = DataFrame.from_dict(data2)  # type: ignore
    return concat([df1, df2])  # type: ignore

function = FunctionWithRequirements.from_callable(  # type: ignore
        template_function,
        ["pandas"],
        [ImportFromModule("pandas", ["DataFrame", "concat"])],
    )

function

FunctionWithRequirements(func=<function template_function at 0x000002ACEFA96FC0>, python_packages=['pandas'], global_imports=[ImportFromModule(module='pandas', imports=('DataFrame', 'concat'))])

functions_module = build_python_functions_file([function])
functions_module

'from pandas import DataFrame, concat\n\ndef template_function() -> DataFrame:  # type: ignore\n    data1 = {\n        "name": ["John", "Anna"],\n        "location": ["New York", "Paris"],\n        "age": [24, 13],\n    }\n    data2 = {\n        "name": ["Peter", "Linda"],\n        "location": ["Berlin", "London"],\n        "age": [53, 33],\n    }\n    df1 = DataFrame.from_dict(data1)  # type: ignore\n    df2 = DataFrame.from_dict(data2)  # type: ignore\n    return concat([df1, df2])  # type: ignore\n\n\n'

assert "from pandas import DataFrame, concat" in functions_module

function2: FunctionWithRequirementsStr = FunctionWithRequirements.from_str(
        textwrap.dedent(
            """
            def template_function2():
                return pd.Series([1, 2])
            """
        ),
        "pandas",
        [Alias("pandas", "pd")],
    )

function2

FunctionWithRequirementsStr(func='\ndef template_function2():\n    return pd.Series([1, 2])\n', compiled_func=<function template_function2 at 0x000002ACA0655A80>, _func_name='template_function2', python_packages='pandas', global_imports=[Alias(name='pandas', alias='pd')])

functions_module2 = build_python_functions_file([function2])
functions_module2

'import pandas as pd\n\n\ndef template_function2():\n    return pd.Series([1, 2])\n\n\n'

assert "import pandas as pd" in functions_module2

代码解释

上面的代码主要展示了如何使用 autogen_core 库中的 FunctionWithRequirements 类来定义和管理具有依赖项的 Python 函数，并将这些函数及其依赖项打包成一个可执行的 Python 模块。

代码分为几个部分：

定义模板函数：
- 定义了一个名为 template_function 的函数，它创建两个字典 data1 和 data2，然后将这些字典转换为 pandas DataFrame 对象，并将它们合并成一个 DataFrame 对象返回。
创建 FunctionWithRequirements 对象：
- 使用 FunctionWithRequirements.from_callable 方法，将 template_function 函数及其依赖项（这里是 pandas 库）封装成一个 FunctionWithRequirements 对象。这个对象包含了函数代码、所需的 Python 包和所需的导入项。
生成 Python 函数文件：
- 使用 build_python_functions_file 函数，将 FunctionWithRequirements 对象转换成一个包含函数代码和必要导入语句的字符串。这个字符串可以被保存为一个 Python 文件，并作为独立模块运行。
断言检查：
- 通过断言确保生成的 Python 模块包含正确的导入语句。
定义第二个模板函数：
- 定义了一个名为 template_function2 的函数，它返回一个 pandas Series 对象。这个函数使用 FunctionWithRequirements.from_str 方法从字符串定义，并使用 pandas 别名 pd。
生成第二个 Python 函数文件：
- 类似于之前的步骤，将 template_function2 函数及其依赖项打包成一个 Python 模块。
断言检查：
- 通过断言确保生成的第二个 Python 模块包含正确的导入语句。