1use dashmap::DashMap;
6use async_trait::async_trait;
7use crate::agent::message::Message;
8use crate::error::Result;
9
10#[async_trait]
12pub trait Cache: Send + Sync {
13 async fn get(&self, messages: &[Message]) -> Result<Option<String>>;
15
16 async fn set(&self, messages: &[Message], response: String) -> Result<()>;
18
19 async fn clear(&self) -> Result<()>;
21}
22
23pub struct InMemoryCache {
28 store: DashMap<String, String>,
29}
30
31impl InMemoryCache {
32 pub fn new() -> Self {
34 Self {
35 store: DashMap::new(),
36 }
37 }
38
39 fn generate_key(&self, messages: &[Message]) -> String {
41 let mut key = String::new();
42 for msg in messages {
43 key.push_str(msg.role.as_str());
44 key.push_str(&msg.text());
45 }
46 use std::collections::hash_map::DefaultHasher;
48 use std::hash::{Hash, Hasher};
49 let mut hasher = DefaultHasher::new();
50 key.hash(&mut hasher);
51 hasher.finish().to_string()
52 }
53}
54
55#[async_trait]
56impl Cache for InMemoryCache {
57 async fn get(&self, messages: &[Message]) -> Result<Option<String>> {
58 let key = self.generate_key(messages);
59 Ok(self.store.get(&key).map(|v| v.value().clone()))
60 }
61
62 async fn set(&self, messages: &[Message], response: String) -> Result<()> {
63 let key = self.generate_key(messages);
64 self.store.insert(key, response);
65 Ok(())
66 }
67
68 async fn clear(&self) -> Result<()> {
69 self.store.clear();
70 Ok(())
71 }
72}
73
74impl Default for InMemoryCache {
75 fn default() -> Self {
76 Self::new()
77 }
78}