85 lines
2.3 KiB
Rust
85 lines
2.3 KiB
Rust
use crate::types::ModelInfo;
|
|
use crate::Result;
|
|
use std::future::Future;
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::RwLock;
|
|
|
|
#[derive(Default, Debug)]
|
|
struct ModelCache {
|
|
models: Vec<ModelInfo>,
|
|
last_refresh: Option<Instant>,
|
|
}
|
|
|
|
/// Caches model listings for improved selection performance
|
|
#[derive(Clone, Debug)]
|
|
pub struct ModelManager {
|
|
cache: Arc<RwLock<ModelCache>>,
|
|
ttl: Duration,
|
|
}
|
|
|
|
impl ModelManager {
|
|
/// Create a new manager with the desired cache TTL
|
|
pub fn new(ttl: Duration) -> Self {
|
|
Self {
|
|
cache: Arc::new(RwLock::new(ModelCache::default())),
|
|
ttl,
|
|
}
|
|
}
|
|
|
|
/// Get cached models, refreshing via the provided fetcher when stale. Returns the up-to-date model list.
|
|
pub async fn get_or_refresh<F, Fut>(
|
|
&self,
|
|
force_refresh: bool,
|
|
fetcher: F,
|
|
) -> Result<Vec<ModelInfo>>
|
|
where
|
|
F: FnOnce() -> Fut,
|
|
Fut: Future<Output = Result<Vec<ModelInfo>>>,
|
|
{
|
|
if !force_refresh {
|
|
if let Some(models) = self.cached_if_fresh().await {
|
|
return Ok(models);
|
|
}
|
|
}
|
|
|
|
let models = fetcher().await?;
|
|
let mut cache = self.cache.write().await;
|
|
cache.models = models.clone();
|
|
cache.last_refresh = Some(Instant::now());
|
|
Ok(models)
|
|
}
|
|
|
|
/// Return cached models without refreshing
|
|
pub async fn cached(&self) -> Vec<ModelInfo> {
|
|
self.cache.read().await.models.clone()
|
|
}
|
|
|
|
/// Drop cached models, forcing next call to refresh
|
|
pub async fn invalidate(&self) {
|
|
let mut cache = self.cache.write().await;
|
|
cache.models.clear();
|
|
cache.last_refresh = None;
|
|
}
|
|
|
|
/// Select a model by id or name from the cache
|
|
pub async fn select(&self, identifier: &str) -> Option<ModelInfo> {
|
|
let cache = self.cache.read().await;
|
|
cache
|
|
.models
|
|
.iter()
|
|
.find(|m| m.id == identifier || m.name == identifier)
|
|
.cloned()
|
|
}
|
|
|
|
async fn cached_if_fresh(&self) -> Option<Vec<ModelInfo>> {
|
|
let cache = self.cache.read().await;
|
|
let fresh = matches!(cache.last_refresh, Some(ts) if ts.elapsed() < self.ttl);
|
|
if fresh && !cache.models.is_empty() {
|
|
Some(cache.models.clone())
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
}
|