mod file_meta;
pub use file_meta::{normalize_path, normalize_path_str, FileMetaStore};
use moka::sync::Cache;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
pub struct EmbeddingCache {
cache: Cache<String, Arc<Vec<f32>>>,
hits: AtomicU64,
misses: AtomicU64,
max_memory_mb: usize,
}
impl EmbeddingCache {
#[allow(dead_code)] pub fn new(max_memory_mb: usize) -> Self {
let avg_embedding_size = 384 * std::mem::size_of::<f32>(); let max_entries = (max_memory_mb * 1024 * 1024) / avg_embedding_size;
let cache = Cache::builder()
.max_capacity(max_entries as u64)
.weigher(|_key: &String, value: &Arc<Vec<f32>>| {
(value.len() * std::mem::size_of::<f32>()) as u32
})
.build();
Self {
cache,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
max_memory_mb,
}
}
#[allow(dead_code)] pub fn get_or_compute<F>(&self, key: &str, compute: F) -> Arc<Vec<f32>>
where
F: FnOnce() -> Vec<f32>,
{
if let Some(value) = self.cache.get(key) {
self.hits.fetch_add(1, Ordering::Relaxed);
value
} else {
self.misses.fetch_add(1, Ordering::Relaxed);
let value = Arc::new(compute());
self.cache.insert(key.to_string(), value.clone());
value
}
}
#[allow(dead_code)] pub fn hit_rate(&self) -> f64 {
let hits = self.hits.load(Ordering::Relaxed);
let misses = self.misses.load(Ordering::Relaxed);
let total = hits + misses;
if total == 0 {
0.0
} else {
hits as f64 / total as f64
}
}
#[allow(dead_code)] pub fn stats(&self) -> CacheStats {
CacheStats {
hits: self.hits.load(Ordering::Relaxed),
misses: self.misses.load(Ordering::Relaxed),
size: self.cache.entry_count(),
max_memory_mb: self.max_memory_mb,
}
}
}
#[derive(Debug, Clone)]
#[allow(dead_code)] pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub size: u64,
pub max_memory_mb: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache() {
let cache = EmbeddingCache::new(100);
let result = cache.get_or_compute("test", || vec![1.0, 2.0, 3.0]);
assert_eq!(*result, vec![1.0, 2.0, 3.0]);
let stats = cache.stats();
assert_eq!(stats.misses, 1);
assert_eq!(stats.hits, 0);
}
}