动手：从 Function Calling 跑到 Research Agent

和《Agent 本质》到《实战》是同一条主线，写法不同：那些篇讲为什么这样设计，这一篇把每一步压成能直接运行的代码，从第一次 Function Calling 写到 Research Agent，附环境搭建和概念速查。

两种读法：先跑通这一篇，再回去补原理；或者读完那些篇，用这一篇对照着敲一遍。

代码基于 LangChain 0.3+ / LangGraph 0.2+，在 LangGraph 1.x 上仍可运行。文中的 set_entry_point("x") 在新代码里更常写成 add_edge(START, "x")，两者等价。

---

《Agent 本质》：Model 加 Harness

1.1 什么是 Agent？

Agent（智能体）不是一个新概念，但在大模型时代它获得了全新的含义。一个AI Agent，本质上是一个能够自主感知环境、做出决策、采取行动并观察结果的系统。

与普通的 LLM 调用相比，Agent 的核心区别在于自主性和闭环性：

维度	普通LLM调用	Agent
交互模式	一问一答，单轮	多轮循环，直到目标达成
工具使用	无	可调用外部工具/API
记忆	仅当前上下文窗口	短期+长期记忆系统
规划	无	任务分解、步骤编排
停止条件	生成结束	目标达成或达到限制

用一个生活化的比喻：普通LLM调用就像一个只能坐在办公室里回答问题的顾问，而Agent则是一个能走出去实地调研、打电话、查资料、写报告的全栈执行者。

1.2 Agent = Model + Harness

这是理解Agent最核心的等式：

Agent = Model + Harness

• Model（模型）：LLM本身，是Agent的"大脑"，负责理解、推理和决策
• Harness（脚手架）：围绕模型构建的基础设施，是Agent的"四肢和感官"，负责工具调用、记忆管理、流程编排、错误处理等

为什么要这样拆分？因为：

1. Model的能力是有边界的 —— LLM只能做文本推理，不能搜索网页、不能执行代码、不能记住昨天的事
2. Harness让模型"动起来" —— 它把模型从被动的文本生成器，变成了主动的任务执行者
3. 两者缺一不可 —— 没有Model的Harness是空壳，没有Harness的Model只是聊天机器人

1.3 Agent 架构全景图

下图展示了 Agent 的完整架构，Model和Harness如何协作：

1.4 为什么 Model 单独不够？

让我们看一个具体的例子。假设用户问："北京今天的天气怎么样？"

只用 Model：

用户：北京今天的天气怎么样？
LLM：我无法获取实时天气信息，我的知识截止于2024年...

Model + Harness (Agent)：

用户：北京今天的天气怎么样？
Agent [Thought]：我需要查询实时天气信息，我应该调用天气API。
Agent [Action]：调用 weather_api(city="北京")
Agent [Observation]：{"temp": 22, "condition": "晴", "humidity": 45%}
Agent [Response]：北京今天天气晴朗，气温22°C，湿度45%，适合外出。

区别一目了然：Model只能基于训练数据回答，而Agent通过Harness获得了与真实世界交互的能力。

1.5 Harness 的四大支柱

Harness 的能力可以归纳为四个核心模块，后面几篇会逐一展开：

1. Tool Use（工具使用）—— 让Agent能"动手"，调用外部API和工具
2. Memory（记忆系统）—— 让Agent能"记住"，跨对话保留信息
3. Planning（规划能力）—— 让Agent能"思考"，分解复杂任务
4. Multi-Agent（多智能体协作）—— 让Agent能"合作"，分工完成大型任务

这四大支柱共同构成了Harness，让Agent从一个简单的聊天机器人进化为真正的智能体。

---

《Model》：推理、提示词和 Function Calling

2.1 LLM 推理的本质：Next-Token Prediction

大语言模型的底层原理非常简单：给定前面的文本，预测下一个token。

# 伪代码示意
next_token = model.predict(context="今天天气")
# 可能输出: "很好"、"不错"、"晴朗"...

但这看似简单的机制，在足够大的规模下，涌现出了惊人的能力：

• 上下文学习（In-Context Learning）：通过示例学会新任务
• 思维链推理（Chain-of-Thought）：逐步推理解决复杂问题
• 指令遵循（Instruction Following）：理解并执行复杂指令

2.2 从单次预测到 Agent 循环

单次 next-token prediction 并不构成 Agent。Agent 的关键在于循环：

while goal_not_achieved:
    1. 观察当前状态 (Observation)
    2. 思考下一步行动 (Thought)
    3. 执行行动 (Action)
    4. 获取行动结果 (Observation)
    5. 判断是否达成目标

这个循环就是著名的 ReAct（Reasoning + Acting）范式，也是 Agent 运行的核心模式。

2.3 System Prompt 的工程化设计

System Prompt 是 Agent 行为的"宪法"，它定义了：

• Agent 的角色和身份
• 可用工具及其使用条件
• 行为约束和安全边界
• 输出格式要求

SYSTEM_PROMPT = """你是一个专业的AI研究助手。

## 角色定义
你能够帮助用户进行学术研究、信息检索和数据分析。

## 可用工具
- search_web: 搜索互联网获取最新信息
- read_paper: 读取学术论文内容
- run_python: 执行Python代码进行数据分析

## 行为规则
1. 在给出结论前，必须先收集足够的证据
2. 如果信息不足，主动使用工具搜索
3. 明确区分事实和你的推测
4. 引用信息来源

## 输出格式
使用Markdown格式，关键信息用加粗标注。
"""

System Prompt 设计的最佳实践：

1. 结构化：用清晰的标题和段落组织，不要写大段文字
2. 具体化：给出具体的规则，而非模糊的描述（"必须先搜索" > "尽量多查资料"）
3. 包含负面约束：明确说什么不能做，比只说能做什么更安全
4. 版本化：像代码一样管理你的System Prompt，记录变更

2.4 Function Calling 协议

Function Calling 是 Model 和 Harness 之间的桥梁。它让 LLM 能够"告诉"系统需要调用什么工具。

协议格式（OpenAI 兼容）：

# 定义工具
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取指定城市的天气信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "城市名称，如'北京'、'上海'"
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "温度单位"
                    }
                },
                "required": ["city"]
            }
        }
    }
]

完整的 Function Calling 流程：

2.5 代码实战：第一次 Function Calling

"""
环境准备：
pip install openai langchain langgraph
"""

from openai import OpenAI

client = OpenAI()  # 默认读取 OPENAI_API_KEY 环境变量

# ========== Step 1: 定义工具 ==========
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取指定城市的当前天气信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "城市名称"
                    }
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "执行数学计算",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "数学表达式，如 '2 + 3 * 4'"
                    }
                },
                "required": ["expression"]
            }
        }
    }
]

# ========== Step 2: 实现工具函数 ==========
def get_weather(city: str) -> dict:
    """模拟天气API（实际项目中替换为真实API）"""
    # 这里用模拟数据，实际项目中调用真实天气API
    mock_data = {
        "北京": {"temp": 22, "condition": "晴", "humidity": 45},
        "上海": {"temp": 25, "condition": "多云", "humidity": 65},
        "深圳": {"temp": 28, "condition": "阵雨", "humidity": 80},
    }
    return mock_data.get(city, {"temp": 20, "condition": "未知", "humidity": 50})

def calculate(expression: str) -> str:
    """执行数学计算"""
    try:
        result = eval(expression)  # 注意：生产环境中需要安全沙箱
        return str(result)
    except Exception as e:
        return f"计算错误: {e}"

# 工具名 -> 函数的映射
tool_map = {
    "get_weather": get_weather,
    "calculate": calculate,
}

# ========== Step 3: Agent 循环 ==========
def run_agent(user_message: str) -> str:
    """完整的 Agent 循环：支持多轮工具调用"""
    messages = [
        {"role": "system", "content": "你是一个有帮助的AI助手。当需要时，使用提供的工具来获取信息。"},
        {"role": "user", "content": user_message}
    ]

    max_iterations = 5  # 防止无限循环

    for i in range(max_iterations):
        print(f"\n--- 迭代 {i+1} ---")

        # 调用LLM
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            tools=tools,
            tool_choice="auto"  # 让模型自己决定是否调用工具
        )

        assistant_message = response.choices[0].message
        messages.append(assistant_message.model_dump())

        # 检查是否有工具调用
        if assistant_message.tool_calls:
            for tool_call in assistant_message.tool_calls:
                func_name = tool_call.function.name
                func_args = json.loads(tool_call.function.arguments)

                print(f"  工具调用: {func_name}({func_args})")

                # 执行工具
                result = tool_map[func_name](**func_args)
                print(f"  工具结果: {result}")

                # 将结果添加到消息列表
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": json.dumps(result, ensure_ascii=False)
                })
        else:
            # 没有工具调用，说明LLM已经给出了最终回答
            print(f"  最终回答: {assistant_message.content}")
            return assistant_message.content

    return "达到最大迭代次数，Agent终止。"

# ========== 运行 ==========
import json

answer = run_agent("北京今天天气怎么样？如果室外温度合适，帮我算算北京和上海的平均温度。")
print(f"\n最终回答: {answer}")

运行输出示例：

--- 迭代 1 ---
  工具调用: get_weather({'city': '北京'})
  工具结果: {"temp": 22, "condition": "晴", "humidity": 45}

--- 迭代 2 ---
  工具调用: get_weather({'city': '上海'})
  工具结果: {"temp": 25, "condition": "多云", "humidity": 65}

--- 迭代 3 ---
  工具调用: calculate({'expression': '(22 + 25) / 2'})
  工具结果: "23.5"

--- 迭代 4 ---
  最终回答: 北京今天天气晴朗，气温22°C，适合户外活动。北京和上海的平均温度为23.5°C。

2.6 关键理解

在这个例子中，你能清楚地看到 Model 和 Harness 的分工：

• **Model（LLM）**负责：理解用户意图、决定调用什么工具、解析工具结果、组织最终回答
• **Harness（我们的代码）**负责：定义工具、执行工具调用、管理消息列表、控制循环

这就是 Agent = Model + Harness 的最简单形态。

---

《Tool Use》：从一次调用到生产级循环

3.1 Tool 的定义、注册与发现

工具是 Agent 与外部世界交互的接口。一个好的工具系统需要解决三个问题：

1. 定义：工具长什么样？输入输出是什么？
2. 注册：如何让 LLM 知道有哪些工具可用？
3. 发现：当有大量工具时，如何让 LLM 高效找到合适的工具？

LangChain 的 Tool 抽象

from langchain_core.tools import tool

# 方式1：使用 @tool 装饰器（推荐）
@tool
def search_web(query: str) -> str:
    """搜索互联网获取信息。当需要查找最新资讯、事实性信息时使用此工具。

    Args:
        query: 搜索关键词
    """
    # 实际项目中替换为真实的搜索API（如 Tavily、SerpAPI）
    import requests
    response = requests.get(f"https://api.search.example.com/search?q={query}")
    return response.json()["results"][0]["content"]

# 方式2：使用 StructuredTool（更灵活）
from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field

class CodeExecutionInput(BaseModel):
    code: str = Field(description="要执行的Python代码")
    language: str = Field(default="python", description="编程语言")

def _run_code(code: str, language: str = "python") -> str:
    """在沙箱中执行代码"""
    import subprocess
    result = subprocess.run(
        ["python3", "-c", code],
        capture_output=True, text=True, timeout=30
    )
    return result.stdout or result.stderr

run_code = StructuredTool.from_function(
    func=_run_code,
    name="run_python",
    description="执行Python代码并返回输出。用于数学计算、数据处理等。",
    args_schema=CodeExecutionInput,
)

# 方式3：从现有函数自动生成
def get_current_time(timezone: str = "Asia/Shanghai") -> str:
    """获取当前时间"""
    from datetime import datetime
    import pytz
    tz = pytz.timezone(timezone)
    return datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")

time_tool = tool(get_current_time)

工具注册中心

当工具数量增多时，需要一个注册中心来管理：

from langchain_core.tools import BaseTool
from typing import Dict, List

class ToolRegistry:
    """工具注册中心 —— 管理所有可用工具"""

    def __init__(self):
        self._tools: Dict[str, BaseTool] = {}

    def register(self, tool: BaseTool):
        self._tools[tool.name] = tool

    def get(self, name: str) -> BaseTool:
        return self._tools[name]

    def get_all(self) -> List[BaseTool]:
        return list(self._tools.values())

    def get_openai_tools_schema(self) -> list:
        """生成OpenAI Function Calling格式的工具定义"""
        return [tool.to_openai_tool() for tool in self._tools.values()]

# 使用
registry = ToolRegistry()
registry.register(search_web)
registry.register(run_code)
registry.register(time_tool)

# 获取所有工具的schema，传给LLM
tools_schema = registry.get_openai_tools_schema()

3.2 ReAct 循环：Thought → Action → Observation

ReAct（Reasoning + Acting）是 Agent 最经典的运行模式。它的核心思想是：先思考，再行动，观察结果，继续思考。

ReAct 的完整工作流程

用户：谁创建了Linux？他还在维护吗？

[Thought 1] 我需要找到Linux的创建者，然后查询他是否还在维护。
[Action 1] search_web("Linux creator history")
[Observation 1] "Linux是由Linus Torvalds在1991年创建的..."
[Thought 2] 我知道了Linux的创建者是Linus Torvalds。现在需要查询他是否还在维护Linux。
[Action 2] search_web("Linus Torvalds Linux maintainer 2024")
[Observation 2] "Linus Torvalds仍然是Linux内核的首席维护者..."
[Thought 3] 我已经有了足够的信息来回答用户的问题。
[Answer] Linux由Linus Torvalds于1991年创建，他至今仍然是Linux内核的首席维护者，继续指导着内核的开发方向。

3.3 用 LangGraph 手写 ReAct Agent

LangGraph 是 LangChain 团队推出的 Agent 编排框架，它的核心概念是状态图（State Graph）：把 Agent 的运行过程建模为一个有限状态机。

"""
pip install langchain langchain-openai langgraph
"""

from typing import Annotated, TypedDict
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

# ========== Step 1: 定义工具 ==========

@tool
def search_web(query: str) -> str:
    """搜索互联网获取最新信息。

    Args:
        query: 搜索关键词
    """
    # 模拟搜索结果（实际项目替换为Tavily/SerpAPI等）
    mock_results = {
        "Linux": "Linux是由Linus Torvalds于1991年创建的开源操作系统内核。",
        "Python": "Python是由Guido van Rossum于1991年发布的编程语言。",
        "LangChain": "LangChain是一个用于构建LLM应用的开源框架，由Harrison Chase创建。",
    }
    for key, value in mock_results.items():
        if key.lower() in query.lower():
            return value
    return f"搜索'{query}'：未找到相关信息，请尝试更具体的关键词。"

@tool
def run_python(code: str) -> str:
    """执行Python代码并返回输出。

    Args:
        code: 要执行的Python代码
    """
    import sys
    from io import StringIO
    old_stdout = sys.stdout
    sys.stdout = StringIO()
    try:
        exec(code, {"__builtins__": {}})
        output = sys.stdout.getvalue()
    except Exception as e:
        output = f"错误: {e}"
    finally:
        sys.stdout = old_stdout
    return output or "(无输出)"

tools = [search_web, run_python]

# ========== Step 2: 定义状态 ==========

class AgentState(TypedDict):
    """Agent的状态：消息列表"""
    messages: Annotated[list, add_messages]

# ========== Step 3: 创建LLM和工具节点 ==========

llm = ChatOpenAI(model="gpt-4o", temperature=0)
llm_with_tools = llm.bind_tools(tools)

# LangGraph 提供的内置工具执行节点
tool_node = ToolNode(tools)

# ========== Step 4: 定义Agent节点 ==========

def agent_node(state: AgentState) -> AgentState:
    """Agent节点：调用LLM，决定下一步是调用工具还是给出最终回答"""
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}

# ========== Step 5: 定义路由逻辑 ==========

def should_continue(state: AgentState) -> str:
    """判断Agent是否需要继续调用工具"""
    last_message = state["messages"][-1]
    if last_message.tool_calls:
        return "tools"  # 还有工具要调用，转到工具节点
    return END  # 没有工具调用了，结束

# ========== Step 6: 构建图 ==========

workflow = StateGraph(AgentState)

# 添加节点
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)

# 设置入口
workflow.set_entry_point("agent")

# 添加边
workflow.add_conditional_edges("agent", should_continue, {"tools": "tools", END: END})
workflow.add_edge("tools", "agent")  # 工具执行后，回到Agent节点

# 编译
app = workflow.compile()

# ========== Step 7: 运行 ==========

result = app.invoke({
    "messages": [{"role": "user", "content": "Linux是谁创建的？帮我算一下从创建到2024年过了多少年。"}]
})

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

运行输出示例：

Linux是由Linus Torvalds于1991年创建的开源操作系统内核。从1991年到2024年，已经过了33年。

3.4 LangGraph 图的可视化

上面的代码构建了一个如下结构的图：

3.5 工具设计的最佳实践

1. 描述要精确：工具的 description 是 LLM 选择工具的唯一依据，要写清楚"何时用"而不仅是"做什么"
2. 参数要明确：用 description 和 enum 约束参数格式，减少 LLM 传错参数的概率
3. 工具粒度要适中：太细的工具让 LLM 难以选择，太粗的工具让 LLM 难以组合
4. 错误信息要可操作：工具返回的错误信息应该帮助 LLM 理解问题并调整策略
5. 控制工具数量：单次调用建议不超过10-20个工具；工具过多时考虑动态检索

---

《Memory》：对话、Checkpoint 和长期记忆

4.1 为什么 Agent 需要记忆？

没有记忆的 Agent 就像一个"金鱼脑"助手——每次对话都从零开始：

用户：我叫小明
Agent：你好小明！

用户：我叫什么名字？
Agent：抱歉，我不知道你的名字。  ← 忘了！

记忆系统让 Agent 能够：

• 记住用户偏好：跨对话记住用户的习惯和偏好
• 保持上下文连贯：长对话中不丢失关键信息
• 积累经验：从历史交互中学习和改进

4.2 记忆的分类

4.3 对话记忆的几种策略

Buffer Memory（缓冲区记忆）

最简单的策略：保留所有对话历史。

from langchain_core.messages import HumanMessage, AIMessage

# 直接保留所有消息
messages = [
    HumanMessage(content="你好，我叫小明"),
    AIMessage(content="你好小明！有什么可以帮你的？"),
    HumanMessage(content="我叫什么？"),
    # LLM 能从历史中找到"小明"
]

问题：对话太长时，token 超限。

Sliding Window Memory（滑动窗口）

只保留最近 K 轮对话：

from collections import deque

class SlidingWindowMemory:
    def __init__(self, max_rounds: int = 5):
        self.max_rounds = max_rounds
        self.messages = deque(maxlen=max_rounds * 2)  # 每轮2条消息

    def add(self, message):
        self.messages.append(message)

    def get(self):
        return list(self.messages)

Summary Memory（摘要记忆）

让 LLM 定期总结历史对话，用摘要替代原文：

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o")

async def summarize_messages(messages: list) -> str:
    """将历史对话压缩为摘要"""
    conversation = "\n".join([
        f"{'用户' if isinstance(m, HumanMessage) else '助手'}: {m.content}"
        for m in messages
    ])

    summary_prompt = f"""请将以下对话历史压缩为一段简洁的摘要，保留关键信息（人名、偏好、重要决定等）：

{conversation}

摘要："""

    response = await llm.ainvoke(summary_prompt)
    return response.content

4.4 长期记忆：向量数据库 + RAG

对于需要长期保存的信息（用户偏好、知识库等），我们使用向量数据库进行语义检索：

"""
pip install langchain-chroma langchain-openai
"""

from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document

# ========== 创建向量存储 ==========
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma(
    collection_name="agent_memory",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)

# ========== 存储记忆 ==========
def save_memory(user_id: str, content: str, metadata: dict = None):
    """将信息存入长期记忆"""
    doc = Document(
        page_content=content,
        metadata={"user_id": user_id, **(metadata or {})}
    )
    vectorstore.add_documents([doc])

# ========== 检索记忆 ==========
def recall_memory(user_id: str, query: str, k: int = 3) -> list:
    """从长期记忆中检索相关信息"""
    results = vectorstore.similarity_search(
        query=query,
        k=k,
        filter={"user_id": user_id}
    )
    return [doc.page_content for doc in results]

# ========== 使用示例 ==========
# 用户告诉Agent自己的偏好
save_memory("user_123", "我喜欢Python编程，偏好函数式编程风格", {"type": "preference"})
save_memory("user_123", "我的项目使用FastAPI和PostgreSQL", {"type": "tech_stack"})

# 后续对话中检索
memories = recall_memory("user_123", "这个用户的技术偏好是什么？")
print(memories)
# ['我喜欢Python编程，偏好函数式编程风格', '我的项目使用FastAPI和PostgreSQL']

4.5 LangGraph 的 Checkpointer 持久化

LangGraph 提供了 Checkpointer 机制，可以自动保存和恢复 Agent 的状态：

"""
pip install langgraph-checkpoint-sqlite
"""

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from typing import Annotated, TypedDict

# ========== 定义状态 ==========
class AgentState(TypedDict):
    messages: Annotated[list, add_messages]

# ========== 创建带记忆的Agent ==========

llm = ChatOpenAI(model="gpt-4o", temperature=0)

def chatbot(state: AgentState) -> AgentState:
    response = llm.invoke(state["messages"])
    return {"messages": [response]}

workflow = StateGraph(AgentState)
workflow.add_node("chatbot", chatbot)
workflow.set_entry_point("chatbot")
workflow.add_edge("chatbot", END)

# ========== 使用 SQLite Checkpointer ==========
# 关键：这行代码让Agent拥有了跨会话记忆
memory = SqliteSaver.from_conn_string("./agent_memory.db")
app = workflow.compile(checkpointer=memory)

# ========== 第一次对话 ==========
config = {"configurable": {"thread_id": "user_123_session_1"}}

result1 = app.invoke(
    {"messages": [{"role": "user", "content": "你好！我叫小明，我是一名Python开发者。"}]},
    config=config
)
print(f"对话1: {result1['messages'][-1].content}")
# 输出: 你好小明！很高兴认识你...

# ========== 第二次对话（新的API调用，但同一个thread） ==========
result2 = app.invoke(
    {"messages": [{"role": "user", "content": "我叫什么名字？我做什么开发？"}]},
    config=config  # 同一个thread_id，能访问之前的对话
)
print(f"对话2: {result2['messages'][-1].content}")
# 输出: 你叫小明，你是一名Python开发者。

4.6 综合记忆 Agent 示例

结合短期记忆和长期记忆的完整 Agent：

"""
带综合记忆的对话Agent
- 短期记忆：LangGraph Checkpointer（对话历史）
- 长期记忆：向量数据库（用户偏好、知识积累）
"""

from typing import Annotated, TypedDict
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document

# ========== 长期记忆（向量数据库）==========
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
long_term_memory = Chroma(
    collection_name="user_long_term_memory",
    embedding_function=embeddings,
    persist_directory="./long_term_memory_db"
)

@tool
def save_to_long_term_memory(content: str, category: str = "general") -> str:
    """将重要信息保存到长期记忆中。用于保存用户偏好、重要事实等。

    Args:
        content: 要保存的信息内容
        category: 信息分类，如"preference"(偏好)、"fact"(事实)、"project"(项目相关)
    """
    doc = Document(page_content=content, metadata={"category": category})
    long_term_memory.add_documents([doc])
    return f"已保存到长期记忆: {content[:50]}..."

@tool
def recall_from_long_term_memory(query: str, category: str = None) -> str:
    """从长期记忆中检索相关信息。

    Args:
        query: 检索关键词
        category: 可选，限制检索的信息分类
    """
    filter_dict = {}
    if category:
        filter_dict["category"] = category
    results = long_term_memory.similarity_search(query=query, k=3, filter=filter_dict or None)
    if not results:
        return "未找到相关记忆。"
    return "\n".join([f"- {doc.page_content}" for doc in results])

# ========== 构建Agent ==========
class AgentState(TypedDict):
    messages: Annotated[list, add_messages]

llm = ChatOpenAI(model="gpt-4o", temperature=0)
tools = [save_to_long_term_memory, recall_from_long_term_memory]
llm_with_tools = llm.bind_tools(tools)

SYSTEM_PROMPT = """你是一个有记忆的AI助手。

## 能力
- 你可以保存重要信息到长期记忆（用户偏好、重要事实等）
- 你可以检索长期记忆中的历史信息

## 行为规则
1. 当用户告诉你关于他们自己的重要信息时（偏好、习惯、项目等），主动保存到长期记忆
2. 在回答问题前，先检索长期记忆看是否有相关信息
3. 不要重复保存相同的信息
"""

from langgraph.prebuilt import ToolNode

def agent_node(state: AgentState) -> AgentState:
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

def should_continue(state: AgentState) -> str:
    last_message = state["messages"][-1]
    return "tools" if last_message.tool_calls else END

workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
workflow.add_node("tools", ToolNode(tools))
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", should_continue, {"tools": "tools", END: END})
workflow.add_edge("tools", "agent")

# 编译，带Checkpointer
memory = SqliteSaver.from_conn_string("./memory_agent.db")
app = workflow.compile(checkpointer=memory)

# ========== 运行 ==========
config = {"configurable": {"thread_id": "demo_user"}}

# 第一轮
r1 = app.invoke(
    {"messages": [{"role": "user", "content": "我喜欢用Python写后端，框架偏好FastAPI。"}]},
    config=config
)
print(f"Agent: {r1['messages'][-1].content}")

# 第二轮（新对话，但同一个thread）
r2 = app.invoke(
    {"messages": [{"role": "user", "content": "帮我推荐一个适合我的技术栈。"}]},
    config=config
)
print(f"Agent: {r2['messages'][-1].content}")
# Agent会检索长期记忆，记住你偏好Python + FastAPI，给出相关推荐

---

《Planning》：这一轮先想完，还是边做边想

5.1 为什么 Agent 需要规划？

简单的 ReAct 循环对于单步任务够用，但面对复杂任务时，Agent 需要先规划、再执行：

用户：帮我写一篇关于量子计算最新进展的调研报告

没有规划的Agent:
[Thought] 我需要搜索量子计算
[Action] search("量子计算")
[Observation] ...
[Thought] 我需要更多信息
[Action] search("量子计算最新")
[Observation] ...
[Thought] 我再搜一下
[Action] search("quantum computing 2024")
... （漫无目的地搜索，效率极低）

有规划的Agent:
[Plan]
1. 搜索量子计算基础概念，确保理解准确
2. 搜索2024年量子计算重大突破
3. 搜索量子计算的商业应用进展
4. 搜索量子计算面临的挑战
5. 综合以上信息，按"背景-进展-应用-挑战-展望"结构撰写报告
[/Plan]
[Execute Plan Step 1] ...
[Execute Plan Step 2] ...
...

5.2 两种规划模式

5.3 用 LangGraph 实现 Plan-and-Execute Agent

"""
Plan-and-Execute Agent
- Planner: 生成任务计划
- Executor: 逐步执行计划中的每一步
- Re-Planner: 根据执行结果调整计划
"""

from typing import Annotated, TypedDict, List
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

# ========== 定义数据结构 ==========

class PlanStep(BaseModel):
    """计划中的一个步骤"""
    step_number: int = Field(description="步骤编号")
    description: str = Field(description="步骤描述")
    tool_to_use: str = Field(default="", description="该步骤需要使用的工具")

class Plan(BaseModel):
    """执行计划"""
    goal: str = Field(description="总体目标")
    steps: List[PlanStep] = Field(description="步骤列表")

# ========== Agent 状态 ==========

class PlanExecuteState(TypedDict):
    """Plan-and-Execute Agent 的状态"""
    goal: str                    # 用户的目标
    plan: Plan                   # 当前计划
    current_step_index: int      # 当前执行到第几步
    step_results: list           # 每一步的执行结果
    final_answer: str            # 最终回答

# ========== 定义工具 ==========

@tool
def search_web(query: str) -> str:
    """搜索互联网获取信息"""
    mock_data = {
        "量子计算": "2024年量子计算领域取得了多项突破：Google发布了新一代量子处理器Willow，IBM推出了1000+量子比特处理器",
        "量子计算应用": "量子计算在药物发现、金融建模、密码学等领域展现出广阔应用前景",
        "量子计算挑战": "量子计算仍面临量子纠错、退相干、可扩展性等核心挑战",
    }
    for key, value in mock_data.items():
        if key in query:
            return value
    return f"关于'{query}'的搜索结果：相关研究正在进行中。"

@tool
def analyze_data(data: str, analysis_type: str = "summary") -> str:
    """分析数据或文本"""
    if analysis_type == "summary":
        return f"分析摘要：{data[:100]}..."
    return f"分析结果({analysis_type})：数据包含关键信息点。"

tools = [search_web, analyze_data]

# ========== Planner 节点 ==========

llm = ChatOpenAI(model="gpt-4o", temperature=0)

def planner_node(state: PlanExecuteState) -> PlanExecuteState:
    """根据目标生成执行计划"""
    goal = state["goal"]

    prompt = f"""你是一个任务规划专家。请为以下目标制定一个详细的执行计划。

目标：{goal}

可用工具：
- search_web: 搜索互联网获取信息
- analyze_data: 分析数据或文本

请以JSON格式输出计划，包含goal和steps字段。每个step包含step_number、description和tool_to_use。
"""

    response = llm.invoke(prompt)
    import json
    try:
        plan_data = json.loads(response.content)
        plan = Plan(**plan_data)
    except:
        # 如果JSON解析失败，创建简单计划
        plan = Plan(
            goal=goal,
            steps=[PlanStep(step_number=1, description=f"搜索关于{goal}的信息", tool_to_use="search_web")]
        )

    return {"plan": plan, "current_step_index": 0, "step_results": []}

# ========== Executor 节点 ==========

def executor_node(state: PlanExecuteState) -> PlanExecuteState:
    """执行当前步骤"""
    plan = state["plan"]
    step_index = state["current_step_index"]

    if step_index >= len(plan.steps):
        return state

    current_step = plan.steps[step_index]

    # 根据步骤的工具选择，调用相应工具
    if "search" in current_step.tool_to_use.lower():
        result = search_web.invoke({"query": current_step.description})
    elif "analyze" in current_step.tool_to_use.lower():
        # 使用之前步骤的结果作为输入
        previous_results = "\n".join(state.get("step_results", []))
        result = analyze_data.invoke({"data": previous_results, "analysis_type": "summary"})
    else:
        # 让LLM直接处理
        result = llm.invoke(current_step.description).content

    new_results = state.get("step_results", []) + [f"步骤{step_index+1}结果: {result}"]
    return {"step_results": new_results, "current_step_index": step_index + 1}

# ========== Re-Planner 节点 ==========

def replanner_node(state: PlanExecuteState) -> PlanExecuteState:
    """根据执行结果决定是否调整计划"""
    plan = state["plan"]
    step_index = state["current_step_index"]

    if step_index >= len(plan.steps):
        # 所有步骤执行完毕，生成最终回答
        results = "\n".join(state["step_results"])
        prompt = f"""基于以下研究结果，回答用户的问题。

目标：{plan.goal}

研究结果：
{results}

请给出完整、有条理的回答："""

        final_answer = llm.invoke(prompt).content
        return {"final_answer": final_answer}

    return state  # 继续执行下一步

# ========== 路由逻辑 ==========

def should_continue_execution(state: PlanExecuteState) -> str:
    """判断是否继续执行"""
    if state.get("final_answer"):
        return END
    if state["current_step_index"] >= len(state["plan"].steps):
        return "replanner"
    return "executor"

# ========== 构建图 ==========

workflow = StateGraph(PlanExecuteState)

workflow.add_node("planner", planner_node)
workflow.add_node("executor", executor_node)
workflow.add_node("replanner", replanner_node)

workflow.set_entry_point("planner")
workflow.add_edge("planner", "executor")
workflow.add_edge("executor", "replanner")
workflow.add_conditional_edges("replanner", should_continue_execution, {
    "executor": "executor",
    "replanner": "replanner",
    END: END
})

app = workflow.compile()

# ========== 运行 ==========
result = app.invoke({
    "goal": "写一篇关于量子计算最新进展的调研报告",
    "plan": Plan(goal="", steps=[]),
    "current_step_index": 0,
    "step_results": [],
    "final_answer": ""
})

print("=" * 60)
print("最终报告：")
print(result["final_answer"])

5.4 自反思与自我纠错

Plan-and-Execute 模式可以进一步增强——让Agent在执行后审视自己的输出，发现不足并改进：

"""
Reflexion 模式：执行 → 反思 → 改进
"""

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

llm = ChatOpenAI(model="gpt-4o", temperature=0)

def reflect_and_improve(task: str, draft: str, max_iterations: int = 3) -> str:
    """反思循环：不断改进输出直到满意"""

    current_draft = draft

    for i in range(max_iterations):
        # Step 1: 反思
        reflection_prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个严格的质量审查员。请审视以下内容，指出其中的不足：
- 事实是否准确？
- 逻辑是否清晰？
- 信息是否充分？
- 是否遗漏了重要方面？

只指出问题，不要给出改进后的版本。"""),
            ("human", "任务: {task}\n\n当前版本:\n{draft}")
        ])

        reflection = llm.invoke(
            reflection_prompt.format_messages(task=task, draft=current_draft)
        ).content

        # Step 2: 判断是否需要改进
        if "没有明显问题" in reflection or "质量良好" in reflection:
            break

        # Step 3: 改进
        improvement_prompt = ChatPromptTemplate.from_messages([
            ("system", "根据审查意见改进你的输出。保持原有优点，修正指出的问题。"),
            ("human", "任务: {task}\n\n当前版本:\n{draft}\n\n审查意见:\n{reflection}\n\n请输出改进后的版本：")
        ])

        current_draft = llm.invoke(
            improvement_prompt.format_messages(
                task=task, draft=current_draft, reflection=reflection
            )
        ).content

    return current_draft

# 使用示例
task = "解释什么是量子纠缠"
draft = "量子纠缠就是两个粒子连在一起，一个动了另一个也会动。"
improved = reflect_and_improve(task, draft)
print(improved)
# 经过反思改进后，输出会变得更加准确和完整

---

《多 Agent》：什么时候值得再开一个循环

6.1 为什么需要多 Agent？

单一 Agent 的能力是有瓶颈的：

• 注意力有限：一个Agent难以同时擅长搜索、编程、写作等多个领域
• 上下文窗口有限：工具越多、记忆越长，LLM越容易"分心"
• 单一视角：不同角色从不同角度看问题，结果更全面

Multi-Agent 的核心思想：让不同的Agent专注于自己擅长的领域，通过协作完成复杂任务。

6.2 四种协作模式

6.3 Supervisor 模式实战：Researcher + Coder + Reviewer

这是最常用也最实用的多Agent模式。一个Supervisor（主管）负责任务分配和整合，三个专业Agent各司其职：

"""
Multi-Agent Supervisor 模式
- Supervisor: 主管Agent，分配任务、收集结果
- Researcher: 研究员Agent，负责搜索和收集信息
- Coder: 编程Agent，负责写代码和数据分析
- Reviewer: 审核Agent，负责质量把关
"""

from typing import Annotated, TypedDict, Literal
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

# ========== 定义工具 ==========

@tool
def search_research(query: str) -> str:
    """搜索学术资料和最新研究"""
    mock_db = {
        "transformer": "Transformer架构由Vaswani等人在2017年提出，基于自注意力机制...",
        "RLHF": "RLHF(基于人类反馈的强化学习)是训练ChatGPT的关键技术...",
        "RAG": "RAG(检索增强生成)通过结合外部知识库提升LLM回答的准确性...",
    }
    for key, value in mock_db.items():
        if key.lower() in query.lower():
            return value
    return f"关于'{query}'的研究：相关领域正在快速发展。"

@tool
def write_code(requirement: str) -> str:
    """根据需求编写Python代码"""
    return f"""```python
# 基于需求生成的代码: {requirement}
import numpy as np

def process_data(data):
    \"\"\"数据处理函数\"\"\"
    result = np.mean(data)
    return result

# 示例使用
if __name__ == "__main__":
    sample_data = [1, 2, 3, 4, 5]
    print(f"处理结果: {{process_data(sample_data)}}")
```"""

@tool
def review_content(content: str, criteria: str = "accuracy") -> str:
    """审核内容的准确性和质量"""
    return f"审核结果({criteria})：内容基本准确，建议补充更多具体数据和引用来源。评分：7/10"

researcher_tools = [search_research]
coder_tools = [write_code]
reviewer_tools = [review_content]

# ========== 创建LLM实例 ==========

llm = ChatOpenAI(model="gpt-4o", temperature=0)

# ========== 定义状态 ==========

class TeamState(TypedDict):
    messages: Annotated[list, add_messages]
    next_agent: str          # 下一个要执行的Agent
    task_description: str    # 当前任务描述
    research_result: str     # 研究结果
    code_result: str         # 代码结果
    review_result: str       # 审核结果
    final_output: str        # 最终输出

# ========== Supervisor Agent ==========

def supervisor_node(state: TeamState) -> TeamState:
    """主管Agent：分析任务，决定分配给谁"""

    supervisor_prompt = """你是一个团队主管，管理三个专业Agent：

1. **Researcher（研究员）**：擅长搜索和收集信息，使用search_research工具
2. **Coder（编程员）**：擅长编写代码和数据分析，使用write_code工具
3. **Reviewer（审核员）**：擅长审核内容质量，使用review_content工具

根据当前任务状态，决定下一步应该让哪个Agent工作。

回复格式只需要Agent名称：Researcher、Coder、Reviewer 或 FINISH（任务完成）
"""

    # 构建上下文
    context = f"任务: {state.get('task_description', '')}\n"
    if state.get('research_result'):
        context += f"研究结果: {state['research_result'][:200]}\n"
    if state.get('code_result'):
        context += f"代码结果: {state['code_result'][:200]}\n"
    if state.get('review_result'):
        context += f"审核结果: {state['review_result'][:200]}\n"

    messages = [
        SystemMessage(content=supervisor_prompt),
        HumanMessage(content=context)
    ]

    response = llm.invoke(messages)

    # 解析决定
    decision = response.content.strip()
    if "FINISH" in decision or "完成" in decision:
        next_agent = "finish"
    elif "Researcher" in decision or "研究" in decision:
        next_agent = "researcher"
    elif "Coder" in decision or "编程" in decision:
        next_agent = "coder"
    elif "Reviewer" in decision or "审核" in decision:
        next_agent = "reviewer"
    else:
        next_agent = "researcher"  # 默认先研究

    return {"next_agent": next_agent}

# ========== Researcher Agent ==========

def researcher_node(state: TeamState) -> TeamState:
    """研究员Agent：搜索和收集信息"""
    researcher_prompt = """你是一个专业研究员。使用search_research工具搜索相关信息。
基于搜索结果，给出简洁的研究摘要。"""

    messages = [
        SystemMessage(content=researcher_prompt),
        HumanMessage(content=state.get('task_description', ''))
    ]

    llm_with_tools = llm.bind_tools(researcher_tools)
    response = llm_with_tools.invoke(messages)

    # 如果有工具调用，执行之
    if response.tool_calls:
        from langgraph.prebuilt import ToolNode
        # 简化处理：直接执行第一个工具调用
        for tool_call in response.tool_calls:
            if tool_call["name"] == "search_research":
                result = search_research.invoke(tool_call["args"])
                # 用结果生成摘要
                summary = llm.invoke([
                    SystemMessage(content="基于搜索结果，给出简洁的研究摘要。"),
                    HumanMessage(content=f"搜索结果: {result}\n\n任务: {state.get('task_description', '')}")
                ]).content
                return {"research_result": summary, "next_agent": "supervisor"}

    return {"research_result": response.content, "next_agent": "supervisor"}

# ========== Coder Agent ==========

def coder_node(state: TeamState) -> TeamState:
    """编程Agent：编写代码"""
    coder_prompt = """你是一个高级Python开发者。根据需求和背景信息，编写高质量的代码。
使用write_code工具生成代码，然后优化和解释代码。"""

    context = state.get('task_description', '')
    if state.get('research_result'):
        context += f"\n\n参考研究: {state['research_result']}"

    messages = [
        SystemMessage(content=coder_prompt),
        HumanMessage(content=context)
    ]

    llm_with_tools = llm.bind_tools(coder_tools)
    response = llm_with_tools.invoke(messages)

    if response.tool_calls:
        for tool_call in response.tool_calls:
            if tool_call["name"] == "write_code":
                code = write_code.invoke(tool_call["args"])
                return {"code_result": code, "next_agent": "supervisor"}

    return {"code_result": response.content, "next_agent": "supervisor"}

# ========== Reviewer Agent ==========

def reviewer_node(state: TeamState) -> TeamState:
    """审核Agent：检查质量"""
    reviewer_prompt = """你是一个严格的质量审核员。审核研究成果和代码质量。
使用review_content工具进行审核，给出改进建议。"""

    content_to_review = ""
    if state.get('research_result'):
        content_to_review += f"研究: {state['research_result']}\n"
    if state.get('code_result'):
        content_to_review += f"代码: {state['code_result']}\n"

    messages = [
        SystemMessage(content=reviewer_prompt),
        HumanMessage(content=f"请审核以下内容:\n{content_to_review}")
    ]

    llm_with_tools = llm.bind_tools(reviewer_tools)
    response = llm_with_tools.invoke(messages)

    if response.tool_calls:
        for tool_call in response.tool_calls:
            if tool_call["name"] == "review_content":
                review = review_content.invoke(tool_call["args"])
                return {"review_result": review, "next_agent": "supervisor"}

    return {"review_result": response.content, "next_agent": "supervisor"}

# ========== Finish 节点 ==========

def finish_node(state: TeamState) -> TeamState:
    """整合所有结果，生成最终输出"""
    prompt = f"""基于以下团队成员的工作成果，生成最终输出：

任务: {state.get('task_description', '')}

研究结果:
{state.get('research_result', '无')}

代码结果:
{state.get('code_result', '无')}

审核意见:
{state.get('review_result', '无')}

请整合以上内容，生成完整的最终报告："""

    final_output = llm.invoke(prompt).content
    return {"final_output": final_output}

# ========== 路由逻辑 ==========

def route_from_supervisor(state: TeamState) -> str:
    return state.get("next_agent", "researcher")

# ========== 构建图 ==========

workflow = StateGraph(TeamState)

# 添加节点
workflow.add_node("supervisor", supervisor_node)
workflow.add_node("researcher", researcher_node)
workflow.add_node("coder", coder_node)
workflow.add_node("reviewer", reviewer_node)
workflow.add_node("finish", finish_node)

# 设置入口
workflow.set_entry_point("supervisor")

# 添加条件边：supervisor -> 各Agent
workflow.add_conditional_edges("supervisor", route_from_supervisor, {
    "researcher": "researcher",
    "coder": "coder",
    "reviewer": "reviewer",
    "finish": "finish",
})

# 各Agent执行后回到supervisor
workflow.add_edge("researcher", "supervisor")
workflow.add_edge("coder", "supervisor")
workflow.add_edge("reviewer", "supervisor")
workflow.add_edge("finish", END)

# 编译
app = workflow.compile()

# ========== 运行 ==========

result = app.invoke({
    "messages": [],
    "next_agent": "",
    "task_description": "研究RAG技术的最新进展，并写一个简单的RAG实现代码",
    "research_result": "",
    "code_result": "",
    "review_result": "",
    "final_output": ""
})

print("=" * 60)
print("最终输出：")
print(result["final_output"])

6.4 LangGraph 的 Sub-Graph 机制

当团队变得更复杂时，可以用 Sub-Graph 把一组 Agent 封装为一个"子团队"：

"""
Sub-Graph 示例：将 Researcher + Coder 封装为"开发子团队"
"""

from langgraph.graph import StateGraph, END

# 1. 定义子团队的State
class DevTeamState(TypedDict):
    task: str
    research_output: str
    code_output: str

# 2. 构建子团队的Graph
dev_team = StateGraph(DevTeamState)

def research_step(state: DevTeamState) -> DevTeamState:
    # 调用researcher agent...
    return {"research_output": "研究完成：..."}

def code_step(state: DevTeamState) -> DevTeamState:
    # 调用coder agent...
    return {"code_output": "代码完成：..."}

dev_team.add_node("research", research_step)
dev_team.add_node("code", code_step)
dev_team.set_entry_point("research")
dev_team.add_edge("research", "code")
dev_team.add_edge("code", END)

dev_team_graph = dev_team.compile()

# 3. 在主Graph中使用子团队
class MainState(TypedDict):
    task: str
    dev_output: str
    review_output: str

main_graph = StateGraph(MainState)

def dev_team_node(state: MainState) -> MainState:
    result = dev_team_graph.invoke({"task": state["task"], "research_output": "", "code_output": ""})
    return {"dev_output": result["code_output"]}

def review_node(state: MainState) -> MainState:
    return {"review_output": "审核通过"}

main_graph.add_node("dev_team", dev_team_node)
main_graph.add_node("review", review_node)
main_graph.set_entry_point("dev_team")
main_graph.add_edge("dev_team", "review")
main_graph.add_edge("review", END)

app = main_graph.compile()

---

《实战》：把各部件收成一个 Research Agent

7.1 项目概述

整合前面所有模块，构建一个完整的 Research Agent，具备以下能力：

• 搜索能力：搜索互联网获取最新信息
• 记忆能力：跨会话记住用户偏好和研究历史
• 规划能力：自动分解研究任务为可执行步骤
• 多Agent协作：Researcher + Writer + Reviewer 分工合作
• 生产级考量：错误处理、成本控制、可观测性

7.2 系统架构

7.3 完整代码

"""
Research Agent —— 端到端实现
整合 Tool Use + Memory + Planning + Multi-Agent

依赖安装:
pip install langchain langchain-openai langgraph langchain-chroma langgraph-checkpoint-sqlite
"""

import json
import time
from typing import Annotated, TypedDict, List, Optional
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_core.documents import Document
from langchain_chroma import Chroma
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.prebuilt import ToolNode

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

@tool
def search_web(query: str) -> str:
    """搜索互联网获取最新信息。用于查找新闻、技术动态、事实性信息等。

    Args:
        query: 搜索关键词
    """
    # 实际项目替换为Tavily/SerpAPI
    mock_db = {
        "RAG": "检索增强生成(RAG)是一种结合外部知识库提升LLM回答准确性的技术。2024年主要进展包括：模态RAG、自适应RAG、GraphRAG等。",
        "Agent": "AI Agent在2024年取得了重大进展，OpenAI、Google、Anthropic等公司都推出了Agent框架。核心趋势是多Agent协作和工具使用的标准化。",
        "Transformer": "Transformer架构持续演进，Mamba等线性注意力模型、混合专家(MoE)模型成为研究热点。",
    }
    for key, value in mock_db.items():
        if key.lower() in query.lower():
            return value
    return f"搜索'{query}'的结果：相关领域正在快速发展，建议查看最新论文和博客。"

@tool
def read_paper(paper_url: str) -> str:
    """读取学术论文内容。提供arXiv链接或论文标题。

    Args:
        paper_url: 论文URL或标题
    """
    return f"论文摘要: {paper_url} - 该论文提出了创新的方法，在多个基准测试上取得了SOTA结果。"

@tool
def run_python_code(code: str) -> str:
    """执行Python代码。用于数据分析、可视化、数学计算等。

    Args:
        code: 要执行的Python代码
    """
    import sys
    from io import StringIO
    old_stdout = sys.stdout
    sys.stdout = StringIO()
    try:
        exec(code, {"__builtins__": {}})
        output = sys.stdout.getvalue()
    except Exception as e:
        output = f"错误: {e}"
    finally:
        sys.stdout = old_stdout
    return output or "(无输出)"

@tool
def save_research_note(content: str, topic: str) -> str:
    """保存研究笔记到长期记忆。

    Args:
        content: 笔记内容
        topic: 研究主题
    """
    # 存入向量数据库
    doc = Document(page_content=content, metadata={"topic": topic, "type": "research_note"})
    vectorstore.add_documents([doc])
    return f"已保存研究笔记: {topic}"

@tool
def recall_research(topic: str, query: str = "") -> str:
    """从长期记忆中检索之前的研究笔记。

    Args:
        topic: 研究主题
        query: 检索关键词（可选）
    """
    search_query = f"{topic} {query}".strip()
    results = vectorstore.similarity_search(query=search_query, k=3)
    if not results:
        return "未找到相关研究笔记。"
    return "\n---\n".join([doc.page_content for doc in results])

# 所有工具
all_tools = [search_web, read_paper, run_python_code, save_research_note, recall_research]

# ==========================================================
# 2. 记忆系统初始化
# ==========================================================

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma(
    collection_name="research_agent_memory",
    embedding_function=embeddings,
    persist_directory="./research_agent_db"
)

# ==========================================================
# 3. Agent 状态定义
# ==========================================================

class ResearchPlan(BaseModel):
    """研究计划"""
    topic: str = Field(description="研究主题")
    sub_questions: List[str] = Field(description="需要回答的子问题")
    approach: str = Field(description="研究方法")

class ResearchAgentState(TypedDict):
    messages: Annotated[list, add_messages]
    task: str                          # 用户的研究任务
    plan: Optional[ResearchPlan]       # 研究计划
    current_step: int                  # 当前执行步骤
    research_notes: list               # 研究笔记
    draft_report: str                  # 草稿报告
    review_feedback: str               # 审核反馈
    revision_count: int                # 修订次数
    final_report: str                  # 最终报告
    token_usage: int                   # Token使用统计

# ==========================================================
# 4. 各Agent节点
# ==========================================================

llm = ChatOpenAI(model="gpt-4o", temperature=0)

# --- Planner Node ---
def planner_node(state: ResearchAgentState) -> dict:
    """制定研究计划"""
    task = state["task"]

    # 先检查是否有历史研究
    history = recall_research.invoke({"topic": task})

    prompt = f"""你是一个研究规划专家。请为以下研究任务制定计划。

任务: {task}

{"相关历史研究:\n" + history if history != "未找到相关研究笔记。" else ""}

请以JSON格式输出计划，包含:
- topic: 研究主题
- sub_questions: 需要回答的3-5个子问题列表
- approach: 研究方法描述
"""
    response = llm.invoke(prompt)
    try:
        plan_data = json.loads(response.content)
        plan = ResearchPlan(**plan_data)
    except:
        plan = ResearchPlan(
            topic=task,
            sub_questions=[f"什么是{task}？", f"{task}的最新进展是什么？", f"{task}面临的挑战是什么？"],
            approach="先搜索基本信息，再深入分析最新进展"
        )

    return {
        "plan": plan,
        "current_step": 0,
        "research_notes": [],
        "revision_count": 0,
        "token_usage": state.get("token_usage", 0)
    }

# --- Researcher Node ---
def researcher_node(state: ResearchAgentState) -> dict:
    """执行研究：搜索和收集信息"""
    plan = state["plan"]
    step = state["current_step"]

    if step >= len(plan.sub_questions):
        return {"current_step": step}

    question = plan.sub_questions[step]

    # 使用LLM + 工具进行搜索
    llm_with_tools = llm.bind_tools([search_web, read_paper, recall_research])
    messages = [
        SystemMessage(content="你是一个专业研究员。使用提供的工具搜索信息，回答研究问题。给出简洁但全面的回答。"),
        HumanMessage(content=f"研究主题: {plan.topic}\n当前问题: {question}")
    ]
    response = llm_with_tools.invoke(messages)

    # 处理工具调用
    note = response.content or ""
    if response.tool_calls:
        for tool_call in response.tool_calls:
            tool_func = {t.name: t for t in all_tools}.get(tool_call["name"])
            if tool_func:
                result = tool_func.invoke(tool_call["args"])
                note += f"\n工具结果: {result}"

    # 保存研究笔记
    save_research_note.invoke({"content": note, "topic": plan.topic})

    new_notes = state.get("research_notes", []) + [f"Q: {question}\nA: {note}"]

    return {
        "research_notes": new_notes,
        "current_step": step + 1
    }

# --- Writer Node ---
def writer_node(state: ResearchAgentState) -> dict:
    """撰写研究报告"""
    plan = state["plan"]
    notes = "\n\n".join(state.get("research_notes", []))

    feedback = state.get("review_feedback", "")

    prompt = f"""你是一个专业学术写作者。基于研究笔记，撰写一份结构清晰的研究报告。

研究主题: {plan.topic}
研究方法: {plan.approach}

研究笔记:
{notes}

{"审核意见（请据此改进）:\n" + feedback if feedback else ""}

请按以下结构撰写报告:
1. 概述
2. 背景与基本概念
3. 最新进展
4. 应用场景
5. 挑战与展望
6. 结论
"""

    draft = llm.invoke(prompt).content
    return {"draft_report": draft}

# --- Reviewer Node ---
def reviewer_node(state: ResearchAgentState) -> dict:
    """审核报告质量"""
    draft = state.get("draft_report", "")

    prompt = f"""你是一个严格的研究审核员。请审核以下研究报告:

{draft}

请从以下维度评分(1-10)并给出具体建议:
1. 准确性 - 信息是否准确？
2. 完整性 - 是否覆盖了所有重要方面？
3. 清晰度 - 逻辑是否清晰？
4. 深度 - 分析是否足够深入？

如果总分>=30，回复"APPROVED"。否则回复"NEEDS_REVISION"并列出改进建议。
"""

    feedback = llm.invoke(prompt).content

    if "APPROVED" in feedback:
        return {"review_feedback": feedback, "final_report": draft}
    else:
        revision_count = state.get("revision_count", 0) + 1
        return {"review_feedback": feedback, "revision_count": revision_count}

# --- Router ---
def route_after_research(state: ResearchAgentState) -> str:
    """研究完成后决定下一步"""
    if state["current_step"] < len(state["plan"].sub_questions):
        return "researcher"  # 继续研究下一个子问题
    return "writer"  # 所有子问题研究完毕，开始写作

def route_after_review(state: ResearchAgentState) -> str:
    """审核后决定下一步"""
    if state.get("final_report"):
        return END  # 审核通过，结束
    if state.get("revision_count", 0) >= 2:
        return END  # 最多修订2次
    return "writer"  # 需要修订

# ==========================================================
# 5. 构建Graph
# ==========================================================

workflow = StateGraph(ResearchAgentState)

# 添加节点
workflow.add_node("planner", planner_node)
workflow.add_node("researcher", researcher_node)
workflow.add_node("writer", writer_node)
workflow.add_node("reviewer", reviewer_node)

# 设置入口和边
workflow.set_entry_point("planner")
workflow.add_edge("planner", "researcher")
workflow.add_conditional_edges("researcher", route_after_research, {
    "researcher": "researcher",
    "writer": "writer"
})
workflow.add_edge("writer", "reviewer")
workflow.add_conditional_edges("reviewer", route_after_review, {
    "writer": "writer",
    END: END
})

# 带记忆编译
memory = SqliteSaver.from_conn_string("./research_agent_checkpoints.db")
app = workflow.compile(checkpointer=memory)

# ==========================================================
# 6. 运行
# ==========================================================

config = {"configurable": {"thread_id": "research_session_1"}}

result = app.invoke({
    "messages": [],
    "task": "RAG（检索增强生成）技术的最新进展研究",
    "plan": None,
    "current_step": 0,
    "research_notes": [],
    "draft_report": "",
    "review_feedback": "",
    "revision_count": 0,
    "final_report": "",
    "token_usage": 0
}, config=config)

print("=" * 70)
print("最终研究报告：")
print("=" * 70)
print(result.get("final_report") or result.get("draft_report", "报告生成失败"))

7.4 生产级考量

错误处理

from tenacity import retry, stop_after_attempt, wait_exponential

# 为工具调用添加重试机制
@retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10))
def robust_tool_call(tool_func, **kwargs):
    """带重试的工具调用"""
    try:
        return tool_func.invoke(kwargs)
    except Exception as e:
        print(f"工具调用失败: {e}, 正在重试...")
        raise

成本控制

class TokenCounter:
    """Token使用量追踪"""
    def __init__(self, max_tokens: int = 100000):
        self.total_tokens = 0
        self.max_tokens = max_tokens

    def count(self, response) -> bool:
        """记录token使用，返回是否还在预算内"""
        if hasattr(response, 'usage_metadata'):
            self.total_tokens += response.usage_metadata.get('total_tokens', 0)
        return self.total_tokens < self.max_tokens

    def report(self) -> str:
        return f"已使用 {self.total_tokens} tokens"

可观测性（LangSmith）

# 设置环境变量即可启用LangSmith追踪
import os

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "research-agent"
os.environ["LANGCHAIN_API_KEY"] = "your-langsmith-api-key"

# 所有LLM调用和工具执行都会被自动记录到LangSmith
# 可以在 https://smith.langchain.com 查看详细追踪

安全护栏

def validate_tool_output(output: str, max_length: int = 5000) -> str:
    """验证工具输出，防止注入攻击"""
    # 截断过长输出
    if len(output) > max_length:
        output = output[:max_length] + "...(截断)"

    # 移除可能的指令注入
    dangerous_patterns = ["ignore previous", "system:", "<|im_start|>"]
    for pattern in dangerous_patterns:
        if pattern.lower() in output.lower():
            output = output.replace(pattern, "[filtered]")

    return output

---

附录

A. 环境搭建

# 创建虚拟环境
python -m venv agent-env
source agent-env/bin/activate  # Windows: agent-env\Scripts\activate

# 安装核心依赖
pip install langchain langchain-openai langgraph
pip install langchain-chroma langgraph-checkpoint-sqlite
pip install pydantic

# 设置API Key
export OPENAI_API_KEY="your-api-key"

B. 推荐阅读

主题	论文/文章	关键词
ReAct	ReAct: Synergizing Reasoning and Acting in Language Models	Thought-Action-Observation
Reflexion	Reflexion: Language Agents with Verbal Reinforcement Learning	自反思、自我改进
Plan-and-Execute	Plan-and-Solve Prompting	任务分解、分步执行
Multi-Agent	ChatDev: Communicative Agents for Software Development	多Agent协作、角色分工
LangGraph	LangGraph Documentation	状态图、Agent编排
Function Calling	OpenAI Function Calling Guide	工具协议、参数定义

C. 核心概念速查

Agent        = Model + Harness
Model        = LLM (推理 + 决策)
Harness      = Tools + Memory + Planning + Orchestration
ReAct        = Thought → Action → Observation (循环)
Plan-Execute = Plan → Execute Steps → Re-Plan (如有需要)
Multi-Agent  = Supervisor → Workers (分工协作)
Checkpointer = LangGraph 的会话持久化机制
Sub-Graph    = 封装一组Agent为可复用组件

D. 从这里出发

1. 动手实践：运行《Model》的Function Calling示例，理解Model和Harness的协作
2. 扩展工具：为《Tool Use》的ReAct Agent添加你自己的工具
3. 添加记忆：用Checkpointer让你的Agent能跨会话记忆
4. 尝试规划：用Plan-and-Execute处理一个多步骤的研究任务
5. 组建团队：用Supervisor模式构建你自己的Multi-Agent系统
6. 走向生产：加上错误处理、成本控制、可观测性，部署你的Agent

记住：Agent = Model + Harness。Model是大脑，Harness是四肢。两者缺一不可，协同才能构建真正的智能体。