parrot Public

Watch 0 Fork 0 Star 0
Go · 3360 bytes Raw Blame History
1 package llm
2
3 import (
4 "crypto/sha256"
5 "encoding/hex"
6 "sync"
7 "time"
8 )
9
10 // ResponseCache provides an LRU cache for LLM responses to avoid redundant calls
11 type ResponseCache struct {
12 mu sync.RWMutex
13 entries map[string]*cacheEntry
14 maxSize int
15 ttl time.Duration
16 }
17
18 type cacheEntry struct {
19 response string
20 timestamp time.Time
21 }
22
23 // Global cache instance
24 var (
25 responseCache *ResponseCache
26 responseCacheOnce sync.Once
27 )
28
29 // GetResponseCache returns the singleton response cache
30 func GetResponseCache() *ResponseCache {
31 responseCacheOnce.Do(func() {
32 responseCache = NewResponseCache(100, 5*time.Minute) // 100 entries, 5 min TTL
33 })
34 return responseCache
35 }
36
37 // NewResponseCache creates a new response cache
38 func NewResponseCache(maxSize int, ttl time.Duration) *ResponseCache {
39 cache := &ResponseCache{
40 entries: make(map[string]*cacheEntry),
41 maxSize: maxSize,
42 ttl: ttl,
43 }
44
45 // Start background cleanup goroutine
46 go cache.cleanupLoop()
47
48 return cache
49 }
50
51 // generateKey creates a cache key from command signature
52 func (c *ResponseCache) generateKey(command, commandType, exitCode, mode string) string {
53 // Create a hash of the command signature
54 h := sha256.New()
55 h.Write([]byte(command))
56 h.Write([]byte("|"))
57 h.Write([]byte(commandType))
58 h.Write([]byte("|"))
59 h.Write([]byte(exitCode))
60 h.Write([]byte("|"))
61 h.Write([]byte(mode))
62 return hex.EncodeToString(h.Sum(nil))[:16] // Use first 16 chars of hash
63 }
64
65 // Get retrieves a cached response if available and not expired
66 func (c *ResponseCache) Get(command, commandType, exitCode, mode string) (string, bool) {
67 key := c.generateKey(command, commandType, exitCode, mode)
68
69 c.mu.RLock()
70 defer c.mu.RUnlock()
71
72 entry, exists := c.entries[key]
73 if !exists {
74 return "", false
75 }
76
77 // Check if expired
78 if time.Since(entry.timestamp) > c.ttl {
79 return "", false
80 }
81
82 return entry.response, true
83 }
84
85 // Set stores a response in the cache
86 func (c *ResponseCache) Set(command, commandType, exitCode, mode, response string) {
87 key := c.generateKey(command, commandType, exitCode, mode)
88
89 c.mu.Lock()
90 defer c.mu.Unlock()
91
92 // Evict oldest entries if at capacity
93 if len(c.entries) >= c.maxSize {
94 c.evictOldest()
95 }
96
97 c.entries[key] = &cacheEntry{
98 response: response,
99 timestamp: time.Now(),
100 }
101 }
102
103 // evictOldest removes the oldest entry (must be called with lock held)
104 func (c *ResponseCache) evictOldest() {
105 var oldestKey string
106 var oldestTime time.Time
107
108 for key, entry := range c.entries {
109 if oldestKey == "" || entry.timestamp.Before(oldestTime) {
110 oldestKey = key
111 oldestTime = entry.timestamp
112 }
113 }
114
115 if oldestKey != "" {
116 delete(c.entries, oldestKey)
117 }
118 }
119
120 // cleanupLoop periodically removes expired entries
121 func (c *ResponseCache) cleanupLoop() {
122 ticker := time.NewTicker(1 * time.Minute)
123 defer ticker.Stop()
124
125 for range ticker.C {
126 c.cleanup()
127 }
128 }
129
130 // cleanup removes all expired entries
131 func (c *ResponseCache) cleanup() {
132 c.mu.Lock()
133 defer c.mu.Unlock()
134
135 now := time.Now()
136 for key, entry := range c.entries {
137 if now.Sub(entry.timestamp) > c.ttl {
138 delete(c.entries, key)
139 }
140 }
141 }
142
143 // Stats returns cache statistics
144 func (c *ResponseCache) Stats() map[string]interface{} {
145 c.mu.RLock()
146 defer c.mu.RUnlock()
147
148 return map[string]interface{}{
149 "size": len(c.entries),
150 "max_size": c.maxSize,
151 "ttl_secs": c.ttl.Seconds(),
152 }
153 }
154