use std::time::Duration;
use crate::config::ModelsConfig;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Limits {
pub max_input_elements: usize,
pub max_output_elements: usize,
pub timeout: Duration,
pub max_concurrency: u32,
}
impl Limits {
pub fn effective(config: &ModelsConfig, model_id: &str) -> Self {
let o = config.override_for(model_id);
let pick_usize =
|ceiling: usize, over: Option<usize>| over.map_or(ceiling, |v| v.min(ceiling));
Self {
max_input_elements: pick_usize(
config.max_input_elements,
o.and_then(|o| o.max_input_elements),
),
max_output_elements: pick_usize(
config.max_output_elements,
o.and_then(|o| o.max_output_elements),
),
timeout: Duration::from_millis(
o.and_then(|o| o.timeout_ms)
.map_or(config.max_timeout_ms, |v| v.min(config.max_timeout_ms)),
),
max_concurrency: o
.and_then(|o| o.max_concurrency)
.map_or(config.max_concurrency_per_model, |v| {
v.min(config.max_concurrency_per_model)
}),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::ModelOverride;
#[test]
fn an_override_narrows_and_never_widens() {
let mut config = ModelsConfig::default();
config.overrides.push(ModelOverride {
id: "ada.c4-tiny".to_string(),
timeout_ms: Some(100),
max_concurrency: Some(4),
max_input_elements: None,
max_output_elements: Some(64),
});
let l = Limits::effective(&config, "ada.c4-tiny");
assert_eq!(l.timeout, Duration::from_millis(100));
assert_eq!(l.max_concurrency, 4);
assert_eq!(l.max_input_elements, config.max_input_elements);
assert_eq!(l.max_output_elements, 64);
let host = Limits::effective(&config, "some.other");
assert_eq!(host.timeout, Duration::from_millis(config.max_timeout_ms));
assert_eq!(host.max_concurrency, config.max_concurrency_per_model);
assert_eq!(host.max_input_elements, config.max_input_elements);
assert_eq!(host.max_output_elements, config.max_output_elements);
config.overrides[0].timeout_ms = Some(u64::MAX);
assert_eq!(
Limits::effective(&config, "ada.c4-tiny").timeout,
Duration::from_millis(config.max_timeout_ms)
);
}
}