如何将AutoGen工具限定到工作目录？优化代码无需手动处理路径

我正在使用AutoGen，为节省资源添加了读写文本文件的工具。我的代理已配置工作目录：

executor = autogen.UserProxyAgent(
    name="executor",
    system_message="Executor. Execute the code written by the Engineer and report the result.",
    human_input_mode="NEVER",
    code_execution_config={
        "last_n_messages": 3,
        "work_dir": WORKING_DIR,
        "use_docker": False,
    },  # Please set use_docker=True if docker is available to run the generated code. Using docker is safer than running the generated code directly.
)

同时编写了以下工具函数：

def read_file(file_name: Annotated[str, "File name has to be json, txt or html"]) -> int:
    if not file_name.endswith(".json") and not file_name.endswith(".txt") and not file_name.endswith(".html"):
        return f"I can read only .json, .txt or .html files you asked for {file_name}. Use python to read other files."
    if not os.path.exists(os.path.join(WORKING_DIR, file_name)):
        return f"File {file_name} does not exist."
    with open(os.path.join(WORKING_DIR, file_name), "r") as f:
        return f.read()


def write_file(file_name: Annotated[str, "File name"], content: Annotated[str, "text or json content"]) -> int:
    # verify that nested folders exists
    if not os.path.exists(f"{WORKING_DIR}/{os.path.dirname(file_name)}"):
        os.makedirs(f"{WORKING_DIR}/{os.path.dirname(file_name)}")

    with open(f"{WORKING_DIR}/{file_name}", "w") as f:
        return f.write(content)

请问如何优化上述代码，无需在工具函数中手动处理工作目录？

可以通过以下几种方式实现工具函数与工作目录的解耦，避免手动拼接路径：

1. 闭包封装工作目录 #

创建工厂函数接收工作目录参数，返回绑定该目录的读写函数，无需在函数内引用全局变量：

import os
from typing import Annotated

def create_file_tools(work_dir):
    def read_file(file_name: Annotated[str, "File name has to be json, txt or html"]) -> str:
        allowed_extensions = (".json", ".txt", ".html")
        if not file_name.endswith(allowed_extensions):
            return f"仅支持读取{allowed_extensions}类型文件，你请求的是{file_name}，请用Python代码读取其他类型文件。"
        full_path = os.path.join(work_dir, file_name)
        if not os.path.exists(full_path):
            return f"文件{file_name}不存在。"
        with open(full_path, "r") as f:
            return f.read()

    def write_file(file_name: Annotated[str, "File name"], content: Annotated[str, "text or json content"]) -> int:
        full_path = os.path.join(work_dir, file_name)
        os.makedirs(os.path.dirname(full_path), exist_ok=True)
        with open(full_path, "w") as f:
            return f.write(content)
    
    return read_file, write_file

# 使用时传入代理的工作目录
read_file, write_file = create_file_tools(WORKING_DIR)

2. 类封装工具逻辑 #

将工作目录作为类的初始化参数，读写方法作为类成员，通过实例调用自动使用绑定的目录：

import os
from typing import Annotated

class FileTools:
    def __init__(self, work_dir):
        self.work_dir = work_dir

    def read_file(self, file_name: Annotated[str, "File name has to be json, txt or html"]) -> str:
        allowed_extensions = (".json", ".txt", ".html")
        if not file_name.endswith(allowed_extensions):
            return f"仅支持读取{allowed_extensions}类型文件，你请求的是{file_name}，请用Python代码读取其他类型文件。"
        full_path = os.path.join(self.work_dir, file_name)
        if not os.path.exists(full_path):
            return f"文件{file_name}不存在。"
        with open(full_path, "r") as f:
            return f.read()

    def write_file(self, file_name: Annotated[str, "File name"], content: Annotated[str, "text or json content"]) -> int:
        full_path = os.path.join(self.work_dir, file_name)
        os.makedirs(os.path.dirname(full_path), exist_ok=True)
        with open(full_path, "w") as f:
            return f.write(content)

# 创建工具实例并注册
file_tools = FileTools(WORKING_DIR)
executor.register_function(
    function_map={
        "read_file": file_tools.read_file,
        "write_file": file_tools.write_file
    }
)

3. AutoGen上下文传递 #

利用AutoGen工具注册时的context参数传递工作目录，工具函数通过上下文获取路径：

import os
from typing import Annotated, Optional

def read_file(
    file_name: Annotated[str, "File name has to be json, txt or html"],
    context: Optional[dict] = None
) -> str:
    work_dir = context.get("work_dir") if context else "."
    allowed_extensions = (".json", ".txt", ".html")
    if not file_name.endswith(allowed_extensions):
        return f"仅支持读取{allowed_extensions}类型文件，你请求的是{file_name}，请用Python代码读取其他类型文件。"
    full_path = os.path.join(work_dir, file_name)
    if not os.path.exists(full_path):
        return f"文件{file_name}不存在。"
    with open(full_path, "r") as f:
        return f.read()

def write_file(
    file_name: Annotated[str, "File name"],
    content: Annotated[str, "text or json content"],
    context: Optional[dict] = None
) -> int:
    work_dir = context.get("work_dir") if context else "."
    full_path = os.path.join(work_dir, file_name)
    os.makedirs(os.path.dirname(full_path), exist_ok=True)
    with open(full_path, "w") as f:
        return f.write(content)

# 注册工具时传入上下文
executor.register_function(
    function_map={
        "read_file": read_file,
        "write_file": write_file
    },
    context={"work_dir": WORKING_DIR}
)

额外细节优化 #

• 统一使用os.path.join拼接路径，避免跨平台路径格式问题
• 用os.makedirs(..., exist_ok=True)替代手动判断目录是否存在的逻辑
• 修正read_file的返回值类型标注（实际返回str而非int）

内容的提问来源于stack exchange，提问作者Vova Bilyachat