swarm_engine_core/learn/record/
mod.rs1mod action;
44mod dependency_graph;
45mod learn_stats;
46mod llm;
47mod strategy_advice;
48mod stream;
49
50use serde::{Deserialize, Serialize};
51
52use crate::events::{ActionEvent, LearningEvent};
53
54pub use action::ActionRecord;
55pub use dependency_graph::DependencyGraphRecord;
56pub use learn_stats::LearnStatsRecord;
57pub use llm::LlmCallRecord;
58pub use strategy_advice::StrategyAdviceRecord;
59pub use stream::RecordStream;
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
70pub enum Record {
71 Action(ActionRecord),
73 Llm(LlmCallRecord),
75 DependencyGraph(DependencyGraphRecord),
77 StrategyAdvice(StrategyAdviceRecord),
79 LearnStats(LearnStatsRecord),
81}
82
83impl Record {
84 pub fn is_action(&self) -> bool {
86 matches!(self, Self::Action(_))
87 }
88
89 pub fn is_llm(&self) -> bool {
91 matches!(self, Self::Llm(_))
92 }
93
94 pub fn is_dependency_graph(&self) -> bool {
96 matches!(self, Self::DependencyGraph(_))
97 }
98
99 pub fn is_strategy_advice(&self) -> bool {
101 matches!(self, Self::StrategyAdvice(_))
102 }
103
104 pub fn is_learn_stats(&self) -> bool {
106 matches!(self, Self::LearnStats(_))
107 }
108
109 pub fn as_action(&self) -> Option<&ActionRecord> {
111 match self {
112 Self::Action(r) => Some(r),
113 _ => None,
114 }
115 }
116
117 pub fn as_llm(&self) -> Option<&LlmCallRecord> {
119 match self {
120 Self::Llm(r) => Some(r),
121 _ => None,
122 }
123 }
124
125 pub fn as_dependency_graph(&self) -> Option<&DependencyGraphRecord> {
127 match self {
128 Self::DependencyGraph(r) => Some(r),
129 _ => None,
130 }
131 }
132
133 pub fn as_strategy_advice(&self) -> Option<&StrategyAdviceRecord> {
135 match self {
136 Self::StrategyAdvice(r) => Some(r),
137 _ => None,
138 }
139 }
140
141 pub fn as_learn_stats(&self) -> Option<&LearnStatsRecord> {
143 match self {
144 Self::LearnStats(r) => Some(r),
145 _ => None,
146 }
147 }
148
149 pub fn worker_id(&self) -> Option<usize> {
151 match self {
152 Self::Action(r) => Some(r.worker_id),
153 Self::Llm(r) => r.worker_id,
154 Self::DependencyGraph(_) => None,
156 Self::StrategyAdvice(_) => None,
157 Self::LearnStats(_) => None,
158 }
159 }
160
161 pub fn timestamp_ms(&self) -> u64 {
163 match self {
164 Self::Action(r) => r.tick,
165 Self::Llm(r) => r.timestamp_ms,
166 Self::DependencyGraph(r) => r.timestamp_ms,
167 Self::StrategyAdvice(r) => r.timestamp_ms,
168 Self::LearnStats(r) => r.timestamp_ms,
169 }
170 }
171}
172
173impl From<ActionRecord> for Record {
174 fn from(record: ActionRecord) -> Self {
175 Self::Action(record)
176 }
177}
178
179impl From<LlmCallRecord> for Record {
180 fn from(record: LlmCallRecord) -> Self {
181 Self::Llm(record)
182 }
183}
184
185impl From<DependencyGraphRecord> for Record {
186 fn from(record: DependencyGraphRecord) -> Self {
187 Self::DependencyGraph(record)
188 }
189}
190
191impl From<StrategyAdviceRecord> for Record {
192 fn from(record: StrategyAdviceRecord) -> Self {
193 Self::StrategyAdvice(record)
194 }
195}
196
197impl From<LearnStatsRecord> for Record {
198 fn from(record: LearnStatsRecord) -> Self {
199 Self::LearnStats(record)
200 }
201}
202
203impl From<&ActionEvent> for Record {
204 fn from(event: &ActionEvent) -> Self {
205 Self::Action(ActionRecord::from(event))
206 }
207}
208
209impl From<&LearningEvent> for Record {
210 fn from(event: &LearningEvent) -> Self {
211 match event {
212 LearningEvent::StrategyAdvice { .. } => {
213 Self::StrategyAdvice(StrategyAdviceRecord::from(event))
214 }
215 LearningEvent::DependencyGraphInference { .. } => {
216 Self::DependencyGraph(DependencyGraphRecord::from(event))
217 }
218 LearningEvent::LearnStatsSnapshot { .. } => {
219 Self::LearnStats(LearnStatsRecord::from(event))
220 }
221 }
222 }
223}
224
225pub trait FromRecord: Sized {
234 fn from_record(record: &Record) -> Option<&Self>;
235}
236
237impl FromRecord for ActionRecord {
238 fn from_record(record: &Record) -> Option<&Self> {
239 record.as_action()
240 }
241}
242
243impl FromRecord for LlmCallRecord {
244 fn from_record(record: &Record) -> Option<&Self> {
245 record.as_llm()
246 }
247}
248
249impl FromRecord for DependencyGraphRecord {
250 fn from_record(record: &Record) -> Option<&Self> {
251 record.as_dependency_graph()
252 }
253}
254
255impl FromRecord for StrategyAdviceRecord {
256 fn from_record(record: &Record) -> Option<&Self> {
257 record.as_strategy_advice()
258 }
259}
260
261impl FromRecord for LearnStatsRecord {
262 fn from_record(record: &Record) -> Option<&Self> {
263 record.as_learn_stats()
264 }
265}
266
267#[cfg(test)]
272mod tests {
273 use super::*;
274
275 #[test]
276 fn test_record_from_action_record() {
277 let action = ActionRecord::new(1, 0, "CheckStatus").success(true);
278 let record = Record::from(action);
279
280 assert!(record.is_action());
281 assert!(!record.is_llm());
282 assert_eq!(record.worker_id(), Some(0));
283 }
284
285 #[test]
286 fn test_record_from_llm_call_record() {
287 let llm = LlmCallRecord::new("decide", "qwen2.5")
288 .worker_id(1)
289 .prompt("test")
290 .response("ok");
291 let record = Record::from(llm);
292
293 assert!(!record.is_action());
294 assert!(record.is_llm());
295 assert_eq!(record.worker_id(), Some(1));
296 }
297
298 #[test]
299 fn test_record_stream_filtering() {
300 let records = vec![
301 Record::from(ActionRecord::new(1, 0, "A").success(true)),
302 Record::from(LlmCallRecord::new("decide", "model").worker_id(0)),
303 Record::from(ActionRecord::new(2, 1, "B").success(true)),
304 ];
305
306 let stream = RecordStream::new(&records);
307
308 assert_eq!(stream.actions().count(), 2);
309 assert_eq!(stream.llm_calls().count(), 1);
310 assert_eq!(stream.by_worker(0).count(), 2);
311 assert_eq!(stream.by_worker(1).count(), 1);
312 }
313}