from typing import TypedDict, Annotated
import operator
from langgraph.graph import StateGraph, START, END


# ===================== 1. 定义全局状态 =====================
class AgentState(TypedDict):
    question: str
    need_tool: bool       # 标记：是否需要调用外部工具
    tool_output: str      # 工具返回结果
    chat_msg: Annotated[list[str], operator.add]
    answer: str


# ===================== 2. 定义节点 =====================
def judge_intent(state: AgentState) -> dict:
    """节点1：意图判断，决定是否要用工具"""
    print("【用户意图判断节点，检测是否需要调用工具】begin:")
    q = state["question"]
    # 模拟规则：包含"查询"则判定需要工具
    print("【用户意图判断节点，检测是否需要调用工具】ongoing:", q)
    need = "查询" in q

    print("【用户意图判断节点，检测是否需要调用工具】done:", need)
    return {
        "need_tool": need,
        "chat_msg": [f"意图判断完成，是否调用工具：{need}"]
    }

def call_external_tool(state: AgentState) -> dict:
    """节点2：工具调用节点"""
    print("【工具调用节点】begin:")
    print("【工具调用节点】ongoing:.....")
    result = "从外部接口拿到业务数据"
    print("【工具调用节点】done:", result)
    return {
        "tool_output": result,
        "chat_msg": ["工具调用成功"]
    }

def direct_answer(state: AgentState) -> dict:
    """节点3：直接回答节点（无需工具）"""
    print("【直接回答节点-大模型调用】begin:")
    print("【直接回答节点-大模型调用】ongoing:")
    answer = "大模型基于常识直接给出的答案";
    print("【直接回答节点-大模型调用】done:", answer)
    return {
        "answer": answer,
        "chat_msg": ["无需工具，直接作答"]
    }

def combine_answer(state: AgentState) -> dict:
    """节点4：结合工具结果生成答案"""
    print("【整合答案节点】begin:")
    print("【整合答案节点】onging:")
    print("【整合答案节点】done:", ["整合工具结果完成作答"])
    return {
        "answer": f"结合工具数据：{state['tool_output']} 生成答案",
        "chat_msg": ["整合工具结果完成作答"]
    }

# ===================== 3. 定义路由函数（分支核心） =====================
def route_tool_or_not(state: AgentState) -> str:
    """
    路由判断函数：
    入参：当前全局状态
    返回：目标节点名称字符串（必须和 path_map 对应）
    """
    print("【路由函数】begin:")
    print("【路由函数】goging:")
    if state["need_tool"]:
        result = "call_tool"
    else:
        result = "direct_ans"
    print("【路由函数】done，返回下个节点名称:", result)
    return result;

# ===================== 4. 构建图 + 配置分支 =====================
builder = StateGraph(AgentState)

# 注册所有节点
builder.add_node("judge", judge_intent)
builder.add_node("call_tool", call_external_tool)
builder.add_node("direct_ans", direct_answer)
builder.add_node("combine_ans", combine_answer)

# 入口流向：judge节点
builder.add_edge(START, "judge")

# 配置条件分支：从source节点开始，到其他节点：path_map，判决的依据：path。
builder.add_conditional_edges(
    source="judge",                # 分支起始节点
    path=route_tool_or_not,        # 路由判断函数，这是条件分支最重要的逻辑，是在节点逻辑执行完，才执行，用于选择下一个节点，返回下一个节点名称
    path_map={                     # 映射：函数返回值 → 真实节点名
        "call_tool": "call_tool",
        "direct_ans": "direct_ans"
    }
)

# 配置后续固定流向
builder.add_edge("call_tool", "combine_ans")  # 工具执行后 → 整合答案
builder.add_edge("direct_ans", "combine_ans")   # 直接回答 → 结束
builder.add_edge("combine_ans", END)          # 整合答案 → 结束

# 编译 & 执行combine_ans
app = builder.compile()

if __name__ == "__main__":
    # 测试1：需要工具
    print("===== 测试1：包含查询，走工具分支 =====")
    # invoke负责触发执行图中的所有节点，包括分支和循环
    res1 = app.invoke({"question": "查询今日订单", "chat_msg": [], "need_tool": False, "tool_output": "", "answer": ""})
    
    print("\n===== 测试2：普通问题，不走工具 =====")
    # 测试2：不需要工具
    res2 = app.invoke({"question": "你好", "chat_msg": [], "need_tool": False, "tool_output": "", "answer": ""})

===== 测试1：包含查询，走工具分支 =====
【用户意图判断节点，检测是否需要调用工具】begin:
【用户意图判断节点，检测是否需要调用工具】ongoing: 查询今日订单
【用户意图判断节点，检测是否需要调用工具】done: True
【路由函数】begin:
【路由函数】goging:
【路由函数】done，返回下个节点名称: call_tool
【工具调用节点】begin:
【工具调用节点】ongoing:.....
【工具调用节点】done: 从外部接口拿到业务数据
【整合答案节点】begin:
【整合答案节点】onging:
【整合答案节点】done: ['整合工具结果完成作答']

===== 测试2：普通问题，不走工具 =====
【用户意图判断节点，检测是否需要调用工具】begin:
【用户意图判断节点，检测是否需要调用工具】ongoing: 你好
【用户意图判断节点，检测是否需要调用工具】done: False
【路由函数】begin:
【路由函数】goging:
【路由函数】done，返回下个节点名称: direct_ans
【直接回答节点-大模型调用】begin:
【直接回答节点-大模型调用】ongoing:
【直接回答节点-大模型调用】done: 大模型基于常识直接给出的答案
【整合答案节点】begin:
【整合答案节点】onging:
【整合答案节点】done: ['整合工具结果完成作答']

一、基本定义

条件分支依托 条件边（Conditional Edge） 实现，是 LangGraph 区别于线性普通边的核心能力。它指：当前节点执行完毕后，框架不再固定跳转下一个节点，而是根据全局 State 中的数据动态判断、选择后续执行路径，以此实现多分支、分流逻辑。

二、核心组成三要素

1. source（源节点） 分支的触发源头，代表该节点执行完成后，立即进入条件判断流程。一个源节点可以对应一套分支逻辑。

2. path（路由函数） 纯判断逻辑函数，也是条件分支的核心。

   • 入参：全局状态 State；
   • 职责：读取状态里的标记、数据、结果等，执行业务规则判断；
   • 出参：固定返回自定义分支标识字符串，用于匹配分支走向。
   • 特点：仅做判断，不执行业务逻辑、不修改状态。

3. path_map（路由映射表） 字典结构，作用是解耦路由函数与真实节点：将路由函数返回的标识字符串，一一映射为图中已注册的节点名 /END 终点。便于统一管理分支、后期维护。

三、执行时序（遵循框架通用规则）

严格遵守 节点先执行，边后执行 的底层逻辑，完整流程：

1. 调度执行源节点（业务函数），完成逻辑处理并更新全局 State；
2. 源节点执行结束，框架自动触发该节点绑定的条件边；
3. 调用路由函数，读取最新 State 并得到分支标识；
4. 通过 path_map 匹配标识，确定下一跳目标节点 /END；
5. 跳转到目标分支节点继续执行；若匹配到 END，则流程终止。

四、与普通边的核心区别

1. 普通边（add_edge） 跳转路径硬编码固定，节点执行完必然走向指定节点，只能实现串行线性流程，无判断能力。
2. 条件边（add_conditional_edges） 跳转路径动态可变，依赖运行时的状态数据做判断，灵活实现分流、分支、前置校验等复杂逻辑。

五、典型应用场景

1. 意图分流：根据用户提问内容，区分问答、查询、闲聊等不同流程；
2. 能力路由：判断是否需要调用外部工具、联网检索、数据库查询，拆分 “工具分支” 和 “直接回答分支”；
3. 结果校验：节点执行完成后，根据输出结果判断是否继续执行、重试或终止；
4. 权限 / 流程管控：结合状态中的权限、审批标记，分流不同权限的执行链路。

六、关键特性与补充说明

1. 多分支支持：path_map 可配置多个映射关系，不止二选一，能实现三路及以上分支；
2. 无状态耦合：路由逻辑、节点逻辑、流向配置相互解耦，路由函数只依赖全局 State，修改规则无需改动业务节点；
3. 可组合使用：条件分支可以和普通边、循环、人在回路搭配，构建多层级复杂流程；
4. 终止兼容：分支目标可以是普通节点，也可以直接指向 END 终点，实现 “部分分支直接结束流程”。

七、总结

条件分支是 LangGraph 实现非线性流程的基础能力。本质是在节点执行完成后，通过一套「状态读取→规则判断→路径映射」的路由逻辑，动态分配后续执行链路，解决了线性串行流程无法适配复杂业务分流的问题，也是智能体、RAG、对话机器人等场景中最常用的基础能力之一。