Skip to main content

vv_llm/
settings.rs

1use crate::{default_chat_backends, BackendType, VvLlmError};
2use serde::{Deserialize, Serialize};
3use serde_json::{Map, Value};
4use std::{collections::HashMap, fs, path::Path};
5
6const DEFAULT_CHAT_CONTEXT_LENGTH: u32 = 32_768;
7
8fn default_true() -> bool {
9    true
10}
11
12fn default_endpoint_rpm() -> u32 {
13    60
14}
15
16fn default_endpoint_tpm() -> u32 {
17    300_000
18}
19
20fn default_endpoint_concurrent_requests() -> u32 {
21    20
22}
23
24fn default_rate_limit_backend() -> String {
25    "memory".to_string()
26}
27
28fn default_rate_limit_rpm() -> u32 {
29    60
30}
31
32fn default_rate_limit_tpm() -> u32 {
33    1_000_000
34}
35
36#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37pub struct EndpointConfig {
38    pub id: String,
39    #[serde(default = "default_true")]
40    pub enabled: bool,
41    #[serde(default)]
42    pub api_base: Option<String>,
43    #[serde(default)]
44    pub api_key: Option<String>,
45    #[serde(default)]
46    pub organization: Option<String>,
47    #[serde(default)]
48    pub response_api: bool,
49    #[serde(default)]
50    pub endpoint_type: Option<String>,
51    #[serde(default)]
52    pub region: Option<String>,
53    #[serde(default)]
54    pub is_azure: bool,
55    #[serde(default)]
56    pub is_bedrock: bool,
57    #[serde(default)]
58    pub is_vertex: bool,
59    #[serde(default)]
60    pub credentials: Value,
61    #[serde(default = "default_endpoint_rpm")]
62    pub rpm: u32,
63    #[serde(default = "default_endpoint_tpm")]
64    pub tpm: u32,
65    #[serde(default = "default_endpoint_concurrent_requests")]
66    pub concurrent_requests: u32,
67    #[serde(default)]
68    pub proxy: Option<String>,
69    #[serde(default)]
70    pub headers: Option<HashMap<String, String>>,
71    #[serde(default)]
72    pub access_token: Option<String>,
73    #[serde(default)]
74    pub access_token_expires_at: Option<f64>,
75    #[serde(default)]
76    pub extra: HashMap<String, Value>,
77}
78
79#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
80pub struct BackendConfig {
81    #[serde(default)]
82    pub models: HashMap<String, ModelConfig>,
83    #[serde(default)]
84    pub default_endpoint: Option<String>,
85    #[serde(flatten)]
86    pub extra: HashMap<String, Value>,
87}
88
89#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
90pub struct ModelConfig {
91    pub id: String,
92    #[serde(default = "default_true")]
93    pub enabled: bool,
94    #[serde(default)]
95    pub endpoints: Vec<EndpointBinding>,
96    #[serde(default)]
97    pub context_length: Option<u32>,
98    #[serde(default)]
99    pub max_output_tokens: Option<u32>,
100    #[serde(default)]
101    pub function_call_available: Option<bool>,
102    #[serde(default)]
103    pub response_format_available: Option<bool>,
104    #[serde(default)]
105    pub native_multimodal: Option<bool>,
106    #[serde(default)]
107    pub protocol: Option<String>,
108    #[serde(default)]
109    pub dimensions: Option<u32>,
110    #[serde(default)]
111    pub default_top_n: Option<u32>,
112    #[serde(default)]
113    pub request_mapping: Option<Value>,
114    #[serde(default)]
115    pub response_mapping: Option<Value>,
116    #[serde(flatten)]
117    pub extra: HashMap<String, Value>,
118}
119
120#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
121pub struct ServerConfig {
122    pub host: String,
123    pub port: u16,
124    #[serde(default)]
125    pub url: Option<String>,
126}
127
128#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
129pub struct RateLimitConfig {
130    #[serde(default)]
131    pub enabled: bool,
132    #[serde(default = "default_rate_limit_backend")]
133    pub backend: String,
134    #[serde(default)]
135    pub redis: Option<Value>,
136    #[serde(default)]
137    pub diskcache: Option<Value>,
138    #[serde(default = "default_rate_limit_rpm")]
139    pub default_rpm: u32,
140    #[serde(default = "default_rate_limit_tpm")]
141    pub default_tpm: u32,
142}
143
144#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
145pub struct LlmSettings {
146    #[serde(rename = "VERSION", default)]
147    pub version: Option<String>,
148    #[serde(default)]
149    pub endpoints: Vec<EndpointConfig>,
150    #[serde(default)]
151    pub token_server: Option<ServerConfig>,
152    #[serde(default)]
153    pub rate_limit: Option<RateLimitConfig>,
154    #[serde(default)]
155    pub backends: HashMap<String, BackendConfig>,
156    #[serde(default)]
157    pub embedding_backends: HashMap<String, BackendConfig>,
158    #[serde(default)]
159    pub rerank_backends: HashMap<String, BackendConfig>,
160    #[serde(flatten)]
161    pub extra: HashMap<String, Value>,
162}
163
164#[derive(Debug, Clone, PartialEq)]
165pub struct ResolvedModelConfig {
166    pub backend: String,
167    pub model: ModelConfig,
168    pub model_id: String,
169    pub endpoint: EndpointConfig,
170}
171
172#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
173#[serde(untagged)]
174pub enum EndpointBinding {
175    Id(String),
176    Config {
177        endpoint_id: String,
178        #[serde(default)]
179        model_id: Option<String>,
180        #[serde(default)]
181        enabled: Option<bool>,
182        #[serde(default)]
183        rpm: Option<u32>,
184        #[serde(default)]
185        tpm: Option<u32>,
186        #[serde(default)]
187        concurrent_requests: Option<u32>,
188        #[serde(flatten)]
189        extra: HashMap<String, Value>,
190    },
191}
192
193impl EndpointBinding {
194    pub fn endpoint_id(&self) -> &str {
195        match self {
196            Self::Id(endpoint_id) => endpoint_id,
197            Self::Config { endpoint_id, .. } => endpoint_id,
198        }
199    }
200
201    pub fn model_id<'a>(&'a self, default_model_id: &'a str) -> &'a str {
202        match self {
203            Self::Id(_) => default_model_id,
204            Self::Config { model_id, .. } => model_id.as_deref().unwrap_or(default_model_id),
205        }
206    }
207
208    pub fn enabled(&self) -> bool {
209        match self {
210            Self::Id(_) => true,
211            Self::Config { enabled, .. } => enabled.unwrap_or(true),
212        }
213    }
214}
215
216impl LlmSettings {
217    pub fn from_json_str(raw: &str) -> Result<Self, VvLlmError> {
218        let mut settings: Self = serde_json::from_str(raw)?;
219        settings.normalize_after_load();
220        Ok(settings)
221    }
222
223    pub fn from_json_file(path: impl AsRef<Path>) -> Result<Self, VvLlmError> {
224        let raw = fs::read_to_string(path.as_ref())
225            .map_err(|error| VvLlmError::Configuration(error.to_string()))?;
226        Self::from_json_str(&raw)
227    }
228
229    pub fn resolve_chat_model(
230        &self,
231        backend: BackendType,
232        model_id: &str,
233    ) -> Result<ResolvedModelConfig, VvLlmError> {
234        self.resolve_model_in_map(&self.backends, backend.as_str(), model_id)
235    }
236
237    pub fn resolve_embedding_model(
238        &self,
239        backend: &str,
240        model_id: &str,
241    ) -> Result<ResolvedModelConfig, VvLlmError> {
242        self.resolve_model_in_map(&self.embedding_backends, backend, model_id)
243    }
244
245    pub fn resolve_rerank_model(
246        &self,
247        backend: &str,
248        model_id: &str,
249    ) -> Result<ResolvedModelConfig, VvLlmError> {
250        self.resolve_model_in_map(&self.rerank_backends, backend, model_id)
251    }
252
253    fn normalize_after_load(&mut self) {
254        for endpoint in &mut self.endpoints {
255            normalize_endpoint_transport_flags(endpoint);
256            normalize_google_openai_base(endpoint);
257            preserve_vertex_cached_token(endpoint);
258        }
259
260        self.merge_default_chat_backends();
261        self.apply_python_chat_model_defaults();
262    }
263
264    fn merge_default_chat_backends(&mut self) {
265        for (backend_name, default_backend) in default_chat_backends() {
266            let backend = match self.backends.remove(&backend_name) {
267                Some(user_backend) => merge_backend_config(default_backend, user_backend),
268                None => default_backend,
269            };
270            self.backends.insert(backend_name, backend);
271        }
272
273        for backend in self.backends.values_mut() {
274            apply_default_endpoint(backend);
275        }
276    }
277
278    fn apply_python_chat_model_defaults(&mut self) {
279        for backend in self.backends.values_mut() {
280            for model in backend.models.values_mut() {
281                model
282                    .context_length
283                    .get_or_insert(DEFAULT_CHAT_CONTEXT_LENGTH);
284                model.function_call_available.get_or_insert(false);
285                model.response_format_available.get_or_insert(false);
286                model.native_multimodal.get_or_insert(false);
287            }
288        }
289    }
290
291    fn resolve_model_in_map(
292        &self,
293        map: &HashMap<String, BackendConfig>,
294        backend: &str,
295        model_id: &str,
296    ) -> Result<ResolvedModelConfig, VvLlmError> {
297        let backend_config = map.get(backend).ok_or_else(|| VvLlmError::ModelNotFound {
298            backend: backend.to_string(),
299            model: model_id.to_string(),
300        })?;
301        let model = backend_config
302            .models
303            .get(model_id)
304            .filter(|model| model.enabled)
305            .or_else(|| {
306                backend_config
307                    .models
308                    .values()
309                    .find(|model| model.enabled && model.id == model_id)
310            })
311            .ok_or_else(|| VvLlmError::ModelNotFound {
312                backend: backend.to_string(),
313                model: model_id.to_string(),
314            })?;
315        let binding = model
316            .endpoints
317            .iter()
318            .find(|binding| binding.enabled())
319            .ok_or_else(|| {
320                VvLlmError::Configuration(format!("model {model_id} has no enabled endpoints"))
321            })?;
322        let endpoint_id = binding.endpoint_id();
323        let endpoint = self
324            .endpoints
325            .iter()
326            .find(|endpoint| endpoint.enabled && endpoint.id == endpoint_id)
327            .ok_or_else(|| VvLlmError::EndpointNotFound(endpoint_id.to_string()))?;
328
329        Ok(ResolvedModelConfig {
330            backend: backend.to_string(),
331            model: model.clone(),
332            model_id: binding.model_id(&model.id).to_string(),
333            endpoint: endpoint.clone(),
334        })
335    }
336}
337
338fn merge_backend_config(
339    mut default_backend: BackendConfig,
340    user_backend: BackendConfig,
341) -> BackendConfig {
342    let default_endpoint = user_backend
343        .default_endpoint
344        .clone()
345        .or(default_backend.default_endpoint.clone());
346
347    for (model_name, user_model) in user_backend.models {
348        let merged_model = match default_backend.models.remove(&model_name) {
349            Some(default_model) => merge_model_config(default_model, user_model),
350            None => user_model,
351        };
352        default_backend.models.insert(model_name, merged_model);
353    }
354
355    default_backend.default_endpoint = default_endpoint;
356    default_backend.extra.extend(user_backend.extra);
357    default_backend
358}
359
360fn merge_model_config(mut default_model: ModelConfig, user_model: ModelConfig) -> ModelConfig {
361    default_model.id = user_model.id;
362    default_model.enabled = user_model.enabled;
363    if !user_model.endpoints.is_empty() {
364        default_model.endpoints = user_model.endpoints;
365    }
366    if user_model.context_length.is_some() {
367        default_model.context_length = user_model.context_length;
368    }
369    if user_model.max_output_tokens.is_some() {
370        default_model.max_output_tokens = user_model.max_output_tokens;
371    }
372    if user_model.function_call_available.is_some() {
373        default_model.function_call_available = user_model.function_call_available;
374    }
375    if user_model.response_format_available.is_some() {
376        default_model.response_format_available = user_model.response_format_available;
377    }
378    if user_model.native_multimodal.is_some() {
379        default_model.native_multimodal = user_model.native_multimodal;
380    }
381    if user_model.protocol.is_some() {
382        default_model.protocol = user_model.protocol;
383    }
384    if user_model.dimensions.is_some() {
385        default_model.dimensions = user_model.dimensions;
386    }
387    if user_model.default_top_n.is_some() {
388        default_model.default_top_n = user_model.default_top_n;
389    }
390    if user_model.request_mapping.is_some() {
391        default_model.request_mapping = user_model.request_mapping;
392    }
393    if user_model.response_mapping.is_some() {
394        default_model.response_mapping = user_model.response_mapping;
395    }
396    default_model.extra.extend(user_model.extra);
397    default_model
398}
399
400fn apply_default_endpoint(backend: &mut BackendConfig) {
401    let Some(default_endpoint) = backend.default_endpoint.clone() else {
402        return;
403    };
404    for model in backend.models.values_mut() {
405        if model.endpoints.is_empty() {
406            model
407                .endpoints
408                .push(EndpointBinding::Id(default_endpoint.clone()));
409        }
410    }
411}
412
413fn normalize_endpoint_transport_flags(endpoint: &mut EndpointConfig) {
414    let endpoint_type = endpoint
415        .endpoint_type
416        .as_deref()
417        .unwrap_or_default()
418        .trim()
419        .to_ascii_lowercase();
420    let normalized = if endpoint_type.is_empty() || endpoint_type == "default" {
421        if endpoint.is_azure {
422            Some("openai_azure")
423        } else if endpoint.is_vertex {
424            Some("anthropic_vertex")
425        } else if endpoint.is_bedrock {
426            Some("anthropic_bedrock")
427        } else if endpoint_type == "default" {
428            Some("default")
429        } else {
430            None
431        }
432    } else {
433        Some(endpoint_type.as_str())
434    };
435
436    if let Some(endpoint_type) = normalized {
437        endpoint.endpoint_type = Some(endpoint_type.to_string());
438        endpoint.is_azure = endpoint_type == "openai_azure";
439        endpoint.is_vertex =
440            endpoint_type == "anthropic_vertex" || endpoint_type == "openai_vertex";
441        endpoint.is_bedrock = endpoint_type == "anthropic_bedrock";
442    }
443}
444
445fn normalize_google_openai_base(endpoint: &mut EndpointConfig) {
446    let Some(api_base) = endpoint.api_base.as_deref() else {
447        return;
448    };
449    if api_base.starts_with("https://generativelanguage.googleapis.com/v1beta")
450        && !api_base.ends_with("openai/")
451    {
452        endpoint.api_base = Some(format!("{}/openai/", api_base.trim_end_matches('/')));
453    }
454}
455
456fn preserve_vertex_cached_token(endpoint: &mut EndpointConfig) {
457    if endpoint.access_token.is_none() && endpoint.access_token_expires_at.is_none() {
458        return;
459    }
460
461    if !endpoint.credentials.is_object() {
462        endpoint.credentials = Value::Object(Map::new());
463    }
464    let credentials = endpoint
465        .credentials
466        .as_object_mut()
467        .expect("credentials was normalized to an object");
468
469    if let Some(access_token) = &endpoint.access_token {
470        credentials
471            .entry("access_token")
472            .or_insert_with(|| Value::String(access_token.clone()));
473    }
474    if let Some(expires_at) = endpoint.access_token_expires_at {
475        if let Some(number) = serde_json::Number::from_f64(expires_at) {
476            credentials
477                .entry("access_token_expires_at")
478                .or_insert_with(|| Value::Number(number));
479        }
480    }
481}