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

lellm_agent/runtime/
typed_state.rs
