1use async_trait::async_trait;
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11use std::collections::HashMap;
12use std::path::Path;
13
14use crate::error::Error;
15use crate::error::Result;
16use crate::message::Message;
17
18use tokio::sync::mpsc;
19
20#[cfg(feature = "anthropic")]
21pub mod anthropic;
22pub mod mock;
23pub mod openai;
24
25#[cfg(feature = "anthropic")]
26pub use anthropic::AnthropicProvider;
27pub use mock::MockProvider;
28pub use openai::OpenAiProvider;
29
30pub type StreamSender = mpsc::UnboundedSender<String>;
33
34#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
36pub struct TokenUsage {
37 pub prompt_tokens: u32,
38 pub completion_tokens: u32,
39 pub total_tokens: u32,
40 pub cache_hit_tokens: u32,
41 pub cache_miss_tokens: u32,
42}
43
44impl TokenUsage {
45 pub fn accumulate(self, other: TokenUsage) -> TokenUsage {
47 TokenUsage {
48 prompt_tokens: self.prompt_tokens.saturating_add(other.prompt_tokens),
49 completion_tokens: self
50 .completion_tokens
51 .saturating_add(other.completion_tokens),
52 total_tokens: self.total_tokens.saturating_add(other.total_tokens),
53 cache_hit_tokens: self.cache_hit_tokens.saturating_add(other.cache_hit_tokens),
54 cache_miss_tokens: self
55 .cache_miss_tokens
56 .saturating_add(other.cache_miss_tokens),
57 }
58 }
59}
60
61#[derive(Debug, Clone, Copy, PartialEq)]
63pub struct ModelPricing {
64 pub input_per_million: f64,
65 pub output_per_million: f64,
66 pub cache_hit_input_per_million: f64,
69}
70
71impl ModelPricing {
72 pub fn cost_usd(&self, usage: TokenUsage) -> f64 {
74 let in_cost = if usage.cache_hit_tokens > 0 {
75 let cache_hit =
77 usage.cache_hit_tokens as f64 * self.cache_hit_input_per_million / 1_000_000.0;
78 let cache_miss = usage.cache_miss_tokens as f64 * self.input_per_million / 1_000_000.0;
82 cache_hit + cache_miss
83 } else {
84 (usage.prompt_tokens as f64) * self.input_per_million / 1_000_000.0
85 };
86 let out_cost = (usage.completion_tokens as f64) * self.output_per_million / 1_000_000.0;
87 in_cost + out_cost
88 }
89}
90
91pub fn pricing_for(model: &str) -> Option<ModelPricing> {
93 match model {
94 "MiniMax-M2" => Some(ModelPricing {
95 input_per_million: 0.30,
96 output_per_million: 1.20,
97 cache_hit_input_per_million: 0.30,
99 }),
100 "deepseek-chat" | "deepseek-v4-flash" => Some(ModelPricing {
101 input_per_million: 0.27,
102 output_per_million: 1.10,
103 cache_hit_input_per_million: 0.027,
105 }),
106 "deepseek-v4-pro" => Some(ModelPricing {
109 input_per_million: 1.89,
110 output_per_million: 7.70,
111 cache_hit_input_per_million: 0.189,
113 }),
114 "glm-4-flash" => Some(ModelPricing {
115 input_per_million: 0.10,
116 output_per_million: 0.10,
117 cache_hit_input_per_million: 0.10,
119 }),
120 "glm-5.1" => Some(ModelPricing {
124 input_per_million: 0.50,
125 output_per_million: 2.00,
126 cache_hit_input_per_million: 0.50,
128 }),
129 _ => {
130 None
134 }
135 }
136}
137
138pub fn load_pricing_from_yaml(path: &Path) -> Result<HashMap<String, ModelPricing>> {
142 use std::fs;
143 use std::io::{self, BufRead};
144
145 let file = fs::File::open(path).map_err(Error::Io)?;
146 let reader = io::BufReader::new(file);
147
148 let mut models: HashMap<String, ModelPricing> = HashMap::new();
149 let mut current_model: Option<String> = None;
150 let mut current_pricing: Option<ModelPricingBuilder> = None;
151 let mut in_models_section = false;
152
153 for line in reader.lines() {
159 let line = line.map_err(Error::Io)?;
160
161 if line.trim().is_empty() || line.trim().starts_with('#') {
163 continue;
164 }
165
166 let leading_spaces = line.len() - line.trim_start().len();
168 let trimmed = line.trim();
169
170 if trimmed == "models:" {
172 in_models_section = true;
173 continue;
174 }
175
176 if in_models_section {
178 if leading_spaces == 2 && trimmed.ends_with(':') {
181 if let (Some(name), Some(builder)) = (current_model.take(), current_pricing.take())
183 {
184 if let Some(pricing) = builder.build() {
185 models.insert(name, pricing);
186 }
187 }
188 let model_name = trimmed.trim_end_matches(':').to_string();
190 current_model = Some(model_name);
191 current_pricing = Some(ModelPricingBuilder::default());
192 continue;
193 }
194
195 if leading_spaces >= 4 && current_model.is_some() {
198 if let Some(ref mut builder) = current_pricing {
199 if let Some((key, value)) = trimmed.split_once(':') {
200 let key = key.trim();
201 let value = value.split('#').next().unwrap_or(value).trim();
203 if let Err(e) = builder.parse_field(key, value) {
204 return Err(Error::Config {
205 message: format!("error parsing {}: {}", path.display(), e),
206 });
207 }
208 }
209 }
210 }
211 }
212 }
213
214 if let (Some(name), Some(builder)) = (current_model, current_pricing) {
216 if let Some(pricing) = builder.build() {
217 models.insert(name, pricing);
218 }
219 }
220
221 Ok(models)
222}
223
224#[derive(Default)]
226struct ModelPricingBuilder {
227 input_per_million: Option<f64>,
228 output_per_million: Option<f64>,
229 cache_hit_input_per_million: Option<f64>,
230}
231
232impl ModelPricingBuilder {
233 fn parse_field(&mut self, key: &str, value: &str) -> Result<(), String> {
234 let value: f64 = value
235 .parse()
236 .map_err(|_| format!("invalid float: {}", value))?;
237 match key {
238 "input_per_million" => self.input_per_million = Some(value),
239 "output_per_million" => self.output_per_million = Some(value),
240 "cache_hit_input_per_million" => self.cache_hit_input_per_million = Some(value),
241 _ => return Err(format!("unknown field: {}", key)),
242 }
243 Ok(())
244 }
245
246 fn build(self) -> Option<ModelPricing> {
247 Some(ModelPricing {
248 input_per_million: self.input_per_million?,
249 output_per_million: self.output_per_million?,
250 cache_hit_input_per_million: self
251 .cache_hit_input_per_million
252 .unwrap_or(self.input_per_million.unwrap_or(0.0)),
253 })
254 }
255}
256
257#[derive(Debug, Clone, Serialize, Deserialize)]
259pub struct ToolSpec {
260 pub name: String,
261 pub description: String,
262 pub parameters: Value,
264}
265
266#[derive(Debug, Clone, Serialize, Deserialize)]
268pub struct ToolCall {
269 pub id: String,
270 pub name: String,
271 pub arguments: Value,
273}
274
275pub struct StructuredRequest {
277 pub messages: Vec<Message>,
278 pub schema: Value,
280 pub schema_name: String,
282}
283
284#[derive(Debug, Clone)]
286pub struct Completion {
287 pub content: String,
288 pub tool_calls: Vec<ToolCall>,
289 pub finish_reason: Option<String>,
290 pub usage: Option<TokenUsage>,
291}
292
293#[async_trait]
294pub trait LlmProvider: Send + Sync {
295 async fn complete(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Completion>;
296
297 async fn complete_structured(&self, _req: StructuredRequest) -> Result<Value> {
301 Err(Error::Config {
302 message: "provider does not support structured output".into(),
303 })
304 }
305
306 async fn stream(
315 &self,
316 messages: &[Message],
317 tools: &[ToolSpec],
318 stream_tx: Option<StreamSender>,
319 ) -> Result<Completion> {
320 let completion = self.complete(messages, tools).await?;
321 if let Some(tx) = stream_tx {
322 if !completion.content.is_empty() {
323 let _ = tx.send(completion.content.clone());
324 }
325 }
326 Ok(completion)
327 }
328}
329
330#[cfg(test)]
331mod tests {
332 use super::*;
333
334 #[tokio::test]
335 async fn mock_structured_returns_default_error() {
336 let provider = MockProvider::new(vec![]).with_structured_responses(vec![]);
339 let req = StructuredRequest {
340 messages: vec![Message::user("hi".to_string())],
341 schema: serde_json::json!({"type": "object", "properties": {"answer": {"type": "string"}}}),
342 schema_name: "test_schema".to_string(),
343 };
344 let result = provider.complete_structured(req).await;
345 assert!(result.is_err());
346 let err = result.unwrap_err();
347 let msg = err.to_string();
348 assert!(
349 msg.contains("no structured responses configured"),
350 "error should mention no structured responses: {msg}"
351 );
352 }
353
354 #[test]
355 fn token_usage_default_is_all_zeros() {
356 let u = TokenUsage::default();
357 assert_eq!(u.prompt_tokens, 0);
358 assert_eq!(u.completion_tokens, 0);
359 assert_eq!(u.total_tokens, 0);
360 }
361
362 #[test]
363 fn token_usage_accumulate_is_saturating() {
364 let u1 = TokenUsage {
365 prompt_tokens: 10,
366 completion_tokens: 5,
367 total_tokens: 15,
368 cache_hit_tokens: 0,
369 cache_miss_tokens: 0,
370 };
371 let u2 = TokenUsage {
372 prompt_tokens: 20,
373 completion_tokens: 30,
374 total_tokens: 50,
375 cache_hit_tokens: 0,
376 cache_miss_tokens: 0,
377 };
378 let acc = u1.accumulate(u2);
379 assert_eq!(acc.prompt_tokens, 30);
380 assert_eq!(acc.completion_tokens, 35);
381 assert_eq!(acc.total_tokens, 65);
382 }
383
384 #[test]
385 fn token_usage_accumulate_is_commutative() {
386 let u1 = TokenUsage {
387 prompt_tokens: 10,
388 completion_tokens: 5,
389 total_tokens: 15,
390 cache_hit_tokens: 0,
391 cache_miss_tokens: 0,
392 };
393 let u2 = TokenUsage {
394 prompt_tokens: 20,
395 completion_tokens: 30,
396 total_tokens: 50,
397 cache_hit_tokens: 0,
398 cache_miss_tokens: 0,
399 };
400 assert_eq!(u1.accumulate(u2), u2.accumulate(u1));
401 }
402
403 #[test]
404 fn token_usage_accumulate_saturates() {
405 let u1 = TokenUsage {
406 prompt_tokens: u32::MAX,
407 completion_tokens: 1,
408 total_tokens: u32::MAX,
409 cache_hit_tokens: 0,
410 cache_miss_tokens: 0,
411 };
412 let u2 = TokenUsage {
413 prompt_tokens: 1,
414 completion_tokens: u32::MAX,
415 total_tokens: u32::MAX,
416 cache_hit_tokens: 0,
417 cache_miss_tokens: 0,
418 };
419 let acc = u1.accumulate(u2);
420 assert_eq!(acc.prompt_tokens, u32::MAX);
421 assert_eq!(acc.completion_tokens, u32::MAX);
422 assert_eq!(acc.total_tokens, u32::MAX);
423 }
424
425 #[test]
426 fn cost_usd_handles_zero_usage() {
427 let pricing = ModelPricing {
428 input_per_million: 1.0,
429 output_per_million: 2.0,
430 cache_hit_input_per_million: 0.1, };
432 let usage = TokenUsage::default();
433 let cost = pricing.cost_usd(usage);
434 assert!((cost - 0.0).abs() < 1e-9);
435 }
436
437 #[test]
438 fn cost_usd_computes_simple_case() {
439 let pricing = ModelPricing {
440 input_per_million: 1.0,
441 output_per_million: 1.0,
442 cache_hit_input_per_million: 0.1,
443 };
444 let usage = TokenUsage {
446 prompt_tokens: 1_000_000,
447 completion_tokens: 0,
448 total_tokens: 1_000_000,
449 cache_hit_tokens: 0,
450 cache_miss_tokens: 0,
451 };
452 let cost = pricing.cost_usd(usage);
453 assert!((cost - 1.0).abs() < 1e-9);
454 }
455
456 #[test]
457 fn cost_usd_mixes_input_and_output() {
458 let pricing = ModelPricing {
459 input_per_million: 1.0,
460 output_per_million: 2.0,
461 cache_hit_input_per_million: 0.1,
462 };
463 let usage = TokenUsage {
465 prompt_tokens: 500_000,
466 completion_tokens: 250_000,
467 total_tokens: 750_000,
468 cache_hit_tokens: 0,
469 cache_miss_tokens: 0,
470 };
471 let cost = pricing.cost_usd(usage);
472 assert!((cost - 1.0).abs() < 1e-9);
474 }
475
476 #[test]
477 fn pricing_for_known_models() {
478 let p1 = pricing_for("MiniMax-M2");
479 assert!(p1.is_some());
480 assert!((p1.unwrap().input_per_million - 0.30).abs() < 1e-9);
481
482 let p2 = pricing_for("deepseek-chat");
483 assert!(p2.is_some());
484 assert!((p2.unwrap().input_per_million - 0.27).abs() < 1e-9);
485 }
486
487 #[test]
488 fn pricing_for_unknown_returns_none() {
489 let p = pricing_for("unknown-model-xyz");
490 assert!(p.is_none());
491 }
492
493 #[test]
494 fn token_usage_accumulate_cache_fields() {
495 let u1 = TokenUsage {
496 prompt_tokens: 100,
497 completion_tokens: 50,
498 total_tokens: 150,
499 cache_hit_tokens: 60,
500 cache_miss_tokens: 40,
501 };
502 let u2 = TokenUsage {
503 prompt_tokens: 200,
504 completion_tokens: 100,
505 total_tokens: 300,
506 cache_hit_tokens: 120,
507 cache_miss_tokens: 80,
508 };
509 let acc = u1.accumulate(u2);
510 assert_eq!(acc.cache_hit_tokens, 180);
511 assert_eq!(acc.cache_miss_tokens, 120);
512 assert_eq!(acc.prompt_tokens, 300);
513 assert_eq!(acc.completion_tokens, 150);
514 assert_eq!(acc.total_tokens, 450);
515 }
516
517 #[test]
519 fn cost_usd_with_no_cache_hit_matches_old_behavior() {
520 let pricing = ModelPricing {
521 input_per_million: 1.0,
522 output_per_million: 2.0,
523 cache_hit_input_per_million: 0.1, };
525 let usage = TokenUsage {
527 prompt_tokens: 1_000_000,
528 completion_tokens: 500_000,
529 total_tokens: 1_500_000,
530 cache_hit_tokens: 0,
531 cache_miss_tokens: 1_000_000,
532 };
533 let cost = pricing.cost_usd(usage);
535 assert!((cost - 2.0).abs() < 1e-9);
536 }
537
538 #[test]
540 fn cost_usd_with_cache_hit_applies_discount() {
541 let pricing = ModelPricing {
543 input_per_million: 0.27,
544 output_per_million: 1.10,
545 cache_hit_input_per_million: 0.027,
546 };
547 let usage = TokenUsage {
549 prompt_tokens: 1_000,
550 completion_tokens: 500,
551 total_tokens: 1_500,
552 cache_hit_tokens: 900,
553 cache_miss_tokens: 100,
554 };
555 let cost = pricing.cost_usd(usage);
556 let expected =
561 900.0 * 0.027 / 1_000_000.0 + 100.0 * 0.27 / 1_000_000.0 + 500.0 * 1.10 / 1_000_000.0;
562 assert!((cost - expected).abs() < 1e-9);
563 }
564
565 #[test]
567 fn pricing_for_deepseek_has_cache_discount() {
568 let pricing = pricing_for("deepseek-chat").expect("deepseek-chat should be known");
569 assert!((pricing.input_per_million - 0.27).abs() < 1e-9);
571 assert!((pricing.cache_hit_input_per_million - 0.027).abs() < 1e-9);
572 }
573
574 #[test]
576 fn pricing_for_unknown_model_returns_none() {
577 let p = pricing_for("unknown-model-xyz");
578 assert!(p.is_none());
579 }
580
581 #[test]
583 fn token_usage_accumulate_preserves_cache_tokens() {
584 let u1 = TokenUsage {
585 prompt_tokens: 1000,
586 completion_tokens: 100,
587 total_tokens: 1100,
588 cache_hit_tokens: 900,
589 cache_miss_tokens: 100,
590 };
591 let u2 = TokenUsage {
592 prompt_tokens: 2000,
593 completion_tokens: 200,
594 total_tokens: 2200,
595 cache_hit_tokens: 1800,
596 cache_miss_tokens: 200,
597 };
598 let acc = u1.accumulate(u2);
599 assert_eq!(acc.cache_hit_tokens, 2700);
600 assert_eq!(acc.cache_miss_tokens, 300);
601 assert_eq!(acc.prompt_tokens, 3000);
602 }
603
604 #[test]
606 fn load_pricing_from_yaml_parses_file() {
607 let temp_dir = std::env::temp_dir();
608 let yaml_path = temp_dir.join("test_pricing.yaml");
609 std::fs::write(
610 &yaml_path,
611 r#"
612models:
613 test-model:
614 input_per_million: 1.0
615 output_per_million: 2.0
616 cache_hit_input_per_million: 0.1
617"#,
618 )
619 .unwrap();
620
621 let result = load_pricing_from_yaml(&yaml_path);
622 assert!(result.is_ok(), "load should succeed: {:?}", result.err());
623 let pricing = result.unwrap();
624 assert!(
625 pricing.contains_key("test-model"),
626 "should contain test-model: {:?}",
627 pricing.keys().collect::<Vec<_>>()
628 );
629 let p = pricing.get("test-model").unwrap();
630 assert!((p.input_per_million - 1.0).abs() < 1e-9);
631 assert!((p.output_per_million - 2.0).abs() < 1e-9);
632 assert!((p.cache_hit_input_per_million - 0.1).abs() < 1e-9);
633
634 std::fs::remove_file(&yaml_path).ok();
635 }
636
637 #[test]
639 fn load_pricing_from_yaml_overrides_hardcoded() {
640 let temp_dir = std::env::temp_dir();
641 let yaml_path = temp_dir.join("test_pricing_override.yaml");
642 std::fs::write(
644 &yaml_path,
645 r#"
646models:
647 deepseek-chat:
648 input_per_million: 99.0
649 output_per_million: 99.0
650 cache_hit_input_per_million: 9.9
651"#,
652 )
653 .unwrap();
654
655 let result = load_pricing_from_yaml(&yaml_path);
656 assert!(result.is_ok());
657 let pricing = result.unwrap();
658 let p = pricing
659 .get("deepseek-chat")
660 .expect("should have deepseek-chat");
661 assert!((p.input_per_million - 99.0).abs() < 1e-9);
663
664 assert!(pricing.contains_key("deepseek-chat"));
666 assert!(!pricing.contains_key("MiniMax-M2"));
668
669 std::fs::remove_file(&yaml_path).ok();
670 }
671
672 #[test]
674 fn load_pricing_from_yaml_missing_model_falls_back() {
675 let temp_dir = std::env::temp_dir();
680 let yaml_path = temp_dir.join("test_pricing_partial.yaml");
681 std::fs::write(
683 &yaml_path,
684 r#"
685models:
686 test-only:
687 input_per_million: 1.0
688 output_per_million: 1.0
689 cache_hit_input_per_million: 0.1
690"#,
691 )
692 .unwrap();
693
694 let result = load_pricing_from_yaml(&yaml_path);
695 assert!(result.is_ok());
696 let external = result.unwrap();
697
698 assert!(
700 external.contains_key("test-only"),
701 "should have test-only: {:?}",
702 external.keys().collect::<Vec<_>>()
703 );
704 let fallback = pricing_for("MiniMax-M2");
706 assert!(fallback.is_some());
707 assert!((fallback.unwrap().input_per_million - 0.30).abs() < 1e-9);
708
709 std::fs::remove_file(&yaml_path).ok();
710 }
711
712 #[test]
714 fn load_pricing_from_yaml_malformed_error() {
715 let temp_dir = std::env::temp_dir();
716 let yaml_path = temp_dir.join("test_pricing_malformed.yaml");
717 std::fs::write(
719 &yaml_path,
720 r#"
721models:
722 bad-model:
723 input_per_million: not_a_number
724"#,
725 )
726 .unwrap();
727
728 let result = load_pricing_from_yaml(&yaml_path);
729 assert!(result.is_err(), "should fail with bad data");
730 let err = result.unwrap_err();
731 let err_str = err.to_string();
733 assert!(
734 err_str.contains("error parsing") || err_str.contains("invalid float"),
735 "error should mention parsing issue: {}",
736 err_str
737 );
738
739 std::fs::remove_file(&yaml_path).ok();
740 }
741}