1use std::collections::HashMap;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4pub enum LayerKind {
5 Input,
6 Intent,
7 Relevance,
8 Compression,
9 Translation,
10 Delivery,
11}
12
13impl LayerKind {
14 pub fn as_str(&self) -> &'static str {
15 match self {
16 Self::Input => "input",
17 Self::Intent => "intent",
18 Self::Relevance => "relevance",
19 Self::Compression => "compression",
20 Self::Translation => "translation",
21 Self::Delivery => "delivery",
22 }
23 }
24
25 pub fn all() -> &'static [LayerKind] {
26 &[
27 Self::Input,
28 Self::Intent,
29 Self::Relevance,
30 Self::Compression,
31 Self::Translation,
32 Self::Delivery,
33 ]
34 }
35}
36
37impl std::fmt::Display for LayerKind {
38 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39 write!(f, "{}", self.as_str())
40 }
41}
42
43#[derive(Debug, Clone)]
44pub struct LayerInput {
45 pub content: String,
46 pub tokens: usize,
47 pub metadata: HashMap<String, String>,
48}
49
50#[derive(Debug, Clone)]
51pub struct LayerOutput {
52 pub content: String,
53 pub tokens: usize,
54 pub metadata: HashMap<String, String>,
55}
56
57#[derive(Debug, Clone)]
58pub struct LayerMetrics {
59 pub layer: LayerKind,
60 pub input_tokens: usize,
61 pub output_tokens: usize,
62 pub duration_us: u64,
63 pub compression_ratio: f64,
64}
65
66impl LayerMetrics {
67 pub fn new(
68 layer: LayerKind,
69 input_tokens: usize,
70 output_tokens: usize,
71 duration_us: u64,
72 ) -> Self {
73 let ratio = if input_tokens == 0 {
74 1.0
75 } else {
76 output_tokens as f64 / input_tokens as f64
77 };
78 Self {
79 layer,
80 input_tokens,
81 output_tokens,
82 duration_us,
83 compression_ratio: ratio,
84 }
85 }
86}
87
88pub trait Layer {
89 fn kind(&self) -> LayerKind;
90 fn process(&self, input: LayerInput) -> LayerOutput;
91}
92
93pub struct Pipeline {
94 layers: Vec<Box<dyn Layer>>,
95}
96
97impl Pipeline {
98 pub fn new() -> Self {
99 Self { layers: Vec::new() }
100 }
101
102 pub fn add_layer(mut self, layer: Box<dyn Layer>) -> Self {
103 self.layers.push(layer);
104 self
105 }
106
107 pub fn execute(&self, input: LayerInput) -> (LayerOutput, Vec<LayerMetrics>) {
108 let mut current = input;
109 let mut metrics = Vec::new();
110
111 for layer in &self.layers {
112 let start = std::time::Instant::now();
113 let input_tokens = current.tokens;
114 let output = layer.process(current);
115 let duration = start.elapsed().as_micros() as u64;
116
117 metrics.push(LayerMetrics::new(
118 layer.kind(),
119 input_tokens,
120 output.tokens,
121 duration,
122 ));
123
124 current = LayerInput {
125 content: output.content,
126 tokens: output.tokens,
127 metadata: output.metadata,
128 };
129 }
130
131 let final_output = LayerOutput {
132 content: current.content,
133 tokens: current.tokens,
134 metadata: current.metadata,
135 };
136
137 (final_output, metrics)
138 }
139
140 pub fn format_metrics(metrics: &[LayerMetrics]) -> String {
141 let mut out = String::from("Pipeline Metrics:\n");
142 let mut total_saved = 0usize;
143 for m in metrics {
144 let saved = m.input_tokens.saturating_sub(m.output_tokens);
145 total_saved += saved;
146 out.push_str(&format!(
147 " {} : {} -> {} tok ({:.0}%, {:.1}ms)\n",
148 m.layer,
149 m.input_tokens,
150 m.output_tokens,
151 m.compression_ratio * 100.0,
152 m.duration_us as f64 / 1000.0,
153 ));
154 }
155 if let (Some(first), Some(last)) = (metrics.first(), metrics.last()) {
156 let total_ratio = if first.input_tokens == 0 {
157 1.0
158 } else {
159 last.output_tokens as f64 / first.input_tokens as f64
160 };
161 out.push_str(&format!(
162 " TOTAL: {} -> {} tok ({:.0}%, saved {})\n",
163 first.input_tokens,
164 last.output_tokens,
165 total_ratio * 100.0,
166 total_saved,
167 ));
168 }
169 out
170 }
171}
172
173#[derive(Debug, Clone, Default)]
174pub struct PipelineStats {
175 pub runs: usize,
176 pub per_layer: HashMap<LayerKind, AggregatedMetrics>,
177}
178
179#[derive(Debug, Clone, Default)]
180pub struct AggregatedMetrics {
181 pub total_input_tokens: usize,
182 pub total_output_tokens: usize,
183 pub total_duration_us: u64,
184 pub count: usize,
185}
186
187impl AggregatedMetrics {
188 pub fn avg_ratio(&self) -> f64 {
189 if self.total_input_tokens == 0 {
190 return 1.0;
191 }
192 self.total_output_tokens as f64 / self.total_input_tokens as f64
193 }
194
195 pub fn avg_duration_ms(&self) -> f64 {
196 if self.count == 0 {
197 return 0.0;
198 }
199 self.total_duration_us as f64 / self.count as f64 / 1000.0
200 }
201}
202
203impl PipelineStats {
204 pub fn record(&mut self, metrics: &[LayerMetrics]) {
205 self.runs += 1;
206 for m in metrics {
207 let agg = self.per_layer.entry(m.layer).or_default();
208 agg.total_input_tokens += m.input_tokens;
209 agg.total_output_tokens += m.output_tokens;
210 agg.total_duration_us += m.duration_us;
211 agg.count += 1;
212 }
213 }
214
215 pub fn total_tokens_saved(&self) -> usize {
216 self.per_layer
217 .values()
218 .map(|a| a.total_input_tokens.saturating_sub(a.total_output_tokens))
219 .sum()
220 }
221
222 pub fn format_summary(&self) -> String {
223 let mut out = format!("Pipeline Stats ({} runs):\n", self.runs);
224 for kind in LayerKind::all() {
225 if let Some(agg) = self.per_layer.get(kind) {
226 out.push_str(&format!(
227 " {}: avg {:.0}% ratio, {:.1}ms, {} invocations\n",
228 kind,
229 agg.avg_ratio() * 100.0,
230 agg.avg_duration_ms(),
231 agg.count,
232 ));
233 }
234 }
235 out.push_str(&format!(" SAVED: {} tokens\n", self.total_tokens_saved()));
236 out
237 }
238}
239
240impl Default for Pipeline {
241 fn default() -> Self {
242 Self::new()
243 }
244}
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249
250 struct PassthroughLayer {
251 kind: LayerKind,
252 }
253
254 impl Layer for PassthroughLayer {
255 fn kind(&self) -> LayerKind {
256 self.kind
257 }
258
259 fn process(&self, input: LayerInput) -> LayerOutput {
260 LayerOutput {
261 content: input.content,
262 tokens: input.tokens,
263 metadata: input.metadata,
264 }
265 }
266 }
267
268 struct CompressionLayer {
269 ratio: f64,
270 }
271
272 impl Layer for CompressionLayer {
273 fn kind(&self) -> LayerKind {
274 LayerKind::Compression
275 }
276
277 fn process(&self, input: LayerInput) -> LayerOutput {
278 let new_tokens = (input.tokens as f64 * self.ratio) as usize;
279 let truncated = if input.content.len() > new_tokens * 4 {
280 input.content[..new_tokens * 4].to_string()
281 } else {
282 input.content
283 };
284 LayerOutput {
285 content: truncated,
286 tokens: new_tokens,
287 metadata: input.metadata,
288 }
289 }
290 }
291
292 #[test]
293 fn layer_kind_all_ordered() {
294 let all = LayerKind::all();
295 assert_eq!(all.len(), 6);
296 assert_eq!(all[0], LayerKind::Input);
297 assert_eq!(all[5], LayerKind::Delivery);
298 }
299
300 #[test]
301 fn passthrough_preserves_content() {
302 let layer = PassthroughLayer {
303 kind: LayerKind::Input,
304 };
305 let input = LayerInput {
306 content: "hello world".to_string(),
307 tokens: 2,
308 metadata: HashMap::new(),
309 };
310 let output = layer.process(input);
311 assert_eq!(output.content, "hello world");
312 assert_eq!(output.tokens, 2);
313 }
314
315 #[test]
316 fn compression_layer_reduces() {
317 let layer = CompressionLayer { ratio: 0.5 };
318 let input = LayerInput {
319 content: "a ".repeat(100),
320 tokens: 100,
321 metadata: HashMap::new(),
322 };
323 let output = layer.process(input);
324 assert_eq!(output.tokens, 50);
325 }
326
327 #[test]
328 fn pipeline_chains_layers() {
329 let pipeline = Pipeline::new()
330 .add_layer(Box::new(PassthroughLayer {
331 kind: LayerKind::Input,
332 }))
333 .add_layer(Box::new(CompressionLayer { ratio: 0.5 }))
334 .add_layer(Box::new(PassthroughLayer {
335 kind: LayerKind::Delivery,
336 }));
337
338 let input = LayerInput {
339 content: "a ".repeat(100),
340 tokens: 100,
341 metadata: HashMap::new(),
342 };
343 let (output, metrics) = pipeline.execute(input);
344 assert_eq!(output.tokens, 50);
345 assert_eq!(metrics.len(), 3);
346 assert_eq!(metrics[0].layer, LayerKind::Input);
347 assert_eq!(metrics[1].layer, LayerKind::Compression);
348 assert_eq!(metrics[2].layer, LayerKind::Delivery);
349 }
350
351 #[test]
352 fn metrics_new_calculates_ratio() {
353 let m = LayerMetrics::new(LayerKind::Compression, 100, 50, 1000);
354 assert!((m.compression_ratio - 0.5).abs() < f64::EPSILON);
355 }
356
357 #[test]
358 fn metrics_format_readable() {
359 let metrics = vec![
360 LayerMetrics::new(LayerKind::Input, 1000, 1000, 100),
361 LayerMetrics::new(LayerKind::Compression, 1000, 300, 5000),
362 LayerMetrics::new(LayerKind::Delivery, 300, 300, 50),
363 ];
364 let formatted = Pipeline::format_metrics(&metrics);
365 assert!(formatted.contains("input"));
366 assert!(formatted.contains("compression"));
367 assert!(formatted.contains("delivery"));
368 assert!(formatted.contains("TOTAL"));
369 }
370
371 #[test]
372 fn empty_pipeline_passes_through() {
373 let pipeline = Pipeline::new();
374 let input = LayerInput {
375 content: "test".to_string(),
376 tokens: 1,
377 metadata: HashMap::new(),
378 };
379 let (output, metrics) = pipeline.execute(input);
380 assert_eq!(output.content, "test");
381 assert!(metrics.is_empty());
382 }
383
384 #[test]
385 fn pipeline_stats_record_and_summarize() {
386 let mut stats = PipelineStats::default();
387 let metrics = vec![
388 LayerMetrics::new(LayerKind::Input, 1000, 1000, 100),
389 LayerMetrics::new(LayerKind::Compression, 1000, 300, 5000),
390 LayerMetrics::new(LayerKind::Delivery, 300, 300, 50),
391 ];
392 stats.record(&metrics);
393 stats.record(&metrics);
394
395 assert_eq!(stats.runs, 2);
396 assert_eq!(stats.total_tokens_saved(), 1400);
397
398 let agg = stats.per_layer.get(&LayerKind::Compression).unwrap();
399 assert_eq!(agg.count, 2);
400 assert_eq!(agg.total_input_tokens, 2000);
401 assert_eq!(agg.total_output_tokens, 600);
402
403 let summary = stats.format_summary();
404 assert!(summary.contains("2 runs"));
405 assert!(summary.contains("SAVED: 1400"));
406 }
407
408 #[test]
409 fn aggregated_metrics_avg() {
410 let mut agg = AggregatedMetrics::default();
411 agg.total_input_tokens = 1000;
412 agg.total_output_tokens = 500;
413 agg.total_duration_us = 10000;
414 agg.count = 2;
415 assert!((agg.avg_ratio() - 0.5).abs() < f64::EPSILON);
416 assert!((agg.avg_duration_ms() - 5.0).abs() < f64::EPSILON);
417 }
418}