一、从DAG到StateGraph

第一篇实现的DAG工作流引擎适合确定性的任务编排——节点和依赖关系在执行前就已确定。但AI Agent推理的本质是不确定的

用户提问 -> LLM判断是否需要工具 -> 如果需要->调用工具->重新思考->...
                                -> 如果不需要->直接回答

这需要一个有状态的图:节点执行后根据状态决定下一步跳到哪里。这正是LangGraph的核心思想。本文用Go从零实现。


二、StateGraph核心概念

+-------------+    条件边     +--------------+
|  llm_think   |------------->|  tool_exec    |
|  (LLM推理)   |              |  (工具执行)    |
+-------------+              +--------------+
       |                            |
       | 条件边                      | 固定边
       v                            v
+-------------+              +--------------+
|   respond    |<-------------|  llm_think    |
|  (最终回答)  |  条件边       |  (继续推理)   |
+-------------+              +--------------+

关键差异:

  • State:在节点间传递的共享状态,节点可读可写
  • 条件边(ConditionalEdge):根据State动态决定下一节点
  • END:特殊终止节点

三、完整实现

package stategraph

import (
    "context"
    "fmt"
    "sync"
)

// State 在节点间传递的共享状态 —— 本质是一个线程安全的map
type State struct {
    mu   sync.RWMutex
    data map[string]interface{}
}

func NewState() *State {
    return &State{data: make(map[string]interface{})}
}

func (s *State) Set(key string, value interface{}) {
    s.mu.Lock()
    defer s.mu.Unlock()
    s.data[key] = value
}

func (s *State) Get(key string) (interface{}, bool) {
    s.mu.RLock()
    defer s.mu.RUnlock()
    v, ok := s.data[key]
    return v, ok
}

func (s *State) GetString(key string) string {
    v, ok := s.Get(key)
    if !ok {
        return ""
    }
    str, _ := v.(string)
    return str
}

func (s *State) GetInt(key string) int {
    v, ok := s.Get(key)
    if !ok {
        return 0
    }
    i, _ := v.(int)
    return i
}

func (s *State) Snapshot() map[string]interface{} {
    s.mu.RLock()
    defer s.mu.RUnlock()
    cp := make(map[string]interface{}, len(s.data))
    for k, v := range s.data {
        cp[k] = v
    }
    return cp
}

// NodeFunc 节点函数:接收State,返回修改后的State(或nil表示不变)和可能的error
type NodeFunc func(ctx context.Context, state *State) (*State, error)

// RouteFunc 条件路由函数:根据State返回下一个节点名
type RouteFunc func(state *State) (string, error)

// node 内部节点结构
type node struct {
    id      string
    handler NodeFunc
}

// StateGraph 有状态图执行引擎
type StateGraph struct {
    mu         sync.RWMutex
    nodes      map[string]*node     // 节点映射
    edges      map[string][]string  // 固定边: from -> [to...]
    condEdges  map[string]RouteFunc // 条件边: from -> RouteFunc
    entryPoint string               // 入口节点
    maxSteps   int                  // 最大步数防无限循环
    visitedMax int                  // 同节点最大访问次数
    visitCount map[string]int       // 节点访问计数
}

// NewStateGraph 创建StateGraph构建器
func NewStateGraph() *StateGraph {
    return &StateGraph{
        nodes:      make(map[string]*node),
        edges:      make(map[string][]string),
        condEdges:  make(map[string]RouteFunc),
        maxSteps:   50,  // 默认最大50步
        visitedMax: 10,  // 同一节点最多访问10次
    }
}

// SetEntryPoint 设置入口节点
func (sg *StateGraph) SetEntryPoint(nodeID string) *StateGraph {
    sg.entryPoint = nodeID
    return sg
}

// AddNode 添加节点
func (sg *StateGraph) AddNode(id string, handler NodeFunc) *StateGraph {
    sg.mu.Lock()
    defer sg.mu.Unlock()
    sg.nodes[id] = &node{id: id, handler: handler}
    return sg
}

// AddEdge 添加固定边 (from -> to)
func (sg *StateGraph) AddEdge(from, to string) *StateGraph {
    sg.mu.Lock()
    defer sg.mu.Unlock()
    sg.edges[from] = append(sg.edges[from], to)
    return sg
}

// AddConditionalEdge 添加条件边 (from -> RouteFunc)
func (sg *StateGraph) AddConditionalEdge(from string, routeFn RouteFunc) *StateGraph {
    sg.mu.Lock()
    defer sg.mu.Unlock()
    sg.condEdges[from] = routeFn
    return sg
}

// End 特殊常量表示终止
const End = "__END__"

// Execute 执行StateGraph
func (sg *StateGraph) Execute(ctx context.Context, initialState map[string]interface{}) (*State, error) {
    if sg.entryPoint == "" {
        return nil, fmt.Errorf("未设置入口节点")
    }

    state := NewState()
    for k, v := range initialState {
        state.Set(k, v)
    }

    sg.visitCount = make(map[string]int)
    current := sg.entryPoint

    for step := 0; step < sg.maxSteps; step++ {
        select {
        case <-ctx.Done():
            return state, ctx.Err()
        default:
        }

        // 1. 访问计数检查(防死循环)
        sg.visitCount[current]++
        if sg.visitCount[current] > sg.visitedMax {
            return state, fmt.Errorf("节点 %s 被访问 %d 次,疑似死循环,已中止",
                current, sg.visitCount[current])
        }

        // 2. 执行当前节点
        nd, ok := sg.nodes[current]
        if !ok {
            return state, fmt.Errorf("节点 %s 未注册", current)
        }

        newState, err := nd.handler(ctx, state)
        if err != nil {
            return state, fmt.Errorf("节点 %s 执行失败: %w", current, err)
        }
        if newState != nil {
            state = newState
        }

        // 3. 决定下一节点
        if routeFn, hasCond := sg.condEdges[current]; hasCond {
            next, err := routeFn(state)
            if err != nil {
                return state, fmt.Errorf("节点 %s 路由决策失败: %w", current, err)
            }
            if next == End || next == "" {
                return state, nil
            }
            current = next
            continue
        }

        if nextNodes, hasFixed := sg.edges[current]; hasFixed && len(nextNodes) > 0 {
            current = nextNodes[0]
            continue
        }

        return state, nil
    }

    return state, fmt.Errorf("超过最大步数 %d,已强制终止", sg.maxSteps)
}

四、实战:用StateGraph实现ReAct Agent

ReAct (Reasoning + Acting) 是AI Agent的经典模式。用StateGraph建模:

package main

import (
    "context"
    "fmt"
    "strings"

    "github.com/yourorg/stategraph"
)

func main() {
    sg := stategraph.NewStateGraph()

    // 节点1: LLM推理
    sg.AddNode("llm_think", llmThinkNode)

    // 节点2: 工具执行
    sg.AddNode("tool_exec", toolExecNode)

    // 节点3: 最终回答
    sg.AddNode("respond", respondNode)

    // 条件边: llm_think -> 需要工具走tool_exec,否则走respond
    sg.AddConditionalEdge("llm_think", routeAfterThink)

    // 固定边: tool_exec -> 回到llm_think继续推理
    sg.AddEdge("tool_exec", "llm_think")

    // 入口
    sg.SetEntryPoint("llm_think")

    // 执行
    ctx := context.Background()
    state, err := sg.Execute(ctx, map[string]interface{}{
        "question": "北京今天的天气怎么样?适合户外运动吗?",
    })
    if err != nil {
        fmt.Printf("Agent执行失败: %v\n", err)
        return
    }

    answer, _ := state.Get("final_answer")
    fmt.Printf("Agent回答: %s\n", answer)
}

// --- 节点实现 ---

func llmThinkNode(ctx context.Context, state *stategraph.State) (*stategraph.State, error) {
    question := state.GetString("question")
    history := state.GetString("conversation_history")

    // 模拟LLM推理(实际项目中这里调用OpenAI/Claude API)
    if strings.Contains(question, "天气") {
        state.Set("llm_action", "tool_call")
        state.Set("tool_name", "get_weather")
        state.Set("tool_args", "北京")
        state.Set("thought", "用户询问天气,需要调用天气API")
    } else {
        state.Set("llm_action", "respond")
        state.Set("thought", "问题不需要工具,直接回答")
    }
    return state, nil
}

func routeAfterThink(state *stategraph.State) (string, error) {
    action := state.GetString("llm_action")
    switch action {
    case "tool_call":
        return "tool_exec", nil
    case "respond":
        return "respond", nil
    default:
        return stategraph.End, nil
    }
}

func toolExecNode(ctx context.Context, state *stategraph.State) (*stategraph.State, error) {
    toolName := state.GetString("tool_name")
    toolArgs := state.GetString("tool_args")

    var result string
    switch toolName {
    case "get_weather":
        result = fmt.Sprintf("%s: 晴,22\u00b0C,微风,适合户外运动", toolArgs)
    case "search":
        result = fmt.Sprintf("搜索结果: %s 相关信息...", toolArgs)
    default:
        result = fmt.Sprintf("未知工具: %s", toolName)
    }

    history := state.GetString("conversation_history")
    history += fmt.Sprintf("\n[工具调用] %s(%s) -> %s", toolName, toolArgs, result)
    state.Set("conversation_history", history)
    state.Set("tool_result", result)

    return state, nil
}

func respondNode(ctx context.Context, state *stategraph.State) (*stategraph.State, error) {
    question := state.GetString("question")
    history := state.GetString("conversation_history")

    var finalAnswer string
    if history != "" {
        finalAnswer = fmt.Sprintf("查询结果显示: %s", state.GetString("tool_result"))
    } else {
        finalAnswer = fmt.Sprintf("关于\"%s\"的回答...", question)
    }
    state.Set("final_answer", finalAnswer)
    return state, nil
}

执行流程可视化

Step 1: llm_think -> action=tool_call -> route_fn返回"tool_exec"
Step 2: tool_exec -> 获取天气->记录历史 -> 固定边回到"llm_think"
Step 3: llm_think -> 看过历史,action=respond -> route_fn返回"respond"
Step 4: respond -> 生成最终答案 -> 无出口,终止

五、多条件路由——支持复杂Agent逻辑

实际场景中一个节点可能需要根据多种条件跳转到不同分支:

// 复杂路由:根据状态字段组合决策
func complexRoute(state *stategraph.State) (string, error) {
    action := state.GetString("action")
    confidence := state.GetInt("confidence")
    retries := state.GetInt("retries")

    // 高置信度且不需要工具 -> 直接回答
    if action == "respond" && confidence > 80 {
        return "final_respond", nil
    }

    // 工具调用失败且未超出重试 -> 重试
    if action == "retry" && retries < 3 {
        return "tool_exec", nil
    }

    // 超出重试 -> 降级处理
    if retries >= 3 {
        return "fallback", nil
    }

    // 默认:让LLM重新思考
    return "llm_think", nil
}

六、循环检测机制

Agent推理可能陷入循环(例如工具反复返回同样结果导致LLM陷入重复思考)。StateGraph内置了双重防护:

// 1. 全局步数上限
sg.maxSteps = 50

// 2. 单节点访问次数上限
sg.visitedMax = 10

当同一节点被访问超过10次时,引擎会强制终止并返回错误,避免token无限消耗。


七、生产环境集成

7.1 与真实LLM集成

type OpenAILLMNode struct {
    client *openai.Client
    model  string
}

func (n *OpenAILLMNode) Think(ctx context.Context, state *State) (*State, error) {
    messages := buildMessages(state)
    resp, err := n.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
        Model:    n.model,
        Messages: messages,
        Tools:    registeredTools,
    })
    if err != nil {
        return state, fmt.Errorf("LLM调用失败: %w", err)
    }

    choice := resp.Choices[0]
    if len(choice.Message.ToolCalls) > 0 {
        state.Set("llm_action", "tool_call")
        state.Set("pending_tool_calls", choice.Message.ToolCalls)
    } else {
        state.Set("llm_action", "respond")
        state.Set("llm_response", choice.Message.Content)
    }
    return state, nil
}

7.2 可观测性

// 每个节点执行时记录trace
type NodeTrace struct {
    NodeID    string        `json:"node_id"`
    Step      int           `json:"step"`
    Duration  time.Duration `json:"duration"`
    StateSize int           `json:"state_size"`
    Error     string        `json:"error,omitempty"`
}

// 在执行循环中收集traces
traces = append(traces, NodeTrace{
    NodeID:    current,
    Step:      step,
    Duration:  elapsed,
    StateSize: len(state.Snapshot()),
})

八、总结

特性 DAG引擎(第一篇) StateGraph(本篇)
节点间流转 固定拓扑 动态路由
状态传递 input map 全局State对象
循环 不允许(检测到报错) 允许(有防护上限)
适用场景 确定性任务编排 不确定性Agent推理
复杂度 O(V+E)拓扑排序 步数上限控制
Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐