use std::collections::HashMap;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelSpec {
pub model_id: String,
pub provider: String,
pub available: bool,
}
impl ModelSpec {
pub fn new(provider: impl Into<String>, model_id: impl Into<String>) -> Self {
Self {
model_id: model_id.into(),
provider: provider.into(),
available: true,
}
}
pub fn is_available(&self) -> bool {
self.available
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct CategoryRouter {
routes: HashMap<String, Vec<ModelSpec>>,
}
impl CategoryRouter {
pub fn new() -> Self {
Self::default()
}
pub fn register(&mut self, category: impl Into<String>, models: Vec<ModelSpec>) {
self.routes.insert(category.into(), models);
}
pub fn resolve<'a>(&'a self, category: &str) -> Option<&'a ModelSpec> {
self.routes
.get(category)
.and_then(|models| models.iter().find(|m| m.is_available()))
}
pub fn mark_unavailable(&mut self, category: &str, model_id: &str) {
if let Some(models) = self.routes.get_mut(category) {
for m in models.iter_mut() {
if m.model_id == model_id {
m.available = false;
}
}
}
}
pub fn mark_available(&mut self, category: &str, model_id: &str) {
if let Some(models) = self.routes.get_mut(category) {
for m in models.iter_mut() {
if m.model_id == model_id {
m.available = true;
}
}
}
}
pub fn categories(&self) -> impl Iterator<Item = &str> {
self.routes.keys().map(|k| k.as_str())
}
}