1use rig::completion::Usage;
2use serde::{Deserialize, Serialize};
3use std::ops::Add;
4
5#[derive(Clone, Debug, Default, Serialize, Deserialize)]
6pub struct LmUsage {
7 pub prompt_tokens: u64,
8 pub completion_tokens: u64,
9 pub total_tokens: u64,
10}
11
12impl From<Usage> for LmUsage {
13 fn from(usage: Usage) -> Self {
14 LmUsage {
15 prompt_tokens: usage.input_tokens,
16 completion_tokens: usage.output_tokens,
17 total_tokens: usage.total_tokens,
18 }
19 }
20}
21
22impl Add for LmUsage {
23 type Output = LmUsage;
24
25 fn add(self, other: LmUsage) -> Self {
26 LmUsage {
27 prompt_tokens: self.prompt_tokens + other.prompt_tokens,
28 completion_tokens: self.completion_tokens + other.completion_tokens,
29 total_tokens: self.total_tokens + other.total_tokens,
30 }
31 }
32}