use std::collections::HashMap;
use crate::driver::{DriverId, ModelSpec};
#[derive(Debug, Clone, PartialEq, Default)]
#[non_exhaustive]
pub struct ModelProfile {
pub context_window: Option<u64>,
pub max_output_tokens: Option<u64>,
pub reasoning_efforts: Vec<String>,
pub thinking: bool,
pub vision: bool,
}
impl ModelProfile {
pub fn new() -> Self {
Self::default()
}
pub fn context_window(mut self, tokens: u64) -> Self {
self.context_window = Some(tokens);
self
}
pub fn max_output_tokens(mut self, tokens: u64) -> Self {
self.max_output_tokens = Some(tokens);
self
}
pub fn reasoning_efforts(
mut self,
efforts: impl IntoIterator<Item = impl Into<String>>,
) -> Self {
self.reasoning_efforts = efforts.into_iter().map(Into::into).collect();
self
}
pub fn thinking(mut self, supported: bool) -> Self {
self.thinking = supported;
self
}
pub fn vision(mut self, supported: bool) -> Self {
self.vision = supported;
self
}
pub fn validate(&self, model: &ModelSpec) -> Result<(), String> {
let Some(reasoning) = &model.reasoning else {
return Ok(());
};
if let Some(effort) = &reasoning.effort {
if self.reasoning_efforts.is_empty() {
return Err(format!(
"model `{}` takes no reasoning effort, but `{effort}` was requested",
model.model
));
}
if !self.reasoning_efforts.iter().any(|known| known == effort) {
return Err(format!(
"model `{}` does not support reasoning effort `{effort}` (supported: {})",
model.model,
self.reasoning_efforts.join(", ")
));
}
}
if reasoning.budget_tokens.is_some() && !self.thinking {
return Err(format!(
"model `{}` does not support extended thinking, but a thinking budget was set",
model.model
));
}
Ok(())
}
}
pub trait ModelCatalog: Send + Sync {
fn profile(&self, driver: &DriverId, model: &str) -> Option<ModelProfile>;
}
#[derive(Debug, Clone, Default)]
pub struct InMemoryModelCatalog {
profiles: HashMap<(String, String), ModelProfile>,
}
impl InMemoryModelCatalog {
pub fn new() -> Self {
Self::default()
}
pub fn with(
mut self,
driver: DriverId,
model: impl Into<String>,
profile: ModelProfile,
) -> Self {
self.profiles
.insert((driver.as_str().to_string(), model.into()), profile);
self
}
}
impl ModelCatalog for InMemoryModelCatalog {
fn profile(&self, driver: &DriverId, model: &str) -> Option<ModelProfile> {
self.profiles
.get(&(driver.as_str().to_string(), model.to_string()))
.cloned()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn gpt() -> ModelProfile {
ModelProfile::new()
.context_window(400_000)
.reasoning_efforts(["low", "high"])
}
#[test]
fn an_empty_profile_only_rejects_what_it_contradicts() {
let empty = ModelProfile::new();
assert!(empty.validate(&ModelSpec::openai("mystery")).is_ok());
assert!(
empty
.validate(&ModelSpec::openai("mystery").reasoning_effort("high"))
.is_err()
);
}
#[test]
fn a_supported_effort_passes_and_an_unsupported_one_names_the_alternatives() {
assert!(
gpt()
.validate(&ModelSpec::openai("gpt").reasoning_effort("high"))
.is_ok()
);
let error = gpt()
.validate(&ModelSpec::openai("gpt").reasoning_effort("ultra"))
.expect_err("unsupported effort");
assert!(error.contains("ultra"), "{error}");
assert!(error.contains("low, high"), "{error}");
}
#[test]
fn a_thinking_budget_needs_a_thinking_model() {
let error = gpt()
.validate(&ModelSpec::openai("gpt").thinking_budget(4_000))
.expect_err("no thinking support");
assert!(error.contains("extended thinking"), "{error}");
let claude = ModelProfile::new().thinking(true);
assert!(
claude
.validate(&ModelSpec::anthropic("claude").thinking_budget(4_000))
.is_ok()
);
}
#[test]
fn a_catalog_is_keyed_by_driver_and_model() {
let catalog = InMemoryModelCatalog::new().with(DriverId::openai(), "gpt", gpt());
assert!(catalog.profile(&DriverId::openai(), "gpt").is_some());
assert!(
catalog.profile(&DriverId::anthropic(), "gpt").is_none(),
"same name on another protocol is another model"
);
assert!(catalog.profile(&DriverId::openai(), "unknown").is_none());
}
}