1
// src/cache/semantic-cache.ts — Redis-backed semantic cache for LLM prompts
2

3
import { createClient, RedisClientType } from "redis";
4
import { BedrockRuntimeClient, InvokeModelCommand } from "@aws-sdk/client-bedrock-runtime";
5

6
interface CacheEntry {
7
  prompt: string;
8
  response: string;
9
  embedding: number[];
10
  modelId: string;
11
  tokensIn: number;
12
  tokensOut: number;
13
  timestamp: number;
14
  hitCount: number;
15
}
16

17
interface CacheConfig {
18
  modelId: string;
19
  similarityThreshold: number;   // 0.0 - 1.0, default ~0.92
20
  maxCacheAgeMs: number;         // TTL per entry
21
  maxEntries: number;            // Evict when exceeded
22
  enabled: boolean;              // Quick toggle
23
}
24

25
const DEFAULT_CONFIGS: Record<string, CacheConfig> = {
26
  "claude-sonnet": {
27
    modelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
28
    similarityThreshold: 0.92,
29
    maxCacheAgeMs: 300_000,       // 5 minutes (matches Bedrock cache window)
30
    maxEntries: 10_000,
31
    enabled: true,
32
  },
33
  "claude-haiku": {
34
    modelId: "anthropic.claude-3-5-haiku-20241022-v1:0",
35
    similarityThreshold: 0.90,
36
    maxCacheAgeMs: 60_000,        // 1 minute (cheap model, short cache)
37
    maxEntries: 5_000,
38
    enabled: true,
39
  },
40
  // You can add custom configs per agent use-case
41
  "code-review": {
42
    modelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
43
    similarityThreshold: 0.85,    // Looser for code review patterns
44
    maxCacheAgeMs: 600_000,       // 10 minutes
45
    maxEntries: 20_000,
46
    enabled: true,
47
  },
48
};
49

50
export class SemanticCache {
51
  private redis: RedisClientType;
52
  private bedrock: BedrockRuntimeClient;
53
  private configs: Record<string, CacheConfig>;
54
  private embeddingModelId: string;
55

56
  constructor(options: {
57
    redisUrl?: string;
58
    region?: string;
59
    configs?: Record<string, CacheConfig>;
60
    embeddingModelId?: string;
61
  }) {
62
    this.redis = createClient({ url: options.redisUrl || process.env.REDIS_URL });
63
    this.bedrock = new BedrockRuntimeClient({ region: options.region || "us-east-1" });
64
    this.configs = options.configs || DEFAULT_CONFIGS;
65
    this.embeddingModelId = options.embeddingModelId || "amazon.titan-embed-text-v2:0";
66
  }
67

68
  async connect(): Promise<void> {
69
    await this.redis.connect();
70
  }
71

72
  /**
73
   * Get cached response for a prompt, or null if no match.
74
   */
75
  async get(
76
    prompt: string,
77
    configKey: string = "claude-sonnet"
78
  ): Promise<{
79
    response: string;
80
    hitCount: number;
81
    similarity: number;
82
  } | null> {
83
    const config = this.configs[configKey];
84
    if (!config || !config.enabled) return null;
85

86
    // Generate embedding for this prompt
87
    const embedding = await this.embedPrompt(prompt);
88

89
    // Search Redis for similar embeddings
90
    const results = await this.redis.ft.search(
91
      "idx:llm-cache",
92
      `@modelId:{${config.modelId}} @embedding:[VECTOR_RANGE 0.08 $vec]=>{$YIELD_DISTANCE:1}`, {
93
        PARAMS: { vec: this.vectorToString(embedding) },
94
        SORTBY: "VECTOR_DISTANCE",
95
        LIMIT: { from: 0, size: 1 },
96
        DIALECT: 2,
97
      }
98
    );
99

100
    if (!results.total) return null;
101

102
    const doc = results.documents[0];
103
    const similarity = 1 - parseFloat(doc.value.distance as string);
104

105
    // Check threshold
106
    if (similarity < config.similarityThreshold) return null;
107

108
    // Check TTL
109
    const age = Date.now() - (doc.value.timestamp as number);
110
    if (age > config.maxCacheAgeMs) {
111
      // Stale entry — delete and treat as miss
112
      await this.redis.del(`cache:${doc.id}`);
113
      return null;
114
    }
115

116
    // Increment hit count
117
    await this.redis.hIncrBy(`cache:${doc.id}`, "hitCount", 1);
118

119
    return {
120
      response: doc.value.response as string,
121
      hitCount: (doc.value.hitCount as number) + 1,
122
      similarity,
123
    };
124
  }
125

126
  /**
127
   * Store a prompt-response pair in the cache.
128
   */
129
  async set(
130
    prompt: string,
131
    response: string,
132
    tokensIn: number,
133
    tokensOut: number,
134
    modelId: string,
135
    configKey?: string
136
  ): Promise<void> {
137
    const key = `cache:${this.hash(prompt)}`;
138
    const embedding = await this.embedPrompt(prompt);
139

140
    // Check if we need to evict
141
    const currentCount = await this.redis.dbSize();
142
    const config = configKey ? this.configs[configKey] : null;
143
    if (config && currentCount >= config.maxEntries) {
144
      await this.evictOldest();
145
    }
146

147
    await this.redis.hSet(key, {
148
      prompt,
149
      response,
150
      embedding: this.vectorToString(embedding),
151
      modelId,
152
      tokensIn,
153
      tokensOut,
154
      timestamp: Date.now(),
155
      hitCount: 1,
156
    });
157

158
    // Add to search index
159
    await this.redis.ft.synUpdate("idx:llm-cache", key, 1, [key]);
160
  }
161

162
  /**
163
   * Invalidate cache entries matching a pattern.
164
   */
165
  async invalidate(pattern: string): Promise<number> {
166
    const keys = await this.redis.keys(`cache:${pattern}*`);
167
    if (keys.length === 0) return 0;
168
    const deleted = await this.redis.del(keys);
169
    return deleted;
170
  }
171

172
  /**
173
   * Clear all cache entries.
174
   */
175
  async clear(): Promise<void> {
176
    const keys = await this.redis.keys("cache:*");
177
    if (keys.length > 0) await this.redis.del(keys);
178
  }
179

180
  /**
181
   * Get cache stats.
182
   */
183
  async getStats(): Promise<{
184
    totalEntries: number;
185
    totalHits: number;
186
    oldestEntry: number;
187
    newestEntry: number;
188
  }> {
189
    const keys = await this.redis.keys("cache:*");
190
    let totalHits = 0;
191
    let oldestEntry = Date.now();
192
    let newestEntry = 0;
193

194
    for (const key of keys) {
195
      const entry = await this.redis.hGetAll(key);
196
      if (entry.hitCount) totalHits += parseInt(entry.hitCount);
197
      if (entry.timestamp) {
198
        const ts = parseInt(entry.timestamp);
199
        if (ts < oldestEntry) oldestEntry = ts;
200
        if (ts > newestEntry) newestEntry = ts;
201
      }
202
    }
203

204
    return {
205
      totalEntries: keys.length,
206
      totalHits,
207
      oldestEntry,
208
      newestEntry,
209
    };
210
  }
211

212
  // ──── Private ────
213

214
  /**
215
   * Generate embedding for a prompt using Amazon Titan Embeddings.
216
   */
217
  private async embedPrompt(prompt: string): Promise<number[]> {
218
    const command = new InvokeModelCommand({
219
      modelId: this.embeddingModelId,
220
      contentType: "application/json",
221
      body: JSON.stringify({
222
        inputText: prompt,
223
        dimensions: 1024,
224
        normalize: true,
225
      }),
226
    });
227

228
    const result = await this.bedrock.send(command);
229
    const body = JSON.parse(new TextDecoder().decode(result.body));
230

231
    return body.embedding;
232
  }
233

234
  private vectorToString(vector: number[]): string {
235
    return `[${vector.join(",")}]`;
236
  }
237

238
  private async evictOldest(): Promise<void> {
239
    // Find and remove the oldest 10% of entries
240
    const keys = await this.redis.keys("cache:*");
241
    const entries = [];
242

243
    for (const key of keys) {
244
      const entry = await this.redis.hGetAll(key);
245
      entries.push({ key, timestamp: parseInt(entry.timestamp || "0") });
246
    }
247

248
    entries.sort((a, b) => a.timestamp - b.timestamp);
249
    const toEvict = entries.slice(0, Math.ceil(entries.length * 0.1));
250

251
    if (toEvict.length > 0) {
252
      await this.redis.del(toEvict.map(e => e.key));
253
    }
254
  }
255

256
  private hash(prompt: string): string {
257
    const crypto = require("crypto");
258
    return crypto.createHash("md5").update(prompt).digest("hex").slice(0, 12);
259
  }
260
}

1
// Emit custom metrics
2
import { CloudWatchClient, PutMetricDataCommand } from "@aws-sdk/client-cloudwatch";
3

4
const cw = new CloudWatchClient({ region: "us-east-1" });
5

6
async function emitCacheMetric(
7
  hit: boolean,
8
  latency: number,
9
  configKey: string,
10
  similarity?: number
11
): Promise<void> {
12
  await cw.send(new PutMetricDataCommand({
13
    Namespace: "Agent/LLMCache",
14
    MetricData: [
15
      {
16
        MetricName: hit ? "CacheHits" : "CacheMisses",
17
        Value: 1,
18
        Unit: "Count",
19
        Dimensions: [{ Name: "ModelConfig", Value: configKey }],
20
      },
21
      {
22
        MetricName: "Latency",
23
        Value: latency,
24
        Unit: "Milliseconds",
25
        Dimensions: [{ Name: "CacheStatus", Value: hit ? "hit" : "miss" }],
26
      },
27
      ...(similarity !== undefined
28
        ? [{
29
            MetricName: "SimilarityScore",
30
            Value: similarity,
31
            Unit: "None",
32
            Dimensions: [{ Name: "ModelConfig", Value: configKey }],
33
          }]
34
        : []),
35
    ],
36
  }));
37
}