use crate::error::PreviewError;
use tt_provider_openai::ClientConfig;
use tt_shared::{ModelPricing, Provider};
#[derive(Debug, Clone)]
pub struct LookupHit {
pub provider: &'static str,
pub input_per_m: f64,
pub output_per_m: f64,
}
fn hit(provider: &'static str, p: &ModelPricing) -> LookupHit {
LookupHit {
provider,
input_per_m: p.input_per_million,
output_per_m: p.output_per_million,
}
}
pub fn lookup(model: &str) -> Result<LookupHit, PreviewError> {
if let Some(p) = tt_provider_anthropic::pricing::pricing_for(model) {
return Ok(hit("anthropic", &p));
}
if let Some(p) = tt_provider_openai::pricing::pricing_for(model) {
return Ok(hit("openai", &p));
}
if let Some(p) = tt_provider_gemini::pricing::pricing_for(model) {
return Ok(hit("gemini", &p));
}
let cfg = ClientConfig::default;
if let Some(p) = tt_provider_groq::GroqProvider::new(cfg()).pricing(model) {
return Ok(hit("groq", &p));
}
if let Some(p) = tt_provider_mistral::MistralProvider::new(cfg()).pricing(model) {
return Ok(hit("mistral", &p));
}
if let Some(p) = tt_provider_together::TogetherProvider::new(cfg()).pricing(model) {
return Ok(hit("together", &p));
}
if let Some(p) = tt_provider_openrouter::OpenRouterProvider::new(cfg()).pricing(model) {
return Ok(hit("openrouter", &p));
}
Err(PreviewError::UnknownModel(model.to_string()))
}
pub fn lookup_with_provider(model: &str, provider: &str) -> Result<LookupHit, PreviewError> {
let cfg = ClientConfig::default;
let found = match provider {
"anthropic" => {
tt_provider_anthropic::pricing::pricing_for(model).map(|p| hit("anthropic", &p))
}
"openai" => tt_provider_openai::pricing::pricing_for(model).map(|p| hit("openai", &p)),
"gemini" => tt_provider_gemini::pricing::pricing_for(model).map(|p| hit("gemini", &p)),
"groq" => tt_provider_groq::GroqProvider::new(cfg())
.pricing(model)
.map(|p| hit("groq", &p)),
"mistral" => tt_provider_mistral::MistralProvider::new(cfg())
.pricing(model)
.map(|p| hit("mistral", &p)),
"together" => tt_provider_together::TogetherProvider::new(cfg())
.pricing(model)
.map(|p| hit("together", &p)),
"openrouter" => tt_provider_openrouter::OpenRouterProvider::new(cfg())
.pricing(model)
.map(|p| hit("openrouter", &p)),
_ => None,
};
found.ok_or_else(|| PreviewError::UnknownModel(model.to_string()))
}
pub fn cost_usd(input_tokens: u32, output_tokens: u32, hit: &LookupHit) -> f64 {
let i = (input_tokens as f64) * hit.input_per_m / 1_000_000.0;
let o = (output_tokens as f64) * hit.output_per_m / 1_000_000.0;
i + o
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cost_math_basics() {
let hit = LookupHit {
provider: "x",
input_per_m: 3.0,
output_per_m: 15.0,
};
let c = cost_usd(1000, 100, &hit);
assert!((c - 0.0045).abs() < 1e-9, "cost = {c}");
}
#[test]
fn lookup_unknown_model_errors() {
let err = lookup("does-not-exist-model").unwrap_err();
assert!(matches!(err, PreviewError::UnknownModel(_)));
}
#[test]
fn lookup_resolves_compat_provider_models() {
let hit = lookup("llama-3.1-8b-instant").expect("groq model should resolve");
assert_eq!(hit.provider, "groq");
assert!(hit.input_per_m > 0.0, "groq pricing should be > 0");
}
#[test]
fn lookup_with_provider_attributes_to_named_provider() {
let hit =
lookup_with_provider("gpt-4o-mini", "openai").expect("openai carries gpt-4o-mini");
assert_eq!(hit.provider, "openai");
let hit = lookup_with_provider("claude-haiku-4-5", "anthropic")
.expect("anthropic carries claude-haiku-4-5");
assert_eq!(hit.provider, "anthropic");
}
#[test]
fn lookup_with_provider_errors_when_provider_lacks_model() {
let err = lookup_with_provider("gpt-4o-mini", "groq").unwrap_err();
assert!(matches!(err, PreviewError::UnknownModel(_)));
let err = lookup_with_provider("gpt-4o-mini", "nope").unwrap_err();
assert!(matches!(err, PreviewError::UnknownModel(_)));
}
}