use std::num::NonZeroUsize;
use lru::LruCache;
use sha2::{Digest, Sha256};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct EmbeddingCacheStats {
pub hit_count: u64,
pub miss_count: u64,
pub len: usize,
}
pub struct EmbeddingCache {
pub inner: LruCache<String, Vec<f32>>,
pub hit_count: u64,
pub miss_count: u64,
}
impl EmbeddingCache {
pub fn new(capacity: usize) -> Self {
let capacity = NonZeroUsize::new(capacity.max(1)).unwrap();
EmbeddingCache {
inner: LruCache::new(capacity),
hit_count: 0,
miss_count: 0,
}
}
pub fn key(text: &str) -> String {
let digest = Sha256::digest(text.as_bytes());
hex::encode(digest)
}
pub fn get(&mut self, text: &str) -> Option<Vec<f32>> {
let key = Self::key(text);
if let Some(vector) = self.inner.get(&key) {
self.hit_count += 1;
Some(vector.clone())
} else {
self.miss_count += 1;
None
}
}
pub fn insert(&mut self, text: &str, vector: Vec<f32>) {
let key = Self::key(text);
self.inner.put(key, vector);
}
pub fn stats(&self) -> EmbeddingCacheStats {
EmbeddingCacheStats {
hit_count: self.hit_count,
miss_count: self.miss_count,
len: self.inner.len(),
}
}
}