Skip to main content

aagt_core/agent/
cache.rs

1//! Semantic caching for LLM responses
2//!
3//! Provides a mechanism to cache and reuse LLM completions based on prompt similarity.
4
5use dashmap::DashMap;
6use async_trait::async_trait;
7use crate::agent::message::Message;
8use crate::error::Result;
9
10/// Trait for semantic caching
11#[async_trait]
12pub trait Cache: Send + Sync {
13    /// Get a cached response for the given messages
14    async fn get(&self, messages: &[Message]) -> Result<Option<String>>;
15    
16    /// Store a response in the cache
17    async fn set(&self, messages: &[Message], response: String) -> Result<()>;
18    
19    /// Clear the cache
20    async fn clear(&self) -> Result<()>;
21}
22
23/// A simple in-memory implementation of the Cache trait
24/// 
25/// Note: This is an exact-match cache for now. Truly 'semantic' caching 
26/// (vector-based) should be implemented using aagt-qmd.
27pub struct InMemoryCache {
28    store: DashMap<String, String>,
29}
30
31impl InMemoryCache {
32    /// Create a new in-memory cache
33    pub fn new() -> Self {
34        Self {
35            store: DashMap::new(),
36        }
37    }
38    
39    /// Generate a simple key based on message content
40    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        // Hash it for a stable fixed-length key if content is huge
47        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}