1
# ===== 第1步：导入LangGraph核心 =====
2
from typing import TypedDict, Literal
3
from langgraph.graph import StateGraph, END
4
from langchain_openai import ChatOpenAI
5
from langchain_community.tools.tavily_search import TavilySearchResults
6

7
# 解释：LangGraph的核心思想是将Agent流程定义为一个状态图（StateGraph）
8
# 状态（State）在节点之间传递，节点是处理逻辑，边决定流转方向
9

10
# ===== 第2步：定义状态（State） =====
11
class AgentState(TypedDict):
12
    messages: list        # 对话历史
13
    next_step: str        # 当前步骤标记
14
    research_result: str  # 研究结果
15
    final_answer: str     # 最终回答
16
# 解释：State是一个字典，用来在各个节点之间传递数据。
17
# LangGraph会自动维护这个状态，你只需在节点函数中读取和修改它。
18

19
# ===== 第3步：初始化LLM和工具 =====
20
llm = ChatOpenAI(model="gpt-4o")
21
search_tool = TavilySearchResults(max_results=3)   # 需要tavily API key，实在没有可以用requests写一个搜索函数代替
22

23
# ===== 第4步：定义节点函数（每个节点是Agent的一步） =====
24

25
# 节点1：分析问题，判断需要搜索什么
26
def analyze_question(state: AgentState) -> AgentState:
27
    """根据用户问题，决定搜索关键词"""
28
    question = state["messages"][-1]["content"]
29
    # 用LLM生成搜索关键词
30
    response = llm.invoke(f"针对'{question}'这个问题，生成3个最合适的搜索关键词，直接返回关键词列表，每行一个。")
31
    keywords = response.content.strip().split("\n")
32
    state["next_step"] = "search"
33
    state["research_result"] = "\n".join(keywords)
34
    return state
35
# 解释：这个节点用LLM做"思考"，把问题转化成可执行的搜索计划
36

37
# 节点2：执行搜索
38
def search(state: AgentState) -> AgentState:
39
    """使用工具搜索信息"""
40
    keywords = state["research_result"]
41
    # 搜索第一个关键词
42
    results = search_tool.invoke(keywords.split("\n")[0])
43
    state["research_result"] = str(results[:2])  # 取前2条结果
44
    state["next_step"] = "synthesize"
45
    return state
46
# 解释：这个节点通过工具（Tool）与外部世界交互，获取实时信息
47

48
# 节点3：综合回答
49
def synthesize(state: AgentState) -> AgentState:
50
    """基于搜索结果生成最终回答"""
51
    question = state["messages"][-1]["content"]
52
    search_data = state["research_result"]
53
    prompt = f"""基于以下搜索结果，回答用户的问题：
54

55
搜索结果：
56
{search_data}
57

58
问题：{question}
59

60
请用中文给出详细、结构化的回答，包含优缺点对比。"""
61
    response = llm.invoke(prompt)
62
    state["final_answer"] = response.content
63
    state["next_step"] = END
64
    return state
65
# 解释：把搜索结果作为上下文，让LLM生成最终答案
66

67
# ===== 第5步：定义条件边（决定流转路径） =====
68
def router(state: AgentState) -> Literal["analyze", "search", "synthesize", "__end__"]:
69
    """根据当前状态决定下一步"""
70
    step = state.get("next_step", "analyze")
71
    if step == "analyze":
72
        return "analyze"
73
    elif step == "search":
74
        return "search"
75
    elif step == "synthesize":
76
        return "synthesize"
77
    else:
78
        return "__end__"
79
# 解释：条件边让Agent可以根据当前状态动态决定下一步做什么
80
# 这里用next_step字段控制，也可以用LLM自己判断
81

82
# ===== 第6步：构建状态图 =====
83
# 创建一个状态图，指定状态结构
84
graph = StateGraph(AgentState)
85

86
# 添加节点
87
graph.add_node("analyze", analyze_question)
88
graph.add_node("search", search)
89
graph.add_node("synthesize", synthesize)
90

91
# 设置入口
92
graph.set_entry_point("analyze")
93

94
# 添加条件边：从"analyze"节点出来后，根据router函数决定去向
95
graph.add_conditional_edges(
96
    "analyze",       # 起始节点
97
    router,          # 路由函数
98
    {                 # 映射：返回值 → 目标节点
99
        "analyze": "analyze",
100
        "search": "search",
101
        "synthesize": "synthesize",
102
        "__end__": END
103
    }
104
)
105
graph.add_conditional_edges("search", router, {
106
    "analyze": "analyze",
107
    "search": "search",
108
    "synthesize": "synthesize",
109
    "__end__": END
110
})
111
graph.add_conditional_edges("synthesize", router, {
112
    "analyze": "analyze",
113
    "search": "search",
114
    "synthesize": "synthesize",
115
    "__end__": END
116
})
117

118
# 编译成可执行的应用
119
app = graph.compile()
120

121
# ===== 第7步：执行 =====
122
if __name__ == "__main__":
123
    # 初始状态
124
    initial_state = {
125
        "messages": [{"role": "user", "content": "对比Transformer和Mamba架构的优缺点"}],
126
        "next_step": "analyze",
127
        "research_result": "",
128
        "final_answer": ""
129
    }
130
    # 运行
131
    result = app.invoke(initial_state)
132
    print("最终回答：")
133
    print(result["final_answer"])