Go语言实现LLM推理图:有状态多步Agent执行引擎
·
一、从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)拓扑排序 | 步数上限控制 |
更多推荐


所有评论(0)