1use crate::models::PerProviderUsage;
10use crate::pricing::{CostSource, ModelPricingMap, resolve_model_cost};
11use crate::usage::{StoredCosts, UsageData};
12use crate::utils::{extract_token_counts, normalize_usage_value};
13use serde::Serialize;
14use serde_json::Value;
15
16#[derive(Debug, Clone, Serialize)]
24pub struct PricedUsageRow {
25 pub cost_usd: f64,
27 #[serde(skip_serializing_if = "Option::is_none")]
29 pub matched_model: Option<String>,
30 pub model: String,
32 pub usage: Value,
34}
35
36pub fn price_usage_data(
46 usage_data: &UsageData,
47 pricing_map: &ModelPricingMap,
48) -> Vec<PricedUsageRow> {
49 let mut rows = Vec::with_capacity(usage_data.models.len());
50
51 for (model, usage) in usage_data.models.iter() {
52 let (cost, matched_model) = resolve_merged_model_cost(
53 model,
54 &usage_data.per_provider,
55 pricing_map,
56 &usage_data.stored_costs,
57 )
58 .unwrap_or_else(|| price_usage_value(model, usage, pricing_map, CostSource::Litellm));
59
60 rows.push(PricedUsageRow {
61 model: model.clone(),
62 usage: normalize_usage_value(usage),
63 cost_usd: cost,
64 matched_model,
65 });
66 }
67
68 rows
69}
70
71pub(crate) fn resolve_merged_model_cost(
80 model: &str,
81 per_provider: &PerProviderUsage,
82 pricing_map: &ModelPricingMap,
83 stored_costs: &StoredCosts,
84) -> Option<(f64, Option<String>)> {
85 let mut total_cost = 0.0;
86 let mut matched_model = None;
87 let mut found = false;
88
89 for usage in [
90 &per_provider.claude,
91 &per_provider.codex,
92 &per_provider.copilot,
93 &per_provider.gemini,
94 ] {
95 if let Some(raw_usage) = usage.get(model) {
96 found = true;
97 let (cost, matched) =
98 price_usage_value(model, raw_usage, pricing_map, CostSource::Litellm);
99 total_cost += cost;
100 if matched_model.is_none() {
101 matched_model = matched;
102 }
103 }
104 }
105
106 let stored =
110 |m: &crate::constants::FastHashMap<String, f64>| m.get(model).copied().unwrap_or(0.0);
111 for (usage, source) in [
112 (&per_provider.grok, CostSource::GrokGauge),
113 (
114 &per_provider.opencode,
115 CostSource::OpenCodeStored(stored(&stored_costs.opencode)),
116 ),
117 (&per_provider.cursor, CostSource::OpenCodeStored(0.0)),
118 (
119 &per_provider.hermes,
120 CostSource::HermesStored(stored(&stored_costs.hermes)),
121 ),
122 ] {
123 if let Some(raw_usage) = usage.get(model) {
124 found = true;
125 let (cost, matched) = price_usage_value(model, raw_usage, pricing_map, source);
126 total_cost += cost;
127 if matched_model.is_none() {
128 matched_model = matched;
129 }
130 }
131 }
132
133 found.then_some((total_cost, matched_model))
134}
135
136pub(crate) fn price_usage_value(
138 model: &str,
139 usage: &Value,
140 pricing_map: &ModelPricingMap,
141 source: CostSource,
142) -> (f64, Option<String>) {
143 let counts = extract_token_counts(usage);
144 resolve_model_cost(model, &counts, pricing_map, source)
145}
146
147#[cfg(test)]
148mod tests {
149 use super::*;
150 use crate::models::{PerProviderUsage, ProviderActiveDays, UsageResult};
151 use crate::pricing::{ModelPricing, clear_pricing_cache};
152 use crate::usage::StoredCosts;
153 use serde_json::json;
154 use std::collections::HashMap;
155
156 #[test]
157 fn priced_rows_include_grok_source_cost() {
158 clear_pricing_cache();
159 let mut raw_pricing = HashMap::new();
160 raw_pricing.insert(
161 "shared-model".to_string(),
162 ModelPricing {
163 input_cost_per_token: 0.01,
164 ..Default::default()
165 },
166 );
167 let pricing_map = ModelPricingMap::new(raw_pricing);
168 let mut models = UsageResult::default();
169 models.insert("shared-model".to_string(), json!({"input_tokens": 200}));
170 let mut per_provider = PerProviderUsage::default();
171 per_provider
172 .claude
173 .insert("shared-model".to_string(), json!({"input_tokens": 100}));
174 per_provider
175 .grok
176 .insert("shared-model".to_string(), json!({"input_tokens": 100}));
177 let usage_data = UsageData {
178 models,
179 per_provider,
180 provider_days: ProviderActiveDays::default(),
181 stored_costs: StoredCosts::default(),
182 };
183
184 let rows = price_usage_data(&usage_data, &pricing_map);
185
186 assert!((rows[0].cost_usd - 2.0).abs() < 1e-9);
187 }
188
189 #[test]
190 fn priced_rows_price_opencode_fallback_only_for_opencode_tokens() {
191 clear_pricing_cache();
192
193 let mut raw_pricing = HashMap::new();
194 raw_pricing.insert(
195 "shared".to_string(),
196 ModelPricing {
197 input_cost_per_token: 0.01,
198 ..Default::default()
199 },
200 );
201 let pricing_map = ModelPricingMap::new(raw_pricing);
202
203 let mut models = UsageResult::default();
204 models.insert("shared-pro".to_string(), json!({"input_tokens": 200}));
205
206 let mut per_provider = PerProviderUsage::default();
207 per_provider
208 .claude
209 .insert("shared-pro".to_string(), json!({"input_tokens": 100}));
210 per_provider
211 .opencode
212 .insert("shared-pro".to_string(), json!({"input_tokens": 100}));
213
214 let mut stored_costs = StoredCosts::default();
215 stored_costs.opencode.insert("shared-pro".to_string(), 7.0);
216
217 let usage_data = UsageData {
218 models,
219 per_provider,
220 provider_days: ProviderActiveDays::default(),
221 stored_costs,
222 };
223
224 let rows = price_usage_data(&usage_data, &pricing_map);
225 assert_eq!(rows.len(), 1);
226 assert_eq!(rows[0].cost_usd, 8.0);
227 assert_eq!(rows[0].matched_model.as_deref(), Some("shared"));
228 }
229}