langchain_rust/llm/cache/
mod.rs

1//! Response caching for any [`LLM`], like langchaingo's `llms/cache`.
2
3use std::{
4    collections::HashMap,
5    pin::Pin,
6    sync::{Arc, Mutex},
7    time::Duration,
8};
9
10use async_trait::async_trait;
11use futures::{Stream, StreamExt};
12use serde_json::{json, Value};
13use tokio::time::Instant;
14
15use crate::{
16    language_models::{
17        box_llm_events, llm::LLM, options::CallOptions, result_events, GenerateResult, LLMError,
18        LLMEvent, LLMEventStream,
19    },
20    schemas::{Message, StreamData},
21};
22
23/// Where [`CachedLLM`] keeps responses. Implement it to cache in Redis, a database, a
24/// file, etc.
25///
26/// Keys are the exact JSON of the request (namespace, messages and call options), so
27/// they never collide but can be long; hash them if your store needs short keys.
28#[async_trait]
29pub trait LLMCache: Send + Sync {
30    /// The cached response for `key`, if any.
31    async fn get(&self, key: &str) -> Option<GenerateResult>;
32    /// Stores a response.
33    async fn put(&self, key: &str, value: GenerateResult);
34}
35
36struct Entry {
37    value: GenerateResult,
38    inserted: Instant,
39    last_used: u64,
40}
41
42#[derive(Default)]
43struct InMemoryState {
44    entries: HashMap<String, Entry>,
45    clock: u64,
46}
47
48/// An in-process [`LLMCache`]. Unbounded by default; [`with_capacity`] evicts the least
49/// recently used response and [`with_ttl`] expires responses after a while.
50///
51/// [`with_capacity`]: InMemoryLLMCache::with_capacity
52/// [`with_ttl`]: InMemoryLLMCache::with_ttl
53#[derive(Default)]
54pub struct InMemoryLLMCache {
55    state: Mutex<InMemoryState>,
56    capacity: Option<usize>,
57    ttl: Option<Duration>,
58}
59
60impl InMemoryLLMCache {
61    pub fn new() -> Self {
62        Self::default()
63    }
64
65    /// Keeps at most `capacity` responses, evicting the least recently used one.
66    pub fn with_capacity(mut self, capacity: usize) -> Self {
67        self.capacity = Some(capacity);
68        self
69    }
70
71    /// Forgets responses older than `ttl`.
72    pub fn with_ttl(mut self, ttl: Duration) -> Self {
73        self.ttl = Some(ttl);
74        self
75    }
76
77    /// Number of stored responses (expired ones included until they are looked up or
78    /// evicted).
79    pub fn len(&self) -> usize {
80        self.state.lock().unwrap().entries.len()
81    }
82
83    pub fn is_empty(&self) -> bool {
84        self.len() == 0
85    }
86
87    pub fn clear(&self) {
88        self.state.lock().unwrap().entries.clear();
89    }
90
91    fn expired(&self, entry: &Entry, now: Instant) -> bool {
92        self.ttl
93            .is_some_and(|ttl| now.duration_since(entry.inserted) >= ttl)
94    }
95}
96
97#[async_trait]
98impl LLMCache for InMemoryLLMCache {
99    async fn get(&self, key: &str) -> Option<GenerateResult> {
100        let now = Instant::now();
101        let mut state = self.state.lock().unwrap();
102        state.clock += 1;
103        let clock = state.clock;
104        let expired = self.expired(state.entries.get(key)?, now);
105        if expired {
106            state.entries.remove(key);
107            return None;
108        }
109        let entry = state.entries.get_mut(key)?;
110        entry.last_used = clock;
111        Some(entry.value.clone())
112    }
113
114    async fn put(&self, key: &str, value: GenerateResult) {
115        let now = Instant::now();
116        let mut state = self.state.lock().unwrap();
117        state.clock += 1;
118        let clock = state.clock;
119        if self.capacity == Some(0) {
120            return;
121        }
122        if let Some(capacity) = self.capacity {
123            if !state.entries.contains_key(key) {
124                // Expired entries go first, then the least recently used ones.
125                state.entries.retain(|_, e| !self.expired(e, now));
126                while state.entries.len() >= capacity {
127                    let oldest = state
128                        .entries
129                        .iter()
130                        .min_by_key(|(_, e)| e.last_used)
131                        .map(|(k, _)| k.clone());
132                    match oldest {
133                        Some(k) => state.entries.remove(&k),
134                        None => break,
135                    };
136                }
137            }
138        }
139        state.entries.insert(
140            key.to_string(),
141            Entry {
142                value,
143                inserted: now,
144                last_used: clock,
145            },
146        );
147    }
148}
149
150/// Wraps an [`LLM`] and answers repeated requests from an [`LLMCache`] instead of calling
151/// the model again.
152///
153/// The cache key is the namespace, the messages and the call options set through
154/// [`add_options`](LLM::add_options) (temperature, stop words, tools, ...). The wrapper
155/// can't see settings inside the model itself, such as the model name, so give each model
156/// configuration its own [namespace](CachedLLM::with_namespace) when they share a cache.
157/// By default the namespace is the wrapped type's name.
158///
159/// - A hit returns the stored result with `tokens: None`: it cost nothing. If a
160///   `streaming_func` is set, it gets the whole text at once.
161/// - Errors are never cached.
162/// - [`generate_stream`](LLM::generate_stream) caches the final result of a streamed
163///   response and replays hits as one `Text` event plus `Done`. The legacy
164///   [`stream`](LLM::stream) replays hits but doesn't store misses.
165/// - Concurrent identical misses each call the model.
166///
167/// ```rust,ignore
168/// let cache = Arc::new(InMemoryLLMCache::new().with_ttl(Duration::from_secs(3600)));
169/// let llm = CachedLLM::new(OpenAI::default(), cache).with_namespace("gpt-4o-mini");
170/// let first = llm.invoke("Say hi").await?; // calls OpenAI
171/// let again = llm.invoke("Say hi").await?; // from the cache
172/// ```
173pub struct CachedLLM {
174    llm: Box<dyn LLM>,
175    cache: Arc<dyn LLMCache>,
176    namespace: String,
177    options: CallOptions,
178}
179
180impl Clone for CachedLLM {
181    fn clone(&self) -> Self {
182        Self {
183            llm: self.llm.clone_box(),
184            cache: self.cache.clone(),
185            namespace: self.namespace.clone(),
186            options: self.options.clone(),
187        }
188    }
189}
190
191impl CachedLLM {
192    pub fn new<L: LLM + 'static>(llm: L, cache: Arc<dyn LLMCache>) -> Self {
193        Self {
194            llm: Box::new(llm),
195            cache,
196            namespace: std::any::type_name::<L>().to_string(),
197            options: CallOptions::new(),
198        }
199    }
200
201    /// Separates this model's entries from others in a shared cache, e.g. the model name.
202    pub fn with_namespace<S: Into<String>>(mut self, namespace: S) -> Self {
203        self.namespace = namespace.into();
204        self
205    }
206
207    /// The key a request is cached under.
208    pub fn cache_key(&self, messages: &[Message]) -> Result<String, LLMError> {
209        let key = json!({
210            "namespace": self.namespace,
211            "messages": messages,
212            "options": options_fingerprint(&self.options),
213        });
214        Ok(serde_json::to_string(&key)?)
215    }
216
217    async fn lookup(&self, key: &str) -> Option<GenerateResult> {
218        let mut hit = self.cache.get(key).await?;
219        hit.tokens = None;
220        Some(hit)
221    }
222}
223
224/// The options that change what the model answers. `streaming_func` and
225/// `stream_usage` only change how it's delivered.
226fn options_fingerprint(o: &CallOptions) -> Value {
227    json!({
228        "candidate_count": o.candidate_count,
229        "max_tokens": o.max_tokens,
230        "temperature": o.temperature,
231        "stop_words": o.stop_words,
232        "top_k": o.top_k,
233        "top_p": o.top_p,
234        "seed": o.seed,
235        "min_length": o.min_length,
236        "max_length": o.max_length,
237        "n": o.n,
238        "repetition_penalty": o.repetition_penalty,
239        "frequency_penalty": o.frequency_penalty,
240        "presence_penalty": o.presence_penalty,
241        "functions": o.functions.as_ref().map(|fs| {
242            fs.iter()
243                .map(|f| json!([f.name, f.description, f.parameters]))
244                .collect::<Vec<_>>()
245        }),
246        "function_call_behavior": o.function_call_behavior.as_ref().map(|b| format!("{b:?}")),
247        "response_format": o.response_format.as_ref().map(|r| format!("{r:?}")),
248    })
249}
250
251#[async_trait]
252impl LLM for CachedLLM {
253    async fn generate(&self, messages: &[Message]) -> Result<GenerateResult, LLMError> {
254        let key = self.cache_key(messages)?;
255        if let Some(hit) = self.lookup(&key).await {
256            if let Some(func) = &self.options.streaming_func {
257                let mut func = func.lock().await;
258                func(hit.generation.clone())
259                    .await
260                    .map_err(|_| LLMError::OtherError("streaming_func failed".into()))?;
261            }
262            return Ok(hit);
263        }
264        let result = self.llm.generate(messages).await?;
265        self.cache.put(&key, result.clone()).await;
266        Ok(result)
267    }
268
269    async fn stream(
270        &self,
271        messages: &[Message],
272    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamData, LLMError>> + Send>>, LLMError> {
273        let key = self.cache_key(messages)?;
274        if let Some(hit) = self.lookup(&key).await {
275            let chunk = StreamData::new(json!(hit.generation), None, hit.generation);
276            return Ok(Box::pin(futures::stream::iter([Ok(chunk)])));
277        }
278        self.llm.stream(messages).await
279    }
280
281    fn generate_stream<'a>(&'a self, messages: &'a [Message]) -> LLMEventStream<'a> {
282        box_llm_events(async_stream::try_stream! {
283            let key = self.cache_key(messages)?;
284            if let Some(hit) = self.lookup(&key).await {
285                for event in result_events(hit) {
286                    yield event;
287                }
288                return;
289            }
290            let mut events = self.llm.generate_stream(messages);
291            while let Some(event) = events.next().await {
292                let event = event?;
293                if let LLMEvent::Done(result) = &event {
294                    self.cache.put(&key, result.clone()).await;
295                }
296                yield event;
297            }
298        })
299    }
300
301    fn add_options(&mut self, options: CallOptions) {
302        self.options.merge_options(options.clone());
303        self.llm.add_options(options);
304    }
305
306    fn messages_to_string(&self, messages: &[Message]) -> String {
307        self.llm.messages_to_string(messages)
308    }
309}
310
311#[cfg(test)]
312mod tests;
langchain_rust/llm/cache/mod.rs

langchain_rust/llm/cache/
mod.rs
