feat: semantic (embedding-based) response cache
- cache.rs: cosine-similarity store (threshold, TTL, ring buffer), partitioned by routing signature; tracks hit rate and model cost saved - pipeline: embed routing text, look up before the model call (route_source=cache on hit), insert successful responses; stats record cache hits - config [cache] + admin settings card; dashboard shows hit-rate and savings - embedding model resolved from the pool (or first kind="embedding")
This commit is contained in:
+12
-3
@@ -8,8 +8,8 @@
|
||||
pub mod assets;
|
||||
|
||||
use crate::config::{
|
||||
AdminCfg, ApiKeyCfg, BreakerCfg, Config, CortiqCfg, LogCfg, ModelCfg, ProtocolsCfg, RouteCfg,
|
||||
RouterCfg, RoutingCfg, RoutingPolicy, StatsCfg, TelemetryCfg, TierTargets,
|
||||
AdminCfg, ApiKeyCfg, BreakerCfg, CacheCfg, Config, CortiqCfg, LogCfg, ModelCfg, ProtocolsCfg,
|
||||
RouteCfg, RouterCfg, RoutingCfg, RoutingPolicy, StatsCfg, TelemetryCfg, TierTargets,
|
||||
};
|
||||
use crate::model::{ChatRequest, GenParams, Message, RequestMeta, RoutingDirective};
|
||||
use crate::state::SharedState;
|
||||
@@ -477,6 +477,7 @@ async fn get_settings(State(state): State<SharedState>) -> ApiResult {
|
||||
"cortiq": c.cortiq,
|
||||
"stats": c.stats,
|
||||
"admin": c.admin,
|
||||
"cache": c.cache,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -491,6 +492,7 @@ struct SettingsBody {
|
||||
cortiq: Option<CortiqCfg>,
|
||||
stats: Option<StatsCfg>,
|
||||
admin: Option<AdminCfg>,
|
||||
cache: Option<CacheCfg>,
|
||||
}
|
||||
|
||||
async fn put_settings(State(state): State<SharedState>, Json(b): Json<SettingsBody>) -> ApiResult {
|
||||
@@ -529,6 +531,11 @@ async fn put_settings(State(state): State<SharedState>, Json(b): Json<SettingsBo
|
||||
}
|
||||
cfg.admin = v;
|
||||
}
|
||||
if let Some(v) = b.cache {
|
||||
// the cache is built at startup; persist now, apply on restart
|
||||
needs_restart = true;
|
||||
cfg.cache = v;
|
||||
}
|
||||
state.reload(cfg)?;
|
||||
ok(json!({ "ok": true, "needs_restart": needs_restart }))
|
||||
}
|
||||
@@ -666,7 +673,9 @@ async fn get_stats(State(state): State<SharedState>, Query(q): Query<StatsQuery>
|
||||
.and_then(parse_duration_secs)
|
||||
.unwrap_or(24 * 3600);
|
||||
let groupby = q.groupby.as_deref().unwrap_or("model");
|
||||
ok(state.stats.snapshot(range_secs, groupby))
|
||||
let mut snap = state.stats.snapshot(range_secs, groupby);
|
||||
snap["cache"] = state.cache.snapshot();
|
||||
ok(snap)
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
|
||||
+171
@@ -0,0 +1,171 @@
|
||||
//! Semantic (embedding-based) response cache.
|
||||
//!
|
||||
//! Each entry stores an L2-normalized embedding of the request plus the response.
|
||||
//! A lookup returns a cached answer when the cosine similarity to a previous request
|
||||
//! (with the same routing signature, within TTL) meets the configured threshold —
|
||||
//! skipping the expensive model call. In-memory ring buffer; counters track hit
|
||||
//! rate and the model cost saved.
|
||||
|
||||
use crate::config::CacheCfg;
|
||||
use crate::stats::parse_duration_secs;
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
struct Entry {
|
||||
sig: String,
|
||||
vec: Vec<f32>, // L2-normalized
|
||||
ts: u64,
|
||||
model_used: String,
|
||||
content: String,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
cost: f64,
|
||||
}
|
||||
|
||||
pub struct CacheHit {
|
||||
pub model_used: String,
|
||||
pub content: String,
|
||||
pub prompt_tokens: u32,
|
||||
pub completion_tokens: u32,
|
||||
pub similarity: f32,
|
||||
}
|
||||
|
||||
pub struct SemanticCache {
|
||||
enabled: bool,
|
||||
threshold: f32,
|
||||
ttl_secs: u64,
|
||||
max_entries: usize,
|
||||
embed_model: Option<String>,
|
||||
entries: Mutex<VecDeque<Entry>>,
|
||||
hits: AtomicU64,
|
||||
misses: AtomicU64,
|
||||
saved_micro_usd: AtomicU64,
|
||||
}
|
||||
|
||||
impl SemanticCache {
|
||||
pub fn new(cfg: &CacheCfg) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
enabled: cfg.enabled,
|
||||
threshold: cfg.threshold,
|
||||
ttl_secs: parse_duration_secs(&cfg.ttl).unwrap_or(3600),
|
||||
max_entries: cfg.max_entries.max(1),
|
||||
embed_model: cfg.embed_model.clone(),
|
||||
entries: Mutex::new(VecDeque::new()),
|
||||
hits: AtomicU64::new(0),
|
||||
misses: AtomicU64::new(0),
|
||||
saved_micro_usd: AtomicU64::new(0),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
pub fn embed_model(&self) -> Option<&str> {
|
||||
self.embed_model.as_deref()
|
||||
}
|
||||
|
||||
/// Look up a cached response by routing signature + query embedding.
|
||||
pub fn lookup(&self, sig: &str, query: &[f32]) -> Option<CacheHit> {
|
||||
if !self.enabled {
|
||||
return None;
|
||||
}
|
||||
let q = normalized(query);
|
||||
let cutoff = now_secs().saturating_sub(self.ttl_secs);
|
||||
let inner = self.entries.lock().unwrap();
|
||||
let mut best: Option<(f32, &Entry)> = None;
|
||||
for e in inner.iter() {
|
||||
if e.sig != sig || e.ts < cutoff {
|
||||
continue;
|
||||
}
|
||||
let sim = dot(&e.vec, &q);
|
||||
if sim >= self.threshold && best.map(|(b, _)| sim > b).unwrap_or(true) {
|
||||
best = Some((sim, e));
|
||||
}
|
||||
}
|
||||
match best {
|
||||
Some((sim, e)) => {
|
||||
self.hits.fetch_add(1, Ordering::Relaxed);
|
||||
self.saved_micro_usd
|
||||
.fetch_add((e.cost * 1_000_000.0) as u64, Ordering::Relaxed);
|
||||
Some(CacheHit {
|
||||
model_used: e.model_used.clone(),
|
||||
content: e.content.clone(),
|
||||
prompt_tokens: e.prompt_tokens,
|
||||
completion_tokens: e.completion_tokens,
|
||||
similarity: sim,
|
||||
})
|
||||
}
|
||||
None => {
|
||||
self.misses.fetch_add(1, Ordering::Relaxed);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Store a response under its routing signature + query embedding.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn insert(
|
||||
&self,
|
||||
sig: String,
|
||||
query: Vec<f32>,
|
||||
model_used: String,
|
||||
content: String,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
cost: f64,
|
||||
) {
|
||||
if !self.enabled || query.is_empty() {
|
||||
return;
|
||||
}
|
||||
let vec = normalized(&query);
|
||||
let mut inner = self.entries.lock().unwrap();
|
||||
inner.push_back(Entry {
|
||||
sig,
|
||||
vec,
|
||||
ts: now_secs(),
|
||||
model_used,
|
||||
content,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
cost,
|
||||
});
|
||||
while inner.len() > self.max_entries {
|
||||
inner.pop_front();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> serde_json::Value {
|
||||
let h = self.hits.load(Ordering::Relaxed);
|
||||
let m = self.misses.load(Ordering::Relaxed);
|
||||
serde_json::json!({
|
||||
"enabled": self.enabled,
|
||||
"hits": h,
|
||||
"misses": m,
|
||||
"hit_rate": if h + m > 0 { h as f64 / (h + m) as f64 } else { 0.0 },
|
||||
"entries": self.entries.lock().unwrap().len(),
|
||||
"saved_usd": self.saved_micro_usd.load(Ordering::Relaxed) as f64 / 1_000_000.0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn dot(a: &[f32], b: &[f32]) -> f32 {
|
||||
a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
|
||||
}
|
||||
|
||||
fn normalized(v: &[f32]) -> Vec<f32> {
|
||||
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if n > 0.0 {
|
||||
v.iter().map(|x| x / n).collect()
|
||||
} else {
|
||||
v.to_vec()
|
||||
}
|
||||
}
|
||||
@@ -31,6 +31,8 @@ pub struct Config {
|
||||
pub admin: AdminCfg,
|
||||
#[serde(default)]
|
||||
pub stats: StatsCfg,
|
||||
#[serde(default)]
|
||||
pub cache: CacheCfg,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
@@ -320,6 +322,46 @@ fn default_stats_retention() -> String {
|
||||
"7d".into()
|
||||
}
|
||||
|
||||
/// Semantic (embedding-based) response cache — returns a cached answer for prompts
|
||||
/// that are semantically near a previous one, skipping the (expensive) model call.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct CacheCfg {
|
||||
#[serde(default)]
|
||||
pub enabled: bool,
|
||||
/// Cosine-similarity threshold for a hit (0..1).
|
||||
#[serde(default = "default_semcache_threshold")]
|
||||
pub threshold: f32,
|
||||
/// Entry time-to-live, e.g. `1h`.
|
||||
#[serde(default = "default_semcache_ttl")]
|
||||
pub ttl: String,
|
||||
/// Maximum number of cached entries (ring buffer).
|
||||
#[serde(default = "default_semcache_max")]
|
||||
pub max_entries: usize,
|
||||
/// Embedding model id (defaults to the first `kind = "embedding"` model).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub embed_model: Option<String>,
|
||||
}
|
||||
impl Default for CacheCfg {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
threshold: default_semcache_threshold(),
|
||||
ttl: default_semcache_ttl(),
|
||||
max_entries: default_semcache_max(),
|
||||
embed_model: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
fn default_semcache_threshold() -> f32 {
|
||||
0.92
|
||||
}
|
||||
fn default_semcache_ttl() -> String {
|
||||
"1h".into()
|
||||
}
|
||||
fn default_semcache_max() -> usize {
|
||||
1000
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Load and validate the config from a file.
|
||||
pub fn load(path: &str) -> anyhow::Result<Self> {
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
#![allow(dead_code)]
|
||||
|
||||
mod admin;
|
||||
mod cache;
|
||||
mod config;
|
||||
mod error;
|
||||
mod model;
|
||||
|
||||
+124
-2
@@ -4,9 +4,9 @@
|
||||
|
||||
use crate::error::{GatewayError, Result};
|
||||
use crate::model::{
|
||||
ChatRequest, ChatResponse, Message, RouteDecision, RouteInfo, RoutingDirective,
|
||||
ChatRequest, ChatResponse, Choice, Message, RouteDecision, RouteInfo, RoutingDirective, Usage,
|
||||
};
|
||||
use crate::state::SharedState;
|
||||
use crate::state::{Live, SharedState};
|
||||
use crate::stats::RequestRecord;
|
||||
use std::time::{Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
@@ -28,6 +28,45 @@ impl Pipeline {
|
||||
}
|
||||
.to_string();
|
||||
|
||||
// --- semantic cache lookup (short-circuit before the model call) ---
|
||||
let cache = &state.cache;
|
||||
let mut cache_key: Option<(String, Vec<f32>)> = None;
|
||||
if cache.enabled() {
|
||||
let live = state.live();
|
||||
let sig = cache_sig(&req);
|
||||
let text = self.extract_text(&req.messages, "concat_all", live.cfg.route.max_chars);
|
||||
let embed_model = cache
|
||||
.embed_model()
|
||||
.map(|s| s.to_string())
|
||||
.or_else(|| first_embedding_model(&live));
|
||||
if let Some(model_id) = embed_model {
|
||||
if let Some(vec) = embed(&live, &model_id, &text).await {
|
||||
if let Some(hit) = cache.lookup(&sig, &vec) {
|
||||
state.stats.record(RequestRecord {
|
||||
ts: now_secs(),
|
||||
account,
|
||||
protocol,
|
||||
directive,
|
||||
task_label: "cache".into(),
|
||||
tier: "cache".into(),
|
||||
score: 0.0,
|
||||
model_id: hit.model_used.clone(),
|
||||
route_source: "cache".into(),
|
||||
prompt_tokens: hit.prompt_tokens,
|
||||
completion_tokens: hit.completion_tokens,
|
||||
cost_usd: 0.0,
|
||||
latency_ms: started.elapsed().as_millis() as u64,
|
||||
outcome: "ok".into(),
|
||||
failover: false,
|
||||
error: None,
|
||||
});
|
||||
return Ok(cached_response(&hit));
|
||||
}
|
||||
cache_key = Some((sig, vec));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let result = self.run_inner(req, state).await;
|
||||
|
||||
// record statistics (for both success and error)
|
||||
@@ -76,6 +115,26 @@ impl Pipeline {
|
||||
};
|
||||
state.stats.record(rec);
|
||||
|
||||
// --- store a successful response in the semantic cache ---
|
||||
if let (Ok(resp), Some((sig, vec))) = (&result, cache_key) {
|
||||
let content = resp
|
||||
.choices
|
||||
.first()
|
||||
.map(|c| c.message.content.clone())
|
||||
.unwrap_or_default();
|
||||
if !content.is_empty() {
|
||||
state.cache.insert(
|
||||
sig,
|
||||
vec,
|
||||
resp.cortiq.selected_model.clone(),
|
||||
content,
|
||||
resp.usage.prompt_tokens,
|
||||
resp.usage.completion_tokens,
|
||||
resp.cortiq.cost_usd,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
@@ -350,6 +409,69 @@ fn now_secs() -> u64 {
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
/// Routing signature used to partition cache entries (so a pinned model's cached
|
||||
/// answer is never returned for a different model/profile).
|
||||
fn cache_sig(req: &ChatRequest) -> String {
|
||||
match &req.routing {
|
||||
RoutingDirective::Auto { profile } => {
|
||||
format!("auto:{}", profile.as_deref().unwrap_or(""))
|
||||
}
|
||||
RoutingDirective::Pinned { model_id } => format!("pinned:{model_id}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn first_embedding_model(live: &Live) -> Option<String> {
|
||||
live.cfg
|
||||
.models
|
||||
.iter()
|
||||
.find(|m| m.kind == "embedding")
|
||||
.map(|m| m.id.clone())
|
||||
}
|
||||
|
||||
/// Embed `text` with the given model; returns the embedding vector (or None on error).
|
||||
async fn embed(live: &Live, model_id: &str, text: &str) -> Option<Vec<f32>> {
|
||||
let provider = live.registry.get(model_id)?;
|
||||
let body = provider.embed(serde_json::json!(text)).await.ok()?;
|
||||
let arr = body["data"][0]["embedding"].as_array()?;
|
||||
Some(
|
||||
arr.iter()
|
||||
.filter_map(|x| x.as_f64())
|
||||
.map(|x| x as f32)
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
fn cached_response(hit: &crate::cache::CacheHit) -> ChatResponse {
|
||||
ChatResponse {
|
||||
id: "cache".to_string(),
|
||||
model_used: hit.model_used.clone(),
|
||||
choices: vec![Choice {
|
||||
index: 0,
|
||||
message: Message {
|
||||
role: "assistant".into(),
|
||||
content: hit.content.clone(),
|
||||
tool_calls: vec![],
|
||||
},
|
||||
finish_reason: "stop".into(),
|
||||
}],
|
||||
usage: Usage {
|
||||
prompt_tokens: hit.prompt_tokens,
|
||||
completion_tokens: hit.completion_tokens,
|
||||
total_tokens: hit.prompt_tokens + hit.completion_tokens,
|
||||
},
|
||||
cortiq: RouteInfo {
|
||||
task_label: "cache".into(),
|
||||
complexity_score: 0.0,
|
||||
complexity_tier: "cache".into(),
|
||||
selected_model: hit.model_used.clone(),
|
||||
route_source: "cache".into(),
|
||||
router_request_id: None,
|
||||
cost_usd: 0.0,
|
||||
failover: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn record_error(
|
||||
state: &SharedState,
|
||||
account: &str,
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
//! a restart: the admin API validates the new config, builds a new `Live`,
|
||||
//! atomically writes the TOML, and swaps the pointer.
|
||||
|
||||
use crate::cache::SemanticCache;
|
||||
use crate::config::Config;
|
||||
use crate::registry::Registry;
|
||||
use crate::router_client::RouterClient;
|
||||
@@ -42,6 +43,7 @@ impl Live {
|
||||
pub struct AppState {
|
||||
pub live: ArcSwap<Live>,
|
||||
pub stats: Arc<Stats>,
|
||||
pub cache: Arc<SemanticCache>,
|
||||
pub secrets: SecretStore,
|
||||
pub config_path: String,
|
||||
pub pipeline: crate::pipeline::Pipeline,
|
||||
@@ -59,10 +61,12 @@ impl SharedState {
|
||||
let secrets_path = sibling(&config_path, "secrets.toml");
|
||||
let secrets = SecretStore::load(&secrets_path);
|
||||
let stats = Stats::new(&cfg.stats);
|
||||
let cache = SemanticCache::new(&cfg.cache);
|
||||
let live = Live::build(cfg, &secrets)?;
|
||||
Ok(SharedState(Arc::new(AppState {
|
||||
live: ArcSwap::from_pointee(live),
|
||||
stats,
|
||||
cache,
|
||||
secrets,
|
||||
config_path,
|
||||
pipeline: crate::pipeline::Pipeline::new(),
|
||||
|
||||
Reference in New Issue
Block a user