状态图

在 Graph 基础上增加条件边（Conditional Edge）：根据当前状态动态选择下一节点。 演示一个情感路由工作流：分类 → 按情感分流到不同处理节点 → 汇总。 对应教程：第8卷《LangGraph》第3章 构建块与第4章 状态记忆。

文件路径: examples/volume-08-langgraph\state_graph.py

---

源代码

"""状态图与条件路由示例。

在 ``Graph`` 基础上增加**条件边（Conditional Edge）**：根据当前状态动态选择下一节点。
演示一个情感路由工作流：分类 → 按情感分流到不同处理节点 → 汇总。

对应教程：第8卷《LangGraph》第3章 构建块与第4章 状态记忆。
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

from hello_graph import END, State

NodeFunc = Callable[[State], State]
ConditionFunc = Callable[[State], str]


@dataclass
class StateGraph:
    """支持条件边的状态图。

    Attributes:
        entry: 入口节点名。
    """

    entry: str | None = None
    _nodes: dict[str, NodeFunc] = field(default_factory=dict)
    _edges: dict[str, str] = field(default_factory=dict)
    _conditional: dict[str, tuple[ConditionFunc, dict[str, str]]] = field(default_factory=dict)

    def add_node(self, name: str, func: NodeFunc) -> StateGraph:
        """添加节点。"""
        self._nodes[name] = func
        return self

    def add_edge(self, from_node: str, to_node: str) -> StateGraph:
        """添加固定边。"""
        self._edges[from_node] = to_node
        return self

    def add_conditional_edges(
        self,
        from_node: str,
        condition: ConditionFunc,
        mapping: dict[str, str],
    ) -> StateGraph:
        """添加条件边。

        Args:
            from_node: 起点节点。
            condition: 条件函数，返回 ``mapping`` 的某个键。
            mapping: 键到目标节点名的映射。

        Returns:
            ``self``。
        """
        self._conditional[from_node] = (condition, mapping)
        return self

    def set_entry_point(self, name: str) -> StateGraph:
        """设置入口。"""
        self.entry = name
        return self

    def _next_node(self, current: str, state: State) -> str:
        if current in self._conditional:
            condition, mapping = self._conditional[current]
            key = condition(state)
            return mapping.get(key, END)
        return self._edges.get(current, END)

    def invoke(self, initial_state: State, max_steps: int = 20) -> State:
        """执行状态图。

        Args:
            initial_state: 初始状态。
            max_steps: 最大步数。

        Returns:
            最终状态。
        """
        if self.entry is None:
            raise ValueError("未设置入口节点")
        state = dict(initial_state)
        current = self.entry
        for _ in range(max_steps):
            if current == END:
                break
            if current not in self._nodes:
                raise ValueError(f"未知节点: {current}")
            update = self._nodes[current](state)
            if update:
                state.update(update)
            current = self._next_node(current, state)
        else:
            raise RuntimeError(f"图执行超过最大步数 {max_steps}")
        return state


def classify(state: State) -> State:
    """示例节点：根据文本中关键词分类情感。"""
    text = state.get("text", "")
    if any(w in text for w in ["好", "棒", "喜欢", "满意"]):
        return {"sentiment": "positive"}
    if any(w in text for w in ["差", "坏", "失望", "糟"]):
        return {"sentiment": "negative"}
    return {"sentiment": "neutral"}


def handle_positive(state: State) -> State:
    return {"reply": f"感谢您的肯定！({state.get('text', '')})"}


def handle_negative(state: State) -> State:
    return {"reply": f"抱歉给您带来不便，已记录并转人工。({state.get('text', '')})"}


def handle_neutral(state: State) -> State:
    return {"reply": f"已收到您的反馈。({state.get('text', '')})"}


def summarize(state: State) -> State:
    return {"final": f"[{state.get('sentiment')}] {state.get('reply', '')}"}


def route_by_sentiment(state: State) -> str:
    """条件函数：按情感路由。"""
    return state.get("sentiment", "neutral")


def build_sentiment_graph() -> StateGraph:
    """构建情感路由状态图。

    流程：classify → (条件) → handle_* → summarize → END
    """
    g = StateGraph()
    g.add_node("classify", classify)
    g.add_node("handle_positive", handle_positive)
    g.add_node("handle_negative", handle_negative)
    g.add_node("handle_neutral", handle_neutral)
    g.add_node("summarize", summarize)

    g.set_entry_point("classify")
    g.add_conditional_edges(
        "classify",
        route_by_sentiment,
        {
            "positive": "handle_positive",
            "negative": "handle_negative",
            "neutral": "handle_neutral",
        },
    )
    for handler in ("handle_positive", "handle_negative", "handle_neutral"):
        g.add_edge(handler, "summarize")
    g.add_edge("summarize", END)
    return g


def main() -> None:
    """演示条件路由状态图。"""
    graph = build_sentiment_graph()
    for text in ["这个产品太棒了，很喜欢！", "质量太差，非常失望", "功能一般"]:
        result = graph.invoke({"text": text})
        print(f"输入: {text}")
        print(f"输出: {result['final']}\n")


if __name__ == "__main__":
    main()