lellm_agent/runtime/
typed_state.rs

1//! Agent Typed State + Mutation — 编译期类型安全的 Agent 状态。
2//!
3//! 替代 `react.rs` 中的 `create_state_from_ctx()` / `sync_state_to_ctx()`
4//! 以及 `runtime.rs` 中的全部 State 辅助函数。
5//!
6//! # 核心原则
7//!
8//! - `AgentState` 是强类型 struct，不是 `HashMap<String, Value>`
9//! - 状态变更通过 `AgentMutation`（领域事件），不是直接修改字段
10//! - 并行合并由 Graph 层的 [`AgentStateMerge`]（`MergeStrategy`）决定
11//! - 零 JSON 序列化开销（节点直接操作 typed state）
12
13use lellm_core::{ChatResponse, Message};
14use lellm_graph::WorkflowState;
15
16use super::context::estimate_tokens;
17use super::event::StopReason;
18
19// ─── AgentState ─────────────────────────────────────────────────
20
21/// Agent 类型化状态。
22///
23/// 替代 `HashMap<String, Value>` 中通过 `StateKey<T>` 存取的模式。
24/// 所有字段编译期可见，零运行时类型检查开销。
25#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
26pub struct AgentState {
27    /// 消息历史（Assistant + Tool + User）
28    pub messages: Vec<Message>,
29    /// 当前迭代轮次
30    pub iterations: usize,
31    /// 累计工具调用总数
32    pub total_tool_calls: usize,
33    /// 累计输出 Token 数（Text，不含 Thinking）
34    pub output_tokens: usize,
35    /// 累计推理 Token 数（Thinking，不含 Text）
36    pub reasoning_tokens: usize,
37    /// 累计压缩次数
38    pub compact_count: usize,
39    /// 停止原因（终态）
40    pub stop_reason: Option<StopReason>,
41    /// 最后一次 LLM 响应
42    pub last_response: Option<ChatResponse>,
43}
44
45impl AgentState {
46    /// 从初始消息列表创建 AgentState。
47    pub fn from_messages(messages: Vec<Message>) -> Self {
48        Self {
49            messages,
50            iterations: 0,
51            total_tool_calls: 0,
52            output_tokens: 0,
53            reasoning_tokens: 0,
54            compact_count: 0,
55            stop_reason: None,
56            last_response: None,
57        }
58    }
59
60    /// 实时派生上下文 Token 数（从 messages 估算）。
61    /// 供 BudgetCondition 判断是否需要压缩。
62    pub fn estimated_context_tokens(&self) -> usize {
63        estimate_tokens(&self.messages)
64    }
65
66    /// 检查是否已达到最大迭代轮次。
67    pub fn reached_max(&self, max_iterations: usize) -> bool {
68        self.iterations >= max_iterations
69    }
70
71    /// 检查是否超过总输出预算。
72    pub fn exceeded_output(&self, max: Option<u32>) -> bool {
73        match max {
74            Some(limit) => self.output_tokens >= limit as usize,
75            None => false,
76        }
77    }
78
79    /// 检查是否超过总推理预算。
80    pub fn exceeded_reasoning(&self, max: Option<u32>) -> bool {
81        match max {
82            Some(limit) => self.reasoning_tokens >= limit as usize,
83            None => false,
84        }
85    }
86
87    /// 检查加上额外 Token 后是否超过总输出预算。
88    ///
89    /// 用于 Mutation 未 apply 时的预判（节点 emit Mutation 之前）。
90    pub fn exceeded_output_with_extra(&self, max: Option<u32>, extra: usize) -> bool {
91        match max {
92            Some(limit) => self.output_tokens + extra >= limit as usize,
93            None => false,
94        }
95    }
96
97    /// 检查加上额外 Token 后是否超过总推理预算。
98    pub fn exceeded_reasoning_with_extra(&self, max: Option<u32>, extra: usize) -> bool {
99        match max {
100            Some(limit) => self.reasoning_tokens + extra >= limit as usize,
101            None => false,
102        }
103    }
104
105    /// 检查是否已终止（有 stop_reason）。
106    pub fn is_terminal(&self) -> bool {
107        self.stop_reason.is_some()
108    }
109
110    /// 获取当前消息引用。
111    pub fn messages_ref(&self) -> &[Message] {
112        &self.messages
113    }
114}
115
116// ─── AgentMutation ────────────────────────────────────────────────
117
118/// Agent 领域事件 — 描述一次状态转换。
119///
120/// 节点通过发射 Mutation 来变更状态，而非直接修改 `AgentState` 字段。
121/// Mutation 是可序列化的、自包含的、不可变的。
122#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
123pub enum AgentMutation {
124    /// 追加一条消息到历史
125    AppendMessage(Message),
126    /// 追加多条消息到历史
127    AppendMessages(Vec<Message>),
128    /// 进入下一轮迭代
129    IncrementIteration,
130    /// 记录工具调用数量
131    AddToolCalls(usize),
132    /// 记录输出 Token
133    AddOutputTokens(usize),
134    /// 记录推理 Token
135    AddReasoningTokens(usize),
136    /// 记录一次压缩
137    IncrementCompactCount,
138    /// 替换消息历史（压缩场景）
139    ReplaceMessages(Vec<Message>),
140    /// 设置停止原因
141    SetStopReason(StopReason),
142    /// 更新最后一次 LLM 响应
143    SetLastResponse(ChatResponse),
144}
145
146impl lellm_graph::workflow_state::StateMutation<AgentState> for AgentMutation {
147    fn apply(self, state: &mut AgentState) {
148        match self {
149            AgentMutation::AppendMessage(msg) => {
150                state.messages.push(msg);
151            }
152            AgentMutation::AppendMessages(msgs) => {
153                state.messages.extend(msgs);
154            }
155            AgentMutation::IncrementIteration => {
156                state.iterations += 1;
157            }
158            AgentMutation::AddToolCalls(n) => {
159                state.total_tool_calls += n;
160            }
161            AgentMutation::AddOutputTokens(n) => {
162                state.output_tokens += n;
163            }
164            AgentMutation::AddReasoningTokens(n) => {
165                state.reasoning_tokens += n;
166            }
167            AgentMutation::IncrementCompactCount => {
168                state.compact_count += 1;
169            }
170            AgentMutation::ReplaceMessages(msgs) => {
171                state.messages = msgs;
172            }
173            AgentMutation::SetStopReason(reason) => {
174                state.stop_reason = Some(reason);
175            }
176            AgentMutation::SetLastResponse(response) => {
177                state.last_response = Some(response);
178            }
179        }
180    }
181}
182
183// ─── WorkflowState for AgentState ───────────────────────────────
184
185impl WorkflowState for AgentState {
186    type Mutation = AgentMutation;
187}
188
189/// AgentState 的默认合并策略。
190///
191/// - messages: 所有分支拼接（chain）
192/// - iterations: 取最大值
193/// - total_tool_calls: 取最大值
194/// - output_tokens: 累加
195/// - reasoning_tokens: 累加
196/// - compact_count: 累加
197/// - stop_reason: 优先取后者
198/// - last_response: 优先取后者
199pub struct AgentStateMerge;
200
201impl lellm_graph::MergeStrategy<AgentState> for AgentStateMerge {
202    fn merge(branches: Vec<AgentState>) -> Result<AgentState, lellm_graph::WorkflowError> {
203        let mut iter = branches.into_iter();
204        let mut merged = iter.next().ok_or_else(|| {
205            lellm_graph::WorkflowError::MergeConflict("no branches to merge".into())
206        })?;
207
208        for branch in iter {
209            merged.messages.extend(branch.messages);
210            merged.iterations = merged.iterations.max(branch.iterations);
211            merged.total_tool_calls = merged.total_tool_calls.max(branch.total_tool_calls);
212            merged.output_tokens += branch.output_tokens;
213            merged.reasoning_tokens += branch.reasoning_tokens;
214            merged.compact_count += branch.compact_count;
215            if merged.stop_reason.is_none() {
216                merged.stop_reason = branch.stop_reason;
217            }
218            if merged.last_response.is_none() {
219                merged.last_response = branch.last_response;
220            }
221        }
222
223        Ok(merged)
224    }
225
226    fn default_instance() -> Self {
227        AgentStateMerge
228    }
229}
230
231// ─── NOTE: 序列化桥接已删除 ─────────────────────────
232// to_value() / from_value() / apply_from_value() / AGENT_STATE_KEY
233// 已删除 — Graph 层只做 Node → Mutation → State 管道，不经过 JSON。
234// 如需 Checkpoint 持久化，由 Checkpoint 层直接序列化 AgentState。
lellm_agent/runtime/typed_state.rs

lellm_agent/runtime/
typed_state.rs
