diff --git a/src/lib.rs b/src/lib.rs index e62adcd..6f63ef9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,7 +3,9 @@ pub mod budget; pub mod client; pub mod persona; +pub mod registry; pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier}; pub use client::{ClientError, ModelClient, OpenRouterClient}; pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier}; +pub use registry::{ModelInfo, ModelPricing, ModelRegistry}; diff --git a/src/registry.rs b/src/registry.rs new file mode 100644 index 0000000..7745470 --- /dev/null +++ b/src/registry.rs @@ -0,0 +1,223 @@ +use crate::persona::{ModelCandidate, Tier}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::RwLock; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ModelPricing { + #[serde(default)] + pub prompt: String, + #[serde(default)] + pub completion: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ModelInfo { + pub id: String, + #[serde(default)] + pub name: Option, + #[serde(default)] + pub context_length: Option, + #[serde(default)] + pub pricing: Option, +} + +#[derive(Debug, Deserialize)] +struct OpenRouterModelsResponse { + data: Vec, +} + +#[derive(Clone)] +pub struct ModelRegistry { + models: Arc>>, + last_refreshed: Arc>>, +} + +impl Default for ModelRegistry { + fn default() -> Self { + Self::new() + } +} + +impl ModelRegistry { + pub fn new() -> Self { + Self { + models: Arc::new(RwLock::new(HashMap::new())), + last_refreshed: Arc::new(RwLock::new(None)), + } + } + + pub async fn model_count(&self) -> usize { + let read = self.models.read().await; + read.len() + } + + pub async fn last_refreshed_at(&self) -> Option { + let read = self.last_refreshed.read().await; + *read + } + + pub async fn update_from_json(&self, json_str: &str) -> Result { + let resp: OpenRouterModelsResponse = serde_json::from_str(json_str) + .map_err(|e| format!("Failed to parse OpenRouter models response: {e}"))?; + + let count = resp.data.len(); + let mut map = HashMap::with_capacity(count); + for m in resp.data { + map.insert(m.id.clone(), m); + } + + { + let mut write = self.models.write().await; + *write = map; + } + { + let mut time_write = self.last_refreshed.write().await; + *time_write = Some(Instant::now()); + } + + Ok(count) + } + + /// Fetches live models catalog from OpenRouter `/v1/models` + pub async fn fetch_live(&self, base_url: &str) -> Result { + let endpoint = format!("{}/models", base_url.trim_end_matches('/')); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(15)) + .build() + .map_err(|e| e.to_string())?; + + let res = client + .get(&endpoint) + .header("HTTP-Referer", "https://github.com/brokecrew/brokecrew") + .header("X-Title", "BrokeCrew") + .send() + .await + .map_err(|e| format!("Network error fetching models catalog: {e}"))?; + + if !res.status().is_success() { + return Err(format!("OpenRouter models endpoint returned HTTP {}", res.status())); + } + + let body = res + .text() + .await + .map_err(|e| format!("Failed to read models response body: {e}"))?; + + self.update_from_json(&body).await + } + + /// Starts a background task refreshing the models registry periodically + pub fn start_background_sync(self: Arc, base_url: String, interval: Duration) { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(interval); + loop { + ticker.tick().await; + if let Err(e) = self.fetch_live(&base_url).await { + tracing::warn!("Periodic models registry sync failed: {}", e); + } else { + tracing::debug!("Periodic models registry sync succeeded"); + } + } + }); + } + + pub async fn is_model_available(&self, model_id: &str) -> bool { + let read = self.models.read().await; + read.contains_key(model_id) + } + + /// Filters and orders candidates according to live availability, stealth permissions, and escalation. + /// If the live registry is empty (e.g. offline mode / uninitialized), candidate presence checks are relaxed. + pub async fn filter_eligible<'a>( + &self, + candidates: &'a [ModelCandidate], + allow_stealth: bool, + require_paid: bool, + ) -> Vec<&'a ModelCandidate> { + let read = self.models.read().await; + let is_empty = read.is_empty(); + + candidates + .iter() + .filter(|c| { + if !allow_stealth && c.is_stealth { + return false; + } + if require_paid && c.tier != Tier::Paid { + return false; + } + // Skip unlisted models if we have a populated registry + if !is_empty && !read.contains_key(&c.id) { + return false; + } + true + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_registry_update_and_filtering() { + let registry = ModelRegistry::new(); + + let json_data = r#"{ + "data": [ + { + "id": "qwen/qwen-2.5-coder-32b-instruct:free", + "name": "Qwen 2.5 Coder 32B", + "context_length": 32768, + "pricing": { "prompt": "0", "completion": "0" } + }, + { + "id": "google/gemini-2.0-flash-001", + "name": "Gemini 2.0 Flash", + "context_length": 1048576, + "pricing": { "prompt": "0.0000001", "completion": "0.0000004" } + } + ] + }"#; + + let count = registry.update_from_json(json_data).await.expect("update failed"); + assert_eq!(count, 2); + + assert!(registry.is_model_available("qwen/qwen-2.5-coder-32b-instruct:free").await); + assert!(!registry.is_model_available("unlisted/model").await); + + let candidates = vec![ + ModelCandidate { + id: "stealth/experiment".into(), + tier: Tier::Free, + is_stealth: true, + }, + ModelCandidate { + id: "qwen/qwen-2.5-coder-32b-instruct:free".into(), + tier: Tier::Free, + is_stealth: false, + }, + ModelCandidate { + id: "google/gemini-2.0-flash-001".into(), + tier: Tier::Paid, + is_stealth: false, + }, + ModelCandidate { + id: "unlisted/dropped-model:free".into(), + tier: Tier::Free, + is_stealth: false, + }, + ]; + + // Eligible without stealth, free-tier + let filtered = registry.filter_eligible(&candidates, false, false).await; + assert_eq!(filtered.len(), 2); + assert_eq!(filtered[0].id, "qwen/qwen-2.5-coder-32b-instruct:free"); + assert_eq!(filtered[1].id, "google/gemini-2.0-flash-001"); + // "unlisted/dropped-model:free" is excluded + } +}