dwara_core/ai/semantic_cache.rs
1//! AI semantic caching (DW-083): embedding-similarity cache for AI
2//! prompts. A paraphrased prompt within the similarity threshold
3//! returns the cached response (no provider call, no token spend).
4//!
5//! Uses an external embedding service (OpenAI-compatible /v1/embeddings
6//! API) to vectorize prompts and `hnsw_rs` (pure Rust HNSW) for
7//! approximate nearest neighbor search.
8//!
9//! Feature-gated behind the `semantic_cache` cargo feature. Without it,
10//! the module compiles to an inert placeholder that always returns
11//! None (the config is accepted but the cache is a no-op).
12//!
13//! # PERF-02 enhancements
14//!
15//! - **Exact-match fast tier**: before calling the embedding service,
16//! a `HashMap<prompt_hash, entry_id>` checks for an exact text match.
17//! A hit skips the embedding call entirely (the common case for
18//! repeated identical prompts — RAG traffic, retry storms, etc.).
19//! - **LRU eviction**: when the cache reaches `max_entries`, the least
20//! recently accessed entry is evicted (not a wholesale reset). Each
21//! entry carries a `last_accessed_ms` timestamp; the eviction scan
22//! is O(n) but runs once per insert, so the amortized cost is O(1).
23//! - **Streaming support**: streaming responses (SSE frame sequences)
24//! are cached and replayed. On a cache miss, the stream is TEE'd:
25//! frames are forwarded to the client AND collected in a bounded
26//! buffer; on stream completion, the collected frames are stored. On
27//! a cache hit, the cached frames are replayed as an SSE response.
28//! The tee buffer is bounded (`stream_cache_max_bytes`, default 1
29//! MiB): a stream exceeding the cap is not cached (the response is
30//! still forwarded to the client unmodified).
31//!
32//! # Lifecycle
33//!
34//! The engine is constructed once at startup and stored on the
35//! dataplane behind an `ArcSwapOption`. It PERSISTS across reloads:
36//! the HNSW index and cached entries survive config refreshes, and a
37//! reload updates the config in place via `update_config` (so a
38//! threshold or TTL change applies to the next lookup with no cache
39//! reset).
40//!
41//! # Request path
42//!
43//! The LOOKUP runs in `serve_ai` AFTER guardrails (the prompt may have
44//! been redacted) and BEFORE model routing + the provider call. A hit
45//! returns the cached response (JSON for non-streaming, SSE frames for
46//! streaming) with no provider call. The lookup is async (it may make
47//! an HTTP call to the embedding service) — acceptable because it runs
48//! in the already-async `serve_ai`. The STORE is fire-and-forget (a
49//! spawned task): the embedding call and the HNSW insert happen AFTER
50//! the response is sent, so they never block the response path.
51
52// -------------------------------------------------------------------------
53// Feature-gated implementation (the `semantic_cache` cargo feature).
54// -------------------------------------------------------------------------
55
56mod enabled {
57 use crate::config::ai::{AiConfig, SemanticCacheConfig};
58 use bytes::Bytes;
59 use hnsw_rs::hnsw::Hnsw;
60 use hnsw_rs::prelude::DistCosine;
61 use http_body_util::{BodyExt as _, Full};
62 use hyper::{Method, Request};
63 use hyper_util::client::legacy::Client;
64 use hyper_util::rt::TokioExecutor;
65 use std::collections::HashMap;
66 use std::sync::atomic::{AtomicUsize, Ordering};
67 use std::sync::{Arc, RwLock};
68 use std::time::{Duration, SystemTime, UNIX_EPOCH};
69
70 /// A cached response: either a complete JSON body (non-streaming)
71 /// or a sequence of SSE frames (streaming). The `frames` variant
72 /// stores the raw `data: ...\n\n` frames as collected by the tee;
73 /// replaying them is a single concatenated body (the client sees
74 /// the same SSE stream the provider emitted).
75 #[derive(Clone)]
76 enum CachedBody {
77 /// Non-streaming: the OpenAI-shaped response JSON.
78 Json(serde_json::Value),
79 /// Streaming: the collected SSE frames (each a `data: ...\n\n`
80 /// string, in order). The `[DONE]` sentinel is included.
81 Frames(Vec<String>),
82 }
83
84 /// One cached response entry.
85 struct CachedEntry {
86 /// The cached response body (JSON or SSE frames).
87 body: CachedBody,
88 /// The model alias the entry was cached for (the cache is
89 /// per-model: a lookup for a different alias is a miss even
90 /// at high similarity).
91 model: String,
92 /// When the entry was stored (epoch millis). Entries older
93 /// than `ttl_secs` are stale and not returned.
94 stored_at_ms: u64,
95 /// When the entry was last accessed (epoch millis). Used for
96 /// LRU eviction (PERF-02): the entry with the smallest
97 /// `last_accessed_ms` is evicted when the cache is full.
98 last_accessed_ms: u64,
99 /// The prompt text (stored for the exact-match fast tier and
100 /// for eviction cleanup). PERF-02: the exact-match map stores
101 /// a hash, but the entry retains the full text so the exact-
102 /// match map can be rebuilt on eviction without re-hashing
103 /// the prompt.
104 prompt_text: String,
105 }
106
107 /// What a cache lookup returns (PERF-02): the cached response,
108 /// either a JSON body (non-streaming) or SSE frames (streaming).
109 #[derive(Clone)]
110 pub enum CachedResponse {
111 /// Non-streaming: the OpenAI-shaped response JSON.
112 Json(serde_json::Value),
113 /// Streaming: the collected SSE frames, ready to replay as
114 /// a single concatenated SSE body.
115 Frames(Vec<String>),
116 }
117
118 /// The semantic cache engine (DW-083). Constructed once at
119 /// startup and stored on the dataplane (persists across reloads
120 /// — the HNSW index and cached entries survive config refreshes).
121 /// Config is updated in place via `update_config`.
122 pub struct SemanticCacheEngine {
123 /// The current config (updated on refresh via RwLock).
124 config: RwLock<SemanticCacheConfig>,
125 /// The HNSW ANN index (replaced on reset). `DistCosine` returns
126 /// cosine DISTANCE (1 - cosine_similarity).
127 hnsw: arc_swap::ArcSwap<Hnsw<'static, f32, DistCosine>>,
128 /// Cached responses keyed by HNSW external id.
129 entries: RwLock<HashMap<usize, CachedEntry>>,
130 /// PERF-02: exact-match fast tier. Maps the prompt text hash
131 /// to the HNSW external id, so an identical prompt skips the
132 /// embedding call entirely. Cleared on reset; updated on store
133 /// and evicted on LRU removal.
134 exact_match: RwLock<HashMap<u64, usize>>,
135 /// Next HNSW external id (monotonic across resets — a stale
136 /// id from a previous index never collides with a live one).
137 next_id: AtomicUsize,
138 /// HTTP client for embedding service calls (constructed once;
139 /// connection-pooled).
140 client: Client<hyper_util::client::legacy::connect::HttpConnector, Full<Bytes>>,
141 }
142
143 impl std::fmt::Debug for SemanticCacheEngine {
144 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145 f.debug_struct("SemanticCacheEngine")
146 .field("enabled", &self.config.read().unwrap().enabled)
147 .field("entry_count", &self.entries.read().unwrap().len())
148 .finish()
149 }
150 }
151
152 impl SemanticCacheEngine {
153 /// Build the engine from a config. Creates a fresh HNSW
154 /// index sized for `max_entries`.
155 pub fn new(config: SemanticCacheConfig) -> Self {
156 let max_entries = config.max_entries.max(1);
157 let hnsw = Arc::new(Hnsw::new(
158 16, // max_nb_connection (M)
159 max_entries, // max_elements (hint)
160 8, // max_layer
161 200, // ef_construction
162 DistCosine, // cosine distance
163 ));
164 let client = Client::builder(TokioExecutor::new()).build_http();
165 SemanticCacheEngine {
166 config: RwLock::new(config),
167 hnsw: arc_swap::ArcSwap::new(hnsw),
168 entries: RwLock::new(HashMap::new()),
169 exact_match: RwLock::new(HashMap::new()),
170 next_id: AtomicUsize::new(0),
171 client,
172 }
173 }
174
175 /// Extract the semantic-cache config from the `ai:` block.
176 /// Returns None when the block or the `semantic_cache` field
177 /// is absent. Does NOT check `enabled` — a disabled engine
178 /// still compiles (so a reload can flip it on); the runtime
179 /// checks `enabled` on every lookup/store.
180 pub fn config_of(cfg: Option<&AiConfig>) -> Option<SemanticCacheConfig> {
181 cfg.and_then(|c| c.semantic_cache.clone())
182 }
183
184 /// Compile from the `ai:` config block. Returns None when
185 /// the `semantic_cache` field is absent. The returned engine
186 /// is live only when `enabled` is true (checked at runtime).
187 pub fn compile(cfg: Option<&AiConfig>) -> Option<Self> {
188 Self::config_of(cfg).map(Self::new)
189 }
190
191 /// Update the config in place (reload path). The HNSW index
192 /// and cached entries PERSIST — only the config (threshold,
193 /// TTL, timeout, etc.) changes. A `max_entries` change does
194 /// NOT rebuild the index immediately; the next reset (when
195 /// the cache fills) sizes the new index to the new value.
196 pub fn update_config(&self, config: SemanticCacheConfig) {
197 *self.config.write().unwrap() = config;
198 }
199
200 /// Whether the cache is enabled (the runtime gate for every
201 /// lookup/store).
202 pub fn is_enabled(&self) -> bool {
203 self.config.read().unwrap().enabled
204 }
205
206 /// Number of cached entries (introspection / tests).
207 pub fn entry_count(&self) -> usize {
208 self.entries.read().unwrap().len()
209 }
210
211 /// Look up a cached response for `prompt_text` + `model`.
212 /// Returns the cached response (JSON or SSE frames) when a
213 /// match is found within the similarity threshold, within TTL,
214 /// and for the same model alias. None otherwise (or when
215 /// disabled, or on any embedding/search error — the cache
216 /// fails open: a miss never blocks the request).
217 ///
218 /// PERF-02: the exact-match fast tier checks a hash map
219 /// BEFORE calling the embedding service. An exact prompt-text
220 /// match skips the embedding call entirely (the common case
221 /// for repeated identical prompts).
222 pub async fn lookup(&self, prompt_text: &str, model: &str) -> Option<CachedResponse> {
223 if !self.is_enabled() {
224 return None;
225 }
226 let cfg = self.config.read().unwrap().clone();
227 let now = now_ms();
228 // PERF-02: exact-match fast tier. Check the hash map
229 // before calling the embedding service. A hit skips the
230 // embedding call entirely.
231 let prompt_hash = hash_prompt(prompt_text);
232 if let Some(id) = self.exact_match.read().unwrap().get(&prompt_hash).copied() {
233 if let Some(entry) = self.validate_entry(id, model, now, &cfg) {
234 tracing::info!(
235 code = "semantic_cache_hit",
236 model = %model,
237 match_type = "exact",
238 "semantic cache exact-match hit; returning cached response"
239 );
240 return Some(entry);
241 }
242 }
243 // Semantic (embedding) lookup.
244 let embedding = match self.embed(prompt_text).await {
245 Ok(v) => v,
246 Err(e) => {
247 tracing::warn!(
248 code = "semantic_cache_embed_failed",
249 "semantic cache lookup embedding failed (failing open as a miss): {e}"
250 );
251 return None;
252 }
253 };
254 let hnsw = self.hnsw.load_full();
255 let neighbors = hnsw.search(&embedding, 1, 200);
256 let neighbor = neighbors.first()?;
257 // DistCosine returns cosine DISTANCE (1 - similarity).
258 let similarity = 1.0 - neighbor.distance as f64;
259 if similarity < cfg.threshold {
260 return None;
261 }
262 if let Some(entry) = self.validate_entry(neighbor.d_id, model, now, &cfg) {
263 tracing::info!(
264 code = "semantic_cache_hit",
265 model = %model,
266 match_type = "semantic",
267 similarity = %similarity,
268 "semantic cache hit; returning cached response"
269 );
270 return Some(entry);
271 }
272 None
273 }
274
275 /// Validate a cached entry by id: check it exists, is within
276 /// TTL, and matches the model. On success, update the
277 /// last-accessed timestamp (LRU) and return the cached body.
278 fn validate_entry(
279 &self,
280 id: usize,
281 model: &str,
282 now: u64,
283 cfg: &SemanticCacheConfig,
284 ) -> Option<CachedResponse> {
285 let mut entries = self.entries.write().unwrap();
286 let entry = entries.get_mut(&id)?;
287 // TTL check.
288 if now.saturating_sub(entry.stored_at_ms) > cfg.ttl_secs * 1000 {
289 return None;
290 }
291 // Model match (the cache is per-model).
292 if entry.model != model {
293 return None;
294 }
295 // PERF-02: update last-accessed for LRU.
296 entry.last_accessed_ms = now;
297 match &entry.body {
298 CachedBody::Json(v) => Some(CachedResponse::Json(v.clone())),
299 CachedBody::Frames(f) => Some(CachedResponse::Frames(f.clone())),
300 }
301 }
302
303 /// Store a non-streaming response in the cache (fire-and-
304 /// forget from the request path). When the cache is full
305 /// (`entry_count >= max_entries`), the least recently used
306 /// entry is evicted (PERF-02 LRU, not a wholesale reset).
307 /// Errors are logged and swallowed.
308 pub async fn store(
309 &self,
310 prompt_text: &str,
311 response_json: &serde_json::Value,
312 model: &str,
313 ) {
314 self.store_body(prompt_text, CachedBody::Json(response_json.clone()), model)
315 .await;
316 }
317
318 /// Store a streaming response (SSE frames) in the cache
319 /// (PERF-02). The frames are the collected SSE `data: ...\n\n`
320 /// strings as tee'd by the stream body. Same LRU eviction as
321 /// the non-streaming store.
322 pub async fn store_streaming(&self, prompt_text: &str, frames: Vec<String>, model: &str) {
323 self.store_body(prompt_text, CachedBody::Frames(frames), model)
324 .await;
325 }
326
327 /// Internal store: inserts a cached body with LRU eviction.
328 async fn store_body(&self, prompt_text: &str, body: CachedBody, model: &str) {
329 if !self.is_enabled() {
330 return;
331 }
332 let cfg = self.config.read().unwrap().clone();
333 // PERF-02: LRU eviction. When the cache is full, evict
334 // the least recently accessed entry (not a wholesale
335 // reset). The scan is O(n) but runs once per insert, so
336 // the amortized cost is O(1).
337 if self.entry_count() >= cfg.max_entries {
338 self.evict_lru();
339 }
340 let embedding = match self.embed(prompt_text).await {
341 Ok(v) => v,
342 Err(e) => {
343 tracing::warn!(
344 code = "semantic_cache_embed_failed",
345 "semantic cache store embedding failed (entry not cached): {e}"
346 );
347 return;
348 }
349 };
350 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
351 let hnsw = self.hnsw.load_full();
352 hnsw.insert((&embedding, id));
353 let now = now_ms();
354 let prompt_hash = hash_prompt(prompt_text);
355 {
356 let mut entries = self.entries.write().unwrap();
357 entries.insert(
358 id,
359 CachedEntry {
360 body,
361 model: model.to_string(),
362 stored_at_ms: now,
363 last_accessed_ms: now,
364 prompt_text: prompt_text.to_string(),
365 },
366 );
367 }
368 self.exact_match.write().unwrap().insert(prompt_hash, id);
369 tracing::info!(
370 code = "semantic_cache_store",
371 model = %model,
372 entries = self.entry_count(),
373 "semantic cache stored a response"
374 );
375 }
376
377 /// PERF-02: evict the least recently used entry. Scans all
378 /// entries for the smallest `last_accessed_ms` and removes it
379 /// from the entries map, the exact-match map, and the HNSW
380 /// index (the HNSW index is not modified — a stale neighbor
381 /// is filtered by the TTL/model check in `validate_entry`).
382 fn evict_lru(&self) {
383 let mut entries = self.entries.write().unwrap();
384 if entries.is_empty() {
385 return;
386 }
387 let (evict_id, _) = entries
388 .iter()
389 .min_by_key(|(_, e)| e.last_accessed_ms)
390 .map(|(id, e)| (*id, e.prompt_text.clone()))
391 .unwrap();
392 let entry = entries.remove(&evict_id);
393 drop(entries);
394 if let Some(entry) = entry {
395 let hash = hash_prompt(&entry.prompt_text);
396 self.exact_match.write().unwrap().remove(&hash);
397 }
398 tracing::info!(
399 code = "semantic_cache_evict",
400 entry_id = evict_id,
401 "semantic cache LRU eviction (max_entries reached)"
402 );
403 }
404
405 /// Reset the cache: a fresh HNSW index replaces the old, and
406 /// all entries are evicted. Called on explicit reset only
407 /// (PERF-02: LRU eviction handles the bounded-memory policy
408 /// at `max_entries`; a wholesale reset is no longer the
409 /// default eviction strategy).
410 pub fn reset(&self) {
411 let cfg = self.config.read().unwrap();
412 let max_entries = cfg.max_entries.max(1);
413 let hnsw = Arc::new(Hnsw::new(16, max_entries, 8, 200, DistCosine));
414 self.hnsw.store(hnsw);
415 self.entries.write().unwrap().clear();
416 self.exact_match.write().unwrap().clear();
417 tracing::info!(
418 code = "semantic_cache_reset",
419 "semantic cache reset; all entries evicted"
420 );
421 }
422
423 /// Call the embedding service: POST
424 /// `{"model": ..., "input": text}` to `embedding_url`, parse
425 /// `{"data": [{"embedding": [...]}]}`. Returns the embedding
426 /// vector on success.
427 async fn embed(&self, text: &str) -> Result<Vec<f32>, String> {
428 let cfg = self.config.read().unwrap().clone();
429 let body = serde_json::json!({
430 "model": cfg.embedding_model,
431 "input": text,
432 });
433 let body_bytes = serde_json::to_vec(&body)
434 .map_err(|e| format!("encode embedding request body: {e}"))?;
435 let mut builder = Request::builder()
436 .method(Method::POST)
437 .uri(&cfg.embedding_url)
438 .header("content-type", "application/json")
439 .header("accept", "application/json");
440 // Optional API key (resolved at compile time; the value
441 // lives only on the wire).
442 if let Some(key) = &cfg.embedding_api_key {
443 let resolved = crate::config::credentials::resolve_configured_secret(key)
444 .map_err(|e| format!("resolve embedding api key: {e}"))?;
445 builder = builder.header("authorization", format!("Bearer {resolved}"));
446 }
447 let req = builder
448 .body(Full::new(Bytes::from(body_bytes)))
449 .map_err(|e| format!("build embedding request: {e}"))?;
450 let timeout = Duration::from_millis(cfg.embedding_timeout_ms);
451 let resp = tokio::time::timeout(timeout, self.client.request(req))
452 .await
453 .map_err(|_| {
454 format!(
455 "embedding service timed out after {} ms",
456 cfg.embedding_timeout_ms
457 )
458 })?
459 .map_err(|e| format!("embedding service request failed: {e}"))?;
460 let status = resp.status();
461 let bytes = resp
462 .into_body()
463 .collect()
464 .await
465 .map_err(|e| format!("read embedding response body: {e}"))?
466 .to_bytes();
467 if !status.is_success() {
468 return Err(format!(
469 "embedding service returned {status}: {}",
470 String::from_utf8_lossy(&bytes)
471 ));
472 }
473 let v: serde_json::Value = serde_json::from_slice(&bytes)
474 .map_err(|e| format!("parse embedding response JSON: {e}"))?;
475 let embedding = v
476 .get("data")
477 .and_then(|d| d.get(0))
478 .and_then(|d| d.get("embedding"))
479 .and_then(|e| e.as_array())
480 .ok_or_else(|| "embedding response missing data[0].embedding".to_string())?;
481 let vec: Vec<f32> = embedding
482 .iter()
483 .map(|x| {
484 x.as_f64()
485 .map(|f| f as f32)
486 .ok_or_else(|| "embedding vector contains a non-number".to_string())
487 })
488 .collect::<Result<Vec<f32>, _>>()?;
489 Ok(vec)
490 }
491 }
492
493 /// Current Unix milliseconds.
494 fn now_ms() -> u64 {
495 SystemTime::now()
496 .duration_since(UNIX_EPOCH)
497 .map(|d| d.as_millis() as u64)
498 .unwrap_or(0)
499 }
500
501 /// Hash a prompt text for the exact-match fast tier (PERF-02).
502 /// Uses the standard library's `DefaultHasher` (a deterministic
503 /// hash within one process — sufficient for an in-memory cache;
504 /// cross-process consistency is not required).
505 fn hash_prompt(text: &str) -> u64 {
506 use std::collections::hash_map::DefaultHasher;
507 use std::hash::{Hash, Hasher};
508 let mut hasher = DefaultHasher::new();
509 text.hash(&mut hasher);
510 hasher.finish()
511 }
512}
513
514pub use enabled::{CachedResponse, SemanticCacheEngine};