crabllm_core/
config.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4/// Per-model token pricing configuration.
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct PricingConfig {
7    /// Cost per million prompt tokens in USD.
8    pub prompt_cost_per_million: f64,
9    /// Cost per million completion tokens in USD.
10    pub completion_cost_per_million: f64,
11}
12
13/// Compute the cost in USD for a given number of prompt and completion tokens.
14pub fn cost(pricing: &PricingConfig, prompt_tokens: u32, completion_tokens: u32) -> f64 {
15    (prompt_tokens as f64 * pricing.prompt_cost_per_million
16        + completion_tokens as f64 * pricing.completion_cost_per_million)
17        / 1_000_000.0
18}
19
20/// Top-level gateway configuration, loaded from TOML.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct GatewayConfig {
23    /// Address to listen on, e.g. "0.0.0.0:8080".
24    pub listen: String,
25    /// Named provider configurations.
26    pub providers: HashMap<String, ProviderConfig>,
27    /// Virtual API keys for client authentication.
28    #[serde(default)]
29    pub keys: Vec<KeyConfig>,
30    /// Extension configurations. Each key is an extension name, value is its config.
31    #[serde(default)]
32    pub extensions: Option<serde_json::Value>,
33    /// Storage backend configuration.
34    #[serde(default)]
35    pub storage: Option<StorageConfig>,
36    /// Model name aliases. Maps friendly names to canonical model names.
37    #[serde(default)]
38    pub aliases: HashMap<String, String>,
39    /// Per-model token pricing for cost tracking and budget enforcement.
40    #[serde(default)]
41    pub pricing: HashMap<String, PricingConfig>,
42    /// Admin API bearer token. If set, enables /v1/admin/* endpoints.
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub admin_token: Option<String>,
45    /// Graceful shutdown timeout in seconds. Default: 30.
46    #[serde(default = "default_shutdown_timeout")]
47    pub shutdown_timeout: u64,
48}
49
50/// Configuration for a single LLM provider.
51#[derive(Debug, Default, Clone, Serialize, Deserialize)]
52pub struct ProviderConfig {
53    /// Provider kind determines the dispatch path.
54    #[serde(
55        default,
56        alias = "standard",
57        skip_serializing_if = "ProviderKind::is_default"
58    )]
59    pub kind: ProviderKind,
60    /// API key (supports `${ENV_VAR}` interpolation).
61    #[serde(default, skip_serializing_if = "Option::is_none")]
62    pub api_key: Option<String>,
63    /// Base URL override. OpenAI-compat providers have sensible defaults.
64    #[serde(default, skip_serializing_if = "Option::is_none")]
65    pub base_url: Option<String>,
66    /// Model names served by this provider.
67    #[serde(default, skip_serializing_if = "Vec::is_empty")]
68    pub models: Vec<String>,
69    /// Routing weight for weighted random selection. Higher = more traffic.
70    #[serde(default, skip_serializing_if = "Option::is_none")]
71    pub weight: Option<u16>,
72    /// Max retries on transient errors before fallback. 0 disables retry.
73    #[serde(default, skip_serializing_if = "Option::is_none")]
74    pub max_retries: Option<u32>,
75    /// API version string, used by Azure OpenAI.
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub api_version: Option<String>,
78    /// Per-request timeout in seconds. Default: 30.
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub timeout: Option<u64>,
81    /// AWS region for Bedrock provider.
82    #[serde(default, skip_serializing_if = "Option::is_none")]
83    pub region: Option<String>,
84    /// AWS access key ID for Bedrock provider.
85    #[serde(default, skip_serializing_if = "Option::is_none")]
86    pub access_key: Option<String>,
87    /// AWS secret access key for Bedrock provider.
88    #[serde(default, skip_serializing)]
89    pub secret_key: Option<String>,
90    /// Path to a GGUF model file for the LlamaCpp provider.
91    #[serde(default, skip_serializing_if = "Option::is_none")]
92    pub model_path: Option<String>,
93    /// Number of GPU layers to offload (LlamaCpp). Default: 0 (CPU only).
94    #[serde(default, skip_serializing_if = "Option::is_none")]
95    pub n_gpu_layers: Option<u32>,
96    /// Context size in tokens (LlamaCpp). Default: 2048.
97    #[serde(default, skip_serializing_if = "Option::is_none")]
98    pub n_ctx: Option<u32>,
99    /// Number of threads for inference (LlamaCpp). Default: system-chosen.
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub n_threads: Option<u32>,
102}
103
104fn default_shutdown_timeout() -> u64 {
105    30
106}
107
108/// Which provider implementation to use.
109#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
110#[serde(rename_all = "snake_case")]
111pub enum ProviderKind {
112    #[default]
113    Openai,
114    Anthropic,
115    Google,
116    Bedrock,
117    Ollama,
118    Azure,
119    #[serde(alias = "llama_cpp")]
120    LlamaCpp,
121}
122
123impl ProviderKind {
124    /// Returns true if this is the default variant (Openai).
125    pub fn is_default(&self) -> bool {
126        *self == Self::Openai
127    }
128}
129
130impl ProviderConfig {
131    /// Resolve the effective provider kind.
132    ///
133    /// Returns `Anthropic` if the field is explicitly set to `Anthropic`,
134    /// or if `base_url` contains "anthropic". Otherwise returns the
135    /// configured kind.
136    pub fn effective_kind(&self) -> ProviderKind {
137        if self.kind == ProviderKind::Anthropic {
138            return ProviderKind::Anthropic;
139        }
140        if let Some(url) = &self.base_url
141            && url.contains("anthropic")
142        {
143            return ProviderKind::Anthropic;
144        }
145        self.kind
146    }
147
148    /// Validate field combinations.
149    pub fn validate(&self, provider_name: &str) -> Result<(), String> {
150        if self.models.is_empty() {
151            return Err(format!("provider '{provider_name}' has no models"));
152        }
153        match self.kind {
154            ProviderKind::Bedrock => {
155                if self.region.is_none() {
156                    return Err(format!(
157                        "provider '{provider_name}' (bedrock) requires region"
158                    ));
159                }
160                if self.access_key.is_none() {
161                    return Err(format!(
162                        "provider '{provider_name}' (bedrock) requires access_key"
163                    ));
164                }
165                if self.secret_key.is_none() {
166                    return Err(format!(
167                        "provider '{provider_name}' (bedrock) requires secret_key"
168                    ));
169                }
170            }
171            ProviderKind::Ollama => {
172                // Ollama doesn't require api_key or base_url.
173            }
174            ProviderKind::LlamaCpp => match &self.model_path {
175                None => {
176                    return Err(format!(
177                        "provider '{provider_name}' (llamacpp) requires model_path"
178                    ));
179                }
180                Some(path) => {
181                    if !std::path::Path::new(path).exists() {
182                        return Err(format!(
183                            "provider '{provider_name}' (llamacpp): model_path '{path}' does not exist"
184                        ));
185                    }
186                }
187            },
188            _ => {
189                if self.api_key.is_none() && self.base_url.is_none() {
190                    return Err(format!(
191                        "provider '{provider_name}' requires api_key or base_url"
192                    ));
193                }
194            }
195        }
196        Ok(())
197    }
198}
199
200/// Virtual API key for client authentication.
201#[derive(Debug, Clone, Serialize, Deserialize)]
202pub struct KeyConfig {
203    /// Human-readable name for this key.
204    pub name: String,
205    /// The key string clients send in Authorization header.
206    pub key: String,
207    /// Which models this key can access. `["*"]` means all.
208    pub models: Vec<String>,
209}
210
211/// Storage backend configuration.
212#[derive(Debug, Clone, Serialize, Deserialize)]
213pub struct StorageConfig {
214    /// Backend kind: "memory" (default) or "sqlite" (requires feature).
215    #[serde(default = "StorageConfig::default_kind")]
216    pub kind: String,
217    /// File path for persistent backends (required for sqlite).
218    #[serde(default)]
219    pub path: Option<String>,
220}
221
222impl StorageConfig {
223    fn default_kind() -> String {
224        "memory".to_string()
225    }
226}
227
228impl GatewayConfig {
229    /// Load config from a TOML file, expanding `${VAR}` patterns in string values.
230    #[cfg(feature = "gateway")]
231    pub fn from_file(path: &std::path::Path) -> Result<Self, Box<dyn std::error::Error>> {
232        let raw = std::fs::read_to_string(path)?;
233        let expanded = expand_env_vars(&raw);
234        let config: GatewayConfig = toml::from_str(&expanded)?;
235        Ok(config)
236    }
237}
238
239/// Expand `${VAR}` patterns in a string using environment variables.
240/// Unknown variables are replaced with empty string.
241#[cfg(feature = "gateway")]
242fn expand_env_vars(input: &str) -> String {
243    let mut result = String::with_capacity(input.len());
244    let mut chars = input.chars().peekable();
245
246    while let Some(c) = chars.next() {
247        if c == '$' && chars.peek() == Some(&'{') {
248            chars.next(); // consume '{'
249            let mut var_name = String::new();
250            for ch in chars.by_ref() {
251                if ch == '}' {
252                    break;
253                }
254                var_name.push(ch);
255            }
256            if let Ok(val) = std::env::var(&var_name) {
257                result.push_str(&val);
258            }
259        } else {
260            result.push(c);
261        }
262    }
263
264    result
265}
crabllm_core/config.rs

crabllm_core/
config.rs
