1use serde::{Deserialize, Serialize};
12
13use crate::budget::Tokens;
14use crate::model::wire::DialectKind;
15
16#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
24pub enum CacheCount {
25 Reported(Tokens),
26 #[default]
27 Unreported,
28}
29
30impl CacheCount {
31 #[must_use]
33 pub const fn reported(self) -> Option<Tokens> {
34 match self {
35 CacheCount::Reported(tokens) => Some(tokens),
36 CacheCount::Unreported => None,
37 }
38 }
39
40 #[must_use]
45 pub const fn or_zero(self) -> Tokens {
46 match self {
47 CacheCount::Reported(tokens) => tokens,
48 CacheCount::Unreported => Tokens::new(0),
49 }
50 }
51}
52
53impl From<Option<Tokens>> for CacheCount {
54 fn from(count: Option<Tokens>) -> CacheCount {
55 match count {
56 Some(tokens) => CacheCount::Reported(tokens),
57 None => CacheCount::Unreported,
58 }
59 }
60}
61
62#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
75#[serde(into = "UsageRow", try_from = "UsageRow")]
76pub struct ModelUsage {
77 pub input_tokens: Tokens,
78 pub output_tokens: Tokens,
79 pub cache_read_tokens: CacheCount,
80 pub cache_write_tokens: CacheCount,
81 pub dialect: Option<DialectKind>,
82}
83
84const USAGE_MEANING: u8 = 1;
88
89#[derive(Serialize, Deserialize)]
92struct UsageRow {
93 input_tokens: Tokens,
94 output_tokens: Tokens,
95 #[serde(default, skip_serializing_if = "Option::is_none")]
96 cache_read_tokens: Option<Tokens>,
97 #[serde(default, skip_serializing_if = "Option::is_none")]
98 cache_write_tokens: Option<Tokens>,
99 #[serde(default, skip_serializing_if = "Option::is_none")]
100 v: Option<u8>,
101 #[serde(default, skip_serializing_if = "Option::is_none")]
102 dialect: Option<DialectKind>,
103}
104
105impl From<ModelUsage> for UsageRow {
106 fn from(usage: ModelUsage) -> UsageRow {
107 UsageRow {
108 input_tokens: usage.input_tokens,
109 output_tokens: usage.output_tokens,
110 cache_read_tokens: usage.cache_read_tokens.reported(),
111 cache_write_tokens: usage.cache_write_tokens.reported(),
112 v: Some(USAGE_MEANING),
113 dialect: usage.dialect,
114 }
115 }
116}
117
118impl TryFrom<UsageRow> for ModelUsage {
119 type Error = String;
120
121 fn try_from(row: UsageRow) -> Result<ModelUsage, String> {
122 let input_tokens = match row.v {
123 Some(USAGE_MEANING) => row.input_tokens,
124 Some(other) => {
125 return Err(format!(
126 "usage row version {other} is newer than this build reads \
127 (it reads {USAGE_MEANING}); read the ledger with a newer sprawling"
128 ));
129 }
130 None => {
131 let cached = CacheCount::from(row.cache_read_tokens)
132 .or_zero()
133 .checked_add(CacheCount::from(row.cache_write_tokens).or_zero())
134 .ok_or("usage row cache counts overflow a token count")?;
135 if cached > row.input_tokens {
136 row.input_tokens
137 .checked_add(cached)
138 .ok_or("usage row counts overflow a token count")?
139 } else {
140 row.input_tokens
141 }
142 }
143 };
144 Ok(ModelUsage {
145 input_tokens,
146 output_tokens: row.output_tokens,
147 cache_read_tokens: CacheCount::from(row.cache_read_tokens),
148 cache_write_tokens: CacheCount::from(row.cache_write_tokens),
149 dialect: row.dialect,
150 })
151 }
152}