use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::path::Path;
use crate::error::Error;
use crate::error::Result;
use crate::message::Message;
use tokio::sync::mpsc;
#[cfg(feature = "anthropic")]
pub mod anthropic;
pub mod mock;
pub mod openai;
#[cfg(feature = "anthropic")]
pub use anthropic::AnthropicProvider;
pub use mock::MockProvider;
pub use openai::OpenAiProvider;
pub type StreamSender = mpsc::UnboundedSender<String>;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TokenUsage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
pub cache_hit_tokens: u32,
pub cache_miss_tokens: u32,
}
impl TokenUsage {
pub fn accumulate(self, other: TokenUsage) -> TokenUsage {
TokenUsage {
prompt_tokens: self.prompt_tokens.saturating_add(other.prompt_tokens),
completion_tokens: self
.completion_tokens
.saturating_add(other.completion_tokens),
total_tokens: self.total_tokens.saturating_add(other.total_tokens),
cache_hit_tokens: self.cache_hit_tokens.saturating_add(other.cache_hit_tokens),
cache_miss_tokens: self
.cache_miss_tokens
.saturating_add(other.cache_miss_tokens),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ModelPricing {
pub input_per_million: f64,
pub output_per_million: f64,
pub cache_hit_input_per_million: f64,
}
impl ModelPricing {
pub fn cost_usd(&self, usage: TokenUsage) -> f64 {
let in_cost = if usage.cache_hit_tokens > 0 {
let cache_hit =
usage.cache_hit_tokens as f64 * self.cache_hit_input_per_million / 1_000_000.0;
let cache_miss = usage.cache_miss_tokens as f64 * self.input_per_million / 1_000_000.0;
cache_hit + cache_miss
} else {
(usage.prompt_tokens as f64) * self.input_per_million / 1_000_000.0
};
let out_cost = (usage.completion_tokens as f64) * self.output_per_million / 1_000_000.0;
in_cost + out_cost
}
}
pub fn pricing_for(model: &str) -> Option<ModelPricing> {
match model {
"MiniMax-M2" => Some(ModelPricing {
input_per_million: 0.30,
output_per_million: 1.20,
cache_hit_input_per_million: 0.30,
}),
"deepseek-chat" | "deepseek-v4-flash" => Some(ModelPricing {
input_per_million: 0.27,
output_per_million: 1.10,
cache_hit_input_per_million: 0.027,
}),
"deepseek-v4-pro" => Some(ModelPricing {
input_per_million: 1.89,
output_per_million: 7.70,
cache_hit_input_per_million: 0.189,
}),
"glm-4-flash" => Some(ModelPricing {
input_per_million: 0.10,
output_per_million: 0.10,
cache_hit_input_per_million: 0.10,
}),
"glm-5.1" => Some(ModelPricing {
input_per_million: 0.50,
output_per_million: 2.00,
cache_hit_input_per_million: 0.50,
}),
_ => {
None
}
}
}
pub fn load_pricing_from_yaml(path: &Path) -> Result<HashMap<String, ModelPricing>> {
use std::fs;
use std::io::{self, BufRead};
let file = fs::File::open(path).map_err(Error::Io)?;
let reader = io::BufReader::new(file);
let mut models: HashMap<String, ModelPricing> = HashMap::new();
let mut current_model: Option<String> = None;
let mut current_pricing: Option<ModelPricingBuilder> = None;
let mut in_models_section = false;
for line in reader.lines() {
let line = line.map_err(Error::Io)?;
if line.trim().is_empty() || line.trim().starts_with('#') {
continue;
}
let leading_spaces = line.len() - line.trim_start().len();
let trimmed = line.trim();
if trimmed == "models:" {
in_models_section = true;
continue;
}
if in_models_section {
if leading_spaces == 2 && trimmed.ends_with(':') {
if let (Some(name), Some(builder)) = (current_model.take(), current_pricing.take())
{
if let Some(pricing) = builder.build() {
models.insert(name, pricing);
}
}
let model_name = trimmed.trim_end_matches(':').to_string();
current_model = Some(model_name);
current_pricing = Some(ModelPricingBuilder::default());
continue;
}
if leading_spaces >= 4 && current_model.is_some() {
if let Some(ref mut builder) = current_pricing {
if let Some((key, value)) = trimmed.split_once(':') {
let key = key.trim();
let value = value.split('#').next().unwrap_or(value).trim();
if let Err(e) = builder.parse_field(key, value) {
return Err(Error::Config {
message: format!("error parsing {}: {}", path.display(), e),
});
}
}
}
}
}
}
if let (Some(name), Some(builder)) = (current_model, current_pricing) {
if let Some(pricing) = builder.build() {
models.insert(name, pricing);
}
}
Ok(models)
}
#[derive(Default)]
struct ModelPricingBuilder {
input_per_million: Option<f64>,
output_per_million: Option<f64>,
cache_hit_input_per_million: Option<f64>,
}
impl ModelPricingBuilder {
fn parse_field(&mut self, key: &str, value: &str) -> Result<(), String> {
let value: f64 = value
.parse()
.map_err(|_| format!("invalid float: {}", value))?;
match key {
"input_per_million" => self.input_per_million = Some(value),
"output_per_million" => self.output_per_million = Some(value),
"cache_hit_input_per_million" => self.cache_hit_input_per_million = Some(value),
_ => return Err(format!("unknown field: {}", key)),
}
Ok(())
}
fn build(self) -> Option<ModelPricing> {
Some(ModelPricing {
input_per_million: self.input_per_million?,
output_per_million: self.output_per_million?,
cache_hit_input_per_million: self
.cache_hit_input_per_million
.unwrap_or(self.input_per_million.unwrap_or(0.0)),
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolSpec {
pub name: String,
pub description: String,
pub parameters: Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String,
pub name: String,
pub arguments: Value,
}
pub struct StructuredRequest {
pub messages: Vec<Message>,
pub schema: Value,
pub schema_name: String,
}
#[derive(Debug, Clone)]
pub struct Completion {
pub content: String,
pub tool_calls: Vec<ToolCall>,
pub finish_reason: Option<String>,
pub usage: Option<TokenUsage>,
}
#[async_trait]
pub trait LlmProvider: Send + Sync {
async fn complete(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Completion>;
async fn complete_structured(&self, _req: StructuredRequest) -> Result<Value> {
Err(Error::Config {
message: "provider does not support structured output".into(),
})
}
async fn stream(
&self,
messages: &[Message],
tools: &[ToolSpec],
stream_tx: Option<StreamSender>,
) -> Result<Completion> {
let completion = self.complete(messages, tools).await?;
if let Some(tx) = stream_tx {
if !completion.content.is_empty() {
let _ = tx.send(completion.content.clone());
}
}
Ok(completion)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn mock_structured_returns_default_error() {
let provider = MockProvider::new(vec![]).with_structured_responses(vec![]);
let req = StructuredRequest {
messages: vec![Message::user("hi".to_string())],
schema: serde_json::json!({"type": "object", "properties": {"answer": {"type": "string"}}}),
schema_name: "test_schema".to_string(),
};
let result = provider.complete_structured(req).await;
assert!(result.is_err());
let err = result.unwrap_err();
let msg = err.to_string();
assert!(
msg.contains("no structured responses configured"),
"error should mention no structured responses: {msg}"
);
}
#[test]
fn token_usage_default_is_all_zeros() {
let u = TokenUsage::default();
assert_eq!(u.prompt_tokens, 0);
assert_eq!(u.completion_tokens, 0);
assert_eq!(u.total_tokens, 0);
}
#[test]
fn token_usage_accumulate_is_saturating() {
let u1 = TokenUsage {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
};
let u2 = TokenUsage {
prompt_tokens: 20,
completion_tokens: 30,
total_tokens: 50,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
};
let acc = u1.accumulate(u2);
assert_eq!(acc.prompt_tokens, 30);
assert_eq!(acc.completion_tokens, 35);
assert_eq!(acc.total_tokens, 65);
}
#[test]
fn token_usage_accumulate_is_commutative() {
let u1 = TokenUsage {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
};
let u2 = TokenUsage {
prompt_tokens: 20,
completion_tokens: 30,
total_tokens: 50,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
};
assert_eq!(u1.accumulate(u2), u2.accumulate(u1));
}
#[test]
fn token_usage_accumulate_saturates() {
let u1 = TokenUsage {
prompt_tokens: u32::MAX,
completion_tokens: 1,
total_tokens: u32::MAX,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
};
let u2 = TokenUsage {
prompt_tokens: 1,
completion_tokens: u32::MAX,
total_tokens: u32::MAX,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
};
let acc = u1.accumulate(u2);
assert_eq!(acc.prompt_tokens, u32::MAX);
assert_eq!(acc.completion_tokens, u32::MAX);
assert_eq!(acc.total_tokens, u32::MAX);
}
#[test]
fn cost_usd_handles_zero_usage() {
let pricing = ModelPricing {
input_per_million: 1.0,
output_per_million: 2.0,
cache_hit_input_per_million: 0.1, };
let usage = TokenUsage::default();
let cost = pricing.cost_usd(usage);
assert!((cost - 0.0).abs() < 1e-9);
}
#[test]
fn cost_usd_computes_simple_case() {
let pricing = ModelPricing {
input_per_million: 1.0,
output_per_million: 1.0,
cache_hit_input_per_million: 0.1,
};
let usage = TokenUsage {
prompt_tokens: 1_000_000,
completion_tokens: 0,
total_tokens: 1_000_000,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
};
let cost = pricing.cost_usd(usage);
assert!((cost - 1.0).abs() < 1e-9);
}
#[test]
fn cost_usd_mixes_input_and_output() {
let pricing = ModelPricing {
input_per_million: 1.0,
output_per_million: 2.0,
cache_hit_input_per_million: 0.1,
};
let usage = TokenUsage {
prompt_tokens: 500_000,
completion_tokens: 250_000,
total_tokens: 750_000,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
};
let cost = pricing.cost_usd(usage);
assert!((cost - 1.0).abs() < 1e-9);
}
#[test]
fn pricing_for_known_models() {
let p1 = pricing_for("MiniMax-M2");
assert!(p1.is_some());
assert!((p1.unwrap().input_per_million - 0.30).abs() < 1e-9);
let p2 = pricing_for("deepseek-chat");
assert!(p2.is_some());
assert!((p2.unwrap().input_per_million - 0.27).abs() < 1e-9);
}
#[test]
fn pricing_for_unknown_returns_none() {
let p = pricing_for("unknown-model-xyz");
assert!(p.is_none());
}
#[test]
fn token_usage_accumulate_cache_fields() {
let u1 = TokenUsage {
prompt_tokens: 100,
completion_tokens: 50,
total_tokens: 150,
cache_hit_tokens: 60,
cache_miss_tokens: 40,
};
let u2 = TokenUsage {
prompt_tokens: 200,
completion_tokens: 100,
total_tokens: 300,
cache_hit_tokens: 120,
cache_miss_tokens: 80,
};
let acc = u1.accumulate(u2);
assert_eq!(acc.cache_hit_tokens, 180);
assert_eq!(acc.cache_miss_tokens, 120);
assert_eq!(acc.prompt_tokens, 300);
assert_eq!(acc.completion_tokens, 150);
assert_eq!(acc.total_tokens, 450);
}
#[test]
fn cost_usd_with_no_cache_hit_matches_old_behavior() {
let pricing = ModelPricing {
input_per_million: 1.0,
output_per_million: 2.0,
cache_hit_input_per_million: 0.1, };
let usage = TokenUsage {
prompt_tokens: 1_000_000,
completion_tokens: 500_000,
total_tokens: 1_500_000,
cache_hit_tokens: 0,
cache_miss_tokens: 1_000_000,
};
let cost = pricing.cost_usd(usage);
assert!((cost - 2.0).abs() < 1e-9);
}
#[test]
fn cost_usd_with_cache_hit_applies_discount() {
let pricing = ModelPricing {
input_per_million: 0.27,
output_per_million: 1.10,
cache_hit_input_per_million: 0.027,
};
let usage = TokenUsage {
prompt_tokens: 1_000,
completion_tokens: 500,
total_tokens: 1_500,
cache_hit_tokens: 900,
cache_miss_tokens: 100,
};
let cost = pricing.cost_usd(usage);
let expected =
900.0 * 0.027 / 1_000_000.0 + 100.0 * 0.27 / 1_000_000.0 + 500.0 * 1.10 / 1_000_000.0;
assert!((cost - expected).abs() < 1e-9);
}
#[test]
fn pricing_for_deepseek_has_cache_discount() {
let pricing = pricing_for("deepseek-chat").expect("deepseek-chat should be known");
assert!((pricing.input_per_million - 0.27).abs() < 1e-9);
assert!((pricing.cache_hit_input_per_million - 0.027).abs() < 1e-9);
}
#[test]
fn pricing_for_unknown_model_returns_none() {
let p = pricing_for("unknown-model-xyz");
assert!(p.is_none());
}
#[test]
fn token_usage_accumulate_preserves_cache_tokens() {
let u1 = TokenUsage {
prompt_tokens: 1000,
completion_tokens: 100,
total_tokens: 1100,
cache_hit_tokens: 900,
cache_miss_tokens: 100,
};
let u2 = TokenUsage {
prompt_tokens: 2000,
completion_tokens: 200,
total_tokens: 2200,
cache_hit_tokens: 1800,
cache_miss_tokens: 200,
};
let acc = u1.accumulate(u2);
assert_eq!(acc.cache_hit_tokens, 2700);
assert_eq!(acc.cache_miss_tokens, 300);
assert_eq!(acc.prompt_tokens, 3000);
}
#[test]
fn load_pricing_from_yaml_parses_file() {
let temp_dir = std::env::temp_dir();
let yaml_path = temp_dir.join("test_pricing.yaml");
std::fs::write(
&yaml_path,
r#"
models:
test-model:
input_per_million: 1.0
output_per_million: 2.0
cache_hit_input_per_million: 0.1
"#,
)
.unwrap();
let result = load_pricing_from_yaml(&yaml_path);
assert!(result.is_ok(), "load should succeed: {:?}", result.err());
let pricing = result.unwrap();
assert!(
pricing.contains_key("test-model"),
"should contain test-model: {:?}",
pricing.keys().collect::<Vec<_>>()
);
let p = pricing.get("test-model").unwrap();
assert!((p.input_per_million - 1.0).abs() < 1e-9);
assert!((p.output_per_million - 2.0).abs() < 1e-9);
assert!((p.cache_hit_input_per_million - 0.1).abs() < 1e-9);
std::fs::remove_file(&yaml_path).ok();
}
#[test]
fn load_pricing_from_yaml_overrides_hardcoded() {
let temp_dir = std::env::temp_dir();
let yaml_path = temp_dir.join("test_pricing_override.yaml");
std::fs::write(
&yaml_path,
r#"
models:
deepseek-chat:
input_per_million: 99.0
output_per_million: 99.0
cache_hit_input_per_million: 9.9
"#,
)
.unwrap();
let result = load_pricing_from_yaml(&yaml_path);
assert!(result.is_ok());
let pricing = result.unwrap();
let p = pricing
.get("deepseek-chat")
.expect("should have deepseek-chat");
assert!((p.input_per_million - 99.0).abs() < 1e-9);
assert!(pricing.contains_key("deepseek-chat"));
assert!(!pricing.contains_key("MiniMax-M2"));
std::fs::remove_file(&yaml_path).ok();
}
#[test]
fn load_pricing_from_yaml_missing_model_falls_back() {
let temp_dir = std::env::temp_dir();
let yaml_path = temp_dir.join("test_pricing_partial.yaml");
std::fs::write(
&yaml_path,
r#"
models:
test-only:
input_per_million: 1.0
output_per_million: 1.0
cache_hit_input_per_million: 0.1
"#,
)
.unwrap();
let result = load_pricing_from_yaml(&yaml_path);
assert!(result.is_ok());
let external = result.unwrap();
assert!(
external.contains_key("test-only"),
"should have test-only: {:?}",
external.keys().collect::<Vec<_>>()
);
let fallback = pricing_for("MiniMax-M2");
assert!(fallback.is_some());
assert!((fallback.unwrap().input_per_million - 0.30).abs() < 1e-9);
std::fs::remove_file(&yaml_path).ok();
}
#[test]
fn load_pricing_from_yaml_malformed_error() {
let temp_dir = std::env::temp_dir();
let yaml_path = temp_dir.join("test_pricing_malformed.yaml");
std::fs::write(
&yaml_path,
r#"
models:
bad-model:
input_per_million: not_a_number
"#,
)
.unwrap();
let result = load_pricing_from_yaml(&yaml_path);
assert!(result.is_err(), "should fail with bad data");
let err = result.unwrap_err();
let err_str = err.to_string();
assert!(
err_str.contains("error parsing") || err_str.contains("invalid float"),
"error should mention parsing issue: {}",
err_str
);
std::fs::remove_file(&yaml_path).ok();
}
}