pi/
provider.rs

1//! LLM provider abstraction layer.
2//!
3//! This module defines the [`Provider`] trait and shared request/response types used by all
4//! backends (Anthropic/OpenAI/Gemini/etc).
5//!
6//! Providers are responsible for:
7//! - Translating [`crate::model::Message`] history into provider-specific HTTP requests.
8//! - Emitting [`StreamEvent`] values as SSE/HTTP chunks arrive.
9//! - Advertising tool schemas to the model (so it can call [`crate::tools`] by name).
10
11pub use crate::model::StreamEvent;
12use crate::model::{Message, ThinkingLevel};
13use async_trait::async_trait;
14use futures::Stream;
15use serde::{Deserialize, Serialize};
16use std::borrow::Cow;
17use std::collections::HashMap;
18use std::pin::Pin;
19
20// ============================================================================
21// Provider Trait
22// ============================================================================
23
24/// An LLM backend capable of streaming assistant output (and tool calls).
25///
26/// A `Provider` is typically configured for a specific API + model and is used by the agent loop
27/// to produce a stream of [`StreamEvent`] updates.
28#[async_trait]
29pub trait Provider: Send + Sync {
30    /// Get the provider name.
31    fn name(&self) -> &str;
32
33    /// Get the API type.
34    fn api(&self) -> &str;
35
36    /// Get the model identifier used by this provider.
37    fn model_id(&self) -> &str;
38
39    /// Start streaming a completion.
40    ///
41    /// Implementations should yield [`StreamEvent`] items as soon as they are decoded, and should
42    /// stop promptly when the request is cancelled.
43    async fn stream(
44        &self,
45        context: &Context<'_>,
46        options: &StreamOptions,
47    ) -> crate::error::Result<Pin<Box<dyn Stream<Item = crate::error::Result<StreamEvent>> + Send>>>;
48}
49
50// ============================================================================
51// Context
52// ============================================================================
53
54/// Inputs to a single completion request.
55///
56/// The agent loop builds a `Context` from the current session state and tool registry, then hands
57/// it to a [`Provider`] implementation to perform provider-specific request encoding.
58///
59/// Uses [`Cow`] for `messages` and `tools` to avoid deep-cloning the full conversation history on
60/// every turn when no mutation is needed (the common case).
61#[derive(Debug, Clone)]
62pub struct Context<'a> {
63    /// Provider-specific system prompt content.
64    ///
65    /// Uses [`Cow`] to borrow from `AgentConfig.system_prompt` on every turn without
66    /// cloning.  Providers that need an owned `String` can call `.into_owned()`.
67    pub system_prompt: Option<Cow<'a, str>>,
68    /// Conversation history (user/assistant/tool results).
69    pub messages: Cow<'a, [Message]>,
70    /// Tool definitions available to the model for this request.
71    pub tools: Cow<'a, [ToolDef]>,
72}
73
74impl Default for Context<'_> {
75    fn default() -> Self {
76        Self {
77            system_prompt: None,
78            messages: Cow::Owned(Vec::new()),
79            tools: Cow::Owned(Vec::new()),
80        }
81    }
82}
83
84impl Context<'_> {
85    /// Create a `Context` with fully-owned data (no borrowing).
86    ///
87    /// Convenient for tests and one-off callers that already have owned vectors.
88    pub fn owned(
89        system_prompt: Option<String>,
90        messages: Vec<Message>,
91        tools: Vec<ToolDef>,
92    ) -> Context<'static> {
93        Context {
94            system_prompt: system_prompt.map(Cow::Owned),
95            messages: Cow::Owned(messages),
96            tools: Cow::Owned(tools),
97        }
98    }
99}
100
101// ============================================================================
102// Tool Definition
103// ============================================================================
104
105/// A tool definition exposed to the model.
106///
107/// Providers translate this struct into the backend's tool/schema representation (typically JSON
108/// Schema) so the model can emit tool calls that the host executes locally.
109#[derive(Debug, Clone)]
110pub struct ToolDef {
111    pub name: String,
112    pub description: String,
113    pub parameters: serde_json::Value, // JSON Schema
114}
115
116// ============================================================================
117// Stream Options
118// ============================================================================
119
120/// Options that control streaming completion behavior.
121///
122/// Most options are passed through to the provider request (temperature, max tokens, headers).
123/// Some fields are Pi-specific conveniences (e.g. `session_id` for logging/correlation).
124#[derive(Debug, Clone, Default)]
125pub struct StreamOptions {
126    pub temperature: Option<f32>,
127    pub max_tokens: Option<u32>,
128    pub api_key: Option<String>,
129    pub cache_retention: CacheRetention,
130    pub session_id: Option<String>,
131    pub headers: HashMap<String, String>,
132    pub thinking_level: Option<ThinkingLevel>,
133    pub thinking_budgets: Option<ThinkingBudgets>,
134}
135
136/// Cache retention policy.
137#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
138pub enum CacheRetention {
139    #[default]
140    None,
141    /// Provider-managed short-lived caching (provider-specific semantics).
142    Short,
143    /// Provider-managed long-lived caching (e.g. ~1 hour TTL on Anthropic).
144    Long,
145}
146
147/// Custom thinking token budgets per level.
148#[derive(Debug, Clone)]
149pub struct ThinkingBudgets {
150    pub minimal: u32,
151    pub low: u32,
152    pub medium: u32,
153    pub high: u32,
154    pub xhigh: u32,
155    pub max: u32,
156}
157
158impl Default for ThinkingBudgets {
159    fn default() -> Self {
160        Self {
161            minimal: 1024,
162            low: 2048,
163            medium: 8192,
164            high: 16384,
165            xhigh: 32768, // Default to double high, or model max? Let's pick a reasonable default.
166            max: 65536,
167        }
168    }
169}
170
171// ============================================================================
172// Model Definition
173// ============================================================================
174
175/// A model definition loaded from the models registry.
176///
177/// This struct is used to drive provider selection, request limits (context window/max tokens),
178/// and cost accounting.
179#[derive(Debug, Clone, Serialize)]
180pub struct Model {
181    pub id: String,
182    pub name: String,
183    pub api: String,
184    pub provider: String,
185    pub base_url: String,
186    pub reasoning: bool,
187    pub input: Vec<InputType>,
188    pub cost: ModelCost,
189    pub context_window: u32,
190    pub max_tokens: u32,
191    pub headers: HashMap<String, String>,
192}
193
194/// Input types supported by a model.
195#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
196#[serde(rename_all = "lowercase")]
197pub enum InputType {
198    Text,
199    Image,
200}
201
202/// Model pricing per million tokens.
203#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
204#[serde(rename_all = "camelCase")]
205pub struct ModelCost {
206    pub input: f64,
207    pub output: f64,
208    pub cache_read: f64,
209    pub cache_write: f64,
210}
211
212impl Model {
213    /// Calculate cost for usage.
214    #[allow(clippy::cast_precision_loss)] // Token counts within practical range won't lose precision
215    pub fn calculate_cost(
216        &self,
217        input: u64,
218        output: u64,
219        cache_read: u64,
220        cache_write: u64,
221    ) -> f64 {
222        let input_cost = (self.cost.input / 1_000_000.0) * input as f64;
223        let output_cost = (self.cost.output / 1_000_000.0) * output as f64;
224        let cache_read_cost = (self.cost.cache_read / 1_000_000.0) * cache_read as f64;
225        let cache_write_cost = (self.cost.cache_write / 1_000_000.0) * cache_write as f64;
226        input_cost + output_cost + cache_read_cost + cache_write_cost
227    }
228}
229
230// ============================================================================
231// Known APIs and Providers
232// ============================================================================
233
234/// Known API types.
235#[derive(Debug, Clone, PartialEq, Eq)]
236pub enum Api {
237    AnthropicMessages,
238    OpenAICompletions,
239    OpenAIResponses,
240    OpenAICodexResponses,
241    AzureOpenAIResponses,
242    BedrockConverseStream,
243    GoogleGenerativeAI,
244    GoogleGeminiCli,
245    GoogleVertex,
246    Custom(String),
247}
248
249impl std::fmt::Display for Api {
250    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
251        match self {
252            Self::AnthropicMessages => write!(f, "anthropic-messages"),
253            Self::OpenAICompletions => write!(f, "openai-completions"),
254            Self::OpenAIResponses => write!(f, "openai-responses"),
255            Self::OpenAICodexResponses => write!(f, "openai-codex-responses"),
256            Self::AzureOpenAIResponses => write!(f, "azure-openai-responses"),
257            Self::BedrockConverseStream => write!(f, "bedrock-converse-stream"),
258            Self::GoogleGenerativeAI => write!(f, "google-generative-ai"),
259            Self::GoogleGeminiCli => write!(f, "google-gemini-cli"),
260            Self::GoogleVertex => write!(f, "google-vertex"),
261            Self::Custom(s) => write!(f, "{s}"),
262        }
263    }
264}
265
266impl std::str::FromStr for Api {
267    type Err = String;
268
269    fn from_str(s: &str) -> Result<Self, Self::Err> {
270        match s {
271            "anthropic-messages" => Ok(Self::AnthropicMessages),
272            "openai-completions" => Ok(Self::OpenAICompletions),
273            "openai-responses" => Ok(Self::OpenAIResponses),
274            "openai-codex-responses" => Ok(Self::OpenAICodexResponses),
275            "azure-openai-responses" => Ok(Self::AzureOpenAIResponses),
276            "bedrock-converse-stream" => Ok(Self::BedrockConverseStream),
277            "google-generative-ai" => Ok(Self::GoogleGenerativeAI),
278            "google-gemini-cli" => Ok(Self::GoogleGeminiCli),
279            "google-vertex" => Ok(Self::GoogleVertex),
280            other if !other.is_empty() => Ok(Self::Custom(other.to_string())),
281            _ => Err("API identifier cannot be empty".to_string()),
282        }
283    }
284}
285
286/// Known providers.
287#[derive(Debug, Clone, PartialEq, Eq)]
288#[allow(clippy::upper_case_acronyms)] // These are proper names/brands
289pub enum KnownProvider {
290    Anthropic,
291    OpenAI,
292    Google,
293    GoogleVertex,
294    AmazonBedrock,
295    AzureOpenAI,
296    GithubCopilot,
297    XAI,
298    Groq,
299    Cerebras,
300    OpenRouter,
301    Mistral,
302    Custom(String),
303}
304
305impl std::fmt::Display for KnownProvider {
306    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
307        match self {
308            Self::Anthropic => write!(f, "anthropic"),
309            Self::OpenAI => write!(f, "openai"),
310            Self::Google => write!(f, "google"),
311            Self::GoogleVertex => write!(f, "google-vertex"),
312            Self::AmazonBedrock => write!(f, "amazon-bedrock"),
313            Self::AzureOpenAI => write!(f, "azure-openai"),
314            Self::GithubCopilot => write!(f, "github-copilot"),
315            Self::XAI => write!(f, "xai"),
316            Self::Groq => write!(f, "groq"),
317            Self::Cerebras => write!(f, "cerebras"),
318            Self::OpenRouter => write!(f, "openrouter"),
319            Self::Mistral => write!(f, "mistral"),
320            Self::Custom(s) => write!(f, "{s}"),
321        }
322    }
323}
324
325impl std::str::FromStr for KnownProvider {
326    type Err = String;
327
328    fn from_str(s: &str) -> Result<Self, Self::Err> {
329        match s {
330            "anthropic" => Ok(Self::Anthropic),
331            "openai" => Ok(Self::OpenAI),
332            "google" => Ok(Self::Google),
333            "google-vertex" => Ok(Self::GoogleVertex),
334            "amazon-bedrock" => Ok(Self::AmazonBedrock),
335            "azure-openai" | "azure" | "azure-cognitive-services" => Ok(Self::AzureOpenAI),
336            "github-copilot" => Ok(Self::GithubCopilot),
337            "xai" => Ok(Self::XAI),
338            "groq" => Ok(Self::Groq),
339            "cerebras" => Ok(Self::Cerebras),
340            "openrouter" => Ok(Self::OpenRouter),
341            "mistral" => Ok(Self::Mistral),
342            other if !other.is_empty() => Ok(Self::Custom(other.to_string())),
343            _ => Err("Provider identifier cannot be empty".to_string()),
344        }
345    }
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351
352    // ========================================================================
353    // Api enum: FromStr + Display round-trips
354    // ========================================================================
355
356    #[test]
357    fn api_from_str_known_variants() {
358        let cases = [
359            ("anthropic-messages", Api::AnthropicMessages),
360            ("openai-completions", Api::OpenAICompletions),
361            ("openai-responses", Api::OpenAIResponses),
362            ("openai-codex-responses", Api::OpenAICodexResponses),
363            ("azure-openai-responses", Api::AzureOpenAIResponses),
364            ("bedrock-converse-stream", Api::BedrockConverseStream),
365            ("google-generative-ai", Api::GoogleGenerativeAI),
366            ("google-gemini-cli", Api::GoogleGeminiCli),
367            ("google-vertex", Api::GoogleVertex),
368        ];
369        for (input, expected) in &cases {
370            let parsed: Api = input.parse().unwrap();
371            assert_eq!(&parsed, expected, "from_str({input})");
372        }
373    }
374
375    #[test]
376    fn api_display_known_variants() {
377        let cases = [
378            (Api::AnthropicMessages, "anthropic-messages"),
379            (Api::OpenAICompletions, "openai-completions"),
380            (Api::OpenAIResponses, "openai-responses"),
381            (Api::OpenAICodexResponses, "openai-codex-responses"),
382            (Api::AzureOpenAIResponses, "azure-openai-responses"),
383            (Api::BedrockConverseStream, "bedrock-converse-stream"),
384            (Api::GoogleGenerativeAI, "google-generative-ai"),
385            (Api::GoogleGeminiCli, "google-gemini-cli"),
386            (Api::GoogleVertex, "google-vertex"),
387        ];
388        for (variant, expected) in &cases {
389            assert_eq!(&variant.to_string(), expected, "display for {variant:?}");
390        }
391    }
392
393    #[test]
394    fn api_round_trip_all_known() {
395        let variants = [
396            Api::AnthropicMessages,
397            Api::OpenAICompletions,
398            Api::OpenAIResponses,
399            Api::OpenAICodexResponses,
400            Api::AzureOpenAIResponses,
401            Api::BedrockConverseStream,
402            Api::GoogleGenerativeAI,
403            Api::GoogleGeminiCli,
404            Api::GoogleVertex,
405        ];
406        for variant in &variants {
407            let s = variant.to_string();
408            let parsed: Api = s.parse().unwrap();
409            assert_eq!(&parsed, variant, "round-trip failed for {variant:?} -> {s}");
410        }
411    }
412
413    #[test]
414    fn api_custom_variant() {
415        let parsed: Api = "my-custom-api".parse().unwrap();
416        assert_eq!(parsed, Api::Custom("my-custom-api".to_string()));
417        assert_eq!(parsed.to_string(), "my-custom-api");
418    }
419
420    #[test]
421    fn api_empty_string_rejected() {
422        let result: Result<Api, _> = "".parse();
423        assert!(result.is_err());
424        assert_eq!(result.unwrap_err(), "API identifier cannot be empty");
425    }
426
427    // ========================================================================
428    // KnownProvider enum: FromStr + Display round-trips
429    // ========================================================================
430
431    #[test]
432    fn provider_from_str_known_variants() {
433        let cases = [
434            ("anthropic", KnownProvider::Anthropic),
435            ("openai", KnownProvider::OpenAI),
436            ("google", KnownProvider::Google),
437            ("google-vertex", KnownProvider::GoogleVertex),
438            ("amazon-bedrock", KnownProvider::AmazonBedrock),
439            ("azure-openai", KnownProvider::AzureOpenAI),
440            ("azure", KnownProvider::AzureOpenAI),
441            ("azure-cognitive-services", KnownProvider::AzureOpenAI),
442            ("github-copilot", KnownProvider::GithubCopilot),
443            ("xai", KnownProvider::XAI),
444            ("groq", KnownProvider::Groq),
445            ("cerebras", KnownProvider::Cerebras),
446            ("openrouter", KnownProvider::OpenRouter),
447            ("mistral", KnownProvider::Mistral),
448        ];
449        for (input, expected) in &cases {
450            let parsed: KnownProvider = input.parse().unwrap();
451            assert_eq!(&parsed, expected, "from_str({input})");
452        }
453    }
454
455    #[test]
456    fn provider_display_known_variants() {
457        let cases = [
458            (KnownProvider::Anthropic, "anthropic"),
459            (KnownProvider::OpenAI, "openai"),
460            (KnownProvider::Google, "google"),
461            (KnownProvider::GoogleVertex, "google-vertex"),
462            (KnownProvider::AmazonBedrock, "amazon-bedrock"),
463            (KnownProvider::AzureOpenAI, "azure-openai"),
464            (KnownProvider::GithubCopilot, "github-copilot"),
465            (KnownProvider::XAI, "xai"),
466            (KnownProvider::Groq, "groq"),
467            (KnownProvider::Cerebras, "cerebras"),
468            (KnownProvider::OpenRouter, "openrouter"),
469            (KnownProvider::Mistral, "mistral"),
470        ];
471        for (variant, expected) in &cases {
472            assert_eq!(&variant.to_string(), expected, "display for {variant:?}");
473        }
474    }
475
476    #[test]
477    fn provider_round_trip_all_known() {
478        let variants = [
479            KnownProvider::Anthropic,
480            KnownProvider::OpenAI,
481            KnownProvider::Google,
482            KnownProvider::GoogleVertex,
483            KnownProvider::AmazonBedrock,
484            KnownProvider::AzureOpenAI,
485            KnownProvider::GithubCopilot,
486            KnownProvider::XAI,
487            KnownProvider::Groq,
488            KnownProvider::Cerebras,
489            KnownProvider::OpenRouter,
490            KnownProvider::Mistral,
491        ];
492        for variant in &variants {
493            let s = variant.to_string();
494            let parsed: KnownProvider = s.parse().unwrap();
495            assert_eq!(&parsed, variant, "round-trip failed for {variant:?} -> {s}");
496        }
497    }
498
499    #[test]
500    fn provider_custom_variant() {
501        let parsed: KnownProvider = "my-custom-provider".parse().unwrap();
502        assert_eq!(
503            parsed,
504            KnownProvider::Custom("my-custom-provider".to_string())
505        );
506        assert_eq!(parsed.to_string(), "my-custom-provider");
507    }
508
509    #[test]
510    fn provider_empty_string_rejected() {
511        let result: Result<KnownProvider, _> = "".parse();
512        assert!(result.is_err());
513        assert_eq!(result.unwrap_err(), "Provider identifier cannot be empty");
514    }
515
516    // ========================================================================
517    // Model::calculate_cost
518    // ========================================================================
519
520    fn test_model() -> Model {
521        Model {
522            id: "test-model".to_string(),
523            name: "Test Model".to_string(),
524            api: "anthropic-messages".to_string(),
525            provider: "anthropic".to_string(),
526            base_url: "https://api.anthropic.com".to_string(),
527            reasoning: false,
528            input: vec![InputType::Text],
529            cost: ModelCost {
530                input: 3.0,   // $3 per million input tokens
531                output: 15.0, // $15 per million output tokens
532                cache_read: 0.3,
533                cache_write: 3.75,
534            },
535            context_window: 200_000,
536            max_tokens: 8192,
537            headers: HashMap::new(),
538        }
539    }
540
541    #[test]
542    fn calculate_cost_basic() {
543        let model = test_model();
544        // 1000 input tokens at $3/M = $0.003
545        // 500 output tokens at $15/M = $0.0075
546        let cost = model.calculate_cost(1000, 500, 0, 0);
547        let input_expected = (3.0 / 1_000_000.0) * 1000.0;
548        let output_expected = (15.0 / 1_000_000.0) * 500.0;
549        let expected = input_expected + output_expected;
550        assert!(
551            (cost - expected).abs() < f64::EPSILON,
552            "expected {expected}, got {cost}"
553        );
554    }
555
556    #[test]
557    fn calculate_cost_with_cache() {
558        let model = test_model();
559        let cost = model.calculate_cost(1000, 500, 2000, 1000);
560        let input_expected = (3.0 / 1_000_000.0) * 1000.0;
561        let output_expected = (15.0 / 1_000_000.0) * 500.0;
562        let cache_read_expected = (0.3 / 1_000_000.0) * 2000.0;
563        let cache_write_expected = (3.75 / 1_000_000.0) * 1000.0;
564        let expected =
565            input_expected + output_expected + cache_read_expected + cache_write_expected;
566        assert!(
567            (cost - expected).abs() < 1e-12,
568            "expected {expected}, got {cost}"
569        );
570    }
571
572    #[test]
573    fn calculate_cost_zero_tokens() {
574        let model = test_model();
575        let cost = model.calculate_cost(0, 0, 0, 0);
576        assert!((cost).abs() < f64::EPSILON, "zero tokens should cost $0");
577    }
578
579    #[test]
580    fn calculate_cost_large_token_count() {
581        let model = test_model();
582        // 1 million tokens each
583        let cost = model.calculate_cost(1_000_000, 1_000_000, 0, 0);
584        let expected = 3.0 + 15.0; // $3 input + $15 output
585        assert!(
586            (cost - expected).abs() < 1e-10,
587            "expected {expected}, got {cost}"
588        );
589    }
590
591    // ========================================================================
592    // Default values
593    // ========================================================================
594
595    #[test]
596    fn thinking_budgets_default() {
597        let budgets = ThinkingBudgets::default();
598        assert_eq!(budgets.minimal, 1024);
599        assert_eq!(budgets.low, 2048);
600        assert_eq!(budgets.medium, 8192);
601        assert_eq!(budgets.high, 16384);
602        assert_eq!(budgets.xhigh, 32768);
603    }
604
605    #[test]
606    fn cache_retention_default_is_none() {
607        assert_eq!(CacheRetention::default(), CacheRetention::None);
608    }
609
610    #[test]
611    fn stream_options_default() {
612        let opts = StreamOptions::default();
613        assert!(opts.temperature.is_none());
614        assert!(opts.max_tokens.is_none());
615        assert!(opts.api_key.is_none());
616        assert_eq!(opts.cache_retention, CacheRetention::None);
617        assert!(opts.session_id.is_none());
618        assert!(opts.headers.is_empty());
619        assert!(opts.thinking_level.is_none());
620        assert!(opts.thinking_budgets.is_none());
621    }
622
623    #[test]
624    fn context_default() {
625        let ctx = Context::default();
626        assert!(ctx.system_prompt.is_none());
627        assert!(ctx.messages.is_empty());
628        assert!(ctx.tools.is_empty());
629    }
630
631    // ========================================================================
632    // InputType serde
633    // ========================================================================
634
635    #[test]
636    fn input_type_serialization() {
637        let text_json = serde_json::to_string(&InputType::Text).unwrap();
638        assert_eq!(text_json, "\"text\"");
639
640        let image_json = serde_json::to_string(&InputType::Image).unwrap();
641        assert_eq!(image_json, "\"image\"");
642
643        let text: InputType = serde_json::from_str("\"text\"").unwrap();
644        assert_eq!(text, InputType::Text);
645
646        let image: InputType = serde_json::from_str("\"image\"").unwrap();
647        assert_eq!(image, InputType::Image);
648    }
649
650    // ========================================================================
651    // ModelCost serde
652    // ========================================================================
653
654    #[test]
655    fn model_cost_camel_case_serialization() {
656        let cost = ModelCost {
657            input: 3.0,
658            output: 15.0,
659            cache_read: 0.3,
660            cache_write: 3.75,
661        };
662        let json = serde_json::to_string(&cost).unwrap();
663        assert!(
664            json.contains("\"cacheRead\""),
665            "should use camelCase: {json}"
666        );
667        assert!(
668            json.contains("\"cacheWrite\""),
669            "should use camelCase: {json}"
670        );
671
672        let parsed: ModelCost = serde_json::from_str(&json).unwrap();
673        assert!((parsed.input - 3.0).abs() < f64::EPSILON);
674        assert!((parsed.cache_read - 0.3).abs() < f64::EPSILON);
675    }
676
677    mod proptests {
678        use super::*;
679        use proptest::prelude::*;
680
681        fn arb_model(rate: f64) -> Model {
682            Model {
683                id: "m".to_string(),
684                name: "m".to_string(),
685                api: "anthropic-messages".to_string(),
686                provider: "test".to_string(),
687                base_url: String::new(),
688                reasoning: false,
689                input: vec![InputType::Text],
690                cost: ModelCost {
691                    input: rate,
692                    output: rate,
693                    cache_read: rate,
694                    cache_write: rate,
695                },
696                context_window: 128_000,
697                max_tokens: 8192,
698                headers: HashMap::new(),
699            }
700        }
701
702        // ====================================================================
703        // calculate_cost
704        // ====================================================================
705
706        proptest! {
707            #[test]
708            fn cost_zero_tokens_is_zero(rate in 0.0f64..1000.0) {
709                let m = arb_model(rate);
710                let cost = m.calculate_cost(0, 0, 0, 0);
711                assert!((cost).abs() < f64::EPSILON);
712            }
713
714            #[test]
715            fn cost_non_negative(
716                rate in 0.0f64..100.0,
717                input in 0u64..10_000_000,
718                output in 0u64..10_000_000,
719                cr in 0u64..10_000_000,
720                cw in 0u64..10_000_000,
721            ) {
722                let m = arb_model(rate);
723                assert!(m.calculate_cost(input, output, cr, cw) >= 0.0);
724            }
725
726            #[test]
727            fn cost_linearity(
728                rate in 0.001f64..50.0,
729                tokens in 1u64..1_000_000,
730            ) {
731                let m = arb_model(rate);
732                let single = m.calculate_cost(tokens, 0, 0, 0);
733                let double = m.calculate_cost(tokens * 2, 0, 0, 0);
734                assert!((double - single * 2.0).abs() < 1e-6,
735                    "doubling tokens should double cost: single={single}, double={double}");
736            }
737
738            #[test]
739            fn cost_additivity(
740                rate in 0.001f64..50.0,
741                input in 0u64..1_000_000,
742                output in 0u64..1_000_000,
743            ) {
744                let m = arb_model(rate);
745                let combined = m.calculate_cost(input, output, 0, 0);
746                let separate = m.calculate_cost(input, 0, 0, 0)
747                    + m.calculate_cost(0, output, 0, 0);
748                assert!((combined - separate).abs() < 1e-10,
749                    "cost should be additive: combined={combined}, separate={separate}");
750            }
751        }
752
753        // ====================================================================
754        // Api FromStr + Display round-trip
755        // ====================================================================
756
757        proptest! {
758            #[test]
759            fn api_custom_round_trip(s in "[a-z][a-z0-9-]{0,20}") {
760                let known = [
761                    "anthropic-messages", "openai-completions", "openai-responses", "openai-codex-responses",
762                    "azure-openai-responses", "bedrock-converse-stream",
763                    "google-generative-ai", "google-gemini-cli", "google-vertex",
764                ];
765                if !known.contains(&s.as_str()) {
766                    let parsed: Api = s.parse().unwrap();
767                    assert_eq!(parsed.to_string(), s);
768                }
769            }
770
771            #[test]
772            fn api_never_panics(s in ".*") {
773                let _ = s.parse::<Api>(); // must not panic
774            }
775
776            #[test]
777            fn api_empty_always_rejected(ws in "[ \t]*") {
778                if ws.is_empty() {
779                    assert!(ws.parse::<Api>().is_err());
780                }
781            }
782        }
783
784        // ====================================================================
785        // KnownProvider FromStr + Display round-trip
786        // ====================================================================
787
788        proptest! {
789            #[test]
790            fn provider_custom_round_trip(s in "[a-z][a-z0-9-]{0,20}") {
791                let known = [
792                    "anthropic", "openai", "google", "google-vertex",
793                    "amazon-bedrock", "azure-openai", "azure",
794                    "azure-cognitive-services", "github-copilot", "xai",
795                    "groq", "cerebras", "openrouter", "mistral",
796                ];
797                if !known.contains(&s.as_str()) {
798                    let parsed: KnownProvider = s.parse().unwrap();
799                    assert_eq!(parsed.to_string(), s);
800                }
801            }
802
803            #[test]
804            fn provider_never_panics(s in ".*") {
805                let _ = s.parse::<KnownProvider>(); // must not panic
806            }
807        }
808
809        // ====================================================================
810        // ModelCost serde round-trip
811        // ====================================================================
812
813        proptest! {
814            #[test]
815            fn model_cost_serde_round_trip(
816                input in 0.0f64..1000.0,
817                output in 0.0f64..1000.0,
818                cr in 0.0f64..1000.0,
819                cw in 0.0f64..1000.0,
820            ) {
821                let cost = ModelCost { input, output, cache_read: cr, cache_write: cw };
822                let json = serde_json::to_string(&cost).unwrap();
823                let parsed: ModelCost = serde_json::from_str(&json).unwrap();
824                assert!((parsed.input - input).abs() < 1e-10);
825                assert!((parsed.output - output).abs() < 1e-10);
826                assert!((parsed.cache_read - cr).abs() < 1e-10);
827                assert!((parsed.cache_write - cw).abs() < 1e-10);
828            }
829        }
830    }
831}
pi/provider.rs

pi/
provider.rs
