llm_optimizer_processor/
core.rs

1//! Core event processing structures and traits
2//!
3//! This module provides the fundamental types for event processing in the stream processor:
4//! - ProcessorEvent: Wrapper around events with metadata
5//! - EventKey: Different types of aggregation keys
6//! - EventTimeExtractor: Trait for extracting event timestamps
7//! - KeyExtractor: Trait for extracting aggregation keys
8
9use chrono::{DateTime, Utc};
10use llm_optimizer_types::events::{CloudEvent, FeedbackEvent};
11use llm_optimizer_types::metrics::MetricPoint;
12use serde::{Deserialize, Serialize};
13use std::fmt;
14
15/// Wrapper around events with metadata for stream processing
16///
17/// This structure wraps any event type with additional metadata needed for
18/// processing, including event time, processing time, and aggregation key.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct ProcessorEvent<T> {
21    /// The actual event payload
22    pub event: T,
23    /// Timestamp when the event occurred (event time semantics)
24    pub event_time: DateTime<Utc>,
25    /// Timestamp when the event was processed (processing time semantics)
26    pub processing_time: DateTime<Utc>,
27    /// Aggregation key for grouping related events
28    pub key: String,
29}
30
31impl<T> ProcessorEvent<T> {
32    /// Create a new ProcessorEvent
33    pub fn new(event: T, event_time: DateTime<Utc>, key: String) -> Self {
34        Self {
35            event,
36            event_time,
37            processing_time: Utc::now(),
38            key,
39        }
40    }
41
42    /// Create a ProcessorEvent with a custom processing time
43    pub fn with_processing_time(
44        event: T,
45        event_time: DateTime<Utc>,
46        processing_time: DateTime<Utc>,
47        key: String,
48    ) -> Self {
49        Self {
50            event,
51            event_time,
52            processing_time,
53            key,
54        }
55    }
56
57    /// Calculate the processing delay (processing_time - event_time)
58    pub fn processing_delay(&self) -> chrono::Duration {
59        self.processing_time - self.event_time
60    }
61
62    /// Check if the event is late (processing delay exceeds threshold)
63    pub fn is_late(&self, threshold: chrono::Duration) -> bool {
64        self.processing_delay() > threshold
65    }
66
67    /// Map the event to a different type while preserving metadata
68    pub fn map<U, F>(self, f: F) -> ProcessorEvent<U>
69    where
70        F: FnOnce(T) -> U,
71    {
72        ProcessorEvent {
73            event: f(self.event),
74            event_time: self.event_time,
75            processing_time: self.processing_time,
76            key: self.key,
77        }
78    }
79
80    /// Get a reference to the event
81    pub fn event_ref(&self) -> &T {
82        &self.event
83    }
84
85    /// Get a mutable reference to the event
86    pub fn event_mut(&mut self) -> &mut T {
87        &mut self.event
88    }
89}
90
91impl<T: fmt::Display> fmt::Display for ProcessorEvent<T> {
92    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
93        write!(
94            f,
95            "ProcessorEvent {{ key: {}, event_time: {}, event: {} }}",
96            self.key, self.event_time, self.event
97        )
98    }
99}
100
101/// Enum representing different types of aggregation keys
102///
103/// Events can be grouped by different dimensions for aggregation:
104/// - By model for model-specific metrics
105/// - By session for user session analysis
106/// - By user for user-specific metrics
107/// - By composite keys for multi-dimensional aggregation
108#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
109#[serde(tag = "type", rename_all = "snake_case")]
110pub enum EventKey {
111    /// Key by model identifier
112    Model(String),
113    /// Key by session identifier
114    Session(String),
115    /// Key by user identifier
116    User(String),
117    /// Composite key with multiple dimensions
118    Composite(Vec<(String, String)>),
119}
120
121impl EventKey {
122    /// Convert the key to a string representation for use as aggregation key
123    pub fn to_key_string(&self) -> String {
124        match self {
125            EventKey::Model(id) => format!("model:{}", id),
126            EventKey::Session(id) => format!("session:{}", id),
127            EventKey::User(id) => format!("user:{}", id),
128            EventKey::Composite(parts) => {
129                let mut key_parts: Vec<String> = parts
130                    .iter()
131                    .map(|(k, v)| format!("{}:{}", k, v))
132                    .collect();
133                key_parts.sort(); // Ensure consistent ordering
134                format!("composite:{}", key_parts.join(","))
135            }
136        }
137    }
138
139    /// Create a composite key from multiple dimensions
140    pub fn composite(parts: Vec<(&str, &str)>) -> Self {
141        EventKey::Composite(
142            parts
143                .into_iter()
144                .map(|(k, v)| (k.to_string(), v.to_string()))
145                .collect(),
146        )
147    }
148
149    /// Extract a specific dimension from a composite key
150    pub fn get_dimension(&self, dimension: &str) -> Option<String> {
151        match self {
152            EventKey::Composite(parts) => parts
153                .iter()
154                .find(|(k, _)| k == dimension)
155                .map(|(_, v)| v.clone()),
156            _ => None,
157        }
158    }
159}
160
161impl fmt::Display for EventKey {
162    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
163        write!(f, "{}", self.to_key_string())
164    }
165}
166
167impl From<EventKey> for String {
168    fn from(key: EventKey) -> String {
169        key.to_key_string()
170    }
171}
172
173/// Trait for extracting event time from events
174///
175/// Implement this trait to define how to extract the timestamp from your event type.
176/// This is used for event time semantics in stream processing.
177pub trait EventTimeExtractor<T> {
178    /// Extract the event timestamp from the event
179    fn extract_event_time(&self, event: &T) -> DateTime<Utc>;
180}
181
182/// Trait for extracting aggregation keys from events
183///
184/// Implement this trait to define how to extract the aggregation key from your event type.
185/// This is used for grouping events in stream processing operations.
186pub trait KeyExtractor<T> {
187    /// Extract the aggregation key from the event
188    fn extract_key(&self, event: &T) -> EventKey;
189}
190
191// ============================================================================
192// Implementations for llm-optimizer-types events
193// ============================================================================
194
195/// Default event time extractor for FeedbackEvent
196#[derive(Debug, Clone, Default)]
197pub struct FeedbackEventTimeExtractor;
198
199impl EventTimeExtractor<FeedbackEvent> for FeedbackEventTimeExtractor {
200    fn extract_event_time(&self, event: &FeedbackEvent) -> DateTime<Utc> {
201        event.timestamp
202    }
203}
204
205/// Default event time extractor for CloudEvent
206#[derive(Debug, Clone, Default)]
207pub struct CloudEventTimeExtractor;
208
209impl EventTimeExtractor<CloudEvent> for CloudEventTimeExtractor {
210    fn extract_event_time(&self, event: &CloudEvent) -> DateTime<Utc> {
211        event.time
212    }
213}
214
215/// Default event time extractor for MetricPoint
216#[derive(Debug, Clone, Default)]
217pub struct MetricPointTimeExtractor;
218
219impl EventTimeExtractor<MetricPoint> for MetricPointTimeExtractor {
220    fn extract_event_time(&self, event: &MetricPoint) -> DateTime<Utc> {
221        event.timestamp
222    }
223}
224
225/// Key extractor for FeedbackEvent based on metadata
226///
227/// Extracts keys in the following priority:
228/// 1. If both model_id and session_id exist, create composite key
229/// 2. If only model_id exists, create Model key
230/// 3. If only session_id exists, create Session key
231/// 4. If user_id exists, create User key
232/// 5. Otherwise, use event ID as key
233#[derive(Debug, Clone, Default)]
234pub struct FeedbackEventKeyExtractor;
235
236impl KeyExtractor<FeedbackEvent> for FeedbackEventKeyExtractor {
237    fn extract_key(&self, event: &FeedbackEvent) -> EventKey {
238        let model_id = event.metadata.get("model_id");
239        let session_id = event.metadata.get("session_id");
240        let user_id = event.metadata.get("user_id");
241
242        match (model_id, session_id, user_id) {
243            (Some(m), Some(s), _) => {
244                EventKey::composite(vec![("model", m.as_str()), ("session", s.as_str())])
245            }
246            (Some(m), None, _) => EventKey::Model(m.clone()),
247            (None, Some(s), _) => EventKey::Session(s.clone()),
248            (None, None, Some(u)) => EventKey::User(u.clone()),
249            _ => EventKey::Session(event.id.to_string()),
250        }
251    }
252}
253
254/// Key extractor for CloudEvent based on source
255#[derive(Debug, Clone, Default)]
256pub struct CloudEventKeyExtractor;
257
258impl KeyExtractor<CloudEvent> for CloudEventKeyExtractor {
259    fn extract_key(&self, event: &CloudEvent) -> EventKey {
260        // Use source as the primary key dimension
261        EventKey::Session(event.source.clone())
262    }
263}
264
265/// Key extractor for MetricPoint based on tags
266///
267/// Extracts keys in the following priority:
268/// 1. If 'model' tag exists, create Model key
269/// 2. If 'session_id' tag exists, create Session key
270/// 3. If 'user_id' tag exists, create User key
271/// 4. Otherwise, use metric name as key
272#[derive(Debug, Clone, Default)]
273pub struct MetricPointKeyExtractor;
274
275impl KeyExtractor<MetricPoint> for MetricPointKeyExtractor {
276    fn extract_key(&self, event: &MetricPoint) -> EventKey {
277        let model = event.tags.get("model");
278        let session_id = event.tags.get("session_id");
279        let user_id = event.tags.get("user_id");
280
281        match (model, session_id, user_id) {
282            (Some(m), Some(s), _) => {
283                EventKey::composite(vec![("model", m.as_str()), ("session", s.as_str())])
284            }
285            (Some(m), None, _) => EventKey::Model(m.clone()),
286            (None, Some(s), _) => EventKey::Session(s.clone()),
287            (None, None, Some(u)) => EventKey::User(u.clone()),
288            _ => EventKey::Session(event.name.clone()),
289        }
290    }
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296    use llm_optimizer_types::events::{EventSource, EventType};
297    use std::collections::HashMap;
298
299    #[test]
300    fn test_processor_event_creation() {
301        let event_time = Utc::now();
302        let event = ProcessorEvent::new(
303            "test_event".to_string(),
304            event_time,
305            "test_key".to_string(),
306        );
307
308        assert_eq!(event.event, "test_event");
309        assert_eq!(event.event_time, event_time);
310        assert_eq!(event.key, "test_key");
311        assert!(event.processing_time >= event_time);
312    }
313
314    #[test]
315    fn test_processor_event_processing_delay() {
316        let event_time = Utc::now() - chrono::Duration::seconds(10);
317        let processing_time = Utc::now();
318
319        let event = ProcessorEvent::with_processing_time(
320            "test".to_string(),
321            event_time,
322            processing_time,
323            "key".to_string(),
324        );
325
326        let delay = event.processing_delay();
327        assert!(delay.num_seconds() >= 9 && delay.num_seconds() <= 11);
328    }
329
330    #[test]
331    fn test_processor_event_is_late() {
332        let event_time = Utc::now() - chrono::Duration::seconds(60);
333        let event = ProcessorEvent::new("test".to_string(), event_time, "key".to_string());
334
335        assert!(event.is_late(chrono::Duration::seconds(30)));
336        assert!(!event.is_late(chrono::Duration::seconds(120)));
337    }
338
339    #[test]
340    fn test_processor_event_map() {
341        let event = ProcessorEvent::new(42, Utc::now(), "key".to_string());
342        let mapped = event.map(|n| n.to_string());
343
344        assert_eq!(mapped.event, "42");
345        assert_eq!(mapped.key, "key");
346    }
347
348    #[test]
349    fn test_event_key_model() {
350        let key = EventKey::Model("claude-3-opus".to_string());
351        assert_eq!(key.to_key_string(), "model:claude-3-opus");
352    }
353
354    #[test]
355    fn test_event_key_session() {
356        let key = EventKey::Session("session-123".to_string());
357        assert_eq!(key.to_key_string(), "session:session-123");
358    }
359
360    #[test]
361    fn test_event_key_user() {
362        let key = EventKey::User("user-456".to_string());
363        assert_eq!(key.to_key_string(), "user:user-456");
364    }
365
366    #[test]
367    fn test_event_key_composite() {
368        let key = EventKey::composite(vec![("model", "claude-3-opus"), ("session", "sess-1")]);
369
370        let key_string = key.to_key_string();
371        // Should contain both parts
372        assert!(key_string.contains("model:claude-3-opus"));
373        assert!(key_string.contains("session:sess-1"));
374        assert!(key_string.starts_with("composite:"));
375    }
376
377    #[test]
378    fn test_event_key_composite_ordering() {
379        // Keys should be consistent regardless of input order
380        let key1 = EventKey::composite(vec![("a", "1"), ("b", "2")]);
381        let key2 = EventKey::composite(vec![("b", "2"), ("a", "1")]);
382
383        assert_eq!(key1.to_key_string(), key2.to_key_string());
384    }
385
386    #[test]
387    fn test_event_key_get_dimension() {
388        let key = EventKey::composite(vec![("model", "claude-3-opus"), ("session", "sess-1")]);
389
390        assert_eq!(
391            key.get_dimension("model"),
392            Some("claude-3-opus".to_string())
393        );
394        assert_eq!(key.get_dimension("session"), Some("sess-1".to_string()));
395        assert_eq!(key.get_dimension("user"), None);
396    }
397
398    #[test]
399    fn test_feedback_event_time_extractor() {
400        let timestamp = Utc::now();
401        let event = FeedbackEvent {
402            id: uuid::Uuid::new_v4(),
403            timestamp,
404            source: EventSource::Observatory,
405            event_type: EventType::Metric,
406            data: serde_json::json!({}),
407            metadata: HashMap::new(),
408        };
409
410        let extractor = FeedbackEventTimeExtractor;
411        assert_eq!(extractor.extract_event_time(&event), timestamp);
412    }
413
414    #[test]
415    fn test_feedback_event_key_extractor_model() {
416        let mut event = FeedbackEvent::new(
417            EventSource::Observatory,
418            EventType::Metric,
419            serde_json::json!({}),
420        );
421        event.metadata.insert("model_id".to_string(), "claude-3-opus".to_string());
422
423        let extractor = FeedbackEventKeyExtractor;
424        let key = extractor.extract_key(&event);
425
426        match key {
427            EventKey::Model(id) => assert_eq!(id, "claude-3-opus"),
428            _ => panic!("Expected Model key"),
429        }
430    }
431
432    #[test]
433    fn test_feedback_event_key_extractor_composite() {
434        let mut event = FeedbackEvent::new(
435            EventSource::Observatory,
436            EventType::Metric,
437            serde_json::json!({}),
438        );
439        event.metadata.insert("model_id".to_string(), "claude-3-opus".to_string());
440        event.metadata.insert("session_id".to_string(), "sess-123".to_string());
441
442        let extractor = FeedbackEventKeyExtractor;
443        let key = extractor.extract_key(&event);
444
445        match key {
446            EventKey::Composite(parts) => {
447                assert_eq!(parts.len(), 2);
448                assert!(parts.contains(&("model".to_string(), "claude-3-opus".to_string())));
449                assert!(parts.contains(&("session".to_string(), "sess-123".to_string())));
450            }
451            _ => panic!("Expected Composite key"),
452        }
453    }
454
455    #[test]
456    fn test_feedback_event_key_extractor_session() {
457        let mut event = FeedbackEvent::new(
458            EventSource::Observatory,
459            EventType::Metric,
460            serde_json::json!({}),
461        );
462        event.metadata.insert("session_id".to_string(), "sess-123".to_string());
463
464        let extractor = FeedbackEventKeyExtractor;
465        let key = extractor.extract_key(&event);
466
467        match key {
468            EventKey::Session(id) => assert_eq!(id, "sess-123"),
469            _ => panic!("Expected Session key"),
470        }
471    }
472
473    #[test]
474    fn test_feedback_event_key_extractor_user() {
475        let mut event = FeedbackEvent::new(
476            EventSource::User,
477            EventType::Feedback,
478            serde_json::json!({}),
479        );
480        event.metadata.insert("user_id".to_string(), "user-789".to_string());
481
482        let extractor = FeedbackEventKeyExtractor;
483        let key = extractor.extract_key(&event);
484
485        match key {
486            EventKey::User(id) => assert_eq!(id, "user-789"),
487            _ => panic!("Expected User key"),
488        }
489    }
490
491    #[test]
492    fn test_feedback_event_key_extractor_fallback() {
493        let event = FeedbackEvent::new(
494            EventSource::System,
495            EventType::PerformanceAlert,
496            serde_json::json!({}),
497        );
498
499        let extractor = FeedbackEventKeyExtractor;
500        let key = extractor.extract_key(&event);
501
502        match key {
503            EventKey::Session(id) => assert_eq!(id, event.id.to_string()),
504            _ => panic!("Expected Session key with event ID"),
505        }
506    }
507
508    #[test]
509    fn test_cloud_event_time_extractor() {
510        let time = Utc::now();
511        let event = CloudEvent {
512            specversion: "1.0".to_string(),
513            event_type: "test.event".to_string(),
514            source: "test-source".to_string(),
515            id: "test-id".to_string(),
516            time,
517            datacontenttype: "application/json".to_string(),
518            data: serde_json::json!({}),
519        };
520
521        let extractor = CloudEventTimeExtractor;
522        assert_eq!(extractor.extract_event_time(&event), time);
523    }
524
525    #[test]
526    fn test_cloud_event_key_extractor() {
527        let event = CloudEvent::new(
528            "test.event",
529            "llm-observatory",
530            serde_json::json!({}),
531        );
532
533        let extractor = CloudEventKeyExtractor;
534        let key = extractor.extract_key(&event);
535
536        match key {
537            EventKey::Session(id) => assert_eq!(id, "llm-observatory"),
538            _ => panic!("Expected Session key"),
539        }
540    }
541
542    #[test]
543    fn test_metric_point_time_extractor() {
544        let timestamp = Utc::now();
545        let point = MetricPoint {
546            name: "latency".to_string(),
547            value: 123.45,
548            unit: "ms".to_string(),
549            timestamp,
550            tags: HashMap::new(),
551            aggregation: None,
552        };
553
554        let extractor = MetricPointTimeExtractor;
555        assert_eq!(extractor.extract_event_time(&point), timestamp);
556    }
557
558    #[test]
559    fn test_metric_point_key_extractor_model() {
560        let mut tags = HashMap::new();
561        tags.insert("model".to_string(), "claude-3-opus".to_string());
562
563        let point = MetricPoint {
564            name: "latency".to_string(),
565            value: 100.0,
566            unit: "ms".to_string(),
567            timestamp: Utc::now(),
568            tags,
569            aggregation: None,
570        };
571
572        let extractor = MetricPointKeyExtractor;
573        let key = extractor.extract_key(&point);
574
575        match key {
576            EventKey::Model(id) => assert_eq!(id, "claude-3-opus"),
577            _ => panic!("Expected Model key"),
578        }
579    }
580
581    #[test]
582    fn test_metric_point_key_extractor_composite() {
583        let mut tags = HashMap::new();
584        tags.insert("model".to_string(), "claude-3-opus".to_string());
585        tags.insert("session_id".to_string(), "sess-123".to_string());
586
587        let point = MetricPoint {
588            name: "latency".to_string(),
589            value: 100.0,
590            unit: "ms".to_string(),
591            timestamp: Utc::now(),
592            tags,
593            aggregation: None,
594        };
595
596        let extractor = MetricPointKeyExtractor;
597        let key = extractor.extract_key(&point);
598
599        match key {
600            EventKey::Composite(parts) => {
601                assert_eq!(parts.len(), 2);
602                assert!(parts.contains(&("model".to_string(), "claude-3-opus".to_string())));
603                assert!(parts.contains(&("session".to_string(), "sess-123".to_string())));
604            }
605            _ => panic!("Expected Composite key"),
606        }
607    }
608
609    #[test]
610    fn test_metric_point_key_extractor_fallback() {
611        let point = MetricPoint::new("latency", 100.0, "ms");
612
613        let extractor = MetricPointKeyExtractor;
614        let key = extractor.extract_key(&point);
615
616        match key {
617            EventKey::Session(id) => assert_eq!(id, "latency"),
618            _ => panic!("Expected Session key with metric name"),
619        }
620    }
621
622    #[test]
623    fn test_processor_event_with_feedback_event() {
624        let feedback = FeedbackEvent::new(
625            EventSource::Observatory,
626            EventType::Metric,
627            serde_json::json!({"latency": 123.45}),
628        );
629
630        let time_extractor = FeedbackEventTimeExtractor;
631        let key_extractor = FeedbackEventKeyExtractor;
632
633        let event_time = time_extractor.extract_event_time(&feedback);
634        let key = key_extractor.extract_key(&feedback);
635
636        let processor_event = ProcessorEvent::new(feedback, event_time, key.to_key_string());
637
638        assert_eq!(processor_event.event_time, event_time);
639        assert!(processor_event.key.contains("session:"));
640    }
641
642    #[test]
643    fn test_processor_event_display() {
644        let event = ProcessorEvent::new("test".to_string(), Utc::now(), "key123".to_string());
645        let display = format!("{}", event);
646
647        assert!(display.contains("key123"));
648        assert!(display.contains("test"));
649    }
650
651    #[test]
652    fn test_event_key_into_string() {
653        let key = EventKey::Model("claude-3-opus".to_string());
654        let key_string: String = key.into();
655
656        assert_eq!(key_string, "model:claude-3-opus");
657    }
658}
llm_optimizer_processor/core.rs

llm_optimizer_processor/
core.rs
