use lru::LruCache;
use std::num::NonZeroUsize;
use std::sync::{LazyLock, Mutex};
const MAX_CACHE_SIZE: usize = 128;
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum TaskType {
System,
Lightweight,
Specialized,
}
pub trait PromptProvider {
fn cache_key(&self) -> String;
fn task_type(&self) -> TaskType;
}
pub struct SystemPromptCache {
entries: Mutex<LruCache<String, String>>,
}
impl Default for SystemPromptCache {
fn default() -> Self {
Self::new()
}
}
impl SystemPromptCache {
pub fn new() -> Self {
let cache_size = NonZeroUsize::new(MAX_CACHE_SIZE).unwrap_or(NonZeroUsize::MIN);
Self {
entries: Mutex::new(LruCache::new(cache_size)),
}
}
pub fn get_or_insert_with<F>(&self, key: &str, builder: F) -> String
where
F: FnOnce() -> String,
{
{
let mut store = self.entries.lock().unwrap_or_else(|p| p.into_inner());
if let Some(value) = store.get(key) {
return value.clone();
}
}
let value = builder();
{
let mut store = self.entries.lock().unwrap_or_else(|p| p.into_inner());
store.put(key.to_string(), value.clone());
}
value
}
pub fn get(&self, key: &str) -> Option<String> {
let mut store = self.entries.lock().unwrap_or_else(|p| p.into_inner());
store.get(key).cloned()
}
pub fn insert(&self, key: String, value: String) {
let mut store = self.entries.lock().unwrap_or_else(|p| p.into_inner());
store.put(key, value);
}
pub fn clear(&self) {
if let Ok(mut store) = self.entries.lock() {
store.clear();
}
}
}
pub static PROMPT_CACHE: LazyLock<SystemPromptCache> = LazyLock::new(SystemPromptCache::new);