《AI Agent 开发入门》12.1 状态机式编排

# minagent/workflow/graph.py
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable

State = dict[str, Any]
NodeFn = Callable[[State], dict[str, Any]]      # 输入 state，返回「增量更新」
CondFn = Callable[[State], bool]                # 输入 state，返回是否走这条边


@dataclass
class Node:
    name: str
    fn: NodeFn


@dataclass
class Edge:
    src: str
    dst: str
    cond: CondFn | None = None                  # None 表示无条件


@dataclass
class Graph:
    nodes: dict[str, Node] = field(default_factory=dict)
    edges: list[Edge] = field(default_factory=list)
    entry: str | None = None

    def add_node(self, name: str, fn: NodeFn) -> "Graph":
        self.nodes[name] = Node(name, fn)
        return self

    def add_edge(self, src: str, dst: str, cond: CondFn | None = None) -> "Graph":
        self.edges.append(Edge(src, dst, cond))
        return self

    def set_entry(self, name: str) -> "Graph":
        self.entry = name
        return self

    def _next(self, name: str, state: State) -> str | None:
        for e in self.edges:
            if e.src == name and (e.cond is None or e.cond(state)):
                return e.dst
        return None

    def run(self, state: State, max_steps: int = 25, on_node=None) -> State:
        cur = self.entry
        steps = 0
        while cur is not None:
            steps += 1
            if steps > max_steps:
                raise RuntimeError(f"超过 max_steps={max_steps}，图未收敛")
            before = dict(state)
            update = self.nodes[cur].fn(state) or {}
            state.update(update)
            if on_node is not None:
                on_node(cur, before, dict(state))
            cur = self._next(cur, state)
        return state

from minagent.fake import FakeModel
from minagent.types import Message
from minagent.workflow.graph import Graph

DOCS = {
    "python": "GIL 是 CPython 的全局解释器锁，同一时刻只允许一个线程执行字节码。",
    "async": "asyncio 用单个事件循环在单线程内并发调度协程。",
}

MODEL = FakeModel([
    "python",                                             # classify 节点：模型判定类别
    "GIL 是 CPython 的全局解释器锁，它限制了多线程的并行度。",  # synthesize 节点：模型给答案
])


def classify(state):
    resp = MODEL.complete([Message(role="user", content=state["question"])], [])
    return {"category": resp.message.content.strip()}


def retrieve(state):
    cat = state["category"]
    return {"docs": [DOCS[cat]] if cat in DOCS else []}


def synthesize(state):
    context = "\n".join(state["docs"])
    resp = MODEL.complete([Message(role="user", content=context)], [])
    return {"answer": resp.message.content.strip()}


g = (
    Graph()
    .add_node("classify", classify)
    .add_node("retrieve", retrieve)
    .add_node("synthesize", synthesize)
    .add_edge("classify", "retrieve")
    .add_edge("retrieve", "synthesize")
    .set_entry("classify")
)


def show(name, before, after):
    print(f"[{name}]")
    print(f"  in : {before}")
    print(f"  out: {after}")


final = g.run({"question": "Python 的 GIL 是什么？"}, on_node=show)
print("answer:", final["answer"])

[classify]
  in : {'question': 'Python 的 GIL 是什么？'}
  out: {'question': 'Python 的 GIL 是什么？', 'category': 'python'}
[retrieve]
  in : {'question': 'Python 的 GIL 是什么？', 'category': 'python'}
  out: {'question': 'Python 的 GIL 是什么？', 'category': 'python', 'docs': ['GIL 是 CPython 的全局解释器锁，同一时刻只允许一个线程执行字节码。']}
[synthesize]
  in : {'question': 'Python 的 GIL 是什么？', 'category': 'python', 'docs': ['GIL 是 CPython 的全局解释器锁，同一时刻只允许一个线程执行字节码。']}
  out: {'question': 'Python 的 GIL 是什么？', 'category': 'python', 'docs': ['GIL 是 CPython 的全局解释器锁，同一时刻只允许一个线程执行字节码。'], 'answer': 'GIL 是 CPython 的全局解释器锁，它限制了多线程的并行度。'}
answer: GIL 是 CPython 的全局解释器锁，它限制了多线程的并行度。