submilli_engine/runtime/
model_calls.rs

1//! Per-run model usage an embedder reads after guest code finishes.
2//!
3//! `submilli:llm` and `submilli:embedding` hand token counts back to the program.
4//! The embedder that has to price the run does not see that value. It owns the
5//! store, so the live dispatch appends a record here. Replay restores a recorded
6//! outcome without running the provider action, and therefore must not append:
7//! recording after `effects::execute` returns would bill the same call again.
8
9use super::call_log::ModelUsage;
10
11/// One live model dispatch in this attempt, in call order.
12///
13/// `model` is the alias the program passed, not a provider id and not a pinned
14/// revision. A missing count is indeterminate, not zero.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct ModelCallUsage {
17    pub model: String,
18    pub usage: AttributedUsage,
19}
20
21/// What the provider reported for one dispatch.
22///
23/// Embeddings are their own variant so a missing completion count — embeddings
24/// have none — is not summed as unknown output.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum AttributedUsage {
27    Llm {
28        input_tokens: Option<u64>,
29        output_tokens: Option<u64>,
30    },
31    Embedding {
32        input_tokens: Option<u64>,
33    },
34}
35
36impl ModelCallUsage {
37    pub(crate) fn llm(model: impl Into<String>, usage: ModelUsage) -> Self {
38        Self {
39            model: model.into(),
40            usage: AttributedUsage::Llm {
41                input_tokens: usage.input_tokens,
42                output_tokens: usage.output_tokens,
43            },
44        }
45    }
46
47    pub(crate) fn embedding(model: impl Into<String>, input_tokens: Option<u64>) -> Self {
48        Self {
49            model: model.into(),
50            usage: AttributedUsage::Embedding { input_tokens },
51        }
52    }
53
54    /// One entry per `(kind, model)`, in the order each group first appeared.
55    ///
56    /// A missing count stays missing. It is not zero, and it is not added to a
57    /// known count. An embedding call never joins an llm group of the same alias.
58    pub fn summed_by_model(calls: &[Self]) -> Vec<Self> {
59        let mut totals = Vec::new();
60        for call in calls {
61            if let Some(total) = totals.iter_mut().find(|total| same_group(total, call)) {
62                add_same_kind(&mut total.usage, &call.usage);
63            } else {
64                totals.push(call.clone());
65            }
66        }
67        totals
68    }
69}
70
71fn same_group(left: &ModelCallUsage, right: &ModelCallUsage) -> bool {
72    left.model == right.model && same_kind(&left.usage, &right.usage)
73}
74
75fn same_kind(left: &AttributedUsage, right: &AttributedUsage) -> bool {
76    matches!(
77        (left, right),
78        (AttributedUsage::Llm { .. }, AttributedUsage::Llm { .. })
79            | (
80                AttributedUsage::Embedding { .. },
81                AttributedUsage::Embedding { .. }
82            )
83    )
84}
85
86/// `same_group` has already required the kinds to match. A mismatch leaves the
87/// total unchanged rather than inventing a cross-kind sum.
88fn add_same_kind(total: &mut AttributedUsage, next: &AttributedUsage) {
89    match (total, next) {
90        (
91            AttributedUsage::Llm {
92                input_tokens,
93                output_tokens,
94            },
95            AttributedUsage::Llm {
96                input_tokens: next_input,
97                output_tokens: next_output,
98            },
99        ) => {
100            *input_tokens = add_count(*input_tokens, *next_input);
101            *output_tokens = add_count(*output_tokens, *next_output);
102        }
103        (
104            AttributedUsage::Embedding { input_tokens },
105            AttributedUsage::Embedding {
106                input_tokens: next_input,
107            },
108        ) => *input_tokens = add_count(*input_tokens, *next_input),
109        _ => {}
110    }
111}
112
113fn add_count(left: Option<u64>, right: Option<u64>) -> Option<u64> {
114    Some(left?.saturating_add(right?))
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    fn llm(model: &str, input: Option<u64>, output: Option<u64>) -> ModelCallUsage {
122        ModelCallUsage::llm(
123            model,
124            ModelUsage {
125                input_tokens: input,
126                output_tokens: output,
127            },
128        )
129    }
130
131    #[test]
132    fn summed_by_model_adds_a_model_and_keeps_a_missing_count() {
133        let calls = [
134            llm("haiku", Some(10), Some(1)),
135            llm("sonnet", Some(3), Some(4)),
136            llm("haiku", None, Some(2)),
137        ];
138        assert_eq!(
139            ModelCallUsage::summed_by_model(&calls),
140            vec![llm("haiku", None, Some(3)), llm("sonnet", Some(3), Some(4)),]
141        );
142    }
143
144    #[test]
145    fn summed_by_model_does_not_merge_llm_and_embedding() {
146        let calls = [
147            llm("haiku", Some(10), Some(1)),
148            ModelCallUsage::embedding("haiku", Some(7)),
149            ModelCallUsage::embedding("haiku", None),
150        ];
151        assert_eq!(
152            ModelCallUsage::summed_by_model(&calls),
153            vec![
154                llm("haiku", Some(10), Some(1)),
155                ModelCallUsage::embedding("haiku", None),
156            ]
157        );
158    }
159}
submilli_engine/runtime/model_calls.rs

submilli_engine/runtime/
model_calls.rs
