Skip to main content

codetether_rlm/context_trace/
trace_summary.rs

1use super::{ContextTrace, ContextTraceSummary};
2use std::collections::HashMap;
3
4impl ContextTrace {
5    /// Get a summary of the trace.
6    pub fn summary(&self) -> ContextTraceSummary {
7        let mut event_counts = HashMap::new();
8        let mut event_tokens = HashMap::new();
9        self.count_events(&mut event_counts, &mut event_tokens);
10        ContextTraceSummary {
11            total_tokens: self.total_tokens,
12            max_tokens: self.max_tokens,
13            budget_used_percent: self.budget_used_percent(),
14            iteration: self.iteration,
15            event_counts,
16            event_tokens,
17            events_len: self.events.len(),
18        }
19    }
20
21    fn count_events(
22        &self,
23        event_counts: &mut HashMap<String, usize>,
24        event_tokens: &mut HashMap<String, usize>,
25    ) {
26        for event in &self.events {
27            let label = event.label().to_string();
28            *event_counts.entry(label.clone()).or_insert(0) += 1;
29            *event_tokens.entry(label).or_insert(0) += event.tokens();
30        }
31    }
32}