273 lines
8.6 KiB
Rust
273 lines
8.6 KiB
Rust
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<String>,
|
|
#[serde(default)]
|
|
pub context_length: Option<u64>,
|
|
#[serde(default)]
|
|
pub pricing: Option<ModelPricing>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct OpenRouterModelsResponse {
|
|
data: Vec<ModelInfo>,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct ModelRegistry {
|
|
models: Arc<RwLock<HashMap<String, ModelInfo>>>,
|
|
last_refreshed: Arc<RwLock<Option<Instant>>>,
|
|
}
|
|
|
|
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<Instant> {
|
|
let read = self.last_refreshed.read().await;
|
|
*read
|
|
}
|
|
|
|
pub async fn update_from_json(&self, json_str: &str) -> Result<usize, String> {
|
|
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<usize, String> {
|
|
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<Self>, 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()
|
|
}
|
|
|
|
/// Discovers active `:free` models in the catalog that best match a persona role hint.
|
|
pub async fn discover_free_candidates_for_role(&self, role_hint: &str) -> Vec<ModelCandidate> {
|
|
let read = self.models.read().await;
|
|
let mut matching = Vec::new();
|
|
let mut generic_free = Vec::new();
|
|
|
|
let hint_lower = role_hint.to_lowercase();
|
|
let keywords: &[&str] = match hint_lower.as_str() {
|
|
"coder" => &["code", "coder", "qwen", "gemma", "north"],
|
|
"reviewer" => &["review", "gemma", "nemotron", "reasoning", "qwen"],
|
|
"architect" => &["ultra", "super", "550b", "120b", "gemma", "nemotron"],
|
|
_ => &["gemma", "qwen", "nemotron", "liquid", "inkling", "ling"],
|
|
};
|
|
|
|
for id in read.keys() {
|
|
if id.ends_with(":free") {
|
|
let id_lower = id.to_lowercase();
|
|
if keywords.iter().any(|&kw| id_lower.contains(kw)) {
|
|
matching.push(ModelCandidate {
|
|
id: id.clone(),
|
|
tier: Tier::Free,
|
|
is_stealth: false,
|
|
});
|
|
} else {
|
|
generic_free.push(ModelCandidate {
|
|
id: id.clone(),
|
|
tier: Tier::Free,
|
|
is_stealth: false,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
matching.sort_by(|a, b| a.id.cmp(&b.id));
|
|
generic_free.sort_by(|a, b| a.id.cmp(&b.id));
|
|
|
|
matching.extend(generic_free);
|
|
matching
|
|
}
|
|
}
|
|
|
|
#[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
|
|
}
|
|
}
|