feat: add CacheStore with xxhash3-128 and sqlx PostgreSQL backend
This commit is contained in:
@@ -0,0 +1,10 @@
|
||||
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);
|
||||
@@ -0,0 +1,200 @@
|
||||
use anyhow::Result;
|
||||
use sqlx::PgPool;
|
||||
use xxhash_rust::xxh3::xxh3_128;
|
||||
|
||||
use crate::types::{ChatCompletionRequest, Message, MessageContent};
|
||||
|
||||
#[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,
|
||||
}
|
||||
|
||||
impl CacheStore {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub async fn init_schema(&self) -> Result<()> {
|
||||
sqlx::query(
|
||||
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?;
|
||||
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)
|
||||
VALUES ($1, $2, $3)
|
||||
ON CONFLICT (hash) DO UPDATE
|
||||
SET hit_count = message_cache.hit_count + 1
|
||||
"#,
|
||||
)
|
||||
.bind(hash)
|
||||
.bind(role)
|
||||
.bind(content_json)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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",
|
||||
)
|
||||
.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)
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_cacheable<'a>(msg: &'a Message) -> (&'a str, Option<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),
|
||||
}
|
||||
}
|
||||
|
||||
fn content_to_json(content: &MessageContent) -> String {
|
||||
serde_json::to_string(content).unwrap_or_default()
|
||||
}
|
||||
|
||||
#[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);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_hash_different_inputs() {
|
||||
let h1 = compute_hash("hello");
|
||||
let h2 = compute_hash("world");
|
||||
assert_ne!(h1, h2);
|
||||
}
|
||||
|
||||
#[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("");
|
||||
assert_eq!(h.len(), 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_cacheable_user() {
|
||||
use crate::types::MessageContent;
|
||||
let msg = Message::User {
|
||||
content: MessageContent::Text("hi".to_string()),
|
||||
name: None,
|
||||
};
|
||||
let (role, content) = extract_cacheable(&msg);
|
||||
assert_eq!(role, "user");
|
||||
assert!(content.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_cacheable_system_returns_none() {
|
||||
use crate::types::MessageContent;
|
||||
let msg = Message::System {
|
||||
content: MessageContent::Text("be helpful".to_string()),
|
||||
name: None,
|
||||
};
|
||||
let (role, content) = extract_cacheable(&msg);
|
||||
assert_eq!(role, "system");
|
||||
assert!(content.is_none(), "system messages should not be cached");
|
||||
}
|
||||
|
||||
#[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(),
|
||||
};
|
||||
let (role, content) = extract_cacheable(&msg);
|
||||
assert_eq!(role, "tool");
|
||||
assert!(content.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hash_consistency_across_messages() {
|
||||
use crate::types::MessageContent;
|
||||
let msg = Message::User {
|
||||
content: MessageContent::Text("same message".to_string()),
|
||||
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");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user