1use crate::utils::{
2 find_pricing_cache_for_date_in, get_pricing_cache_path_in, list_pricing_cache_files_in,
3};
4use anyhow::{Context, Result};
5use chrono::Utc;
6use serde::{Deserialize, Serialize};
7use serde_json::{Map, Value};
8use std::collections::HashMap;
9use std::fs;
10use std::path::Path;
11
12#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
25pub struct ThresholdTier {
26 pub threshold_tokens: i64,
28 #[serde(default)]
30 pub input_cost_per_token: f64,
31 #[serde(default)]
33 pub output_cost_per_token: f64,
34 #[serde(default)]
36 pub cache_read_input_token_cost: f64,
37 #[serde(default)]
39 pub cache_creation_input_token_cost: f64,
40 #[serde(default)]
42 pub cache_creation_input_token_cost_above_1hr: f64,
43}
44
45#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
51pub struct TierRange {
52 pub min_tokens: i64,
54 pub max_tokens: i64,
56 #[serde(default)]
58 pub input_cost_per_token: f64,
59 #[serde(default)]
61 pub output_cost_per_token: f64,
62 #[serde(default)]
64 pub cache_read_input_token_cost: f64,
65 #[serde(default)]
67 pub output_cost_per_reasoning_token: f64,
68}
69
70#[derive(Debug, Clone, Default, Serialize, Deserialize)]
88pub struct ModelPricing {
89 #[serde(default)]
91 pub input_cost_per_token: f64,
92 #[serde(default)]
94 pub output_cost_per_token: f64,
95 #[serde(default)]
97 pub cache_read_input_token_cost: f64,
98 #[serde(default)]
100 pub cache_creation_input_token_cost: f64,
101
102 #[serde(default)]
106 pub cache_creation_input_token_cost_above_1hr: f64,
107
108 #[serde(default)]
115 pub output_cost_per_reasoning_token: f64,
116
117 #[serde(default)]
124 pub web_search_cost_per_query: f64,
125
126 #[serde(default, skip_serializing_if = "Vec::is_empty")]
128 pub tiers: Vec<ThresholdTier>,
129
130 #[serde(default, skip_serializing_if = "Option::is_none")]
132 pub ranges: Option<Vec<TierRange>>,
133}
134
135fn parse_threshold_suffix(suffix: &str) -> Option<i64> {
139 let without_tokens = suffix.strip_suffix("_tokens")?;
140 let num_part = without_tokens.strip_suffix('k')?;
141 num_part.parse::<i64>().ok().map(|n| n * 1000)
142}
143
144fn parse_tier_range(value: &serde_json::Value) -> Option<TierRange> {
149 let obj = value.as_object()?;
150 let range = obj.get("range")?.as_array()?;
151 if range.len() != 2 {
152 return None;
153 }
154 let min = range[0].as_f64()? as i64;
155 let max = range[1].as_f64()? as i64;
156 let f = |k: &str| obj.get(k).and_then(|v| v.as_f64()).unwrap_or(0.0);
157 Some(TierRange {
158 min_tokens: min,
159 max_tokens: max,
160 input_cost_per_token: f("input_cost_per_token"),
161 output_cost_per_token: f("output_cost_per_token"),
162 cache_read_input_token_cost: f("cache_read_input_token_cost"),
163 output_cost_per_reasoning_token: f("output_cost_per_reasoning_token"),
164 })
165}
166
167pub fn parse_litellm_entry(value: &serde_json::Value) -> ModelPricing {
175 let obj = match value.as_object() {
176 Some(o) => o,
177 None => return ModelPricing::default(),
178 };
179
180 let mut pricing = ModelPricing::default();
181 let mut tier_input: HashMap<i64, f64> = HashMap::new();
182 let mut tier_output: HashMap<i64, f64> = HashMap::new();
183 let mut tier_cache_read: HashMap<i64, f64> = HashMap::new();
184 let mut tier_cache_creation: HashMap<i64, f64> = HashMap::new();
185 let mut tier_cache_creation_1hr: HashMap<i64, f64> = HashMap::new();
187
188 for (key, raw_val) in obj {
189 if key == "tiered_pricing" {
190 if let Some(arr) = raw_val.as_array() {
191 let ranges: Vec<TierRange> = arr.iter().filter_map(parse_tier_range).collect();
192 if !ranges.is_empty() {
193 pricing.ranges = Some(ranges);
194 }
195 }
196 continue;
197 }
198
199 if key == "search_context_cost_per_query" {
200 if let Some(sc) = raw_val.as_object() {
204 let pick = |k: &str| sc.get(k).and_then(|v| v.as_f64());
205 pricing.web_search_cost_per_query = pick("search_context_size_medium")
206 .or_else(|| pick("search_context_size_low"))
207 .or_else(|| pick("search_context_size_high"))
208 .unwrap_or(0.0);
209 }
210 continue;
211 }
212
213 let num_value = match raw_val.as_f64() {
214 Some(n) => n,
215 None => continue,
216 };
217
218 match key.as_str() {
219 "input_cost_per_token" => pricing.input_cost_per_token = num_value,
220 "output_cost_per_token" => pricing.output_cost_per_token = num_value,
221 "cache_read_input_token_cost" => pricing.cache_read_input_token_cost = num_value,
222 "cache_creation_input_token_cost" => {
223 pricing.cache_creation_input_token_cost = num_value
224 }
225 "cache_creation_input_token_cost_above_1hr" => {
226 pricing.cache_creation_input_token_cost_above_1hr = num_value;
228 }
229 "output_cost_per_reasoning_token" => {
230 pricing.output_cost_per_reasoning_token = num_value;
231 }
232 _ => {
233 if let Some(suffix) = key.strip_prefix("input_cost_per_token_above_") {
234 if let Some(th) = parse_threshold_suffix(suffix) {
235 tier_input.insert(th, num_value);
236 }
237 } else if let Some(suffix) = key.strip_prefix("output_cost_per_token_above_") {
238 if let Some(th) = parse_threshold_suffix(suffix) {
239 tier_output.insert(th, num_value);
240 }
241 } else if let Some(suffix) = key.strip_prefix("cache_read_input_token_cost_above_")
242 {
243 if let Some(th) = parse_threshold_suffix(suffix) {
244 tier_cache_read.insert(th, num_value);
245 }
246 } else if let Some(suffix) =
247 key.strip_prefix("cache_creation_input_token_cost_above_")
248 {
249 if let Some(inner) = suffix.strip_prefix("1hr_above_") {
253 if let Some(th) = parse_threshold_suffix(inner) {
254 tier_cache_creation_1hr.insert(th, num_value);
255 }
256 } else if !suffix.starts_with("1hr")
257 && let Some(th) = parse_threshold_suffix(suffix)
258 {
259 tier_cache_creation.insert(th, num_value);
260 }
261 }
262 }
263 }
264 }
265
266 let mut thresholds: Vec<i64> = tier_input
267 .keys()
268 .chain(tier_output.keys())
269 .chain(tier_cache_read.keys())
270 .chain(tier_cache_creation.keys())
271 .chain(tier_cache_creation_1hr.keys())
272 .copied()
273 .collect();
274 thresholds.sort();
275 thresholds.dedup();
276
277 pricing.tiers = thresholds
278 .into_iter()
279 .map(|th| ThresholdTier {
280 threshold_tokens: th,
281 input_cost_per_token: *tier_input.get(&th).unwrap_or(&pricing.input_cost_per_token),
282 output_cost_per_token: *tier_output
283 .get(&th)
284 .unwrap_or(&pricing.output_cost_per_token),
285 cache_read_input_token_cost: *tier_cache_read
286 .get(&th)
287 .unwrap_or(&pricing.cache_read_input_token_cost),
288 cache_creation_input_token_cost: *tier_cache_creation
289 .get(&th)
290 .unwrap_or(&pricing.cache_creation_input_token_cost),
291 cache_creation_input_token_cost_above_1hr: tier_cache_creation_1hr
298 .get(&th)
299 .copied()
300 .unwrap_or(0.0),
301 })
302 .collect();
303
304 if let Some(ranges) = pricing.ranges.as_mut() {
308 ranges.sort_by_key(|r| r.min_tokens);
309 }
310
311 pricing
312}
313
314pub fn parse_litellm_pricing_map(raw: serde_json::Value) -> HashMap<String, ModelPricing> {
316 let obj = match raw.as_object() {
317 Some(o) => o,
318 None => return HashMap::new(),
319 };
320 obj.iter()
321 .filter(|(_, v)| v.is_object())
322 .map(|(k, v)| (k.clone(), parse_litellm_entry(v)))
323 .collect()
324}
325
326pub fn filter_cost_fields(value: &Value) -> Option<Value> {
346 let obj = value.as_object()?;
347 let mut filtered = Map::with_capacity(obj.len());
348 for (k, v) in obj {
349 if k.contains("cost") || k == "tiered_pricing" {
350 filtered.insert(k.clone(), v.clone());
351 }
352 }
353 if filtered.is_empty() {
354 None
355 } else {
356 Some(Value::Object(filtered))
357 }
358}
359
360pub fn build_filtered_cost_json(raw: &Value) -> Value {
365 let obj = match raw.as_object() {
366 Some(o) => o,
367 None => return Value::Object(Map::new()),
368 };
369 let mut filtered_map = Map::with_capacity(obj.len());
370 for (model, entry) in obj {
371 if let Some(filtered) = filter_cost_fields(entry) {
372 filtered_map.insert(model.clone(), filtered);
373 }
374 }
375 Value::Object(filtered_map)
376}
377
378pub(crate) fn pricing_cache_date() -> String {
379 Utc::now().date_naive().format("%Y-%m-%d").to_string()
380}
381
382pub fn cleanup_old_cache_in(dir: &Path) {
387 let today = pricing_cache_date();
388
389 for (filename, path) in list_pricing_cache_files_in(dir) {
390 if !filename.contains(&today) {
391 let _ = fs::remove_file(&path);
392 log::debug!("Removed old cache file: {:?}", path);
393 }
394 }
395}
396
397pub fn load_from_cache_in(dir: &Path) -> Result<HashMap<String, ModelPricing>> {
421 let today = pricing_cache_date();
422 let cache_path = find_pricing_cache_for_date_in(dir, &today)
423 .ok_or_else(|| anyhow::anyhow!("No cache file found for today"))?;
424
425 let content = fs::read_to_string(&cache_path).context("Failed to read cached pricing file")?;
426 let raw: Value =
427 serde_json::from_str(&content).context("Failed to parse cached pricing JSON")?;
428
429 if !raw.is_object() {
430 anyhow::bail!("cached pricing JSON must be an object");
431 }
432
433 if looks_like_legacy_pricing_cache(&raw) {
434 log::warn!(
435 "Detected pre-Phase-2 pricing cache format at {:?}; refetching to avoid silent tier/range data loss",
436 cache_path
437 );
438 anyhow::bail!("legacy pricing cache format detected, forcing refetch");
439 }
440
441 let parsed = parse_litellm_pricing_map(raw);
442 if !pricing_rates_are_valid(&parsed) {
443 anyhow::bail!("cached pricing JSON contains a negative or non-finite price");
444 }
445 let pricing = normalize_pricing(parsed);
446 if pricing.is_empty() {
447 anyhow::bail!("cached pricing JSON has no priced models");
448 }
449 Ok(pricing)
450}
451
452fn looks_like_legacy_pricing_cache(raw: &Value) -> bool {
462 let Some(obj) = raw.as_object() else {
463 return false;
464 };
465 obj.values()
466 .filter_map(|v| v.as_object())
467 .any(|entry| entry.contains_key("tiers") || entry.contains_key("ranges"))
468}
469
470pub fn save_to_cache_in(dir: &Path, filtered_raw: &Value) -> Result<()> {
485 let today = pricing_cache_date();
486 let cache_path = get_pricing_cache_path_in(dir, &today);
487
488 crate::utils::write_json_atomic_pretty(&cache_path, filtered_raw)
489 .context("Failed to write pricing cache file")?;
490
491 cleanup_old_cache_in(dir);
492 Ok(())
493}
494
495pub fn normalize_pricing(
508 mut pricing: HashMap<String, ModelPricing>,
509) -> HashMap<String, ModelPricing> {
510 pricing.retain(|_name, p| {
511 let has_base = p.input_cost_per_token > 0.0
512 || p.output_cost_per_token > 0.0
513 || p.cache_read_input_token_cost > 0.0
514 || p.cache_creation_input_token_cost > 0.0
515 || p.cache_creation_input_token_cost_above_1hr > 0.0
516 || p.output_cost_per_reasoning_token > 0.0
517 || p.web_search_cost_per_query > 0.0;
518 let has_positive_tier = p.tiers.iter().any(|t| {
519 t.input_cost_per_token > 0.0
520 || t.output_cost_per_token > 0.0
521 || t.cache_read_input_token_cost > 0.0
522 || t.cache_creation_input_token_cost > 0.0
523 || t.cache_creation_input_token_cost_above_1hr > 0.0
524 });
525 let has_positive_range = p
526 .ranges
527 .as_ref()
528 .map(|rs| {
529 rs.iter().any(|r| {
530 r.input_cost_per_token > 0.0
531 || r.output_cost_per_token > 0.0
532 || r.cache_read_input_token_cost > 0.0
533 || r.output_cost_per_reasoning_token > 0.0
534 })
535 })
536 .unwrap_or(false);
537 model_pricing_rates_are_valid(p) && (has_base || has_positive_tier || has_positive_range)
538 });
539 pricing
540}
541
542pub(crate) fn pricing_rates_are_valid(pricing: &HashMap<String, ModelPricing>) -> bool {
543 pricing.values().all(model_pricing_rates_are_valid)
544}
545
546fn model_pricing_rates_are_valid(pricing: &ModelPricing) -> bool {
547 let valid = |rate: f64| rate.is_finite() && rate >= 0.0;
548 valid(pricing.input_cost_per_token)
549 && valid(pricing.output_cost_per_token)
550 && valid(pricing.cache_read_input_token_cost)
551 && valid(pricing.cache_creation_input_token_cost)
552 && valid(pricing.cache_creation_input_token_cost_above_1hr)
553 && valid(pricing.output_cost_per_reasoning_token)
554 && valid(pricing.web_search_cost_per_query)
555 && pricing.tiers.iter().all(|tier| {
556 valid(tier.input_cost_per_token)
557 && valid(tier.output_cost_per_token)
558 && valid(tier.cache_read_input_token_cost)
559 && valid(tier.cache_creation_input_token_cost)
560 && valid(tier.cache_creation_input_token_cost_above_1hr)
561 })
562 && pricing.ranges.as_ref().is_none_or(|ranges| {
563 ranges.iter().all(|range| {
564 valid(range.input_cost_per_token)
565 && valid(range.output_cost_per_token)
566 && valid(range.cache_read_input_token_cost)
567 && valid(range.output_cost_per_reasoning_token)
568 })
569 })
570}
571
572#[cfg(test)]
573mod parser_tests {
574 use super::*;
575 use serde_json::json;
576
577 #[test]
578 fn parses_flat_model_no_tiers() {
579 let raw = json!({
581 "input_cost_per_token": 5e-6,
582 "output_cost_per_token": 2.5e-5,
583 "cache_read_input_token_cost": 5e-7,
584 "cache_creation_input_token_cost": 6.25e-6,
585 "cache_creation_input_token_cost_above_1hr": 1e-5,
586 "max_input_tokens": 200000
587 });
588 let p = parse_litellm_entry(&raw);
589 assert_eq!(p.input_cost_per_token, 5e-6);
590 assert_eq!(p.output_cost_per_token, 2.5e-5);
591 assert_eq!(p.cache_read_input_token_cost, 5e-7);
592 assert_eq!(p.cache_creation_input_token_cost, 6.25e-6);
593 assert!(p.tiers.is_empty());
595 assert!(p.ranges.is_none());
596 }
597
598 #[test]
599 fn parses_web_search_cost_per_query() {
600 let raw = json!({
605 "input_cost_per_token": 5e-6,
606 "output_cost_per_token": 2.5e-5,
607 "search_context_cost_per_query": {
608 "search_context_size_low": 0.01,
609 "search_context_size_medium": 0.01,
610 "search_context_size_high": 0.01
611 }
612 });
613 let p = parse_litellm_entry(&raw);
614 assert_eq!(p.web_search_cost_per_query, 0.01);
615 assert_eq!(p.input_cost_per_token, 5e-6);
616 }
617
618 #[test]
619 fn web_search_cost_absent_defaults_to_zero() {
620 let raw = json!({ "input_cost_per_token": 5e-6 });
621 assert_eq!(parse_litellm_entry(&raw).web_search_cost_per_query, 0.0);
622 }
623
624 #[test]
625 fn parses_sonnet_like_with_200k_tier() {
626 let raw = json!({
627 "input_cost_per_token": 3e-6,
628 "output_cost_per_token": 1.5e-5,
629 "cache_read_input_token_cost": 3e-7,
630 "cache_creation_input_token_cost": 3.75e-6,
631 "input_cost_per_token_above_200k_tokens": 6e-6,
632 "output_cost_per_token_above_200k_tokens": 2.25e-5,
633 "cache_read_input_token_cost_above_200k_tokens": 6e-7,
634 "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6
635 });
636 let p = parse_litellm_entry(&raw);
637 assert_eq!(p.tiers.len(), 1);
638 let t = &p.tiers[0];
639 assert_eq!(t.threshold_tokens, 200_000);
640 assert_eq!(t.input_cost_per_token, 6e-6);
641 assert_eq!(t.output_cost_per_token, 2.25e-5);
642 assert_eq!(t.cache_read_input_token_cost, 6e-7);
643 assert_eq!(t.cache_creation_input_token_cost, 7.5e-6);
644 }
645
646 #[test]
647 fn parses_multiple_thresholds_sorted() {
648 let raw = json!({
650 "input_cost_per_token": 1e-6,
651 "output_cost_per_token": 2e-6,
652 "input_cost_per_token_above_272k_tokens": 4e-6,
653 "output_cost_per_token_above_272k_tokens": 8e-6,
654 "input_cost_per_token_above_128k_tokens": 2e-6,
655 "output_cost_per_token_above_128k_tokens": 4e-6
656 });
657 let p = parse_litellm_entry(&raw);
658 assert_eq!(p.tiers.len(), 2);
659 assert_eq!(p.tiers[0].threshold_tokens, 128_000);
661 assert_eq!(p.tiers[1].threshold_tokens, 272_000);
662 assert_eq!(p.tiers[0].input_cost_per_token, 2e-6);
663 assert_eq!(p.tiers[1].input_cost_per_token, 4e-6);
664 }
665
666 #[test]
667 fn missing_tier_fields_fall_back_to_base() {
668 let raw = json!({
670 "input_cost_per_token": 1e-6,
671 "output_cost_per_token": 2e-6,
672 "cache_read_input_token_cost": 1e-7,
673 "input_cost_per_token_above_200k_tokens": 2e-6
674 });
675 let p = parse_litellm_entry(&raw);
676 assert_eq!(p.tiers.len(), 1);
677 let t = &p.tiers[0];
678 assert_eq!(t.input_cost_per_token, 2e-6);
679 assert_eq!(t.output_cost_per_token, 2e-6); assert_eq!(t.cache_read_input_token_cost, 1e-7); }
682
683 #[test]
684 fn parses_tiered_pricing_ranges() {
685 let raw = json!({
687 "tiered_pricing": [
688 {
689 "range": [0, 32000],
690 "input_cost_per_token": 1e-6,
691 "output_cost_per_token": 5e-6,
692 "cache_read_input_token_cost": 1e-7
693 },
694 {
695 "range": [32000, 128000],
696 "input_cost_per_token": 1.8e-6,
697 "output_cost_per_token": 9e-6
698 },
699 {
700 "range": [256000, 1000000],
701 "input_cost_per_token": 6e-6,
702 "output_cost_per_token": 6e-5
703 }
704 ]
705 });
706 let p = parse_litellm_entry(&raw);
707 let ranges = p.ranges.expect("ranges should be parsed");
708 assert_eq!(ranges.len(), 3);
709 assert_eq!(ranges[0].min_tokens, 0);
710 assert_eq!(ranges[0].max_tokens, 32_000);
711 assert_eq!(ranges[0].input_cost_per_token, 1e-6);
712 assert_eq!(ranges[1].input_cost_per_token, 1.8e-6);
713 assert_eq!(ranges[2].max_tokens, 1_000_000);
714 }
715
716 #[test]
717 fn skips_non_token_tiered_pricing() {
718 let raw = json!({
720 "tiered_pricing": [
721 {"max_results_range": [0, 25], "input_cost_per_query": 0.005}
722 ]
723 });
724 let p = parse_litellm_entry(&raw);
725 assert!(p.ranges.is_none());
726 }
727
728 #[test]
729 fn parses_combined_1hr_plus_200k_tier() {
730 let raw = json!({
734 "input_cost_per_token": 3e-6,
735 "output_cost_per_token": 1.5e-5,
736 "cache_creation_input_token_cost": 3.75e-6,
737 "cache_creation_input_token_cost_above_1hr": 7.5e-6,
738 "input_cost_per_token_above_200k_tokens": 6e-6,
739 "output_cost_per_token_above_200k_tokens": 2.25e-5,
740 "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6,
741 "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.5e-5
742 });
743 let p = parse_litellm_entry(&raw);
744 assert_eq!(p.cache_creation_input_token_cost_above_1hr, 7.5e-6);
745 assert_eq!(p.tiers.len(), 1);
746 let t = &p.tiers[0];
747 assert_eq!(t.threshold_tokens, 200_000);
748 assert_eq!(t.cache_creation_input_token_cost, 7.5e-6);
749 assert_eq!(t.cache_creation_input_token_cost_above_1hr, 1.5e-5);
752 }
753
754 #[test]
755 fn tier_1hr_left_zero_when_missing_so_calculate_cost_can_fall_back() {
756 let raw = json!({
759 "input_cost_per_token": 3e-6,
760 "cache_creation_input_token_cost": 3.75e-6,
761 "cache_creation_input_token_cost_above_1hr": 6e-6,
762 "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6
763 });
764 let p = parse_litellm_entry(&raw);
765 assert_eq!(p.cache_creation_input_token_cost_above_1hr, 6e-6);
766 let t = &p.tiers[0];
767 assert_eq!(t.cache_creation_input_token_cost, 7.5e-6);
768 assert_eq!(t.cache_creation_input_token_cost_above_1hr, 0.0);
769 }
770
771 #[test]
772 fn cache_reload_reconstructs_tiers_from_raw_keys() {
773 let raw = json!({
778 "input_cost_per_token": 3e-6,
779 "output_cost_per_token": 1.5e-5,
780 "cache_read_input_token_cost": 3e-7,
781 "cache_creation_input_token_cost": 3.75e-6,
782 "input_cost_per_token_above_200k_tokens": 6e-6,
783 "output_cost_per_token_above_200k_tokens": 2.25e-5,
784 "cache_read_input_token_cost_above_200k_tokens": 6e-7,
785 "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6
786 });
787 let p = parse_litellm_entry(&raw);
788 assert_eq!(p.tiers.len(), 1, "tier must be rebuilt on cache reload");
789 assert_eq!(p.tiers[0].threshold_tokens, 200_000);
790 }
791
792 #[test]
793 fn parses_output_cost_per_reasoning_token() {
794 let raw = json!({
798 "input_cost_per_token": 3e-7,
799 "output_cost_per_token": 2.5e-6,
800 "output_cost_per_reasoning_token": 2.5e-6
801 });
802 let p = parse_litellm_entry(&raw);
803 assert_eq!(p.output_cost_per_reasoning_token, 2.5e-6);
804 }
805
806 #[test]
807 fn filter_cost_fields_keeps_only_cost_keys() {
808 let raw = json!({
809 "input_cost_per_token": 3e-6,
810 "output_cost_per_token": 1.5e-5,
811 "cache_creation_input_token_cost_above_1hr": 6e-6,
812 "max_input_tokens": 200_000,
813 "supports_vision": true,
814 "litellm_provider": "anthropic",
815 "search_context_cost_per_query": {"search_context_size_high": 0.01}
816 });
817 let filtered = filter_cost_fields(&raw).expect("has cost keys");
818 let obj = filtered.as_object().unwrap();
819 assert!(obj.contains_key("input_cost_per_token"));
820 assert!(obj.contains_key("output_cost_per_token"));
821 assert!(obj.contains_key("cache_creation_input_token_cost_above_1hr"));
822 assert!(
823 obj.contains_key("search_context_cost_per_query"),
824 "nested cost objects must survive the filter"
825 );
826 assert!(!obj.contains_key("max_input_tokens"));
827 assert!(!obj.contains_key("supports_vision"));
828 assert!(!obj.contains_key("litellm_provider"));
829 }
830
831 #[test]
832 fn filter_cost_fields_returns_none_for_non_cost_entries() {
833 let raw = json!({
837 "max_input_tokens": 8192,
838 "litellm_provider": "azure"
839 });
840 assert!(filter_cost_fields(&raw).is_none());
841 }
842
843 #[test]
844 fn build_filtered_cost_json_skips_entries_without_cost_keys() {
845 let raw = json!({
846 "model-a": {
847 "input_cost_per_token": 1e-6,
848 "max_input_tokens": 8192
849 },
850 "model-b": {
851 "max_input_tokens": 16384
852 }
853 });
854 let filtered = build_filtered_cost_json(&raw);
855 let obj = filtered.as_object().unwrap();
856 assert!(obj.contains_key("model-a"));
857 assert!(!obj.contains_key("model-b"));
858 let a = obj["model-a"].as_object().unwrap();
859 assert!(a.contains_key("input_cost_per_token"));
860 assert!(!a.contains_key("max_input_tokens"));
861 }
862
863 #[test]
864 fn ranges_are_sorted_by_min_tokens_after_parse() {
865 let raw = json!({
867 "tiered_pricing": [
868 {"range": [128_000, 256_000], "input_cost_per_token": 3e-6},
869 {"range": [0, 32_000], "input_cost_per_token": 1e-6},
870 {"range": [32_000, 128_000], "input_cost_per_token": 2e-6}
871 ]
872 });
873 let p = parse_litellm_entry(&raw);
874 let ranges = p.ranges.expect("ranges");
875 assert_eq!(ranges.len(), 3);
876 assert_eq!(ranges[0].min_tokens, 0);
877 assert_eq!(ranges[1].min_tokens, 32_000);
878 assert_eq!(ranges[2].min_tokens, 128_000);
879 }
880
881 #[test]
882 fn ignores_unknown_fields() {
883 let raw = json!({
884 "input_cost_per_token": 1e-6,
885 "output_cost_per_token": 2e-6,
886 "input_cost_per_token_priority": 5e-6,
887 "input_cost_per_token_batches": 5e-7,
888 "output_cost_per_reasoning_token": 3e-6,
889 "supports_vision": true,
890 "litellm_provider": "anthropic"
891 });
892 let p = parse_litellm_entry(&raw);
893 assert_eq!(p.input_cost_per_token, 1e-6);
894 assert_eq!(p.output_cost_per_token, 2e-6);
895 assert!(p.tiers.is_empty());
896 assert!(p.ranges.is_none());
897 }
898
899 #[test]
900 fn filter_cost_fields_preserves_tiered_pricing() {
901 let raw = json!({
906 "input_cost_per_token": 0.0,
907 "tiered_pricing": [
908 {
909 "range": [0, 32000],
910 "input_cost_per_token": 1e-6,
911 "output_cost_per_token": 5e-6
912 },
913 {
914 "range": [32000, 128000],
915 "input_cost_per_token": 1.8e-6,
916 "output_cost_per_token": 9e-6
917 }
918 ],
919 "max_input_tokens": 1_000_000,
920 "litellm_provider": "dashscope"
921 });
922 let filtered = filter_cost_fields(&raw).expect("has cost keys");
923 let obj = filtered.as_object().unwrap();
924 assert!(
925 obj.contains_key("tiered_pricing"),
926 "tiered_pricing must survive the filter — it carries range-based pricing data"
927 );
928 let ranges = obj["tiered_pricing"].as_array().expect("array preserved");
929 assert_eq!(ranges.len(), 2);
930 assert!(!obj.contains_key("max_input_tokens"));
931 assert!(!obj.contains_key("litellm_provider"));
932 }
933
934 #[test]
935 fn cache_roundtrip_preserves_range_priced_models() {
936 let upstream = json!({
943 "qwen3-coder-plus": {
944 "tiered_pricing": [
945 {
946 "range": [0, 32000],
947 "input_cost_per_token": 1e-6,
948 "output_cost_per_token": 5e-6
949 },
950 {
951 "range": [32000, 128000],
952 "input_cost_per_token": 1.8e-6,
953 "output_cost_per_token": 9e-6
954 }
955 ],
956 "max_input_tokens": 1_000_000,
957 "litellm_provider": "dashscope"
958 }
959 });
960
961 let filtered = build_filtered_cost_json(&upstream);
962 let reloaded = parse_litellm_pricing_map(filtered);
963
964 let p = reloaded
965 .get("qwen3-coder-plus")
966 .expect("model must survive roundtrip");
967 let ranges = p.ranges.as_ref().expect("ranges must be rebuilt on reload");
968 assert_eq!(ranges.len(), 2);
969 assert_eq!(ranges[0].min_tokens, 0);
970 assert_eq!(ranges[0].max_tokens, 32_000);
971 assert_eq!(ranges[0].input_cost_per_token, 1e-6);
972 assert_eq!(ranges[1].min_tokens, 32_000);
973 assert_eq!(ranges[1].input_cost_per_token, 1.8e-6);
974 }
975
976 #[test]
977 fn looks_like_legacy_pricing_cache_flags_tiers_array() {
978 let legacy = json!({
983 "claude-sonnet-4-6": {
984 "input_cost_per_token": 3e-6,
985 "output_cost_per_token": 1.5e-5,
986 "tiers": [
987 {
988 "threshold_tokens": 200_000,
989 "input_cost_per_token": 6e-6,
990 "output_cost_per_token": 2.25e-5
991 }
992 ]
993 }
994 });
995 assert!(looks_like_legacy_pricing_cache(&legacy));
996 }
997
998 #[test]
999 fn looks_like_legacy_pricing_cache_flags_ranges_field() {
1000 let legacy = json!({
1001 "qwen-plus": {
1002 "ranges": [
1003 {
1004 "min_tokens": 0,
1005 "max_tokens": 32_000,
1006 "input_cost_per_token": 1e-6
1007 }
1008 ]
1009 }
1010 });
1011 assert!(looks_like_legacy_pricing_cache(&legacy));
1012 }
1013
1014 #[test]
1015 fn looks_like_legacy_pricing_cache_accepts_new_format() {
1016 let new_format = json!({
1020 "claude-sonnet-4-6": {
1021 "input_cost_per_token": 3e-6,
1022 "output_cost_per_token": 1.5e-5,
1023 "input_cost_per_token_above_200k_tokens": 6e-6,
1024 "output_cost_per_token_above_200k_tokens": 2.25e-5
1025 },
1026 "qwen3-coder-plus": {
1027 "tiered_pricing": [
1028 {
1029 "range": [0, 32000],
1030 "input_cost_per_token": 1e-6
1031 }
1032 ]
1033 }
1034 });
1035 assert!(!looks_like_legacy_pricing_cache(&new_format));
1036 }
1037}
1038
1039#[cfg(test)]
1040mod serialization_tests {
1041 use super::*;
1042 use std::collections::HashMap;
1043
1044 #[test]
1045 fn test_model_pricing_default() {
1046 let pricing = ModelPricing::default();
1048
1049 assert_eq!(pricing.input_cost_per_token, 0.0);
1050 assert_eq!(pricing.output_cost_per_token, 0.0);
1051 assert_eq!(pricing.cache_read_input_token_cost, 0.0);
1052 assert_eq!(pricing.cache_creation_input_token_cost, 0.0);
1053 assert!(pricing.tiers.is_empty());
1054 assert!(pricing.ranges.is_none());
1055 }
1056
1057 #[test]
1058 fn test_model_pricing_serialization() {
1059 let pricing = ModelPricing {
1061 input_cost_per_token: 0.000001,
1062 output_cost_per_token: 0.000002,
1063 cache_read_input_token_cost: 0.0000001,
1064 cache_creation_input_token_cost: 0.0000005,
1065 tiers: vec![ThresholdTier {
1066 threshold_tokens: 200_000,
1067 input_cost_per_token: 0.000002,
1068 output_cost_per_token: 0.000004,
1069 cache_read_input_token_cost: 0.0000002,
1070 cache_creation_input_token_cost: 0.000001,
1071 ..Default::default()
1072 }],
1073 ranges: None,
1074 ..Default::default()
1075 };
1076
1077 let json = serde_json::to_string(&pricing).unwrap();
1078 let deserialized: ModelPricing = serde_json::from_str(&json).unwrap();
1079
1080 assert_eq!(
1081 deserialized.input_cost_per_token,
1082 pricing.input_cost_per_token
1083 );
1084 assert_eq!(deserialized.tiers.len(), 1);
1085 assert_eq!(deserialized.tiers[0].threshold_tokens, 200_000);
1086 assert_eq!(deserialized.tiers[0].input_cost_per_token, 0.000002);
1087 }
1088
1089 #[test]
1090 fn test_model_pricing_clone() {
1091 let pricing1 = ModelPricing {
1093 input_cost_per_token: 0.000001,
1094 output_cost_per_token: 0.000002,
1095 ..Default::default()
1096 };
1097
1098 let pricing2 = pricing1.clone();
1099
1100 assert_eq!(pricing1.input_cost_per_token, pricing2.input_cost_per_token);
1101 assert_eq!(
1102 pricing1.output_cost_per_token,
1103 pricing2.output_cost_per_token
1104 );
1105 }
1106
1107 #[test]
1108 fn test_model_pricing_debug() {
1109 let pricing = ModelPricing::default();
1111 let debug_str = format!("{:?}", pricing);
1112
1113 assert!(debug_str.contains("ModelPricing"));
1114 }
1115
1116 #[test]
1117 fn test_model_pricing_with_partial_data() {
1118 let json = r#"{"input_cost_per_token": 0.000001}"#;
1120 let pricing: ModelPricing = serde_json::from_str(json).unwrap();
1121
1122 assert_eq!(pricing.input_cost_per_token, 0.000001);
1123 assert_eq!(pricing.output_cost_per_token, 0.0); }
1125
1126 #[test]
1127 fn test_model_pricing_empty_json() {
1128 let json = "{}";
1130 let pricing: ModelPricing = serde_json::from_str(json).unwrap();
1131
1132 assert_eq!(pricing.input_cost_per_token, 0.0);
1133 assert_eq!(pricing.output_cost_per_token, 0.0);
1134 }
1135
1136 #[test]
1137 fn test_model_pricing_hashmap_serialization() {
1138 let mut pricing_map = HashMap::new();
1140 pricing_map.insert(
1141 "gpt-4".to_string(),
1142 ModelPricing {
1143 input_cost_per_token: 0.000030,
1144 output_cost_per_token: 0.000060,
1145 ..Default::default()
1146 },
1147 );
1148 pricing_map.insert(
1149 "claude-3".to_string(),
1150 ModelPricing {
1151 input_cost_per_token: 0.000015,
1152 output_cost_per_token: 0.000075,
1153 ..Default::default()
1154 },
1155 );
1156
1157 let json = serde_json::to_string(&pricing_map).unwrap();
1158 let deserialized: HashMap<String, ModelPricing> = serde_json::from_str(&json).unwrap();
1159
1160 assert_eq!(deserialized.len(), 2);
1161 assert!(deserialized.contains_key("gpt-4"));
1162 assert!(deserialized.contains_key("claude-3"));
1163 }
1164
1165 #[test]
1166 fn test_model_pricing_all_fields() {
1167 let pricing = ModelPricing {
1169 input_cost_per_token: 1.0,
1170 output_cost_per_token: 2.0,
1171 cache_read_input_token_cost: 3.0,
1172 cache_creation_input_token_cost: 4.0,
1173 output_cost_per_reasoning_token: 11.0,
1174 web_search_cost_per_query: 0.01,
1175 tiers: vec![ThresholdTier {
1176 threshold_tokens: 200_000,
1177 input_cost_per_token: 5.0,
1178 output_cost_per_token: 6.0,
1179 cache_read_input_token_cost: 7.0,
1180 cache_creation_input_token_cost: 8.0,
1181 cache_creation_input_token_cost_above_1hr: 12.0,
1182 }],
1183 cache_creation_input_token_cost_above_1hr: 10.0,
1184 ranges: Some(vec![TierRange {
1185 min_tokens: 0,
1186 max_tokens: 32_000,
1187 input_cost_per_token: 0.1,
1188 output_cost_per_token: 0.2,
1189 cache_read_input_token_cost: 0.01,
1190 output_cost_per_reasoning_token: 0.5,
1191 }]),
1192 };
1193
1194 let json = serde_json::to_string(&pricing).unwrap();
1195 let deserialized: ModelPricing = serde_json::from_str(&json).unwrap();
1196
1197 assert_eq!(deserialized.input_cost_per_token, 1.0);
1198 assert_eq!(deserialized.output_cost_per_token, 2.0);
1199 assert_eq!(deserialized.cache_read_input_token_cost, 3.0);
1200 assert_eq!(deserialized.cache_creation_input_token_cost, 4.0);
1201 assert_eq!(deserialized.output_cost_per_reasoning_token, 11.0);
1202 assert_eq!(deserialized.web_search_cost_per_query, 0.01);
1203 assert_eq!(deserialized.tiers.len(), 1);
1204 assert_eq!(deserialized.tiers[0].threshold_tokens, 200_000);
1205 assert_eq!(deserialized.tiers[0].input_cost_per_token, 5.0);
1206 assert_eq!(deserialized.tiers[0].output_cost_per_token, 6.0);
1207 assert_eq!(deserialized.tiers[0].cache_read_input_token_cost, 7.0);
1208 assert_eq!(deserialized.tiers[0].cache_creation_input_token_cost, 8.0);
1209 let ranges = deserialized.ranges.unwrap();
1210 assert_eq!(ranges.len(), 1);
1211 assert_eq!(ranges[0].max_tokens, 32_000);
1212 assert_eq!(ranges[0].output_cost_per_reasoning_token, 0.5);
1213 }
1214
1215 #[test]
1216 fn test_model_pricing_zero_values() {
1217 let pricing = ModelPricing::default();
1219 let json = serde_json::to_string(&pricing).unwrap();
1220 let deserialized: ModelPricing = serde_json::from_str(&json).unwrap();
1221
1222 assert_eq!(deserialized.input_cost_per_token, 0.0);
1224 assert_eq!(deserialized.output_cost_per_token, 0.0);
1225 }
1226
1227 #[test]
1228 fn test_model_pricing_negative_values() {
1229 let pricing = ModelPricing {
1231 input_cost_per_token: -0.000001,
1232 output_cost_per_token: -0.000002,
1233 ..Default::default()
1234 };
1235
1236 let json = serde_json::to_string(&pricing).unwrap();
1237 let deserialized: ModelPricing = serde_json::from_str(&json).unwrap();
1238
1239 assert_eq!(deserialized.input_cost_per_token, -0.000001);
1240 assert_eq!(deserialized.output_cost_per_token, -0.000002);
1241 }
1242
1243 #[test]
1244 fn test_model_pricing_very_small_values() {
1245 let pricing = ModelPricing {
1247 input_cost_per_token: 1e-10,
1248 output_cost_per_token: 1e-15,
1249 ..Default::default()
1250 };
1251
1252 let json = serde_json::to_string(&pricing).unwrap();
1253 let deserialized: ModelPricing = serde_json::from_str(&json).unwrap();
1254
1255 assert!((deserialized.input_cost_per_token - 1e-10).abs() < 1e-20);
1256 assert!((deserialized.output_cost_per_token - 1e-15).abs() < 1e-25);
1257 }
1258
1259 #[test]
1260 fn test_model_pricing_large_values() {
1261 let pricing = ModelPricing {
1263 input_cost_per_token: 1000000.0,
1264 output_cost_per_token: 9999999.99,
1265 ..Default::default()
1266 };
1267
1268 let json = serde_json::to_string(&pricing).unwrap();
1269 let deserialized: ModelPricing = serde_json::from_str(&json).unwrap();
1270
1271 assert_eq!(deserialized.input_cost_per_token, 1000000.0);
1272 assert_eq!(deserialized.output_cost_per_token, 9999999.99);
1273 }
1274}