mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
clean models manager (#9168)
Have only the following Methods: - `list_models`: getting current available models - `try_list_models`: sync version no refresh for tui use - `get_default_model`: get the default model (should be tightened to core and received on session configuration) - `get_model_info`: get `ModelInfo` for a specific model (should be tightened to core but used in tests) - `refresh_if_new_etag`: trigger refresh on different etags Also move the cache to its own struct
This commit is contained in:
@@ -5,9 +5,95 @@ use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use tokio::fs;
|
||||
use tracing::error;
|
||||
|
||||
/// Manages loading and saving of models cache to disk.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct ModelsCacheManager {
|
||||
cache_path: PathBuf,
|
||||
cache_ttl: Duration,
|
||||
}
|
||||
|
||||
impl ModelsCacheManager {
|
||||
/// Create a new cache manager with the given path and TTL.
|
||||
pub(crate) fn new(cache_path: PathBuf, cache_ttl: Duration) -> Self {
|
||||
Self {
|
||||
cache_path,
|
||||
cache_ttl,
|
||||
}
|
||||
}
|
||||
|
||||
/// Attempt to load a fresh cache entry. Returns `None` if the cache doesn't exist or is stale.
|
||||
pub(crate) async fn load_fresh(&self) -> Option<ModelsCache> {
|
||||
let cache = match self.load().await {
|
||||
Ok(cache) => cache?,
|
||||
Err(err) => {
|
||||
error!("failed to load models cache: {err}");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if !cache.is_fresh(self.cache_ttl) {
|
||||
return None;
|
||||
}
|
||||
Some(cache)
|
||||
}
|
||||
|
||||
/// Persist the cache to disk, creating parent directories as needed.
|
||||
pub(crate) async fn persist_cache(&self, models: &[ModelInfo], etag: Option<String>) {
|
||||
let cache = ModelsCache {
|
||||
fetched_at: Utc::now(),
|
||||
etag,
|
||||
models: models.to_vec(),
|
||||
};
|
||||
if let Err(err) = self.save_internal(&cache).await {
|
||||
error!("failed to write models cache: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
async fn load(&self) -> io::Result<Option<ModelsCache>> {
|
||||
match fs::read(&self.cache_path).await {
|
||||
Ok(contents) => {
|
||||
let cache = serde_json::from_slice(&contents)
|
||||
.map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?;
|
||||
Ok(Some(cache))
|
||||
}
|
||||
Err(err) if err.kind() == ErrorKind::NotFound => Ok(None),
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
async fn save_internal(&self, cache: &ModelsCache) -> io::Result<()> {
|
||||
if let Some(parent) = self.cache_path.parent() {
|
||||
fs::create_dir_all(parent).await?;
|
||||
}
|
||||
let json = serde_json::to_vec_pretty(cache)
|
||||
.map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?;
|
||||
fs::write(&self.cache_path, json).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
/// Set the cache TTL.
|
||||
pub(crate) fn set_ttl(&mut self, ttl: Duration) {
|
||||
self.cache_ttl = ttl;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
/// Manipulate cache file for testing. Allows setting a custom fetched_at timestamp.
|
||||
pub(crate) async fn manipulate_cache_for_test<F>(&self, f: F) -> io::Result<()>
|
||||
where
|
||||
F: FnOnce(&mut DateTime<Utc>),
|
||||
{
|
||||
let mut cache = match self.load().await? {
|
||||
Some(cache) => cache,
|
||||
None => return Err(io::Error::new(ErrorKind::NotFound, "cache not found")),
|
||||
};
|
||||
f(&mut cache.fetched_at);
|
||||
self.save_internal(&cache).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Serialized snapshot of models and metadata cached on disk.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -20,7 +106,7 @@ pub(crate) struct ModelsCache {
|
||||
|
||||
impl ModelsCache {
|
||||
/// Returns `true` when the cache entry has not exceeded the configured TTL.
|
||||
pub(crate) fn is_fresh(&self, ttl: Duration) -> bool {
|
||||
fn is_fresh(&self, ttl: Duration) -> bool {
|
||||
if ttl.is_zero() {
|
||||
return false;
|
||||
}
|
||||
@@ -31,26 +117,3 @@ impl ModelsCache {
|
||||
age <= ttl_duration
|
||||
}
|
||||
}
|
||||
|
||||
/// Read and deserialize the cache file if it exists.
|
||||
pub(crate) async fn load_cache(path: &Path) -> io::Result<Option<ModelsCache>> {
|
||||
match fs::read(path).await {
|
||||
Ok(contents) => {
|
||||
let cache = serde_json::from_slice(&contents)
|
||||
.map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?;
|
||||
Ok(Some(cache))
|
||||
}
|
||||
Err(err) if err.kind() == ErrorKind::NotFound => Ok(None),
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
/// Persist the cache contents to disk, creating parent directories as needed.
|
||||
pub(crate) async fn save_cache(path: &Path, cache: &ModelsCache) -> io::Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).await?;
|
||||
}
|
||||
let json = serde_json::to_vec_pretty(cache)
|
||||
.map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?;
|
||||
fs::write(path, json).await
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user