mistralrs_server_core/
cached_responses.rs

1//! ## Response caching functionality for the Responses API.
2
3use anyhow::Result;
4use std::collections::HashMap;
5use std::sync::LazyLock;
6use std::sync::{Arc, RwLock};
7
8use crate::openai::Message;
9use crate::responses_types::ResponseResource;
10
11/// Trait for caching responses
12pub trait ResponseCache: Send + Sync {
13    /// Store a response object with the given ID
14    fn store_response(&self, id: String, response: ResponseResource) -> Result<()>;
15
16    /// Retrieve a response object by ID
17    fn get_response(&self, id: &str) -> Result<Option<ResponseResource>>;
18
19    /// Delete a response object by ID
20    fn delete_response(&self, id: &str) -> Result<bool>;
21
22    /// Store conversation history for a response
23    fn store_conversation_history(&self, id: String, messages: Vec<Message>) -> Result<()>;
24
25    /// Retrieve conversation history for a response
26    fn get_conversation_history(&self, id: &str) -> Result<Option<Vec<Message>>>;
27}
28
29/// In-memory implementation of ResponseCache
30pub struct InMemoryResponseCache {
31    responses: Arc<RwLock<HashMap<String, ResponseResource>>>,
32    conversation_histories: Arc<RwLock<HashMap<String, Vec<Message>>>>,
33}
34
35impl InMemoryResponseCache {
36    /// Create a new in-memory cache
37    pub fn new() -> Self {
38        Self {
39            responses: Arc::new(RwLock::new(HashMap::new())),
40            conversation_histories: Arc::new(RwLock::new(HashMap::new())),
41        }
42    }
43}
44
45impl Default for InMemoryResponseCache {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl ResponseCache for InMemoryResponseCache {
52    fn store_response(&self, id: String, response: ResponseResource) -> Result<()> {
53        let mut responses = self.responses.write().unwrap();
54        responses.insert(id, response);
55        Ok(())
56    }
57
58    fn get_response(&self, id: &str) -> Result<Option<ResponseResource>> {
59        let responses = self.responses.read().unwrap();
60        Ok(responses.get(id).cloned())
61    }
62
63    fn delete_response(&self, id: &str) -> Result<bool> {
64        // IMPORTANT: Lock ordering must be maintained to prevent deadlocks.
65        // Order: responses -> conversation_histories
66        // All methods that acquire multiple locks must follow this order.
67        //
68        // We acquire all locks before any modifications to ensure atomicity.
69        // The locks are released in reverse order when dropped at end of scope.
70        let mut responses = self.responses.write().unwrap();
71        let mut histories = self.conversation_histories.write().unwrap();
72
73        let response_removed = responses.remove(id).is_some();
74        let history_removed = histories.remove(id).is_some();
75
76        Ok(response_removed || history_removed)
77    }
78
79    fn store_conversation_history(&self, id: String, messages: Vec<Message>) -> Result<()> {
80        let mut histories = self.conversation_histories.write().unwrap();
81        histories.insert(id, messages);
82        Ok(())
83    }
84
85    fn get_conversation_history(&self, id: &str) -> Result<Option<Vec<Message>>> {
86        let histories = self.conversation_histories.read().unwrap();
87        Ok(histories.get(id).cloned())
88    }
89}
90
91/// Global response cache instance
92pub static RESPONSE_CACHE: LazyLock<Arc<dyn ResponseCache>> =
93    LazyLock::new(|| Arc::new(InMemoryResponseCache::new()));
94
95/// Helper function to get the global cache instance
96pub fn get_response_cache() -> Arc<dyn ResponseCache> {
97    RESPONSE_CACHE.clone()
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::responses_types::{ItemStatus, OutputContent, OutputItem, ResponseStatus};
104
105    #[test]
106    fn test_in_memory_cache() {
107        let cache = InMemoryResponseCache::new();
108
109        // Create a test response
110        let response =
111            ResponseResource::new("test-id".to_string(), "test-model".to_string(), 1234567890)
112                .with_status(ResponseStatus::Completed)
113                .with_output(vec![OutputItem::message(
114                    "msg-1".to_string(),
115                    vec![OutputContent::text("Hello".to_string())],
116                    ItemStatus::Completed,
117                )]);
118
119        // Store and retrieve
120        cache
121            .store_response("test-id".to_string(), response.clone())
122            .unwrap();
123        let retrieved = cache.get_response("test-id").unwrap();
124        assert!(retrieved.is_some());
125        assert_eq!(retrieved.unwrap().id, "test-id");
126
127        // Delete
128        let deleted = cache.delete_response("test-id").unwrap();
129        assert!(deleted);
130        let retrieved = cache.get_response("test-id").unwrap();
131        assert!(retrieved.is_none());
132    }
133
134    #[test]
135    fn test_conversation_history() {
136        let cache = InMemoryResponseCache::new();
137
138        let messages = vec![Message {
139            content: Some(crate::openai::MessageContent::from_text(
140                "Hello".to_string(),
141            )),
142            role: "user".to_string(),
143            name: None,
144            tool_calls: None,
145            tool_call_id: None,
146        }];
147
148        cache
149            .store_conversation_history("test-id".to_string(), messages.clone())
150            .unwrap();
151
152        let retrieved = cache.get_conversation_history("test-id").unwrap();
153        assert!(retrieved.is_some());
154        assert_eq!(retrieved.unwrap().len(), 1);
155    }
156}
mistralrs_server_core/cached_responses.rs

mistralrs_server_core/
cached_responses.rs
