lc_core/cache/
llm_cache.rs

1// src/core/cache/llm_cache.rs
2//! LLM 调用缓存实现
3//!
4//! 基于内存的 LRU 缓存，缓存重复的 LLM 调用结果。
5//! 支持可选的 TTL 过期和最大条目限制。
6
7use crate::language_models::LLMResult;
8use lc_schema::Message;
9use std::collections::HashMap;
10use std::time::{Duration, Instant};
11use tokio::sync::RwLock;
12
13/// 缓存的 LLM 结果，包含过期时间
14#[derive(Debug, Clone)]
15pub struct CachedLLMResult {
16    /// LLM 返回结果
17    pub result: LLMResult,
18    /// 缓存时间戳
19    pub cached_at: Instant,
20}
21
22/// 缓存配置
23#[derive(Debug, Clone)]
24pub struct CacheConfig {
25    /// 最大缓存条目数（0 表示不限制）
26    pub max_entries: usize,
27    /// TTL 过期时间（None 表示永不过期）
28    pub ttl: Option<Duration>,
29    /// 是否启用
30    pub enabled: bool,
31}
32
33impl Default for CacheConfig {
34    fn default() -> Self {
35        Self {
36            max_entries: 1000,
37            ttl: Some(Duration::from_secs(3600)), // 默认 1 小时过期
38            enabled: true,
39        }
40    }
41}
42
43impl CacheConfig {
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    /// 禁用 TTL（永不过期）
49    pub fn no_ttl(mut self) -> Self {
50        self.ttl = None;
51        self
52    }
53
54    /// 设置 TTL
55    pub fn with_ttl(mut self, ttl: Duration) -> Self {
56        self.ttl = Some(ttl);
57        self
58    }
59
60    /// 设置最大条目数
61    pub fn with_max_entries(mut self, max: usize) -> Self {
62        self.max_entries = max;
63        self
64    }
65
66    /// 禁用缓存
67    pub fn disabled(mut self) -> Self {
68        self.enabled = false;
69        self
70    }
71}
72
73/// LLM 调用缓存
74///
75/// 缓存 LLM 调用的输入输出，避免相同请求重复调用 API。
76///
77/// # 示例
78/// ```ignore
79/// use langchainrust::core::cache::LLMCache;
80///
81/// let cache = LLMCache::new();
82/// cache.put("key", llm_result).await;
83///
84/// if let Some(cached) = cache.get("key").await {
85///     println!("缓存命中: {}", cached.result.content);
86/// }
87/// ```
88pub struct LLMCache {
89    config: CacheConfig,
90    store: RwLock<HashMap<String, CachedLLMResult>>,
91}
92
93impl LLMCache {
94    pub fn new() -> Self {
95        Self::with_config(CacheConfig::default())
96    }
97
98    pub fn with_config(config: CacheConfig) -> Self {
99        Self {
100            config,
101            store: RwLock::new(HashMap::new()),
102        }
103    }
104
105    /// 从消息列表生成缓存键
106    ///
107    /// 将消息列表序列化为 JSON 字符串作为键。
108    /// 包含 model 名称以确保不同模型的调用不互相影响。
109    /// 如果序列化失败，返回错误而非回退到空字符串（M34）。
110    pub fn build_key(messages: &[Message], model: &str) -> Result<String, String> {
111        let serialized = serde_json::to_string(messages)
112            .map_err(|e| format!("cache key serialization failed: {}", e))?;
113        Ok(format!("{}:{}", model, serialized))
114    }
115
116    /// 获取缓存结果
117    ///
118    /// 如果发现过期条目，会立即删除（H36）。
119    pub async fn get(&self, key: &str) -> Option<CachedLLMResult> {
120        if !self.config.enabled {
121            return None;
122        }
123
124        let store = self.store.read().await;
125        if let Some(entry) = store.get(key) {
126            // 检查 TTL
127            if let Some(ttl) = self.config.ttl {
128                if entry.cached_at.elapsed() > ttl {
129                    // H36: 过期条目需要删除，先释放读锁再获取写锁
130                    drop(store);
131                    let mut store = self.store.write().await;
132                    // Double-check after acquiring write lock
133                    if let Some(entry) = store.get(key) {
134                        if entry.cached_at.elapsed() > ttl {
135                            store.remove(key);
136                        }
137                    }
138                    return None;
139                }
140            }
141            Some(entry.clone())
142        } else {
143            None
144        }
145    }
146
147    /// 存入缓存结果
148    pub async fn put(&self, key: String, result: LLMResult) {
149        if !self.config.enabled {
150            return;
151        }
152
153        let mut store = self.store.write().await;
154
155        // 检查是否需要淘汰
156        if self.config.max_entries > 0 && store.len() >= self.config.max_entries {
157            // 移除最早的一条
158            if let Some(oldest_key) = store
159                .iter()
160                .min_by_key(|(_, v)| v.cached_at)
161                .map(|(k, _)| k.clone())
162            {
163                store.remove(&oldest_key);
164            }
165        }
166
167        store.insert(
168            key,
169            CachedLLMResult {
170                result,
171                cached_at: Instant::now(),
172            },
173        );
174    }
175
176    /// 清除缓存
177    pub async fn clear(&self) {
178        let mut store = self.store.write().await;
179        store.clear();
180    }
181
182    /// 获取缓存大小
183    pub async fn len(&self) -> usize {
184        let store = self.store.read().await;
185        store.len()
186    }
187
188    /// 缓存是否为空
189    pub async fn is_empty(&self) -> bool {
190        self.len().await == 0
191    }
192
193    /// 移除过期条目
194    pub async fn evict_expired(&self) -> usize {
195        if let Some(ttl) = self.config.ttl {
196            let mut store = self.store.write().await;
197            let before = store.len();
198            store.retain(|_, v| v.cached_at.elapsed() <= ttl);
199            before - store.len()
200        } else {
201            0
202        }
203    }
204}
205
206impl Default for LLMCache {
207    fn default() -> Self {
208        Self::new()
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215    use crate::language_models::TokenUsage;
216
217    fn make_result(content: &str) -> LLMResult {
218        LLMResult {
219            content: content.to_string(),
220            model: "test-model".to_string(),
221            token_usage: Some(TokenUsage {
222                prompt_tokens: 10,
223                completion_tokens: 5,
224                total_tokens: 15,
225            }),
226            tool_calls: None,
227            thinking_content: None,
228        }
229    }
230
231    #[tokio::test]
232    async fn test_cache_put_and_get() {
233        let cache = LLMCache::new();
234        let key = "test-key";
235        let result = make_result("Hello, world!");
236
237        cache.put(key.to_string(), result.clone()).await;
238        let cached = cache.get(key).await;
239
240        assert!(cached.is_some());
241        assert_eq!(cached.unwrap().result.content, "Hello, world!");
242    }
243
244    #[tokio::test]
245    async fn test_cache_miss() {
246        let cache = LLMCache::new();
247        let cached = cache.get("non-existent").await;
248        assert!(cached.is_none());
249    }
250
251    #[tokio::test]
252    async fn test_cache_clear() {
253        let cache = LLMCache::new();
254        cache.put("k1".to_string(), make_result("r1")).await;
255        cache.put("k2".to_string(), make_result("r2")).await;
256        assert_eq!(cache.len().await, 2);
257
258        cache.clear().await;
259        assert_eq!(cache.len().await, 0);
260    }
261
262    #[tokio::test]
263    async fn test_cache_disabled() {
264        let config = CacheConfig::new().disabled();
265        let cache = LLMCache::with_config(config);
266
267        cache.put("key".to_string(), make_result("test")).await;
268        let cached = cache.get("key").await;
269        assert!(cached.is_none());
270    }
271
272    #[tokio::test]
273    async fn test_cache_ttl_expiry() {
274        let config = CacheConfig::new().with_ttl(Duration::from_millis(10));
275        let cache = LLMCache::with_config(config);
276
277        cache.put("key".to_string(), make_result("test")).await;
278        assert!(cache.get("key").await.is_some());
279
280        // 等待过期
281        tokio::time::sleep(Duration::from_millis(20)).await;
282        assert!(cache.get("key").await.is_none());
283    }
284
285    #[tokio::test]
286    async fn test_cache_max_entries() {
287        let config = CacheConfig::new().with_max_entries(3).no_ttl();
288        let cache = LLMCache::with_config(config);
289
290        cache.put("a".to_string(), make_result("1")).await;
291        cache.put("b".to_string(), make_result("2")).await;
292        cache.put("c".to_string(), make_result("3")).await;
293        assert_eq!(cache.len().await, 3);
294
295        // 超过限制，淘汰最早的一条
296        cache.put("d".to_string(), make_result("4")).await;
297        assert_eq!(cache.len().await, 3);
298        // a 应该被淘汰
299        assert!(cache.get("a").await.is_none());
300    }
301
302    #[tokio::test]
303    async fn test_cache_no_ttl() {
304        let config = CacheConfig::new().no_ttl();
305        let cache = LLMCache::with_config(config);
306
307        cache.put("key".to_string(), make_result("persist")).await;
308
309        // 即使等待许久也不应过期
310        tokio::time::sleep(Duration::from_millis(10)).await;
311        assert!(cache.get("key").await.is_some());
312    }
313
314    #[tokio::test]
315    async fn test_cache_evict_expired() {
316        // 用 0 TTL 确保立即过期
317        let config = CacheConfig::new().with_ttl(Duration::from_millis(0));
318        let cache = LLMCache::with_config(config);
319
320        cache.put("key".to_string(), make_result("test")).await;
321        tokio::time::sleep(Duration::from_millis(1)).await;
322
323        let evicted = cache.evict_expired().await;
324        assert_eq!(evicted, 1);
325        assert!(cache.is_empty().await);
326    }
327
328    #[tokio::test]
329    async fn test_cache_build_key() {
330        let messages = vec![Message::human("Hello"), Message::ai("Hi!")];
331        let key = LLMCache::build_key(&messages, "gpt-4").unwrap();
332        assert!(key.contains("gpt-4"));
333        assert!(key.contains("Hello"));
334    }
335
336    #[tokio::test]
337    async fn test_cache_is_empty() {
338        let cache = LLMCache::new();
339        assert!(cache.is_empty().await);
340
341        cache.put("key".to_string(), make_result("test")).await;
342        assert!(!cache.is_empty().await);
343    }
344}
lc_core/cache/llm_cache.rs

lc_core/cache/
llm_cache.rs
