1use std::collections::HashMap;
4
5#[derive(Debug, Clone)]
7pub struct Action {
8 pub timestamp: f64,
9 pub action_type: String,
10 pub payload: Vec<u8>,
11 pub duration: f64,
12 pub success: bool,
13}
14
15impl Action {
16 pub fn new(ts: f64, action_type: &str) -> Self {
17 Self {
18 timestamp: ts,
19 action_type: action_type.to_string(),
20 payload: vec![],
21 duration: 0.0,
22 success: true,
23 }
24 }
25 pub fn with_duration(mut self, d: f64) -> Self {
26 self.duration = d;
27 self
28 }
29 pub fn with_success(mut self, s: bool) -> Self {
30 self.success = s;
31 self
32 }
33}
34
35#[derive(Debug, Clone)]
37pub struct Trace {
38 pub agent_id: String,
39 pub actions: Vec<Action>,
40 pub start_time: f64,
41 pub end_time: f64,
42}
43
44impl Trace {
45 pub fn new(agent_id: &str) -> Self {
46 Self {
47 agent_id: agent_id.to_string(),
48 actions: vec![],
49 start_time: 0.0,
50 end_time: 0.0,
51 }
52 }
53
54 pub fn record(&mut self, action: Action) {
55 if self.actions.is_empty() {
56 self.start_time = action.timestamp;
57 }
58 self.end_time = action.timestamp;
59 self.actions.push(action);
60 }
61
62 pub fn duration(&self) -> f64 {
63 self.end_time - self.start_time
64 }
65 pub fn action_count(&self) -> usize {
66 self.actions.len()
67 }
68
69 pub fn success_rate(&self) -> f64 {
70 if self.actions.is_empty() {
71 return 0.0;
72 }
73 self.actions.iter().filter(|a| a.success).count() as f64 / self.actions.len() as f64
74 }
75
76 pub fn action_types(&self) -> HashMap<String, usize> {
77 let mut counts = HashMap::new();
78 for a in &self.actions {
79 *counts.entry(a.action_type.clone()).or_insert(0) += 1;
80 }
81 counts
82 }
83
84 pub fn filter_by_type(&self, action_type: &str) -> Vec<&Action> {
85 self.actions
86 .iter()
87 .filter(|a| a.action_type == action_type)
88 .collect()
89 }
90
91 pub fn avg_duration(&self) -> f64 {
92 if self.actions.is_empty() {
93 return 0.0;
94 }
95 self.actions.iter().map(|a| a.duration).sum::<f64>() / self.actions.len() as f64
96 }
97}
98
99#[derive(Debug, Clone)]
101pub struct Shadow {
102 pub traces: HashMap<String, Trace>,
103 pub max_traces: usize,
104}
105
106impl Default for Shadow {
107 fn default() -> Self {
108 Self::new()
109 }
110}
111
112impl Shadow {
113 pub fn new() -> Self {
114 Self {
115 traces: HashMap::new(),
116 max_traces: 100,
117 }
118 }
119 pub fn with_max_traces(mut self, n: usize) -> Self {
120 self.max_traces = n;
121 self
122 }
123
124 pub fn record(&mut self, agent_id: &str, action: Action) {
125 let trace = self
126 .traces
127 .entry(agent_id.to_string())
128 .or_insert_with(|| Trace::new(agent_id));
129 trace.record(action);
130 if self.traces.len() > self.max_traces {
131 if let Some(oldest) = self
133 .traces
134 .iter()
135 .min_by(|a, b| a.1.start_time.partial_cmp(&b.1.start_time).unwrap())
136 .map(|(k, _)| k.clone())
137 {
138 self.traces.remove(&oldest);
139 }
140 }
141 }
142
143 pub fn get_trace(&self, agent_id: &str) -> Option<&Trace> {
144 self.traces.get(agent_id)
145 }
146
147 pub fn compare(&self, agent_a: &str, agent_b: &str) -> Option<ComparisonResult> {
148 let a = self.traces.get(agent_a)?;
149 let b = self.traces.get(agent_b)?;
150 Some(ComparisonResult {
151 action_diff: a.action_count() as i64 - b.action_count() as i64,
152 success_diff: a.success_rate() - b.success_rate(),
153 duration_diff: a.avg_duration() - b.avg_duration(),
154 })
155 }
156}
157
158#[derive(Debug, Clone, Copy)]
159pub struct ComparisonResult {
160 pub action_diff: i64,
161 pub success_diff: f64,
162 pub duration_diff: f64,
163}
164
165#[cfg(test)]
166mod tests {
167 use super::*;
168
169 #[test]
170 fn test_trace_recording() {
171 let mut trace = Trace::new("agent-1");
172 trace.record(Action::new(0.0, "search").with_duration(0.5));
173 trace.record(Action::new(1.0, "fetch").with_duration(0.3));
174 assert_eq!(trace.action_count(), 2);
175 assert!((trace.duration() - 1.0).abs() < 0.001);
176 }
177
178 #[test]
179 fn test_success_rate() {
180 let mut trace = Trace::new("agent-1");
181 trace.record(Action::new(0.0, "a").with_success(true));
182 trace.record(Action::new(1.0, "b").with_success(false));
183 trace.record(Action::new(2.0, "c").with_success(true));
184 assert!((trace.success_rate() - 0.667).abs() < 0.01);
185 }
186
187 #[test]
188 fn test_action_types() {
189 let mut trace = Trace::new("agent-1");
190 trace.record(Action::new(0.0, "search"));
191 trace.record(Action::new(1.0, "search"));
192 trace.record(Action::new(2.0, "fetch"));
193 let types = trace.action_types();
194 assert_eq!(types.get("search"), Some(&2));
195 assert_eq!(types.get("fetch"), Some(&1));
196 }
197
198 #[test]
199 fn test_shadow_record() {
200 let mut shadow = Shadow::new();
201 shadow.record("a", Action::new(0.0, "x"));
202 shadow.record("a", Action::new(1.0, "y"));
203 shadow.record("b", Action::new(0.5, "z"));
204 assert_eq!(shadow.get_trace("a").unwrap().action_count(), 2);
205 assert_eq!(shadow.get_trace("b").unwrap().action_count(), 1);
206 }
207
208 #[test]
209 fn test_compare() {
210 let mut shadow = Shadow::new();
211 shadow.record(
212 "a",
213 Action::new(0.0, "x").with_success(true).with_duration(1.0),
214 );
215 shadow.record(
216 "a",
217 Action::new(1.0, "y").with_success(false).with_duration(2.0),
218 );
219 shadow.record(
220 "b",
221 Action::new(0.0, "x").with_success(true).with_duration(0.5),
222 );
223 let cmp = shadow.compare("a", "b").unwrap();
224 assert_eq!(cmp.action_diff, 1);
225 }
226
227 #[test]
228 fn test_filter() {
229 let mut trace = Trace::new("agent-1");
230 trace.record(Action::new(0.0, "search"));
231 trace.record(Action::new(1.0, "fetch"));
232 trace.record(Action::new(2.0, "search"));
233 assert_eq!(trace.filter_by_type("search").len(), 2);
234 }
235}