feat: add Merkle-chain conversation ID and conversation response cache
This commit is contained in:
+459
-118
@@ -1,200 +1,541 @@
|
||||
use anyhow::Result;
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
use sqlx::PgPool;
|
||||
use tokio::sync::mpsc;
|
||||
use xxhash_rust::xxh3::xxh3_128;
|
||||
|
||||
use crate::types::{ChatCompletionRequest, Message, MessageContent};
|
||||
|
||||
// ── Constants ─────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Max entries in per-node DRAM cache. At avg 32 KiB each ≈ 320 MiB worst-case.
|
||||
const L1_CAP: usize = 10_000;
|
||||
/// Max conversation responses in per-node DRAM cache. At avg 4 KiB ≈ 40 MiB worst-case.
|
||||
const CONV_L1_CAP: usize = 10_000;
|
||||
/// Redis TTL in seconds (24 h).
|
||||
const REDIS_TTL: usize = 86_400;
|
||||
/// Redis key prefix for individual message cache.
|
||||
const REDIS_PFX: &str = "mc:";
|
||||
/// Redis key prefix for conversation response cache.
|
||||
const CONV_REDIS_PFX: &str = "cv:";
|
||||
|
||||
// ── Public types ──────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CachedEntry {
|
||||
pub role: Arc<str>,
|
||||
pub content_json: Arc<str>,
|
||||
}
|
||||
|
||||
// ── Internal flush messages ───────────────────────────────────────────────────
|
||||
|
||||
struct Flush {
|
||||
hash_hex: String,
|
||||
role: String,
|
||||
content_json: String,
|
||||
}
|
||||
|
||||
struct ConvFlush {
|
||||
conv_id_hex: String,
|
||||
response_json: Arc<[u8]>,
|
||||
}
|
||||
|
||||
// ── FIFO eviction ring ────────────────────────────────────────────────────────
|
||||
|
||||
/// Single-writer FIFO ring that tracks insertion order for L1 eviction.
|
||||
/// Lives entirely inside the background flush task — zero contention on reads.
|
||||
struct FifoRing {
|
||||
buf: VecDeque<u128>,
|
||||
cap: usize,
|
||||
}
|
||||
|
||||
impl FifoRing {
|
||||
fn new(cap: usize) -> Self {
|
||||
Self { buf: VecDeque::with_capacity(cap + 1), cap }
|
||||
}
|
||||
/// Returns the hash to evict if over capacity, else None.
|
||||
fn push(&mut self, h: u128) -> Option<u128> {
|
||||
self.buf.push_back(h);
|
||||
if self.buf.len() > self.cap { self.buf.pop_front() } else { None }
|
||||
}
|
||||
}
|
||||
|
||||
// ── CacheStore ────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Three-tier write-through cache:
|
||||
/// L1 DashMap (DRAM, this node, FIFO eviction) ~50 ns
|
||||
/// L2 Redis (shared all nodes, TTL 24 h, optional) ~200 µs
|
||||
/// L3 Postgres (durable, async background, analytics) ~2 ms — NEVER on hot path
|
||||
///
|
||||
/// Clone is O(1) — all state is behind Arc.
|
||||
#[derive(Clone)]
|
||||
pub struct CacheStore {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
pub fn compute_hash(content: &str) -> String {
|
||||
let hash = xxh3_128(content.as_bytes());
|
||||
format!("{:032x}", hash)
|
||||
}
|
||||
|
||||
#[derive(Debug, sqlx::FromRow)]
|
||||
pub struct CachedMessage {
|
||||
pub hash: String,
|
||||
pub role: String,
|
||||
pub content_json: String,
|
||||
pub hit_count: i64,
|
||||
// ── message-content cache (existing) ──
|
||||
l1: Arc<DashMap<u128, CachedEntry>>,
|
||||
flush_tx: mpsc::UnboundedSender<Flush>,
|
||||
// ── conversation-response cache (new) ──
|
||||
/// conv_id → full serialised response bytes
|
||||
conv_l1: Arc<DashMap<u128, Arc<[u8]>>>,
|
||||
conv_flush_tx: mpsc::UnboundedSender<ConvFlush>,
|
||||
redis: Option<deadpool_redis::Pool>,
|
||||
/// Kept for init_schema only; L3 writes go through flush channels.
|
||||
pg: PgPool,
|
||||
}
|
||||
|
||||
impl CacheStore {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
pub fn new(pg: PgPool, redis: Option<deadpool_redis::Pool>) -> Self {
|
||||
let l1: Arc<DashMap<u128, CachedEntry>> = Arc::new(DashMap::with_capacity(L1_CAP));
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
tokio::spawn(flush_task(rx, l1.clone(), pg.clone()));
|
||||
|
||||
let conv_l1: Arc<DashMap<u128, Arc<[u8]>>> =
|
||||
Arc::new(DashMap::with_capacity(CONV_L1_CAP));
|
||||
let (conv_tx, conv_rx) = mpsc::unbounded_channel();
|
||||
tokio::spawn(conv_flush_task(conv_rx, conv_l1.clone(), pg.clone()));
|
||||
|
||||
Self { l1, flush_tx: tx, conv_l1, conv_flush_tx: conv_tx, redis, pg }
|
||||
}
|
||||
|
||||
pub async fn init_schema(&self) -> Result<()> {
|
||||
pub async fn init_schema(&self) -> anyhow::Result<()> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS message_cache (
|
||||
r#"CREATE TABLE IF NOT EXISTS message_cache (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
hash TEXT NOT NULL UNIQUE,
|
||||
role TEXT NOT NULL,
|
||||
content_json TEXT NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
hit_count BIGINT NOT NULL DEFAULT 0
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_message_cache_hash ON message_cache (hash);
|
||||
"#,
|
||||
)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
)"#,
|
||||
).execute(&self.pg).await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS idx_message_cache_hash ON message_cache (hash)",
|
||||
).execute(&self.pg).await?;
|
||||
// conversation_cache stores full response bytes keyed by conv_id
|
||||
sqlx::query(
|
||||
r#"CREATE TABLE IF NOT EXISTS conversation_cache (
|
||||
conv_id TEXT NOT NULL PRIMARY KEY,
|
||||
response BYTEA NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
hit_count BIGINT NOT NULL DEFAULT 0
|
||||
)"#,
|
||||
).execute(&self.pg).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn upsert_message(&self, hash: &str, role: &str, content_json: &str) -> Result<()> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO message_cache (hash, role, content_json)
|
||||
/// Look up a hash. L1 hit → 50 ns. L2 hit → 200 µs, backfills L1.
|
||||
/// Returns None only on complete miss (write path still caches on put).
|
||||
pub async fn get(&self, hash: u128) -> Option<CachedEntry> {
|
||||
if let Some(e) = self.l1.get(&hash) {
|
||||
return Some(e.clone());
|
||||
}
|
||||
if let Some(e) = self.redis_get(hash).await {
|
||||
self.l1.insert(hash, e.clone());
|
||||
return Some(e);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Write through all three tiers. L1+L2 are synchronous (fast).
|
||||
/// L3 (Postgres) is fire-and-forget via unbounded channel.
|
||||
pub async fn put(&self, hash: u128, role: &str, content_json: &str) {
|
||||
let entry = CachedEntry {
|
||||
role: Arc::from(role),
|
||||
content_json: Arc::from(content_json),
|
||||
};
|
||||
self.l1.insert(hash, entry);
|
||||
self.redis_put(hash, role, content_json).await;
|
||||
// Non-blocking: L3 write queued to background task
|
||||
let _ = self.flush_tx.send(Flush {
|
||||
hash_hex: hash_to_hex(hash),
|
||||
role: role.to_string(),
|
||||
content_json: content_json.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
/// Cache all messages from a request. Infallible — errors are absorbed.
|
||||
pub async fn cache_messages(&self, req: &ChatCompletionRequest) {
|
||||
for msg in &req.messages {
|
||||
if let Some((role, cj)) = extract_cacheable(msg) {
|
||||
self.put(xxh3_128(cj.as_bytes()), role, &cj).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cache an assistant reply. Returns hex hash.
|
||||
pub async fn cache_assistant_reply(&self, content: &str) -> String {
|
||||
let cj = serde_json::to_string(content).unwrap_or_else(|_| content.to_string());
|
||||
let h = xxh3_128(cj.as_bytes());
|
||||
self.put(h, "assistant", &cj).await;
|
||||
hash_to_hex(h)
|
||||
}
|
||||
|
||||
// ── Redis helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
async fn redis_get(&self, hash: u128) -> Option<CachedEntry> {
|
||||
let pool = self.redis.as_ref()?;
|
||||
let mut conn = pool.get().await.ok()?;
|
||||
let key = format!("{}{}", REDIS_PFX, hash_to_hex(hash));
|
||||
let raw: Option<String> = redis::cmd("GET")
|
||||
.arg(&key)
|
||||
.query_async(&mut *conn)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
let raw = raw?;
|
||||
let (role, cj) = raw.split_once('\n')?;
|
||||
Some(CachedEntry { role: Arc::from(role), content_json: Arc::from(cj) })
|
||||
}
|
||||
|
||||
async fn redis_put(&self, hash: u128, role: &str, content_json: &str) {
|
||||
let pool = match &self.redis { Some(p) => p, None => return };
|
||||
let mut conn = match pool.get().await { Ok(c) => c, Err(_) => return };
|
||||
let key = format!("{}{}", REDIS_PFX, hash_to_hex(hash));
|
||||
let val = format!("{}\n{}", role, content_json);
|
||||
let _: Result<(), _> = redis::cmd("SETEX")
|
||||
.arg(&key).arg(REDIS_TTL).arg(&val)
|
||||
.query_async(&mut *conn).await;
|
||||
}
|
||||
|
||||
// ── Conversation-response cache ───────────────────────────────────────────
|
||||
|
||||
/// Look up a full response by conversation ID.
|
||||
/// L1 hit → 50 ns. L2 (Redis) hit → 200 µs, backfills L1.
|
||||
pub async fn get_conv(&self, conv_id: u128) -> Option<Arc<[u8]>> {
|
||||
if let Some(e) = self.conv_l1.get(&conv_id) {
|
||||
return Some(e.clone());
|
||||
}
|
||||
if let Some(bytes) = self.conv_redis_get(conv_id).await {
|
||||
self.conv_l1.insert(conv_id, bytes.clone());
|
||||
return Some(bytes);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Store a full response keyed by conversation ID.
|
||||
/// L1+L2 synchronous; L3 (Postgres) fire-and-forget.
|
||||
pub async fn put_conv(&self, conv_id: u128, response: &[u8]) {
|
||||
let arc: Arc<[u8]> = Arc::from(response);
|
||||
self.conv_l1.insert(conv_id, arc.clone());
|
||||
self.conv_redis_put(conv_id, response).await;
|
||||
let _ = self.conv_flush_tx.send(ConvFlush {
|
||||
conv_id_hex: hash_to_hex(conv_id),
|
||||
response_json: arc,
|
||||
});
|
||||
}
|
||||
|
||||
async fn conv_redis_get(&self, conv_id: u128) -> Option<Arc<[u8]>> {
|
||||
let pool = self.redis.as_ref()?;
|
||||
let mut conn = pool.get().await.ok()?;
|
||||
let key = format!("{}{}", CONV_REDIS_PFX, hash_to_hex(conv_id));
|
||||
let raw: Option<Vec<u8>> = redis::cmd("GET")
|
||||
.arg(&key)
|
||||
.query_async(&mut *conn)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
raw.map(|v| Arc::from(v.as_slice()))
|
||||
}
|
||||
|
||||
async fn conv_redis_put(&self, conv_id: u128, response: &[u8]) {
|
||||
let pool = match &self.redis { Some(p) => p, None => return };
|
||||
let mut conn = match pool.get().await { Ok(c) => c, Err(_) => return };
|
||||
let key = format!("{}{}", CONV_REDIS_PFX, hash_to_hex(conv_id));
|
||||
let _: Result<(), _> = redis::cmd("SETEX")
|
||||
.arg(&key).arg(REDIS_TTL).arg(response)
|
||||
.query_async(&mut *conn).await;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Background flush task ─────────────────────────────────────────────────────
|
||||
|
||||
/// Drains the flush channel, evicts L1 FIFO, writes Postgres.
|
||||
/// Single task — the FifoRing has zero contention.
|
||||
async fn flush_task(
|
||||
mut rx: mpsc::UnboundedReceiver<Flush>,
|
||||
l1: Arc<DashMap<u128, CachedEntry>>,
|
||||
pg: PgPool,
|
||||
) {
|
||||
let mut ring = FifoRing::new(L1_CAP);
|
||||
while let Some(msg) = rx.recv().await {
|
||||
let h = u128::from_str_radix(&msg.hash_hex, 16).unwrap_or(0);
|
||||
if let Some(evict) = ring.push(h) {
|
||||
l1.remove(&evict);
|
||||
}
|
||||
// Best-effort Postgres upsert; never blocks the hot path
|
||||
let _ = sqlx::query(
|
||||
r#"INSERT INTO message_cache (hash, role, content_json)
|
||||
VALUES ($1, $2, $3)
|
||||
ON CONFLICT (hash) DO UPDATE
|
||||
SET hit_count = message_cache.hit_count + 1
|
||||
"#,
|
||||
SET hit_count = message_cache.hit_count + 1"#,
|
||||
)
|
||||
.bind(hash)
|
||||
.bind(role)
|
||||
.bind(content_json)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
.bind(&msg.hash_hex)
|
||||
.bind(&msg.role)
|
||||
.bind(&msg.content_json)
|
||||
.execute(&pg)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_by_hash(&self, hash: &str) -> Result<Option<CachedMessage>> {
|
||||
let row = sqlx::query_as::<_, CachedMessage>(
|
||||
"SELECT hash, role, content_json, hit_count FROM message_cache WHERE hash = $1",
|
||||
/// Drains the conv flush channel, evicts conv L1 FIFO, writes Postgres.
|
||||
async fn conv_flush_task(
|
||||
mut rx: mpsc::UnboundedReceiver<ConvFlush>,
|
||||
conv_l1: Arc<DashMap<u128, Arc<[u8]>>>,
|
||||
pg: PgPool,
|
||||
) {
|
||||
let mut ring = FifoRing::new(CONV_L1_CAP);
|
||||
while let Some(msg) = rx.recv().await {
|
||||
let h = u128::from_str_radix(&msg.conv_id_hex, 16).unwrap_or(0);
|
||||
if let Some(evict) = ring.push(h) {
|
||||
conv_l1.remove(&evict);
|
||||
}
|
||||
let _ = sqlx::query(
|
||||
r#"INSERT INTO conversation_cache (conv_id, response)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT (conv_id) DO UPDATE
|
||||
SET hit_count = conversation_cache.hit_count + 1"#,
|
||||
)
|
||||
.bind(hash)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
Ok(row)
|
||||
}
|
||||
|
||||
pub async fn cache_messages(&self, req: &ChatCompletionRequest) -> Result<()> {
|
||||
for msg in &req.messages {
|
||||
let (role, content_json) = extract_cacheable(msg);
|
||||
if let Some(json) = content_json {
|
||||
let hash = compute_hash(&json);
|
||||
self.upsert_message(&hash, role, &json).await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn cache_assistant_reply(&self, content: &str) -> Result<String> {
|
||||
let content_json = serde_json::to_string(content).unwrap_or_else(|_| content.to_string());
|
||||
let hash = compute_hash(&content_json);
|
||||
self.upsert_message(&hash, "assistant", &content_json).await?;
|
||||
Ok(hash)
|
||||
.bind(&msg.conv_id_hex)
|
||||
.bind(msg.response_json.as_ref())
|
||||
.execute(&pg)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_cacheable<'a>(msg: &'a Message) -> (&'a str, Option<String>) {
|
||||
// ── Public helpers ────────────────────────────────────────────────────────────
|
||||
|
||||
pub fn compute_hash(content: &str) -> u128 {
|
||||
xxh3_128(content.as_bytes())
|
||||
}
|
||||
|
||||
pub fn hash_to_hex(h: u128) -> String {
|
||||
format!("{:032x}", h)
|
||||
}
|
||||
|
||||
/// Compute a Merkle-chain conversation ID from an ordered message slice.
|
||||
///
|
||||
/// The ID is derived by chaining: `running = xxh3(running_bytes ‖ msg_hash_bytes)`
|
||||
/// for each cacheable message in sequence. System and Developer messages are
|
||||
/// excluded — they are configuration, not conversation turns, so two requests
|
||||
/// that differ only in their system prompt share the same conversation chain.
|
||||
///
|
||||
/// Properties:
|
||||
/// - Two requests with identical message history produce identical conv_ids.
|
||||
/// - A single added message produces a new ID dependent on the full prior chain.
|
||||
/// - O(N) in message count, single allocation-free pass.
|
||||
/// - Returns 0 for an empty / all-system message list (no cacheable turns).
|
||||
pub fn compute_conv_id(messages: &[crate::types::Message]) -> u128 {
|
||||
let mut running = 0u128;
|
||||
for msg in messages {
|
||||
if let Some((_, cj)) = extract_cacheable(msg) {
|
||||
let msg_hash = xxh3_128(cj.as_bytes());
|
||||
let mut buf = [0u8; 32];
|
||||
buf[..16].copy_from_slice(&running.to_le_bytes());
|
||||
buf[16..].copy_from_slice(&msg_hash.to_le_bytes());
|
||||
running = xxh3_128(&buf);
|
||||
}
|
||||
}
|
||||
running
|
||||
}
|
||||
|
||||
fn extract_cacheable(msg: &Message) -> Option<(&'static str, String)> {
|
||||
match msg {
|
||||
Message::User { content, .. } => {
|
||||
("user", Some(content_to_json(content)))
|
||||
}
|
||||
Message::Assistant { .. } => {
|
||||
let json = serde_json::to_string(msg).ok();
|
||||
("assistant", json)
|
||||
}
|
||||
Message::Tool { .. } => {
|
||||
let json = serde_json::to_string(msg).ok();
|
||||
("tool", json)
|
||||
}
|
||||
Message::Function { .. } => {
|
||||
let json = serde_json::to_string(msg).ok();
|
||||
("function", json)
|
||||
}
|
||||
Message::System { .. } | Message::Developer { .. } => ("system", None),
|
||||
Message::User { content, .. } =>
|
||||
Some(("user", serde_json::to_string(content).unwrap_or_default())),
|
||||
Message::Assistant { .. } =>
|
||||
Some(("assistant", serde_json::to_string(msg).unwrap_or_default())),
|
||||
Message::Tool { .. } =>
|
||||
Some(("tool", serde_json::to_string(msg).unwrap_or_default())),
|
||||
Message::Function { .. } =>
|
||||
Some(("function", serde_json::to_string(msg).unwrap_or_default())),
|
||||
Message::System { .. } | Message::Developer { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn content_to_json(content: &MessageContent) -> String {
|
||||
serde_json::to_string(content).unwrap_or_default()
|
||||
}
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_compute_hash_deterministic() {
|
||||
let h1 = compute_hash("hello world");
|
||||
let h2 = compute_hash("hello world");
|
||||
assert_eq!(h1, h2);
|
||||
assert_eq!(compute_hash("hello world"), compute_hash("hello world"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_hash_different_inputs() {
|
||||
let h1 = compute_hash("hello");
|
||||
let h2 = compute_hash("world");
|
||||
assert_ne!(h1, h2);
|
||||
assert_ne!(compute_hash("hello"), compute_hash("world"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_hash_is_32_hex_chars() {
|
||||
let h = compute_hash("test input");
|
||||
assert_eq!(h.len(), 32, "hash should be 32 hex chars: {}", h);
|
||||
assert!(h.chars().all(|c| c.is_ascii_hexdigit()), "hash should be hex: {}", h);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_hash_empty_string() {
|
||||
let h = compute_hash("");
|
||||
fn test_hash_to_hex_32_chars() {
|
||||
let h = hash_to_hex(compute_hash("test input"));
|
||||
assert_eq!(h.len(), 32);
|
||||
assert!(h.chars().all(|c| c.is_ascii_hexdigit()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hash_to_hex_empty() {
|
||||
assert_eq!(hash_to_hex(compute_hash("")), format!("{:032x}", xxh3_128(b"")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fifo_ring_no_eviction_below_cap() {
|
||||
let mut r = FifoRing::new(3);
|
||||
assert!(r.push(1).is_none());
|
||||
assert!(r.push(2).is_none());
|
||||
assert!(r.push(3).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fifo_ring_evicts_oldest() {
|
||||
let mut r = FifoRing::new(3);
|
||||
r.push(1); r.push(2); r.push(3);
|
||||
assert_eq!(r.push(4), Some(1));
|
||||
assert_eq!(r.push(5), Some(2));
|
||||
assert_eq!(r.push(6), Some(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_cacheable_user() {
|
||||
use crate::types::MessageContent;
|
||||
let msg = Message::User {
|
||||
content: MessageContent::Text("hi".to_string()),
|
||||
content: MessageContent::Text("hi".into()),
|
||||
name: None,
|
||||
};
|
||||
let (role, content) = extract_cacheable(&msg);
|
||||
let (role, _) = extract_cacheable(&msg).unwrap();
|
||||
assert_eq!(role, "user");
|
||||
assert!(content.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_cacheable_system_returns_none() {
|
||||
use crate::types::MessageContent;
|
||||
fn test_extract_cacheable_system_none() {
|
||||
let msg = Message::System {
|
||||
content: MessageContent::Text("be helpful".to_string()),
|
||||
content: MessageContent::Text("be helpful".into()),
|
||||
name: None,
|
||||
};
|
||||
let (role, content) = extract_cacheable(&msg);
|
||||
assert_eq!(role, "system");
|
||||
assert!(content.is_none(), "system messages should not be cached");
|
||||
assert!(extract_cacheable(&msg).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_cacheable_tool() {
|
||||
use crate::types::MessageContent;
|
||||
let msg = Message::Tool {
|
||||
content: MessageContent::Text("result".to_string()),
|
||||
tool_call_id: "call_123".to_string(),
|
||||
content: MessageContent::Text("result".into()),
|
||||
tool_call_id: "call_1".into(),
|
||||
};
|
||||
let (role, content) = extract_cacheable(&msg);
|
||||
let (role, _) = extract_cacheable(&msg).unwrap();
|
||||
assert_eq!(role, "tool");
|
||||
assert!(content.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hash_consistency_across_messages() {
|
||||
use crate::types::MessageContent;
|
||||
fn test_same_content_same_hash() {
|
||||
let msg = Message::User {
|
||||
content: MessageContent::Text("same message".to_string()),
|
||||
content: MessageContent::Text("same".into()),
|
||||
name: None,
|
||||
};
|
||||
let (_, c1) = extract_cacheable(&msg);
|
||||
let (_, c2) = extract_cacheable(&msg);
|
||||
let h1 = compute_hash(&c1.unwrap());
|
||||
let h2 = compute_hash(&c2.unwrap());
|
||||
assert_eq!(h1, h2, "same message must always produce same hash");
|
||||
let (_, c1) = extract_cacheable(&msg).unwrap();
|
||||
let (_, c2) = extract_cacheable(&msg).unwrap();
|
||||
assert_eq!(compute_hash(&c1), compute_hash(&c2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_redis_value_roundtrip() {
|
||||
let role = "user";
|
||||
let cj = r#""hello\nworld""#;
|
||||
let val = format!("{}\n{}", role, cj);
|
||||
let (r, c) = val.split_once('\n').unwrap();
|
||||
assert_eq!(r, role);
|
||||
assert_eq!(c, cj);
|
||||
}
|
||||
|
||||
// ── compute_conv_id tests ─────────────────────────────────────────────────
|
||||
|
||||
fn user(text: &str) -> Message {
|
||||
Message::User { content: MessageContent::Text(text.into()), name: None }
|
||||
}
|
||||
fn assistant(text: &str) -> Message {
|
||||
Message::Assistant {
|
||||
content: Some(MessageContent::Text(text.into())),
|
||||
name: None, tool_calls: None, function_call: None,
|
||||
refusal: None, audio: None,
|
||||
}
|
||||
}
|
||||
fn system(text: &str) -> Message {
|
||||
Message::System { content: MessageContent::Text(text.into()), name: None }
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conv_id_empty_is_zero() {
|
||||
assert_eq!(compute_conv_id(&[]), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conv_id_system_only_is_zero() {
|
||||
// System messages don't advance the chain
|
||||
assert_eq!(compute_conv_id(&[system("be helpful")]), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conv_id_deterministic() {
|
||||
let msgs = vec![user("hello"), assistant("hi there"), user("what's 2+2?")];
|
||||
assert_eq!(compute_conv_id(&msgs), compute_conv_id(&msgs));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conv_id_different_content_differs() {
|
||||
let a = vec![user("hello")];
|
||||
let b = vec![user("goodbye")];
|
||||
assert_ne!(compute_conv_id(&a), compute_conv_id(&b));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conv_id_order_matters() {
|
||||
let a = vec![user("hello"), user("world")];
|
||||
let b = vec![user("world"), user("hello")];
|
||||
assert_ne!(compute_conv_id(&a), compute_conv_id(&b));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conv_id_prefix_shared() {
|
||||
// First two turns are shared; only the third differs
|
||||
let prefix = vec![user("hello"), assistant("hi")];
|
||||
let id_at_prefix = compute_conv_id(&prefix);
|
||||
|
||||
let mut branch_a = prefix.clone();
|
||||
branch_a.push(user("what's 2+2?"));
|
||||
let mut branch_b = prefix.clone();
|
||||
branch_b.push(user("what's the time?"));
|
||||
|
||||
// Both branches produce IDs different from the shared prefix ID
|
||||
assert_ne!(compute_conv_id(&branch_a), id_at_prefix);
|
||||
assert_ne!(compute_conv_id(&branch_b), id_at_prefix);
|
||||
// And different from each other
|
||||
assert_ne!(compute_conv_id(&branch_a), compute_conv_id(&branch_b));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conv_id_system_ignored_in_chain() {
|
||||
// Changing only the system prompt must not change the conv_id
|
||||
let a = vec![system("be helpful"), user("hello"), assistant("hi")];
|
||||
let b = vec![system("be terse"), user("hello"), assistant("hi")];
|
||||
assert_eq!(compute_conv_id(&a), compute_conv_id(&b));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conv_id_extends_parent() {
|
||||
// ID after turn N must differ from ID after turn N-1
|
||||
let m0 = user("turn 0");
|
||||
let m1 = user("turn 1");
|
||||
|
||||
let id0 = compute_conv_id(&[m0.clone()]);
|
||||
let id1 = compute_conv_id(&[m0.clone(), m1.clone()]);
|
||||
assert_ne!(id0, id1);
|
||||
assert_ne!(id0, 0);
|
||||
assert_ne!(id1, 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ use reqwest::Client;
|
||||
use serde_json::Value;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::cache::CacheStore;
|
||||
use crate::cache::{self, CacheStore};
|
||||
use crate::middleware::Pipeline;
|
||||
use crate::router::{self, Provider};
|
||||
use crate::types::ChatCompletionRequest;
|
||||
@@ -59,12 +59,27 @@ impl ProxyClient {
|
||||
debug!(provider = ?target.provider, model = %target.model, "routing");
|
||||
req.model = target.model.clone();
|
||||
|
||||
// Compute conv_id from the original messages (before middleware mutation).
|
||||
// Non-streaming only — streaming responses can't be cached.
|
||||
let is_streaming = req.stream.unwrap_or(false);
|
||||
let conv_id = if !is_streaming {
|
||||
let id = cache::compute_conv_id(&req.messages);
|
||||
if id != 0 {
|
||||
if let Some(cached) = self.0.cache.get_conv(id).await {
|
||||
debug!(conv_id = %cache::hash_to_hex(id), "conv cache hit");
|
||||
return Ok(ProxyResponse::Buffered(Bytes::from(cached.to_vec())));
|
||||
}
|
||||
}
|
||||
id
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
// Pre-send: middleware may modify req and auth_header
|
||||
self.0.pipeline.run_pre(&mut req, &mut auth_header);
|
||||
|
||||
self.0.cache.cache_messages(&req).await;
|
||||
|
||||
let is_streaming = req.stream.unwrap_or(false);
|
||||
let url = format!("{}/chat/completions", target.base_url);
|
||||
|
||||
let mut builder = self
|
||||
@@ -111,6 +126,10 @@ impl ProxyClient {
|
||||
}
|
||||
|
||||
let modified = serde_json::to_vec(&json).context("failed to re-serialize response")?;
|
||||
// Store full response in conv cache for future identical requests
|
||||
if conv_id != 0 {
|
||||
self.0.cache.put_conv(conv_id, &modified).await;
|
||||
}
|
||||
return Ok(ProxyResponse::Buffered(Bytes::from(modified)));
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user