atman_runtime/
cost.rs

1use std::collections::HashMap;
2
3use crate::event::{Event, LlmCallStatus};
4use crate::provider::TokenUsage;
5
6#[derive(Debug, Default, Clone)]
7pub struct CostSummary {
8    pub calls: u64,
9    pub failures: u64,
10    pub usage: TokenUsage,
11    pub wallclock_ms: u64,
12}
13
14impl CostSummary {
15    fn accumulate(&mut self, usage: &TokenUsage, wallclock_ms: u64, failed: bool) {
16        self.calls += 1;
17        if failed {
18            self.failures += 1;
19        }
20        self.usage.input = self.usage.input.saturating_add(usage.input);
21        self.usage.cached_input = self.usage.cached_input.saturating_add(usage.cached_input);
22        self.usage.output = self.usage.output.saturating_add(usage.output);
23        self.usage.cache_write = self.usage.cache_write.saturating_add(usage.cache_write);
24        self.wallclock_ms = self.wallclock_ms.saturating_add(wallclock_ms);
25    }
26}
27
28pub fn summarize_by_model(events: &[Event]) -> HashMap<String, CostSummary> {
29    let mut out: HashMap<String, CostSummary> = HashMap::new();
30    for e in events {
31        if let Event::LlmCall {
32            model,
33            usage,
34            wallclock_ms,
35            status,
36            ..
37        } = e
38        {
39            let entry = out.entry(model.clone()).or_default();
40            entry.accumulate(
41                usage,
42                *wallclock_ms,
43                matches!(status, LlmCallStatus::Errored { .. }),
44            );
45        }
46    }
47    out
48}
49
50pub fn summarize_by_provider(events: &[Event]) -> HashMap<String, CostSummary> {
51    let mut out: HashMap<String, CostSummary> = HashMap::new();
52    for e in events {
53        if let Event::LlmCall {
54            provider,
55            usage,
56            wallclock_ms,
57            status,
58            ..
59        } = e
60        {
61            let entry = out.entry(provider.clone()).or_default();
62            entry.accumulate(
63                usage,
64                *wallclock_ms,
65                matches!(status, LlmCallStatus::Errored { .. }),
66            );
67        }
68    }
69    out
70}
71
72pub fn total(events: &[Event]) -> CostSummary {
73    let mut acc = CostSummary::default();
74    for e in events {
75        if let Event::LlmCall {
76            usage,
77            wallclock_ms,
78            status,
79            ..
80        } = e
81        {
82            acc.accumulate(
83                usage,
84                *wallclock_ms,
85                matches!(status, LlmCallStatus::Errored { .. }),
86            );
87        }
88    }
89    acc
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    fn ok_call(model: &str, provider: &str, in_tok: u64, out_tok: u64) -> Event {
97        Event::LlmCall {
98            model: model.into(),
99            provider: provider.into(),
100            usage: TokenUsage {
101                input: in_tok,
102                output: out_tok,
103                ..Default::default()
104            },
105            wallclock_ms: 100,
106            status: LlmCallStatus::Ok,
107            ttft_ms: None,
108            tokens_per_second: None,
109            run_id: None,
110            node_id: None,
111        }
112    }
113
114    fn err_call(model: &str) -> Event {
115        Event::LlmCall {
116            model: model.into(),
117            provider: "p".into(),
118            usage: TokenUsage {
119                input: 5,
120                ..Default::default()
121            },
122            wallclock_ms: 50,
123            status: LlmCallStatus::Errored {
124                message: "boom".into(),
125            },
126            ttft_ms: None,
127            tokens_per_second: None,
128            run_id: None,
129            node_id: None,
130        }
131    }
132
133    #[test]
134    fn summarize_by_model_groups_multiple_calls() {
135        let events = vec![
136            ok_call("opus", "anthropic", 10, 20),
137            ok_call("opus", "anthropic", 5, 8),
138            ok_call("mini", "openai", 3, 4),
139        ];
140        let s = summarize_by_model(&events);
141        assert_eq!(s.get("opus").unwrap().calls, 2);
142        assert_eq!(s.get("opus").unwrap().usage.input, 15);
143        assert_eq!(s.get("opus").unwrap().usage.output, 28);
144        assert_eq!(s.get("mini").unwrap().calls, 1);
145    }
146
147    #[test]
148    fn failures_counted_separately() {
149        let events = vec![ok_call("m", "p", 1, 2), err_call("m"), err_call("m")];
150        let s = summarize_by_model(&events);
151        assert_eq!(s.get("m").unwrap().calls, 3);
152        assert_eq!(s.get("m").unwrap().failures, 2);
153    }
154
155    #[test]
156    fn total_sums_wallclock() {
157        let events = vec![ok_call("a", "p", 1, 1), ok_call("b", "p", 2, 2)];
158        let t = total(&events);
159        assert_eq!(t.calls, 2);
160        assert_eq!(t.wallclock_ms, 200);
161    }
162
163    #[test]
164    fn non_llm_events_ignored() {
165        use crate::event::{FlowRunId, FlowStatus};
166        let run_id = FlowRunId::now();
167        let events = vec![
168            Event::FlowStart {
169                run_id: run_id.clone(),
170                flow_name: "t".into(),
171                parent_run_id: None,
172                parent_node_id: None,
173                spawned: false,
174            },
175            ok_call("m", "p", 5, 5),
176            Event::FlowEnd {
177                run_id,
178                flow_name: "t".into(),
179                status: FlowStatus::Ok,
180            },
181        ];
182        let t = total(&events);
183        assert_eq!(t.calls, 1);
184        assert_eq!(t.usage.input, 5);
185    }
186}
atman_runtime/cost.rs

atman_runtime/
cost.rs
