1
import os
2
from pathlib import Path
3
import subprocess
4

5
from dotenv import load_dotenv
6
from langchain.chat_models import init_chat_model
7
from langchain.tools import tool
8

9
load_dotenv(override=True)
10

11
model = init_chat_model(
12
    model="deepseek-v4-flash",
13
    model_provider="openai",
14
    api_key=os.getenv("DEEPSEEK_API_KEY"),
15
    base_url=os.getenv("DEEPSEEK_BASE_URL"),
16
)
17

18
# 工作区：相对"运行脚本时的当前目录"（课程里脚本放在子目录，所以用了 ..）
19
WORKSPACE = Path("../todo_workspace")
20

21
@tool
22
def list_files(path: str = ".") -> str:
23
    """
24
    列出工作区指定目录下的文件和子目录。path 只能是相对路径。
25
    Args:
26
        path: 工作区下的相对路径，一定指向目录，默认为.，表示工作区根路径，不能访问工作区外的目录
27
    """
28
    target = (WORKSPACE / path).resolve()
29
    workspace_root = WORKSPACE.resolve()
30
    if not str(target).startswith(str(workspace_root)):
31
        return "错误：只允许访问工作区内的目录。"
32
    if not target.exists():
33
        return f"错误：目录不存在: {path}"
34
    if not target.is_dir():
35
        return f"错误：不是目录: {path}"
36
    items = sorted(target.iterdir(), key=lambda p: (p.is_file(), p.name.lower()))
37
    if not items:
38
        return f"目录为空: {path}"
39
    lines = []
40
    for item in items:
41
        rel = item.relative_to(workspace_root)
42
        kind = "[DIR]" if item.is_dir() else "[FILE]"
43
        lines.append(f"{kind} {rel.as_posix()}")
44
    return "\n".join(lines)
45

46
@tool
47
def read_file(path: str) -> str:
48
    """
49
    读取工作区中的文本文件内容。path 只能是相对路径。
50
    Args:
51
        path: 工作区内的文件名
52
    """
53
    file_path = (WORKSPACE / path).resolve()
54
    if not str(file_path).startswith(str(WORKSPACE.resolve())):
55
        return "错误：只允许读取工作区内的文件。"
56
    if not file_path.exists():
57
        return f"错误：文件不存在: {path}"
58
    return file_path.read_text(encoding="utf-8")
59

60
@tool
61
def write_file(path: str, content: str) -> str:
62
    """
63
    写入工作区中的文本文件。path 只能是相对路径。
64
    Args:
65
        path: 工作区内的文件名
66
        content: 写入文件的内容
67
    """
68
    file_path = (WORKSPACE / path).resolve()
69
    if not str(file_path).startswith(str(WORKSPACE.resolve())):
70
        return "错误：只允许写入工作区内的文件。"
71
    file_path.write_text(content, encoding="utf-8")
72
    return f"已写入文件: {path}"
73

74
@tool
75
def run_tests() -> str:
76
    """
77
    在工作区运行 pytest -q，并返回输出。
78
    不接收任何参数，返回格式为
79
    returncode=0|1
80
    STDOUT:
81
    STDERR:
82
    """
83
    try:
84
        result = subprocess.run(
85
            ["pytest", "-q"],
86
            cwd=str(WORKSPACE),
87
            capture_output=True,
88
            text=True,
89
            timeout=20,
90
        )
91
        return (
92
            f"returncode={result.returncode}\n\n"
93
            f"STDOUT:\n{result.stdout}\n\n"
94
            f"STDERR:\n{result.stderr}"
95
        )
96
    except Exception as e:
97
        return f"运行测试失败: {e}"

1
from langchain.agents import create_agent
2
from langchain.agents.middleware import TodoListMiddleware
3
from langchain.messages import HumanMessage
4
from rich import print as rprint
5

6
# 1. 初始化 Agent
7
agent = create_agent(
8
    model=model,
9
    # 除 write_todos 等工具外，TodoListMiddleware 需要配合这些业务工具使用
10
    tools=[list_files, read_file, write_file, run_tests],
11
    # 引入 Todo 列表中间件
12
    middleware=[TodoListMiddleware()],
13
    system_prompt=(
14
        "你是一个代码修复助手。遇到多步骤任务时，先使用 write_todos 制定待办事项；"
15
        "然后读取文件、修复代码并运行测试。工作全部在工作区下进行。"
16
    ),
17
)
18

19
# 2. 同步调用：让 Agent 自己改到测试通过
20
print("正在执行 Agent 任务...")
21
final_state = agent.invoke(
22
    {"messages": [HumanMessage(content="请测试并修复工作区下 my_add.py 文件中的代码")]}
23
)
24

25
# 3. 直观展示中间件产生的数据结果
26
rprint(final_state)
27
# TodoListMiddleware 运行期间，会自动把规划好的步骤注入到 state 的 "todos" 字段
28
todos = final_state.get("todos", [])
29
if todos:
30
    for i, item in enumerate(todos, 1):
31
        print(f"{i}. [{item.get('status', 'unknown')}] {item.get('content')}")
32
else:
33
    print("未检测到待办事项（可能 Agent 认为不需要规划，或未触发 write_todos 工具）")