53 lines
1.4 KiB
Rust
53 lines
1.4 KiB
Rust
|
use crate::store::StoreOps;
|
||
|
|
use serde_json::json;
|
||
|
|
|
||
|
|
const KEY_SELECTED_MODEL: &str = "selected_model";
|
||
|
|
|
||
|
|
pub fn get_model_preference(store: &dyn StoreOps) -> Result<Option<String>, String> {
|
||
|
|
if let Some(model) = store
|
||
|
|
.get(KEY_SELECTED_MODEL)
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|val| val.as_str())
|
||
|
|
{
|
||
|
|
return Ok(Some(model.to_string()));
|
||
|
|
}
|
||
|
|
Ok(None)
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn set_model_preference(model: String, store: &dyn StoreOps) -> Result<(), String> {
|
||
|
|
store.set(KEY_SELECTED_MODEL, json!(model));
|
||
|
|
store.save()?;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
use crate::store::JsonFileStore;
|
||
|
|
use tempfile::tempdir;
|
||
|
|
|
||
|
|
fn make_store(dir: &tempfile::TempDir) -> JsonFileStore {
|
||
|
|
JsonFileStore::new(dir.path().join("test_store.json")).unwrap()
|
||
|
|
}
|
||
|
|
|
||
|
|
// --- model preference ---
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn model_preference_none_by_default() {
|
||
|
|
let dir = tempdir().unwrap();
|
||
|
|
let store = make_store(&dir);
|
||
|
|
assert!(get_model_preference(&store).unwrap().is_none());
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn set_and_get_model_preference() {
|
||
|
|
let dir = tempdir().unwrap();
|
||
|
|
let store = make_store(&dir);
|
||
|
|
set_model_preference("claude-3-sonnet".to_string(), &store).unwrap();
|
||
|
|
assert_eq!(
|
||
|
|
get_model_preference(&store).unwrap(),
|
||
|
|
Some("claude-3-sonnet".to_string())
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|