⚙️ DeepAgents 中间件

"""
TodoListMiddleware - 任务规划示例
"""
from deepagents import create_deep_agent
from deepagents.middleware import TodoListMiddleware

agent = create_deep_agent(
    model="gpt-4o",
    middleware=[
        TodoListMiddleware(
            system_prompt="""当遇到多步骤任务时，
            使用 write_todos 工具分解任务、跟踪进度。"""
        )
    ]
)

# 使用 Agent
result = agent.invoke({
    "messages": [{
        "role": "user",
        "content": "帮我研究 AI Agent 技术栈并写一份报告"
    }]
})

# Agent 会自动：
# 1. 使用 write_todos 分解任务：
#    - 研究 Agent 框架
#    - 收集代码示例
#    - 分析最佳实践
#    - 撰写报告
# 2. 逐步执行并更新待办事项状态
# 3. 完成所有任务后生成报告

"""
FilesystemMiddleware - 记忆管理示例
"""
from deepagents import create_deep_agent
from deepagents.middleware import FilesystemMiddleware
from deepagents.backends import CompositeBackend, StateBackend, StoreBackend
from langgraph.store.memory import InMemoryStore

# 配置混合存储后端
def make_backend(runtime):
    return CompositeBackend(
        default=StateBackend(runtime),  # 临时文件
        routes={
            "/memories/": StoreBackend(runtime)  # 长期记忆
        }
    )

agent = create_deep_agent(
    model="claude-sonnet-4-5-20250929",
    store=InMemoryStore(),
    middleware=[
        FilesystemMiddleware(
            backend=make_backend,
            custom_tool_descriptions={
                "write_file": "保存搜索结果、中间数据到文件系统",
                "read_file": "读取已保存的文件内容",
            }
        )
    ],
    system_prompt="""你可以使用文件系统管理数据：

- 临时数据：保存到 /tmp/ 下
- 长期记忆：保存到 /memories/ 下（跨对话持久化）

示例：
- /tmp/search_results.txt（临时）
- /memories/user_preferences.txt（持久化）
"""
)

# Agent 可以自动管理记忆
result = agent.invoke({
    "messages": [{
        "role": "user",
        "content": "搜索 LangChain 最新功能并保存到长期记忆"
    }]
})

# Agent 会：
# 1. 搜索信息
# 2. 使用 write_file 保存到 /memories/langchain_features.txt
# 3. 下次对话可以用 read_file 读取

"""
SubAgentMiddleware - 任务委派示例
"""
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.middleware import SubAgentMiddleware

@tool
def internet_search(query: str) -> dict:
    """搜索网络信息"""
    # 实际搜索逻辑
    return {"results": [...]}

agent = create_deep_agent(
    model="claude-sonnet-4-5-20250929",
    middleware=[
        SubAgentMiddleware(
            subagents=[
                {
                    "name": "research-agent",
                    "description": "深入研究特定主题，综合多个来源",
                    "system_prompt": "你是研究专家...",
                    "tools": [internet_search],
                    "model": "gpt-4o",  # 子代理可使用不同模型
                }
            ],
            default_model="claude-sonnet-4-5-20250929"
        )
    ]
)

# 使用 Agent
result = agent.invoke({
    "messages": [{
        "role": "user",
        "content": "深入研究量子计算的最新进展"
    }]
})

# Agent 会：
# 1. 识别这是复杂研究任务
# 2. 使用 task 工具委派给 research-agent
# 3. 子代理执行详细研究（上下文隔离）
# 4. 子代理返回简洁摘要给主 Agent
# 5. 主 Agent 整合结果返回用户

"""
禁用特定中间件
"""
from deepagents import create_deep_agent
from deepagents.middleware import FilesystemMiddleware

# 只启用 FilesystemMiddleware，禁用其他
agent = create_deep_agent(
    model="gpt-4o",
    middleware=[
        FilesystemMiddleware()  # 仅文件系统中间件
    ]
)

# 此 Agent：
# ✅ 有文件系统工具（ls, read_file, write_file, edit_file）
# ❌ 无待办事项工具（write_todos）
# ❌ 无子代理工具（task）

"""
子代理独立中间件配置
"""
from deepagents import create_deep_agent
from deepagents.middleware import (
    TodoListMiddleware,
    FilesystemMiddleware,
    SubAgentMiddleware
)

agent = create_deep_agent(
    model="claude-sonnet-4-5-20250929",

    # 主 Agent 的中间件
    middleware=[
        TodoListMiddleware(),
        SubAgentMiddleware(
            subagents=[{
                "name": "data-processor",
                "description": "处理大量数据文件",
                "tools": [...],

                # 子代理的独立中间件配置
                "middleware": [
                    FilesystemMiddleware()  # 仅文件系统
                ]
            }]
        )
    ]
)

# 配置结果：
# 主 Agent：有 TodoList + SubAgent 能力
# 子代理 data-processor：仅有 Filesystem 能力

# 用途：
# - 主 Agent 专注协调和规划
# - 子代理专注数据处理，无需规划工具

"""
完整实战示例 - 研究助手（综合使用三种中间件）
"""
import os
from typing import Literal
from tavily import TavilyClient
from deepagents import create_deep_agent
from deepagents.middleware import (
    TodoListMiddleware,
    FilesystemMiddleware,
    SubAgentMiddleware
)
from deepagents.backends import CompositeBackend, StateBackend, StoreBackend
from langgraph.store.memory import InMemoryStore

# ========== 1. 定义工具 ==========

tavily_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])

def internet_search(query: str, max_results: int = 5) -> dict:
    """搜索网络信息"""
    return tavily_client.search(query, max_results=max_results)

def analyze_data(data: str) -> dict:
    """分析数据并生成洞察"""
    # 简化示例：实际会执行数据分析
    return {
        "insights": ["发现 1", "发现 2"],
        "summary": "数据分析摘要..."
    }

# ========== 2. 配置存储后端 ==========

store = InMemoryStore()

def make_backend(runtime):
    return CompositeBackend(
        default=StateBackend(runtime),
        routes={
            "/memories/": StoreBackend(runtime),  # 长期记忆
            "/tmp/": StateBackend(runtime),        # 临时数据
        }
    )

# ========== 3. 配置中间件 ==========

middleware = [
    # 中间件 1: 任务规划
    TodoListMiddleware(
        system_prompt="""当处理复杂研究任务时：

1. 使用 write_todos 分解任务为具体步骤
2. 按顺序执行每个步骤
3. 完成后更新待办事项状态
"""
    ),

    # 中间件 2: 文件系统和记忆
    FilesystemMiddleware(
        backend=make_backend,
        custom_tool_descriptions={
            "write_file": "保存研究结果到文件（临时或长期）",
            "read_file": "读取已保存的研究笔记",
            "ls": "列出已保存的文件",
            "edit_file": "更新既有的研究笔记",
        }
    ),

    # 中间件 3: 子代理
    SubAgentMiddleware(
        subagents=[
            {
                "name": "web-researcher",
                "description": "专门搜索和分析网络信息",
                "system_prompt": """你是网络研究专家。

任务：
1. 使用 internet_search 搜索相关信息
2. 分析搜索结果质量
3. 综合信息并生成结构化摘要

输出格式：
- 核心发现（3-5 点）
- 关键数据
- 信息来源

保持输出在 500 字以内。""",
                "tools": [internet_search],
                "model": "gpt-4o",
            },
            {
                "name": "data-analyst",
                "description": "分析数据并生成洞察报告",
                "system_prompt": """你是数据分析专家。

任务：
1. 使用 analyze_data 分析数据
2. 识别模式和趋势
3. 生成可操作的建议

输出格式：
- 关键洞察
- 数据支撑
- 建议行动

保持输出在 400 字以内。""",
                "tools": [analyze_data],
                "model": "claude-sonnet-4-5-20250929",
            }
        ],
        default_model="claude-sonnet-4-5-20250929"
    )
]

# ========== 4. 创建研究助手 ==========

research_assistant_prompt = """你是一位高效的研究助手。你的核心能力包括：

## 1. 任务规划能力
- 使用 write_todos 分解复杂任务
- 跟踪进度并适应性调整

## 2. 记忆管理能力
- 临时数据：保存到 /tmp/ 下（对话结束后清除）
- 长期记忆：保存到 /memories/ 下（跨对话持久化）

文件组织：
- /memories/research/[主题]/notes.txt：研究笔记
- /memories/research/[主题]/sources.txt：信息来源
- /tmp/temp_data.txt：临时工作数据

## 3. 任务委派能力
- **web-researcher**：网络搜索和信息综合
- **data-analyst**：数据分析和洞察生成

委派策略：
- 需要搜索多个来源 → 委派给 web-researcher
- 需要分析数据 → 委派给 data-analyst
- 复杂任务可以连续委派多个子代理

## 工作流程

1. **接收任务** → 使用 write_todos 分解
2. **执行研究** → 委派给相应子代理
3. **保存数据** → 使用 write_file 持久化
4. **生成报告** → 整合所有信息
5. **完成任务** → 更新待办事项状态

始终保持主上下文清洁，将详细工作委派给子代理。
"""

agent = create_deep_agent(
    model="claude-sonnet-4-5-20250929",
    middleware=middleware,
    store=store,
    system_prompt=research_assistant_prompt
)

# ========== 5. 使用研究助手 ==========

result = agent.invoke({
    "messages": [{
        "role": "user",
        "content": """请研究 AI Agent 的主流技术栈：

1. 搜索当前流行的 Agent 框架
2. 分析它们的特点和优势
3. 生成对比报告并保存到长期记忆
"""
    }]
})

print(result["messages"][-1].content)

# Agent 执行流程：
#
# 1. 【TodoList】分解任务：
#    - 搜索 Agent 框架
#    - 分析框架特点
#    - 生成对比报告
#    - 保存到长期记忆
#
# 2. 【SubAgent】委派给 web-researcher：
#    - 搜索 LangChain, AutoGPT, CrewAI 等框架
#    - 返回简洁摘要（< 500 字）
#
# 3. 【SubAgent】委派给 data-analyst：
#    - 分析框架对比数据
#    - 返回洞察报告（< 400 字）
#
# 4. 【Filesystem】保存结果：
#    - write_file("/memories/research/ai_agents/report.md")
#    - write_file("/memories/research/ai_agents/sources.txt")
#
# 5. 【TodoList】更新状态：
#    - 标记所有任务为已完成
#
# 6. 返回最终报告给用户

# 优势：
# - 主 Agent 上下文始终清洁（子代理隔离了详细工作）
# - 任务进度透明（TodoList 跟踪）
# - 数据持久化（下次对话可以继续）
# - 专门化处理（不同子代理用最合适的模型）

# ✅ 好的配置：子代理仅需文件系统
subagent = {
    "name": "data-processor",
    "tools": [...],
    "middleware": [FilesystemMiddleware()]  # 仅文件系统
}

# ❌ 不好的配置：子代理有不必要的中间件
subagent = {
    "name": "data-processor",
    "tools": [...],
    "middleware": [
        TodoListMiddleware(),  # 不需要规划
        SubAgentMiddleware()   # 不需要再委派
    ]
}

agent = create_deep_agent(
    model="gpt-4o",
    middleware=[]  # 禁用所有中间件
)

# 此 Agent 只有你显式提供的 tools，无内置工具

subagent = {
    "name": "my-subagent",
    "middleware": [FilesystemMiddleware()]  # 必须显式指定
}

SubAgentMiddleware(
    subagents=[
        {"name": "agent1"},  # 使用 default_model
        {"name": "agent2", "model": "gpt-4o"},  # 使用显式模型
    ],
    default_model="claude-sonnet-4-5-20250929"
)

from deepagents import CompiledSubAgent
from langgraph.graph import StateGraph, END

# 构建 LangGraph 工作流
workflow = StateGraph(...)
# ... 添加节点和边 ...
custom_graph = workflow.compile()

# 包装为子代理
custom_subagent = CompiledSubAgent(
    name="custom-workflow",
    description="执行复杂的多步骤工作流",
    runnable=custom_graph
)

agent = create_deep_agent(
    subagents=[custom_subagent]
)

import os

# 启用 LangSmith 追踪
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "your-api-key"

# 所有中间件工具调用会被记录
result = agent.invoke({...})

# 在 LangSmith UI 中查看：
# - write_todos 调用
# - write_file / read_file 调用
# - task 子代理调用
# - 每个工具的输入输出