1use serde::{Deserialize, Serialize};
12
13use crate::budget::Tokens;
14use crate::model::wire::DialectKind;
15
16#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
28#[serde(into = "UsageRow", try_from = "UsageRow")]
29pub struct ModelUsage {
30 pub input_tokens: Tokens,
31 pub output_tokens: Tokens,
32 pub cache_read_tokens: Tokens,
33 pub cache_write_tokens: Tokens,
34 pub dialect: Option<DialectKind>,
35}
36
37const USAGE_MEANING: u8 = 1;
41
42#[derive(Serialize, Deserialize)]
45struct UsageRow {
46 input_tokens: Tokens,
47 output_tokens: Tokens,
48 cache_read_tokens: Tokens,
49 cache_write_tokens: Tokens,
50 #[serde(default, skip_serializing_if = "Option::is_none")]
51 v: Option<u8>,
52 #[serde(default, skip_serializing_if = "Option::is_none")]
53 dialect: Option<DialectKind>,
54}
55
56impl From<ModelUsage> for UsageRow {
57 fn from(usage: ModelUsage) -> UsageRow {
58 UsageRow {
59 input_tokens: usage.input_tokens,
60 output_tokens: usage.output_tokens,
61 cache_read_tokens: usage.cache_read_tokens,
62 cache_write_tokens: usage.cache_write_tokens,
63 v: Some(USAGE_MEANING),
64 dialect: usage.dialect,
65 }
66 }
67}
68
69impl TryFrom<UsageRow> for ModelUsage {
70 type Error = String;
71
72 fn try_from(row: UsageRow) -> Result<ModelUsage, String> {
73 let input_tokens = match row.v {
74 Some(USAGE_MEANING) => row.input_tokens,
75 Some(other) => {
76 return Err(format!(
77 "usage row version {other} is newer than this build reads \
78 (it reads {USAGE_MEANING}); read the ledger with a newer sprawling"
79 ));
80 }
81 None => {
82 let cached = row
83 .cache_read_tokens
84 .checked_add(row.cache_write_tokens)
85 .ok_or("usage row cache counts overflow a token count")?;
86 if cached > row.input_tokens {
87 row.input_tokens
88 .checked_add(cached)
89 .ok_or("usage row counts overflow a token count")?
90 } else {
91 row.input_tokens
92 }
93 }
94 };
95 Ok(ModelUsage {
96 input_tokens,
97 output_tokens: row.output_tokens,
98 cache_read_tokens: row.cache_read_tokens,
99 cache_write_tokens: row.cache_write_tokens,
100 dialect: row.dialect,
101 })
102 }
103}