llm_trait/
types.rs

1//! Usage info types.
2
3/// Token usage information for an LLM call.
4#[derive(Clone, Debug, Default)]
5pub struct UsageInfo {
6    pub prompt_tokens: Option<u32>,
7    pub completion_tokens: Option<u32>,
8    pub total_tokens: Option<u32>,
9    /// Tokens used for reasoning / thinking (DeepSeek, OpenAI o-series, etc.).
10    /// Always `None` on the Anthropic protocol: its usage object has no
11    /// thinking breakdown — thinking tokens are folded into
12    /// `completion_tokens`.
13    pub reasoning_tokens: Option<u32>,
14}
15
16impl UsageInfo {
17    /// Merge a partial usage event into `self`; the incoming `Some` field
18    /// wins over the existing value, a `None` leaves it untouched.
19    ///
20    /// Streaming protocols may split usage across several events (Anthropic
21    /// sends `prompt_tokens` in `message_start` and the final
22    /// `completion_tokens` in `message_delta`). Replacing the whole struct
23    /// on each event would silently zero out fields the last event omits,
24    /// so consumers must fold events with this method instead.
25    pub fn merge(&mut self, other: &UsageInfo) {
26        if other.prompt_tokens.is_some() {
27            self.prompt_tokens = other.prompt_tokens;
28        }
29        if other.completion_tokens.is_some() {
30            self.completion_tokens = other.completion_tokens;
31        }
32        if other.total_tokens.is_some() {
33            self.total_tokens = other.total_tokens;
34        }
35        if other.reasoning_tokens.is_some() {
36            self.reasoning_tokens = other.reasoning_tokens;
37        }
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44
45    #[test]
46    fn merge_keeps_existing_when_incoming_none() {
47        let mut acc = UsageInfo {
48            prompt_tokens: Some(5),
49            completion_tokens: Some(2),
50            total_tokens: Some(7),
51            reasoning_tokens: None,
52        };
53        // Anthropic `message_delta`: only the final output count arrives.
54        acc.merge(&UsageInfo {
55            prompt_tokens: None,
56            completion_tokens: Some(26),
57            total_tokens: None,
58            reasoning_tokens: None,
59        });
60        assert_eq!(
61            acc.prompt_tokens,
62            Some(5),
63            "prompt survives the partial event"
64        );
65        assert_eq!(acc.completion_tokens, Some(26), "later Some wins");
66        assert_eq!(acc.total_tokens, Some(7));
67    }
68
69    #[test]
70    fn merge_lets_later_some_override() {
71        // DashScope's anthropic endpoint reports the full `input_tokens` in
72        // `message_delta` (context so far), superseding the `message_start`
73        // value (latest message only) — the newer number is more useful.
74        let mut acc = UsageInfo {
75            prompt_tokens: Some(2),
76            ..Default::default()
77        };
78        acc.merge(&UsageInfo {
79            prompt_tokens: Some(63),
80            ..Default::default()
81        });
82        assert_eq!(acc.prompt_tokens, Some(63));
83    }
84}
llm_trait/types.rs

llm_trait/
types.rs
