mojentic/tracer/
tracer_system.rs

1//! Tracer system for coordinating tracer events
2//!
3//! This module provides the central system for recording, filtering, and querying
4//! tracer events. It coordinates with the event store and provides convenience methods
5//! for recording different types of events.
6
7use super::event_store::EventStore;
8use super::tracer_events::*;
9use std::collections::HashMap;
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::Arc;
12use std::time::{SystemTime, UNIX_EPOCH};
13
14/// Central system for capturing and querying tracer events
15///
16/// The TracerSystem is responsible for recording events related to LLM calls,
17/// tool usage, and agent interactions, providing a way to trace through the
18/// major events of the system.
19pub struct TracerSystem {
20    event_store: Arc<EventStore>,
21    enabled: Arc<AtomicBool>,
22}
23
24impl TracerSystem {
25    /// Create a new tracer system
26    ///
27    /// # Arguments
28    ///
29    /// * `event_store` - Optional event store to use. If None, a new one will be created.
30    /// * `enabled` - Whether the tracer system is enabled (default: true)
31    pub fn new(event_store: Option<Arc<EventStore>>, enabled: bool) -> Self {
32        Self {
33            event_store: event_store.unwrap_or_else(|| Arc::new(EventStore::default())),
34            enabled: Arc::new(AtomicBool::new(enabled)),
35        }
36    }
37
38    /// Check if the tracer is enabled
39    pub fn is_enabled(&self) -> bool {
40        self.enabled.load(Ordering::SeqCst)
41    }
42
43    /// Enable the tracer system
44    pub fn enable(&self) {
45        self.enabled.store(true, Ordering::SeqCst);
46    }
47
48    /// Disable the tracer system
49    pub fn disable(&self) {
50        self.enabled.store(false, Ordering::SeqCst);
51    }
52
53    /// Record a tracer event in the event store
54    ///
55    /// # Arguments
56    ///
57    /// * `event` - The tracer event to record
58    pub fn record_event(&self, event: Box<dyn TracerEvent>) {
59        if !self.is_enabled() {
60            return;
61        }
62        self.event_store.store(event);
63    }
64
65    /// Record an LLM call event
66    ///
67    /// # Arguments
68    ///
69    /// * `model` - The name of the LLM model being called
70    /// * `messages` - The messages sent to the LLM (simplified representation)
71    /// * `temperature` - The temperature setting for the LLM call
72    /// * `tools` - The tools available to the LLM, if any
73    /// * `source` - The source of the event
74    /// * `correlation_id` - UUID string for tracing related events
75    pub fn record_llm_call(
76        &self,
77        model: impl Into<String>,
78        messages: Vec<HashMap<String, serde_json::Value>>,
79        temperature: f64,
80        tools: Option<Vec<HashMap<String, serde_json::Value>>>,
81        source: impl Into<String>,
82        correlation_id: impl Into<String>,
83    ) {
84        if !self.is_enabled() {
85            return;
86        }
87
88        let event = Box::new(LlmCallTracerEvent {
89            timestamp: current_timestamp(),
90            correlation_id: correlation_id.into(),
91            source: source.into(),
92            model: model.into(),
93            messages,
94            temperature,
95            tools,
96        });
97
98        self.event_store.store(event);
99    }
100
101    /// Record an LLM response event
102    ///
103    /// # Arguments
104    ///
105    /// * `model` - The name of the LLM model that responded
106    /// * `content` - The content of the LLM response
107    /// * `tool_calls` - Any tool calls made by the LLM in its response
108    /// * `call_duration_ms` - The duration of the LLM call in milliseconds
109    /// * `source` - The source of the event
110    /// * `correlation_id` - UUID string for tracing related events
111    pub fn record_llm_response(
112        &self,
113        model: impl Into<String>,
114        content: impl Into<String>,
115        tool_calls: Option<Vec<HashMap<String, serde_json::Value>>>,
116        call_duration_ms: Option<f64>,
117        source: impl Into<String>,
118        correlation_id: impl Into<String>,
119    ) {
120        if !self.is_enabled() {
121            return;
122        }
123
124        let event = Box::new(LlmResponseTracerEvent {
125            timestamp: current_timestamp(),
126            correlation_id: correlation_id.into(),
127            source: source.into(),
128            model: model.into(),
129            content: content.into(),
130            tool_calls,
131            call_duration_ms,
132        });
133
134        self.event_store.store(event);
135    }
136
137    /// Record a tool call event
138    ///
139    /// # Arguments
140    ///
141    /// * `tool_name` - The name of the tool being called
142    /// * `arguments` - The arguments provided to the tool
143    /// * `result` - The result returned by the tool
144    /// * `caller` - The name of the agent or component calling the tool
145    /// * `call_duration_ms` - The duration of the tool call in milliseconds
146    /// * `source` - The source of the event
147    /// * `correlation_id` - UUID string for tracing related events
148    #[allow(clippy::too_many_arguments)]
149    pub fn record_tool_call(
150        &self,
151        tool_name: impl Into<String>,
152        arguments: HashMap<String, serde_json::Value>,
153        result: serde_json::Value,
154        caller: Option<String>,
155        call_duration_ms: Option<f64>,
156        source: impl Into<String>,
157        correlation_id: impl Into<String>,
158    ) {
159        if !self.is_enabled() {
160            return;
161        }
162
163        let event = Box::new(ToolCallTracerEvent {
164            timestamp: current_timestamp(),
165            correlation_id: correlation_id.into(),
166            source: source.into(),
167            tool_name: tool_name.into(),
168            arguments,
169            result,
170            caller,
171            call_duration_ms,
172        });
173
174        self.event_store.store(event);
175    }
176
177    /// Record an agent interaction event
178    ///
179    /// # Arguments
180    ///
181    /// * `from_agent` - The name of the agent sending the event
182    /// * `to_agent` - The name of the agent receiving the event
183    /// * `event_type` - The type of event being processed
184    /// * `event_id` - A unique identifier for the event
185    /// * `source` - The source of the event
186    /// * `correlation_id` - UUID string for tracing related events
187    pub fn record_agent_interaction(
188        &self,
189        from_agent: impl Into<String>,
190        to_agent: impl Into<String>,
191        event_type: impl Into<String>,
192        event_id: Option<String>,
193        source: impl Into<String>,
194        correlation_id: impl Into<String>,
195    ) {
196        if !self.is_enabled() {
197            return;
198        }
199
200        let event = Box::new(AgentInteractionTracerEvent {
201            timestamp: current_timestamp(),
202            correlation_id: correlation_id.into(),
203            source: source.into(),
204            from_agent: from_agent.into(),
205            to_agent: to_agent.into(),
206            event_type: event_type.into(),
207            event_id,
208        });
209
210        self.event_store.store(event);
211    }
212
213    /// Get event summaries from the store, optionally filtered
214    ///
215    /// # Arguments
216    ///
217    /// * `start_time` - Include events with timestamp >= start_time
218    /// * `end_time` - Include events with timestamp <= end_time
219    /// * `filter_func` - Custom filter function to apply to events
220    ///
221    /// # Returns
222    ///
223    /// Vector of event summaries matching the filter criteria
224    pub fn get_event_summaries(
225        &self,
226        start_time: Option<f64>,
227        end_time: Option<f64>,
228        filter_func: Option<&dyn super::EventFilterFn>,
229    ) -> Vec<String> {
230        self.event_store.get_event_summaries(start_time, end_time, filter_func)
231    }
232
233    /// Get the last N event summaries, optionally filtered
234    ///
235    /// # Arguments
236    ///
237    /// * `n` - Number of events to return
238    /// * `filter_func` - Optional custom filter function
239    ///
240    /// # Returns
241    ///
242    /// Vector of the last N event summaries matching the filter criteria
243    pub fn get_last_n_summaries(
244        &self,
245        n: usize,
246        filter_func: Option<&dyn super::EventFilterFn>,
247    ) -> Vec<String> {
248        self.event_store.get_last_n_summaries(n, filter_func)
249    }
250
251    /// Count events matching filters
252    ///
253    /// # Arguments
254    ///
255    /// * `start_time` - Include events with timestamp >= start_time
256    /// * `end_time` - Include events with timestamp <= end_time
257    /// * `filter_func` - Custom filter function to apply to events
258    ///
259    /// # Returns
260    ///
261    /// Number of events matching the filter criteria
262    pub fn count_events(
263        &self,
264        start_time: Option<f64>,
265        end_time: Option<f64>,
266        filter_func: Option<&dyn super::EventFilterFn>,
267    ) -> usize {
268        self.event_store.count_events(start_time, end_time, filter_func)
269    }
270
271    /// Clear all events from the event store
272    pub fn clear(&self) {
273        self.event_store.clear();
274    }
275
276    /// Get the total number of events in the store
277    pub fn len(&self) -> usize {
278        self.event_store.len()
279    }
280
281    /// Check if the event store is empty
282    pub fn is_empty(&self) -> bool {
283        self.event_store.is_empty()
284    }
285}
286
287impl Default for TracerSystem {
288    fn default() -> Self {
289        Self::new(None, true)
290    }
291}
292
293/// Get current timestamp as Unix timestamp (seconds since epoch)
294fn current_timestamp() -> f64 {
295    SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs_f64()
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    #[test]
303    fn test_new_tracer_system() {
304        let tracer = TracerSystem::default();
305        assert!(tracer.is_enabled());
306        assert_eq!(tracer.len(), 0);
307    }
308
309    #[test]
310    fn test_enable_disable() {
311        let tracer = TracerSystem::default();
312        assert!(tracer.is_enabled());
313
314        tracer.disable();
315        assert!(!tracer.is_enabled());
316
317        tracer.enable();
318        assert!(tracer.is_enabled());
319    }
320
321    #[test]
322    fn test_record_llm_call() {
323        let tracer = TracerSystem::default();
324
325        tracer.record_llm_call("llama3.2", vec![], 0.7, None, "test", "corr-123");
326
327        assert_eq!(tracer.len(), 1);
328    }
329
330    #[test]
331    fn test_record_llm_response() {
332        let tracer = TracerSystem::default();
333
334        tracer.record_llm_response(
335            "llama3.2",
336            "Hello, world!",
337            None,
338            Some(150.5),
339            "test",
340            "corr-456",
341        );
342
343        assert_eq!(tracer.len(), 1);
344    }
345
346    #[test]
347    fn test_record_tool_call() {
348        let tracer = TracerSystem::default();
349        let mut args = HashMap::new();
350        args.insert("input".to_string(), serde_json::json!("test"));
351
352        tracer.record_tool_call(
353            "example_tool",
354            args,
355            serde_json::json!({"output": "result"}),
356            Some("agent1".to_string()),
357            Some(25.0),
358            "test",
359            "corr-789",
360        );
361
362        assert_eq!(tracer.len(), 1);
363    }
364
365    #[test]
366    fn test_record_agent_interaction() {
367        let tracer = TracerSystem::default();
368
369        tracer.record_agent_interaction(
370            "agent1",
371            "agent2",
372            "message",
373            Some("evt-123".to_string()),
374            "test",
375            "corr-abc",
376        );
377
378        assert_eq!(tracer.len(), 1);
379    }
380
381    #[test]
382    fn test_disabled_tracer_doesnt_record() {
383        let tracer = TracerSystem::new(None, false);
384        assert!(!tracer.is_enabled());
385
386        tracer.record_llm_call("llama3.2", vec![], 1.0, None, "test", "corr-123");
387
388        assert_eq!(tracer.len(), 0);
389    }
390
391    #[test]
392    fn test_clear() {
393        let tracer = TracerSystem::default();
394
395        tracer.record_llm_call("llama3.2", vec![], 1.0, None, "test", "corr-123");
396
397        assert_eq!(tracer.len(), 1);
398
399        tracer.clear();
400        assert_eq!(tracer.len(), 0);
401        assert!(tracer.is_empty());
402    }
403
404    #[test]
405    fn test_multiple_events() {
406        let tracer = TracerSystem::default();
407
408        for i in 0..5 {
409            tracer.record_llm_call("llama3.2", vec![], 1.0, None, "test", format!("corr-{}", i));
410        }
411
412        assert_eq!(tracer.len(), 5);
413    }
414}
mojentic/tracer/tracer_system.rs

mojentic/tracer/
tracer_system.rs
