oxirs_vec/reranking/
cache.rs1use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::sync::RwLock;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct RerankingCacheConfig {
9 pub max_size: usize,
10}
11
12pub struct RerankingCache {
13 cache: RwLock<HashMap<String, f32>>,
14 max_size: usize,
15}
16
17impl RerankingCache {
18 pub fn new(max_size: usize) -> Self {
19 Self {
20 cache: RwLock::new(HashMap::new()),
21 max_size,
22 }
23 }
24
25 pub fn get(&self, key: &str) -> Option<f32> {
26 self.cache.read().ok()?.get(key).copied()
27 }
28
29 pub fn put(&self, key: String, value: f32) {
30 if let Ok(mut cache) = self.cache.write() {
31 if cache.len() >= self.max_size {
32 if let Some(first_key) = cache.keys().next().cloned() {
34 cache.remove(&first_key);
35 }
36 }
37 cache.insert(key, value);
38 }
39 }
40
41 pub fn clear(&self) {
42 if let Ok(mut cache) = self.cache.write() {
43 cache.clear();
44 }
45 }
46
47 pub fn stats(&self) -> (usize, usize) {
48 if let Ok(cache) = self.cache.read() {
49 (cache.len(), self.max_size)
50 } else {
51 (0, self.max_size)
52 }
53 }
54}