From a111aa3b66baf65d674df84e315a09683a14c473 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 7 Jul 2026 02:39:30 +0000 Subject: [PATCH 1/6] feat(rag): add encrypted embeddings foundation --- src/main.rs | 9 +- src/models/mod.rs | 1 + src/models/user_embeddings.rs | 85 +++ src/rag.rs | 1141 +++++++++++++++++++++++++++++++++ src/web/mod.rs | 3 + src/web/openai.rs | 220 +++++-- src/web/rag.rs | 216 +++++++ 7 files changed, 1603 insertions(+), 72 deletions(-) create mode 100644 src/models/user_embeddings.rs create mode 100644 src/rag.rs create mode 100644 src/web/rag.rs diff --git a/src/main.rs b/src/main.rs index 86cc5ff6..c58fe4aa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,7 +21,7 @@ use crate::web::openai_auth::validate_openai_auth; use crate::web::platform_login_routes; use crate::web::{ conversation_projects_routes, conversations_routes, health_routes_with_state, - instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, + instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, rag_routes, responses_routes, }; use crate::{attestation_routes::SessionState, web::platform_routes}; @@ -98,6 +98,7 @@ mod os_flags; mod private_key; mod provider_routing; mod proxy_config; +mod rag; #[cfg(test)] mod security_invariants; mod seed_wrapping; @@ -465,6 +466,7 @@ pub struct AppState { enclave_key: Vec, proxy_router: Arc, provider_router: Arc, + rag_cache: Arc>, resend_api_key: Option, ephemeral_keys: Arc>>, session_states: Arc>>, @@ -768,6 +770,7 @@ impl AppStateBuilder { enclave_key, proxy_router, provider_router, + rag_cache: Arc::new(tokio::sync::Mutex::new(rag::RagCache::default())), resend_api_key: self.resend_api_key, ephemeral_keys: Arc::new(RwLock::new(HashMap::new())), session_states: Arc::new(tokio::sync::RwLock::new(HashMap::new())), @@ -3211,6 +3214,10 @@ async fn main() -> Result<(), Error> { instructions_routes(app_state.clone()) .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), ) + .merge( + rag_routes(app_state.clone()) + .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), + ) .merge(attestation_routes::router(app_state.clone())) .merge(oauth_routes(app_state.clone())) .merge(platform_login_routes(app_state.clone())) diff --git a/src/models/mod.rs b/src/models/mod.rs index 396d4da3..8d92d1a4 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -18,6 +18,7 @@ pub mod responses; pub mod schema; pub mod token_usage; pub mod user_api_keys; +pub mod user_embeddings; pub mod user_kv; pub mod user_seed_wrappings; pub mod users; diff --git a/src/models/user_embeddings.rs b/src/models/user_embeddings.rs new file mode 100644 index 00000000..804cb470 --- /dev/null +++ b/src/models/user_embeddings.rs @@ -0,0 +1,85 @@ +use crate::models::schema::user_embeddings; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +#[derive(Error, Debug)] +pub enum UserEmbeddingError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, AsChangeset, Serialize, Deserialize, Clone, Debug)] +#[diesel(table_name = user_embeddings)] +pub struct UserEmbedding { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub source_type: String, + pub user_message_id: Option, + pub assistant_message_id: Option, + pub conversation_id: Option, + pub vector_enc: Vec, + pub embedding_model: String, + pub vector_dim: i32, + pub content_enc: Vec, + pub metadata_enc: Option>, + pub tags_enc: Vec>, + pub token_count: i32, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = user_embeddings)] +pub struct NewUserEmbedding { + pub uuid: Uuid, + pub user_id: Uuid, + pub source_type: String, + pub user_message_id: Option, + pub assistant_message_id: Option, + pub conversation_id: Option, + pub vector_enc: Vec, + pub embedding_model: String, + pub vector_dim: i32, + pub content_enc: Vec, + pub metadata_enc: Option>, + pub tags_enc: Vec>, + pub token_count: i32, +} + +impl NewUserEmbedding { + pub fn insert(&self, conn: &mut PgConnection) -> Result { + diesel::insert_into(user_embeddings::table) + .values(self) + .get_result::(conn) + .map_err(UserEmbeddingError::DatabaseError) + } +} + +impl UserEmbedding { + pub fn delete_all_for_user( + conn: &mut PgConnection, + user_uuid: Uuid, + ) -> Result { + diesel::delete(user_embeddings::table.filter(user_embeddings::user_id.eq(user_uuid))) + .execute(conn) + .map_err(UserEmbeddingError::DatabaseError) + } + + pub fn delete_by_uuid_for_user( + conn: &mut PgConnection, + user_uuid: Uuid, + embedding_uuid: Uuid, + ) -> Result { + diesel::delete( + user_embeddings::table + .filter(user_embeddings::user_id.eq(user_uuid)) + .filter(user_embeddings::uuid.eq(embedding_uuid)), + ) + .execute(conn) + .map_err(UserEmbeddingError::DatabaseError) + } +} diff --git a/src/rag.rs b/src/rag.rs new file mode 100644 index 00000000..71922fa2 --- /dev/null +++ b/src/rag.rs @@ -0,0 +1,1141 @@ +use crate::encrypt::{decrypt_with_key, encrypt_key_deterministic, encrypt_with_key}; +use crate::models::schema::user_embeddings; +use crate::models::user_embeddings::NewUserEmbedding; +use crate::models::users::User; +use crate::web::openai_auth::AuthMethod; +use crate::{ApiError, AppState}; +use base64::{engine::general_purpose::STANDARD as B64_STANDARD, Engine as _}; +use diesel::prelude::*; +use secp256k1::SecretKey; +use serde::Serialize; +use std::cmp::Ordering; +use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tracing::{debug, error}; +use uuid::Uuid; + +pub const DEFAULT_EMBEDDING_MODEL: &str = "nomic-embed-text"; +pub const DEFAULT_EMBEDDING_DIM: i32 = 768; + +const CACHE_MAX_USERS: usize = 100; +const CACHE_TTL: Duration = Duration::from_secs(5 * 60); + +const DB_SCAN_BATCH_SIZE: i64 = 1000; + +#[allow(dead_code)] +pub const SOURCE_TYPE_MESSAGE: &str = "message"; +pub const SOURCE_TYPE_ARCHIVAL: &str = "archival"; + +#[derive(Debug, Clone, Serialize)] +pub struct RagSearchResult { + pub content: String, + pub score: f32, + pub token_count: i32, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RagEmbeddingsStatus { + pub total_embeddings: i64, + pub by_model: HashMap, + pub stale_count: i64, +} + +#[derive(Debug, Clone)] +pub struct CachedEmbedding { + pub source_type: String, + pub conversation_id: Option, + pub vector: Vec, + pub content_enc: Vec, + pub token_count: i32, +} + +#[derive(Debug, Clone)] +struct CacheEntry { + loaded_at: Instant, + embeddings: Arc>, +} + +#[derive(Debug)] +pub struct RagCache { + max_users: usize, + ttl: Duration, + entries: HashMap, + lru: VecDeque, +} + +impl Default for RagCache { + fn default() -> Self { + Self::new(CACHE_MAX_USERS, CACHE_TTL) + } +} + +impl RagCache { + pub fn new(max_users: usize, ttl: Duration) -> Self { + Self { + max_users, + ttl, + entries: HashMap::new(), + lru: VecDeque::new(), + } + } + + pub fn evict_user(&mut self, user_id: Uuid) { + self.entries.remove(&user_id); + self.lru.retain(|u| *u != user_id); + } + + pub fn get(&mut self, user_id: Uuid) -> Option>> { + self.evict_expired(); + + let (loaded_at, embeddings) = { + let entry = self.entries.get(&user_id)?; + (entry.loaded_at, entry.embeddings.clone()) + }; + + if loaded_at.elapsed() > self.ttl { + self.evict_user(user_id); + return None; + } + + self.touch(user_id); + Some(embeddings) + } + + pub fn put(&mut self, user_id: Uuid, embeddings: Arc>) { + self.entries.insert( + user_id, + CacheEntry { + loaded_at: Instant::now(), + embeddings, + }, + ); + self.touch(user_id); + + while self.entries.len() > self.max_users { + if let Some(lru_user) = self.lru.pop_back() { + self.entries.remove(&lru_user); + } else { + break; + } + } + } + + fn touch(&mut self, user_id: Uuid) { + self.lru.retain(|u| *u != user_id); + self.lru.push_front(user_id); + } + + fn evict_expired(&mut self) { + let ttl = self.ttl; + let expired: Vec = self + .entries + .iter() + .filter_map(|(user_id, entry)| { + if entry.loaded_at.elapsed() > ttl { + Some(*user_id) + } else { + None + } + }) + .collect(); + + for user_id in expired { + self.evict_user(user_id); + } + } +} + +pub fn serialize_f32_le(values: &[f32]) -> Vec { + let mut out = Vec::with_capacity(values.len() * 4); + for v in values { + out.extend_from_slice(&v.to_le_bytes()); + } + out +} + +pub fn deserialize_f32_le(bytes: &[u8]) -> Result, ApiError> { + if !bytes.len().is_multiple_of(4) { + return Err(ApiError::BadRequest); + } + + let mut out = Vec::with_capacity(bytes.len() / 4); + for chunk in bytes.chunks_exact(4) { + let arr: [u8; 4] = chunk.try_into().map_err(|_| ApiError::BadRequest)?; + out.push(f32::from_le_bytes(arr)); + } + Ok(out) +} + +pub fn cosine_similarity(a: &[f32], b: &[f32]) -> Result { + if a.len() != b.len() { + return Err(ApiError::BadRequest); + } + + let mut dot = 0.0f32; + let mut norm_a = 0.0f32; + let mut norm_b = 0.0f32; + + for (x, y) in a.iter().zip(b.iter()) { + dot += x * y; + norm_a += x * x; + norm_b += y * y; + } + + if norm_a == 0.0 || norm_b == 0.0 { + return Ok(0.0); + } + + Ok(dot / (norm_a.sqrt() * norm_b.sqrt())) +} + +#[derive(Debug, Clone)] +struct HeapItem { + score: f32, + token_count: i32, + content_enc: Vec, +} + +impl Eq for HeapItem {} + +impl PartialEq for HeapItem { + fn eq(&self, other: &Self) -> bool { + self.score.to_bits() == other.score.to_bits() && self.token_count == other.token_count + } +} + +impl Ord for HeapItem { + fn cmp(&self, other: &Self) -> Ordering { + self.score + .total_cmp(&other.score) + .then_with(|| other.token_count.cmp(&self.token_count)) + } +} + +impl PartialOrd for HeapItem { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +fn top_k_candidates( + query: &[f32], + embeddings: &[CachedEmbedding], + top_k: usize, + source_types: Option<&[String]>, + conversation_id: Option, +) -> Result, ApiError> { + let allowed_source_types: Option> = source_types.map(|v| { + v.iter() + .map(|s| s.as_str()) + .collect::>() + }); + + let mut heap: BinaryHeap> = BinaryHeap::new(); + + for e in embeddings { + if let Some(allowed) = &allowed_source_types { + if !allowed.contains(e.source_type.as_str()) { + continue; + } + } + + if let Some(conv_id) = conversation_id { + if e.conversation_id != Some(conv_id) { + continue; + } + } + + if e.vector.len() != query.len() { + continue; + } + + if heap.len() < top_k { + let score = cosine_similarity(query, &e.vector)?; + let item = HeapItem { + score, + token_count: e.token_count, + content_enc: e.content_enc.clone(), + }; + heap.push(std::cmp::Reverse(item)); + continue; + } + + let score = cosine_similarity(query, &e.vector)?; + if let Some(std::cmp::Reverse(min)) = heap.peek() { + let ordering = score + .total_cmp(&min.score) + .then_with(|| min.token_count.cmp(&e.token_count)); + if ordering == Ordering::Greater { + heap.pop(); + let item = HeapItem { + score, + token_count: e.token_count, + content_enc: e.content_enc.clone(), + }; + heap.push(std::cmp::Reverse(item)); + } + } + } + + let mut out: Vec = heap.into_iter().map(|r| r.0).collect(); + out.sort_by(|a, b| b.cmp(a)); + Ok(out) +} + +fn apply_token_budget(results: Vec, budget: i32) -> Vec { + let mut total: i32 = 0; + let mut limited: Vec = Vec::new(); + for r in results { + if total + r.token_count > budget { + break; + } + total += r.token_count; + limited.push(r); + } + limited +} + +async fn embed_text_via_tinfoil( + state: &Arc, + user: &User, + auth_method: AuthMethod, + text: &str, +) -> Result<(Vec, i32), ApiError> { + crate::web::get_embedding_vector( + state, + user, + auth_method, + DEFAULT_EMBEDDING_MODEL, + text, + Some(DEFAULT_EMBEDDING_DIM), + ) + .await +} + +fn normalize_tags<'a>(tags: impl Iterator) -> Vec { + let mut seen: HashSet = HashSet::new(); + let mut out: Vec = Vec::new(); + + for tag in tags { + let normalized = tag.trim().to_lowercase(); + if normalized.is_empty() { + continue; + } + + if seen.insert(normalized.clone()) { + out.push(normalized); + } + } + + out +} + +fn encrypt_tags_b64(user_key: &SecretKey, tags: &[String]) -> Vec> { + tags.iter() + .map(|tag| { + let ciphertext = encrypt_key_deterministic(user_key, tag.as_bytes()); + Some(B64_STANDARD.encode(ciphertext)) + }) + .collect() +} + +fn extract_tags_from_metadata(metadata: Option<&serde_json::Value>) -> Vec { + let Some(metadata) = metadata else { + return Vec::new(); + }; + + let Some(tags) = metadata.get("tags") else { + return Vec::new(); + }; + + match tags { + serde_json::Value::Array(arr) => normalize_tags(arr.iter().filter_map(|v| v.as_str())), + serde_json::Value::String(s) => normalize_tags(s.split(',')), + _ => Vec::new(), + } +} + +pub async fn insert_archival_embedding( + state: &Arc, + user: &User, + auth_method: AuthMethod, + user_key: &SecretKey, + text: &str, + metadata: Option<&serde_json::Value>, +) -> Result { + let text = text.trim(); + if text.is_empty() { + return Err(ApiError::BadRequest); + } + + let user_id = user.uuid; + let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; + + let vector_bytes = serialize_f32_le(&vector); + let vector_enc = encrypt_with_key(user_key, &vector_bytes).await; + let content_enc = encrypt_with_key(user_key, text.as_bytes()).await; + + let metadata_enc = if let Some(m) = metadata { + let m_bytes = serde_json::to_vec(m).map_err(|_| ApiError::BadRequest)?; + Some(encrypt_with_key(user_key, &m_bytes).await) + } else { + None + }; + + let tags = extract_tags_from_metadata(metadata); + let tags_enc = encrypt_tags_b64(user_key, &tags); + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let inserted = NewUserEmbedding { + uuid: Uuid::new_v4(), + user_id, + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + user_message_id: None, + assistant_message_id: None, + conversation_id: None, + vector_enc, + embedding_model: DEFAULT_EMBEDDING_MODEL.to_string(), + vector_dim: DEFAULT_EMBEDDING_DIM, + content_enc, + metadata_enc, + tags_enc, + token_count, + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to insert archival embedding: {:?}", e); + ApiError::InternalServerError + })?; + + state.rag_cache.lock().await.evict_user(user_id); + Ok(inserted) +} + +#[allow(dead_code)] +#[allow(clippy::too_many_arguments)] +pub async fn insert_message_embedding( + state: &Arc, + user: &User, + auth_method: AuthMethod, + user_key: &SecretKey, + text: &str, + conversation_id: i64, + user_message_id: Option, + assistant_message_id: Option, +) -> Result { + let user_id = user.uuid; + let text = text.trim(); + if text.is_empty() { + return Err(ApiError::BadRequest); + } + + let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; + + let vector_bytes = serialize_f32_le(&vector); + let vector_enc = encrypt_with_key(user_key, &vector_bytes).await; + let content_enc = encrypt_with_key(user_key, text.as_bytes()).await; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let inserted = NewUserEmbedding { + uuid: Uuid::new_v4(), + user_id, + source_type: SOURCE_TYPE_MESSAGE.to_string(), + user_message_id, + assistant_message_id, + conversation_id: Some(conversation_id), + vector_enc, + embedding_model: DEFAULT_EMBEDDING_MODEL.to_string(), + vector_dim: DEFAULT_EMBEDDING_DIM, + content_enc, + metadata_enc: None, + tags_enc: Vec::new(), + token_count, + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to insert message embedding: {:?}", e); + ApiError::InternalServerError + })?; + + state.rag_cache.lock().await.evict_user(user_id); + Ok(inserted) +} + +async fn load_all_user_embeddings( + state: &AppState, + user_id: Uuid, + user_key: &SecretKey, +) -> Result>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let mut last_id: i64 = 0; + let mut out: Vec = Vec::new(); + + #[derive(Queryable)] + struct EmbeddingScanRow { + source_type: String, + conversation_id: Option, + vector_enc: Vec, + content_enc: Vec, + token_count: i32, + vector_dim: i32, + id: i64, + } + + loop { + let rows: Vec = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::embedding_model.eq(DEFAULT_EMBEDDING_MODEL)) + .filter(user_embeddings::id.gt(last_id)) + .order(user_embeddings::id.asc()) + .select(( + user_embeddings::source_type, + user_embeddings::conversation_id, + user_embeddings::vector_enc, + user_embeddings::content_enc, + user_embeddings::token_count, + user_embeddings::vector_dim, + user_embeddings::id, + )) + .limit(DB_SCAN_BATCH_SIZE) + .load(&mut conn) + .map_err(|e| { + error!( + "Failed to load embeddings batch for user={} after id={}: {:?}", + user_id, last_id, e + ); + ApiError::InternalServerError + })?; + + if rows.is_empty() { + break; + } + + for row in rows { + let EmbeddingScanRow { + source_type, + conversation_id, + vector_enc, + content_enc, + token_count, + vector_dim, + id, + } = row; + last_id = id; + + let vector_bytes = decrypt_with_key(user_key, &vector_enc) + .map_err(|_| ApiError::InternalServerError)?; + let vector = deserialize_f32_le(&vector_bytes)?; + + if vector.len() != vector_dim as usize { + debug!( + "Skipping embedding id={} for user={} due to dim mismatch (expected={}, got={})", + id, + user_id, + vector_dim, + vector.len() + ); + continue; + } + + out.push(CachedEmbedding { + source_type, + conversation_id, + vector, + content_enc, + token_count, + }); + } + } + + Ok(Arc::new(out)) +} + +async fn load_user_embeddings_by_tags( + state: &AppState, + user_id: Uuid, + user_key: &SecretKey, + source_types: Option<&[String]>, + conversation_id: Option, + tags_enc_filter: &[Option], +) -> Result>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let tags_enc_filter = tags_enc_filter.to_vec(); + + let mut last_id: i64 = 0; + let mut out: Vec = Vec::new(); + + #[derive(Queryable)] + struct EmbeddingScanRow { + source_type: String, + conversation_id: Option, + vector_enc: Vec, + content_enc: Vec, + token_count: i32, + vector_dim: i32, + id: i64, + } + + loop { + let mut query = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::embedding_model.eq(DEFAULT_EMBEDDING_MODEL)) + .filter(user_embeddings::id.gt(last_id)) + .filter(user_embeddings::tags_enc.overlaps_with(tags_enc_filter.clone())) + .into_boxed(); + + if let Some(source_types) = source_types { + query = query.filter(user_embeddings::source_type.eq_any(source_types)); + } + + if let Some(conversation_id) = conversation_id { + query = query.filter(user_embeddings::conversation_id.eq(Some(conversation_id))); + } + + let rows: Vec = query + .order(user_embeddings::id.asc()) + .select(( + user_embeddings::source_type, + user_embeddings::conversation_id, + user_embeddings::vector_enc, + user_embeddings::content_enc, + user_embeddings::token_count, + user_embeddings::vector_dim, + user_embeddings::id, + )) + .limit(DB_SCAN_BATCH_SIZE) + .load(&mut conn) + .map_err(|e| { + error!( + "Failed to load tag-filtered embeddings batch for user={} after id={}: {:?}", + user_id, last_id, e + ); + ApiError::InternalServerError + })?; + + if rows.is_empty() { + break; + } + + for row in rows { + let EmbeddingScanRow { + source_type, + conversation_id, + vector_enc, + content_enc, + token_count, + vector_dim, + id, + } = row; + last_id = id; + + let vector_bytes = decrypt_with_key(user_key, &vector_enc) + .map_err(|_| ApiError::InternalServerError)?; + let vector = deserialize_f32_le(&vector_bytes)?; + + if vector.len() != vector_dim as usize { + debug!( + "Skipping embedding id={} for user={} due to dim mismatch (expected={}, got={})", + id, + user_id, + vector_dim, + vector.len() + ); + continue; + } + + out.push(CachedEmbedding { + source_type, + conversation_id, + vector, + content_enc, + token_count, + }); + } + } + + Ok(Arc::new(out)) +} + +#[allow(clippy::too_many_arguments)] +pub async fn search_user_embeddings( + state: &Arc, + user: &User, + auth_method: AuthMethod, + user_key: &SecretKey, + query: &str, + top_k: usize, + max_tokens: Option, + source_types: Option<&[String]>, + conversation_id: Option, + tags: Option<&[String]>, +) -> Result, ApiError> { + let top_k = top_k.clamp(1, 20); + + let user_id = user.uuid; + + let (query_vec, _query_tokens) = + embed_text_via_tinfoil(state, user, auth_method, query).await?; + + let tags_enc_filter = tags + .map(|t| normalize_tags(t.iter().map(|s| s.as_str()))) + .filter(|t| !t.is_empty()) + .map(|t| encrypt_tags_b64(user_key, &t)); + + let embeddings = if let Some(tags_enc_filter) = tags_enc_filter.as_ref() { + load_user_embeddings_by_tags( + state, + user_id, + user_key, + source_types, + conversation_id, + tags_enc_filter, + ) + .await? + } else { + let cached = { + let mut cache = state.rag_cache.lock().await; + cache.get(user_id) + }; + + if let Some(hit) = cached { + hit + } else { + let loaded = load_all_user_embeddings(state, user_id, user_key).await?; + state.rag_cache.lock().await.put(user_id, loaded.clone()); + loaded + } + }; + + let candidates = top_k_candidates( + &query_vec, + &embeddings, + top_k, + source_types, + conversation_id, + )?; + + let mut results: Vec = Vec::with_capacity(candidates.len()); + for c in candidates { + let plaintext = decrypt_with_key(user_key, &c.content_enc) + .map_err(|_| ApiError::InternalServerError)?; + let content = String::from_utf8(plaintext).map_err(|_| ApiError::InternalServerError)?; + results.push(RagSearchResult { + content, + score: c.score, + token_count: c.token_count, + }); + } + + if let Some(budget) = max_tokens { + results = apply_token_budget(results, budget); + } + + Ok(results) +} + +pub async fn delete_all_user_embeddings(state: &AppState, user_id: Uuid) -> Result<(), ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + diesel::delete(user_embeddings::table.filter(user_embeddings::user_id.eq(user_id))) + .execute(&mut conn) + .map_err(|e| { + error!( + "Failed to delete all embeddings for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + + state.rag_cache.lock().await.evict_user(user_id); + Ok(()) +} + +pub async fn delete_user_embedding_by_uuid( + state: &AppState, + user_id: Uuid, + embedding_uuid: Uuid, +) -> Result<(), ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let affected = diesel::delete( + user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::uuid.eq(embedding_uuid)), + ) + .execute(&mut conn) + .map_err(|e| { + error!( + "Failed to delete embedding user={} uuid={}: {:?}", + user_id, embedding_uuid, e + ); + ApiError::InternalServerError + })?; + + if affected == 0 { + return Err(ApiError::NotFound); + } + + state.rag_cache.lock().await.evict_user(user_id); + Ok(()) +} + +pub async fn embeddings_status( + state: &AppState, + user_id: Uuid, +) -> Result { + use diesel::dsl::count_star; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let total_embeddings: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .select(count_star()) + .first(&mut conn) + .map_err(|e| { + error!( + "Failed to count embeddings for user={} (total): {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + + let grouped: Vec<(String, i64)> = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .group_by(user_embeddings::embedding_model) + .select((user_embeddings::embedding_model, count_star())) + .load(&mut conn) + .map_err(|e| { + error!( + "Failed to group embeddings by model for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + + let mut by_model: HashMap = HashMap::new(); + for (model, count) in grouped { + by_model.insert(model, count); + } + + let stale_count: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::embedding_model.ne(DEFAULT_EMBEDDING_MODEL)) + .select(count_star()) + .first(&mut conn) + .map_err(|e| { + error!( + "Failed to count stale embeddings for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + + Ok(RagEmbeddingsStatus { + total_embeddings, + by_model, + stale_count, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn serialize_deserialize_roundtrip() { + let v = vec![0.0f32, 1.5, -2.25, 42.0]; + let bytes = serialize_f32_le(&v); + let decoded = deserialize_f32_le(&bytes).unwrap(); + assert_eq!(v, decoded); + } + + #[tokio::test] + async fn vector_encrypt_roundtrip() { + let key = SecretKey::from_slice(&[7u8; 32]).unwrap(); + let v = vec![0.0f32, 1.5, -2.25, 42.0]; + + let bytes = serialize_f32_le(&v); + let enc = encrypt_with_key(&key, &bytes).await; + let dec = decrypt_with_key(&key, &enc).unwrap(); + let decoded = deserialize_f32_le(&dec).unwrap(); + assert_eq!(v, decoded); + } + + #[test] + fn cosine_similarity_known_values() { + let a = vec![1.0f32, 0.0, 0.0]; + let b = vec![1.0f32, 0.0, 0.0]; + let c = vec![0.0f32, 1.0, 0.0]; + + assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 1e-6); + assert!((cosine_similarity(&a, &c).unwrap() - 0.0).abs() < 1e-6); + } + + #[test] + fn cosine_similarity_mismatched_dimensions_errors() { + let a = vec![1.0f32, 0.0]; + let b = vec![1.0f32]; + assert!(matches!( + cosine_similarity(&a, &b), + Err(ApiError::BadRequest) + )); + } + + #[test] + fn cosine_similarity_zero_vectors_return_zero() { + let zero = vec![0.0f32, 0.0, 0.0]; + let nonzero = vec![1.0f32, 2.0, 3.0]; + + assert_eq!(cosine_similarity(&zero, &nonzero).unwrap(), 0.0); + assert_eq!(cosine_similarity(&nonzero, &zero).unwrap(), 0.0); + assert_eq!(cosine_similarity(&zero, &zero).unwrap(), 0.0); + } + + #[test] + fn deserialize_invalid_byte_length_errors() { + let bytes = vec![0u8; 3]; + assert!(matches!( + deserialize_f32_le(&bytes), + Err(ApiError::BadRequest) + )); + } + + #[test] + fn apply_token_budget_is_prefix() { + let results = vec![ + RagSearchResult { + content: "a".to_string(), + score: 1.0, + token_count: 6, + }, + RagSearchResult { + content: "b".to_string(), + score: 0.9, + token_count: 6, + }, + RagSearchResult { + content: "c".to_string(), + score: 0.8, + token_count: 1, + }, + ]; + + let limited = apply_token_budget(results, 10); + assert_eq!(limited.len(), 1); + assert_eq!(limited[0].content, "a"); + } + + #[test] + fn apply_token_budget_exact_fit_keeps_all() { + let results = vec![ + RagSearchResult { + content: "a".to_string(), + score: 1.0, + token_count: 3, + }, + RagSearchResult { + content: "b".to_string(), + score: 0.9, + token_count: 7, + }, + ]; + + let limited = apply_token_budget(results, 10); + assert_eq!(limited.len(), 2); + } + + #[test] + fn apply_token_budget_zero_budget_returns_empty() { + let results = vec![RagSearchResult { + content: "a".to_string(), + score: 1.0, + token_count: 1, + }]; + + let limited = apply_token_budget(results, 0); + assert!(limited.is_empty()); + } + + #[test] + fn top_k_heap_ranking_and_filters() { + let query = vec![1.0f32, 0.0]; + + let embeddings = vec![ + CachedEmbedding { + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + conversation_id: None, + vector: vec![1.0, 0.0], + content_enc: b"a".to_vec(), + token_count: 5, + }, + CachedEmbedding { + source_type: SOURCE_TYPE_MESSAGE.to_string(), + conversation_id: Some(123), + vector: vec![0.0, 1.0], + content_enc: b"b".to_vec(), + token_count: 7, + }, + CachedEmbedding { + source_type: SOURCE_TYPE_MESSAGE.to_string(), + conversation_id: Some(123), + vector: vec![0.8, 0.2], + content_enc: b"c".to_vec(), + token_count: 9, + }, + ]; + + let items = top_k_candidates( + &query, + &embeddings, + 2, + Some(&[SOURCE_TYPE_MESSAGE.to_string()]), + Some(123), + ) + .unwrap(); + + assert_eq!(items.len(), 2); + assert!(items[0].score >= items[1].score); + assert_eq!(items[0].content_enc, b"c"); + assert_eq!(items[1].content_enc, b"b"); + } + + #[test] + fn top_k_tie_break_prefers_fewer_tokens() { + let query = vec![1.0f32, 0.0]; + + let embeddings = vec![ + CachedEmbedding { + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + conversation_id: None, + vector: vec![1.0, 0.0], + content_enc: b"a".to_vec(), + token_count: 10, + }, + CachedEmbedding { + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + conversation_id: None, + vector: vec![1.0, 0.0], + content_enc: b"b".to_vec(), + token_count: 5, + }, + ]; + + let items = top_k_candidates(&query, &embeddings, 1, None, None).unwrap(); + assert_eq!(items.len(), 1); + assert_eq!(items[0].content_enc, b"b"); + } + + #[test] + fn top_k_when_k_gt_embeddings_returns_all() { + let query = vec![1.0f32, 0.0]; + + let embeddings = vec![ + CachedEmbedding { + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + conversation_id: None, + vector: vec![1.0, 0.0], + content_enc: b"a".to_vec(), + token_count: 5, + }, + CachedEmbedding { + source_type: SOURCE_TYPE_MESSAGE.to_string(), + conversation_id: Some(123), + vector: vec![0.0, 1.0], + content_enc: b"b".to_vec(), + token_count: 7, + }, + ]; + + let total = embeddings.len(); + let items = top_k_candidates(&query, &embeddings, 10, None, None).unwrap(); + assert_eq!(items.len(), total); + } + + #[test] + fn top_k_with_empty_embeddings_returns_empty() { + let query = vec![1.0f32, 0.0]; + let embeddings: Vec = vec![]; + let items = top_k_candidates(&query, &embeddings, 10, None, None).unwrap(); + assert!(items.is_empty()); + } + + #[test] + fn rag_cache_evict_user_removes_entry() { + let mut cache = RagCache::new(10, Duration::from_secs(60)); + let user_id = Uuid::new_v4(); + + cache.put(user_id, Arc::new(vec![])); + assert!(cache.entries.contains_key(&user_id)); + + cache.evict_user(user_id); + assert!(!cache.entries.contains_key(&user_id)); + assert!(!cache.lru.contains(&user_id)); + assert!(cache.get(user_id).is_none()); + } + + #[tokio::test] + async fn rag_cache_lru_eviction() { + let mut cache = RagCache::new(2, Duration::from_secs(60)); + + let v1 = Arc::new(vec![]); + let v2 = Arc::new(vec![]); + let v3 = Arc::new(vec![]); + + let u1 = Uuid::new_v4(); + let u2 = Uuid::new_v4(); + let u3 = Uuid::new_v4(); + + cache.put(u1, v1); + cache.put(u2, v2); + // touch u1 so u2 becomes LRU + cache.get(u1); + cache.put(u3, v3); + + assert!(cache.entries.contains_key(&u1)); + assert!(!cache.entries.contains_key(&u2)); + assert!(cache.entries.contains_key(&u3)); + } + + #[tokio::test] + async fn rag_cache_ttl_expiration() { + let mut cache = RagCache::new(10, Duration::from_millis(5)); + let user = Uuid::new_v4(); + + cache.put(user, Arc::new(vec![])); + tokio::time::sleep(Duration::from_millis(10)).await; + assert!(cache.get(user).is_none()); + } +} diff --git a/src/web/mod.rs b/src/web/mod.rs index dd9a5cc0..5e7eff65 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -8,14 +8,17 @@ mod openai; pub mod openai_auth; pub mod platform; pub mod protected_routes; +pub mod rag; pub mod responses; pub use health_routes::router_with_state as health_routes_with_state; pub use login_routes::router as login_routes; pub use oauth_routes::router as oauth_routes; +pub use openai::get_embedding_vector; pub use openai::router as openai_routes; pub use platform::router as platform_routes; pub use protected_routes::router as protected_routes; +pub use rag::router as rag_routes; pub use responses::conversation_projects_router as conversation_projects_routes; pub use responses::conversations_router as conversations_routes; pub use responses::instructions_router as instructions_routes; diff --git a/src/web/openai.rs b/src/web/openai.rs index fec095c5..23e6cc2b 100644 --- a/src/web/openai.rs +++ b/src/web/openai.rs @@ -2,7 +2,7 @@ use crate::model_config::{model_catalog_response, openai_models_response}; use crate::models::token_usage::NewTokenUsage; use crate::models::users::User; use crate::provider_routing::ProviderRoutingError; -use crate::proxy_config::ProxyConfig; +use crate::proxy_config::{canonicalize_tinfoil_model, ProxyConfig}; use crate::sqs::UsageEvent; use crate::web::audio_utils::{merge_transcriptions, AudioSplitter, TINFOIL_MAX_SIZE}; use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; @@ -1989,53 +1989,10 @@ async fn proxy_tts( encrypt_response(&state, &session_id, &audio_response).await } -async fn proxy_embeddings( - State(state): State>, - _headers: HeaderMap, - axum::Extension(session_id): axum::Extension, - axum::Extension(user): axum::Extension, - axum::Extension(_auth_method): axum::Extension, - axum::Extension(embedding_request): axum::Extension, -) -> Result>, ApiError> { - // Check if guest user is allowed (paid guests are allowed, free guests are not) - if user.is_guest() { - if let Some(billing_client) = &state.billing_client { - match billing_client.is_user_paid(user.uuid).await { - Ok(true) => { - debug!("Paid guest user allowed for embeddings: {}", user.uuid); - } - Ok(false) => { - error!( - "Free guest user attempted to use embeddings feature: {}", - user.uuid - ); - return Err(ApiError::Unauthorized); - } - Err(e) => { - error!("Billing check failed for guest user {}: {}", user.uuid, e); - return Err(ApiError::Unauthorized); - } - } - } else { - error!( - "Guest user attempted to use embeddings without billing client: {}", - user.uuid - ); - return Err(ApiError::Unauthorized); - } - } - - // Validate input is not empty - let is_empty = match &embedding_request.input { - Value::String(s) => s.trim().is_empty(), - Value::Array(arr) => arr.is_empty(), - _ => true, - }; - if is_empty { - error!("Input is empty or invalid"); - return Err(ApiError::BadRequest); - } - +async fn request_embeddings_from_primary( + state: &AppState, + embedding_request: &EmbeddingRequest, +) -> Result<(Value, String, String), ApiError> { let proxy_config = state.proxy_router.get_tinfoil_proxy(); // Create a new hyper client @@ -2045,8 +2002,14 @@ async fn proxy_embeddings( .pool_max_idle_per_host(10) .build::<_, HyperBody>(https); + let mut provider_request = embedding_request.clone(); + if proxy_config.provider_name == "tinfoil" { + provider_request.model = canonicalize_tinfoil_model(&embedding_request.model); + } + let effective_model = provider_request.model.clone(); + // Build request body - let request_body = serde_json::to_string(&embedding_request).map_err(|e| { + let request_body = serde_json::to_string(&provider_request).map_err(|e| { error!("Failed to serialize embedding request: {:?}", e); ApiError::InternalServerError })?; @@ -2111,32 +2074,147 @@ async fn proxy_embeddings( ApiError::InternalServerError })?; + Ok((response_json, proxy_config.provider_name, effective_model)) +} + +async fn request_embeddings_with_billing( + state: &Arc, + user: &User, + auth_method: AuthMethod, + embedding_request: &EmbeddingRequest, +) -> Result<(Value, String, i32), ApiError> { + // Check if guest user is allowed (paid guests are allowed, free guests are not) + if user.is_guest() { + if let Some(billing_client) = &state.billing_client { + match billing_client.is_user_paid(user.uuid).await { + Ok(true) => { + debug!("Paid guest user allowed for embeddings: {}", user.uuid); + } + Ok(false) => { + error!( + "Free guest user attempted to use embeddings feature: {}", + user.uuid + ); + return Err(ApiError::Unauthorized); + } + Err(e) => { + error!("Billing check failed for guest user {}: {}", user.uuid, e); + return Err(ApiError::Unauthorized); + } + } + } else { + error!( + "Guest user attempted to use embeddings without billing client: {}", + user.uuid + ); + return Err(ApiError::Unauthorized); + } + } + + let (response_json, provider_name, effective_model) = + request_embeddings_from_primary(state, embedding_request).await?; + + let prompt_tokens: i32 = response_json + .get("usage") + .and_then(|u| u.get("prompt_tokens")) + .and_then(|v| v.as_i64()) + .and_then(|n| i32::try_from(n).ok()) + .filter(|n| *n >= 0) + .unwrap_or(0); + // Handle billing - embeddings only have prompt_tokens (no completion_tokens) - if let Some(usage) = response_json.get("usage") { - let prompt_tokens = usage - .get("prompt_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0) as i32; - - if prompt_tokens > 0 { - let billing_context = - BillingContext::new(_auth_method, embedding_request.model.clone()); - let embedding_usage = CompletionUsage { - prompt_tokens, - completion_tokens: 0, // Embeddings don't have completion tokens - cached_prompt_tokens: None, - }; - publish_usage_event_internal( - &state, - &user, - &billing_context, - embedding_usage, - &proxy_config.provider_name, - ) - .await; + if prompt_tokens > 0 { + let billing_context = BillingContext::new(auth_method, effective_model); + let embedding_usage = CompletionUsage { + prompt_tokens, + completion_tokens: 0, + cached_prompt_tokens: None, + }; + publish_usage_event_internal( + state, + user, + &billing_context, + embedding_usage, + &provider_name, + ) + .await; + } + + Ok((response_json, provider_name, prompt_tokens)) +} + +pub async fn get_embedding_vector( + state: &Arc, + user: &User, + auth_method: AuthMethod, + model: &str, + input: &str, + dimensions: Option, +) -> Result<(Vec, i32), ApiError> { + let request = EmbeddingRequest { + input: Value::String(input.to_string()), + model: model.to_string(), + encoding_format: Some("float".to_string()), + dimensions, + user: None, + }; + + let (response_json, _provider, prompt_tokens) = + request_embeddings_with_billing(state, user, auth_method, &request).await?; + + let embedding = response_json + .get("data") + .and_then(|v| v.as_array()) + .and_then(|arr| arr.first()) + .and_then(|obj| obj.get("embedding")) + .and_then(|v| v.as_array()) + .ok_or(ApiError::InternalServerError)?; + + let mut vector: Vec = Vec::with_capacity(embedding.len()); + for v in embedding { + let f = v.as_f64().ok_or(ApiError::InternalServerError)?; + vector.push(f as f32); + } + + if let Some(dim) = dimensions { + if vector.len() != dim as usize { + error!( + "Embeddings response dim mismatch: expected={}, got={}", + dim, + vector.len() + ); + return Err(ApiError::InternalServerError); } } + Ok((vector, prompt_tokens)) +} + +async fn proxy_embeddings( + State(state): State>, + _headers: HeaderMap, + axum::Extension(session_id): axum::Extension, + axum::Extension(user): axum::Extension, + axum::Extension(auth_method): axum::Extension, + axum::Extension(embedding_request): axum::Extension, +) -> Result>, ApiError> { + debug!("Entering proxy_embeddings function"); + + // Validate input is not empty + let is_empty = match &embedding_request.input { + Value::String(s) => s.trim().is_empty(), + Value::Array(arr) => arr.is_empty(), + _ => true, + }; + if is_empty { + error!("Input is empty or invalid"); + return Err(ApiError::BadRequest); + } + + let (response_json, _provider_name, _prompt_tokens) = + request_embeddings_with_billing(&state, &user, auth_method, &embedding_request).await?; + + debug!("Exiting proxy_embeddings function"); // Encrypt and return the response encrypt_response(&state, &session_id, &response_json).await } diff --git a/src/web/rag.rs b/src/web/rag.rs new file mode 100644 index 00000000..657a3e92 --- /dev/null +++ b/src/web/rag.rs @@ -0,0 +1,216 @@ +use axum::{ + extract::{Path, State}, + middleware::from_fn_with_state, + routing::{delete, get, post}, + Extension, Json, Router, +}; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::sync::Arc; +use uuid::Uuid; + +use crate::jwt::AuthContext; +use crate::models::users::User; +use crate::rag; +use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; +use crate::web::openai_auth::AuthMethod; +use crate::web::responses::error_mapping; +use crate::{ApiError, AppMode, AppState}; + +#[derive(Debug, Clone, Deserialize)] +struct InsertEmbeddingRequest { + text: String, + metadata: Option, +} + +#[derive(Debug, Clone, Serialize)] +struct InsertEmbeddingResponse { + id: Uuid, + source_type: String, + embedding_model: String, + token_count: i32, + created_at: DateTime, +} + +#[derive(Debug, Clone, Deserialize)] +struct SearchRequest { + query: String, + top_k: Option, + max_tokens: Option, + source_types: Option>, + conversation_id: Option, +} + +#[derive(Debug, Clone, Serialize)] +struct SearchResponse { + results: Vec, +} + +pub fn router(app_state: Arc) -> Router<()> { + // Experimental endpoints: only enabled in Local/Dev. + if !matches!(app_state.app_mode, AppMode::Local | AppMode::Dev) { + return Router::new().with_state(app_state); + } + + Router::new() + .route( + "/v1/rag/embeddings", + post(insert_archival_embedding).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/rag/embeddings", + delete(delete_all_embeddings) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/rag/embeddings/:id", + delete(delete_embedding) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/rag/search", + post(search).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/rag/embeddings/status", + get(status).layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .with_state(app_state) +} + +async fn insert_archival_embedding( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result { + if body.text.trim().is_empty() { + return Err(ApiError::BadRequest); + } + if let Some(m) = &body.metadata { + if !m.is_object() { + return Err(ApiError::BadRequest); + } + } + + let user_key = state + .get_user_key(&user, &auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let inserted = rag::insert_archival_embedding( + &state, + &user, + AuthMethod::Jwt, + &user_key, + &body.text, + body.metadata.as_ref(), + ) + .await?; + + let response = InsertEmbeddingResponse { + id: inserted.uuid, + source_type: inserted.source_type, + embedding_model: inserted.embedding_model, + token_count: inserted.token_count, + created_at: inserted.created_at, + }; + + let encrypted = encrypt_response(&state, &session_id, &response).await?; + Ok((axum::http::StatusCode::CREATED, encrypted)) +} + +async fn search( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + if body.query.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + let top_k = body.top_k.unwrap_or(5); + if top_k == 0 || top_k > 20 { + return Err(ApiError::BadRequest); + } + + if let Some(max_tokens) = body.max_tokens { + if max_tokens <= 0 { + return Err(ApiError::BadRequest); + } + } + + if let Some(source_types) = &body.source_types { + if source_types.is_empty() { + return Err(ApiError::BadRequest); + } + } + + let conversation_internal_id = if let Some(conversation_uuid) = body.conversation_id { + let conversation = state + .db + .get_conversation_by_uuid_and_user(conversation_uuid, user.uuid) + .map_err(error_mapping::map_conversation_error)?; + Some(conversation.id) + } else { + None + }; + + let user_key = state + .get_user_key(&user, &auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let results = rag::search_user_embeddings( + &state, + &user, + AuthMethod::Jwt, + &user_key, + &body.query, + top_k, + body.max_tokens, + body.source_types.as_deref(), + conversation_internal_id, + None, + ) + .await?; + + let response = SearchResponse { results }; + encrypt_response(&state, &session_id, &response).await +} + +async fn delete_all_embeddings( + State(state): State>, + Extension(user): Extension, +) -> Result { + rag::delete_all_user_embeddings(&state, user.uuid).await?; + Ok(axum::http::StatusCode::NO_CONTENT) +} + +async fn delete_embedding( + State(state): State>, + Path(embedding_id): Path, + Extension(user): Extension, +) -> Result { + rag::delete_user_embedding_by_uuid(&state, user.uuid, embedding_id).await?; + Ok(axum::http::StatusCode::NO_CONTENT) +} + +async fn status( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + let response = rag::embeddings_status(&state, user.uuid).await?; + encrypt_response(&state, &session_id, &response).await +} From 712a3dbe68068469702391ce6a300e1158d1f526 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 7 Jul 2026 04:25:06 +0000 Subject: [PATCH 2/6] fix(rag): address foundation review feedback --- .../down.sql | 3 + .../up.sql | 9 + src/main.rs | 8 +- src/rag.rs | 1495 +++++++++++++---- src/web/mod.rs | 2 - src/web/rag.rs | 216 --- 6 files changed, 1208 insertions(+), 525 deletions(-) create mode 100644 migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql create mode 100644 migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql delete mode 100644 src/web/rag.rs diff --git a/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql new file mode 100644 index 00000000..c0057325 --- /dev/null +++ b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql @@ -0,0 +1,3 @@ +DROP INDEX IF EXISTS idx_user_embeddings_user_model_conversation_created_id; +DROP INDEX IF EXISTS idx_user_embeddings_user_model_source_created_id; +DROP INDEX IF EXISTS idx_user_embeddings_user_model_created_id; diff --git a/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql new file mode 100644 index 00000000..b857135b --- /dev/null +++ b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql @@ -0,0 +1,9 @@ +-- Hot RAG retrieval paths: active embedding model for a user, newest first. +CREATE INDEX idx_user_embeddings_user_model_created_id + ON user_embeddings(user_id, embedding_model, created_at DESC, id DESC); + +CREATE INDEX idx_user_embeddings_user_model_source_created_id + ON user_embeddings(user_id, embedding_model, source_type, created_at DESC, id DESC); + +CREATE INDEX idx_user_embeddings_user_model_conversation_created_id + ON user_embeddings(user_id, embedding_model, conversation_id, created_at DESC, id DESC); diff --git a/src/main.rs b/src/main.rs index c58fe4aa..baaa0206 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,7 +21,7 @@ use crate::web::openai_auth::validate_openai_auth; use crate::web::platform_login_routes; use crate::web::{ conversation_projects_routes, conversations_routes, health_routes_with_state, - instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, rag_routes, + instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, responses_routes, }; use crate::{attestation_routes::SessionState, web::platform_routes}; @@ -1709,6 +1709,7 @@ impl AppState { encrypted_password, new_wrapping, )?; + self.rag_cache.lock().await.evict_user(user.uuid); // Send confirmation email in the background let app_state = self.clone(); @@ -2099,6 +2100,7 @@ impl AppState { // Secret verification succeeded, proceed with account deletion // Use the transaction-based method for atomicity self.db.mark_and_delete_user(&user, &deletion_request)?; + self.rag_cache.lock().await.evict_user(user.uuid); // Send confirmation email in the background if the user has an email if let Some(email) = user.email.clone() { @@ -3214,10 +3216,6 @@ async fn main() -> Result<(), Error> { instructions_routes(app_state.clone()) .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), ) - .merge( - rag_routes(app_state.clone()) - .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), - ) .merge(attestation_routes::router(app_state.clone())) .merge(oauth_routes(app_state.clone())) .merge(platform_login_routes(app_state.clone())) diff --git a/src/rag.rs b/src/rag.rs index 71922fa2..dc5a2b65 100644 --- a/src/rag.rs +++ b/src/rag.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] + use crate::encrypt::{decrypt_with_key, encrypt_key_deterministic, encrypt_with_key}; use crate::models::schema::user_embeddings; use crate::models::user_embeddings::NewUserEmbedding; @@ -5,23 +7,41 @@ use crate::models::users::User; use crate::web::openai_auth::AuthMethod; use crate::{ApiError, AppState}; use base64::{engine::general_purpose::STANDARD as B64_STANDARD, Engine as _}; +use chrono::{DateTime, Utc}; use diesel::prelude::*; +use diesel::QueryableByName; use secp256k1::SecretKey; use serde::Serialize; use std::cmp::Ordering; use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque}; use std::sync::Arc; use std::time::{Duration, Instant}; -use tracing::{debug, error}; +use tokio::sync::Notify; +use tokio::time::{sleep, timeout}; +use tracing::{debug, error, info, warn}; use uuid::Uuid; pub const DEFAULT_EMBEDDING_MODEL: &str = "nomic-embed-text"; pub const DEFAULT_EMBEDDING_DIM: i32 = 768; -const CACHE_MAX_USERS: usize = 100; +const DEFAULT_CACHE_MAX_BYTES: usize = 4 * 1024 * 1024 * 1024; +const DEFAULT_CACHE_MAX_USER_BYTES: usize = 128 * 1024 * 1024; const CACHE_TTL: Duration = Duration::from_secs(5 * 60); -const DB_SCAN_BATCH_SIZE: i64 = 1000; +const DEFAULT_SCAN_LIMIT: i64 = 25_000; +const DEFAULT_MAX_INSERT_TEXT_BYTES: usize = 16 * 1024; +const DEFAULT_MAX_SEARCH_QUERY_BYTES: usize = 4 * 1024; +const DEFAULT_MAX_USER_EMBEDDINGS: i64 = 250_000; +const DEFAULT_MAX_USER_STORED_BYTES: i64 = 2 * 1024 * 1024 * 1024; +const DEFAULT_MAX_INSERTS_PER_USER_PER_HOUR: i64 = 10_000; +const DEFAULT_MAX_PROJECT_EMBEDDINGS: i64 = 2_500_000; +const DEFAULT_MAX_PROJECT_STORED_BYTES: i64 = 20 * 1024 * 1024 * 1024; + +const CACHE_LOAD_WAIT_TIMEOUT: Duration = Duration::from_secs(60); +const CACHE_LOAD_TIMEOUT_BACKOFF: Duration = Duration::from_millis(250); +const FINAL_RESULT_OVERFETCH_MULTIPLIER: usize = 3; +const FINAL_RESULT_OVERFETCH_MIN_EXTRA: usize = 10; +const FINAL_RESULT_OVERFETCH_MAX: usize = 100; #[allow(dead_code)] pub const SOURCE_TYPE_MESSAGE: &str = "message"; @@ -34,6 +54,41 @@ pub struct RagSearchResult { pub token_count: i32, } +#[derive(Debug, Clone, Serialize)] +pub struct RagSearchOutcome { + pub results: Vec, + pub feedback: Option, + pub scan_limit_hit: bool, + pub scanned_rows: usize, + pub skipped_rows: usize, +} + +#[derive(Debug, Clone, Default)] +pub struct RagSearchFilters { + pub source_types: Option>, + pub conversation_id: Option, + pub tags: Option>, + pub begin_date: Option>, + pub end_date: Option>, +} + +#[derive(Debug, Clone)] +pub struct RagSearchOptions { + pub limit: usize, + pub max_tokens: Option, + pub filters: RagSearchFilters, +} + +impl Default for RagSearchOptions { + fn default() -> Self { + Self { + limit: 5, + max_tokens: None, + filters: RagSearchFilters::default(), + } + } +} + #[derive(Debug, Clone, Serialize)] pub struct RagEmbeddingsStatus { pub total_embeddings: i64, @@ -43,81 +98,266 @@ pub struct RagEmbeddingsStatus { #[derive(Debug, Clone)] pub struct CachedEmbedding { + pub id: i64, + pub uuid: Uuid, pub source_type: String, pub conversation_id: Option, pub vector: Vec, - pub content_enc: Vec, pub token_count: i32, + pub created_at: DateTime, + pub updated_at: DateTime, } #[derive(Debug, Clone)] struct CacheEntry { loaded_at: Instant, embeddings: Arc>, + bytes: usize, + scan_limit_hit: bool, +} + +#[derive(Debug)] +struct InFlightLoad { + notify: Arc, + timeout_duplicate_started: bool, +} + +#[derive(Debug, Clone)] +enum CacheLoadPermit { + Start, + Wait(Arc), +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CacheAppendResult { + Appended, + Missing, + EvictedOverLimit, } #[derive(Debug)] pub struct RagCache { - max_users: usize, + max_bytes: usize, + max_user_bytes: usize, ttl: Duration, + total_bytes: usize, entries: HashMap, lru: VecDeque, + in_flight_loads: HashMap, } impl Default for RagCache { fn default() -> Self { - Self::new(CACHE_MAX_USERS, CACHE_TTL) + Self::from_env() } } impl RagCache { - pub fn new(max_users: usize, ttl: Duration) -> Self { + pub fn from_env() -> Self { + Self::new( + read_env_usize("RAG_CACHE_MAX_BYTES", DEFAULT_CACHE_MAX_BYTES), + read_env_usize("RAG_CACHE_MAX_USER_BYTES", DEFAULT_CACHE_MAX_USER_BYTES), + CACHE_TTL, + ) + } + + pub fn new(max_bytes: usize, max_user_bytes: usize, ttl: Duration) -> Self { Self { - max_users, + max_bytes, + max_user_bytes, ttl, + total_bytes: 0, entries: HashMap::new(), lru: VecDeque::new(), + in_flight_loads: HashMap::new(), } } pub fn evict_user(&mut self, user_id: Uuid) { - self.entries.remove(&user_id); + self.evict_user_with_reason(user_id, "manual"); + } + + fn evict_user_with_reason(&mut self, user_id: Uuid, reason: &'static str) { + if let Some(entry) = self.entries.remove(&user_id) { + self.total_bytes = self.total_bytes.saturating_sub(entry.bytes); + info!( + target: "rag", + user_id = %user_id, + bytes = entry.bytes, + total_cache_bytes = self.total_bytes, + reason, + "rag_cache_evict" + ); + } self.lru.retain(|u| *u != user_id); } - pub fn get(&mut self, user_id: Uuid) -> Option>> { + pub fn get(&mut self, user_id: Uuid) -> Option<(Arc>, bool)> { self.evict_expired(); - let (loaded_at, embeddings) = { + let (loaded_at, embeddings, scan_limit_hit) = { let entry = self.entries.get(&user_id)?; - (entry.loaded_at, entry.embeddings.clone()) + ( + entry.loaded_at, + entry.embeddings.clone(), + entry.scan_limit_hit, + ) }; if loaded_at.elapsed() > self.ttl { - self.evict_user(user_id); + self.evict_user_with_reason(user_id, "ttl"); return None; } self.touch(user_id); - Some(embeddings) + Some((embeddings, scan_limit_hit)) } - pub fn put(&mut self, user_id: Uuid, embeddings: Arc>) { + pub fn put( + &mut self, + user_id: Uuid, + embeddings: Arc>, + scan_limit_hit: bool, + ) -> bool { + let bytes = cached_embeddings_bytes(&embeddings); + if bytes > self.max_user_bytes || bytes > self.max_bytes { + info!( + target: "rag", + user_id = %user_id, + bytes, + max_user_bytes = self.max_user_bytes, + max_cache_bytes = self.max_bytes, + reason = "entry_over_limit", + "rag_cache_skip" + ); + self.evict_user_with_reason(user_id, "entry_over_limit"); + return false; + } + + if self.entries.contains_key(&user_id) { + self.evict_user_with_reason(user_id, "replace"); + } + + while self.total_bytes.saturating_add(bytes) > self.max_bytes { + if let Some(lru_user) = self.lru.pop_back() { + self.evict_user_with_reason(lru_user, "global_byte_limit"); + } else { + break; + } + } + + if self.total_bytes.saturating_add(bytes) > self.max_bytes { + info!( + target: "rag", + user_id = %user_id, + bytes, + total_cache_bytes = self.total_bytes, + max_cache_bytes = self.max_bytes, + reason = "global_byte_limit", + "rag_cache_skip" + ); + return false; + } + self.entries.insert( user_id, CacheEntry { loaded_at: Instant::now(), embeddings, + bytes, + scan_limit_hit, }, ); + self.total_bytes = self.total_bytes.saturating_add(bytes); self.touch(user_id); + true + } - while self.entries.len() > self.max_users { - if let Some(lru_user) = self.lru.pop_back() { - self.entries.remove(&lru_user); - } else { - break; - } + pub fn append(&mut self, user_id: Uuid, embedding: CachedEmbedding) -> CacheAppendResult { + self.evict_expired(); + let Some(entry) = self.entries.get(&user_id) else { + return CacheAppendResult::Missing; + }; + + if entry.loaded_at.elapsed() > self.ttl { + self.evict_user_with_reason(user_id, "ttl"); + return CacheAppendResult::Missing; + } + + let row_bytes = embedding.estimated_cache_bytes(); + let new_entry_bytes = entry.bytes.saturating_add(row_bytes); + let new_total_bytes = self.total_bytes.saturating_add(row_bytes); + if new_entry_bytes > self.max_user_bytes || new_total_bytes > self.max_bytes { + self.evict_user_with_reason(user_id, "append_over_limit"); + return CacheAppendResult::EvictedOverLimit; + } + + if let Some(entry) = self.entries.get_mut(&user_id) { + Arc::make_mut(&mut entry.embeddings).push(embedding); + entry.bytes = new_entry_bytes; + self.total_bytes = new_total_bytes; + } + self.touch(user_id); + CacheAppendResult::Appended + } + + pub fn remove_embedding_by_uuid(&mut self, user_id: Uuid, embedding_uuid: Uuid) { + self.evict_expired(); + let Some(entry) = self.entries.get(&user_id) else { + return; + }; + + let before_len = entry.embeddings.len(); + let old_bytes = entry.bytes; + let Some(entry) = self.entries.get_mut(&user_id) else { + return; + }; + Arc::make_mut(&mut entry.embeddings).retain(|e| e.uuid != embedding_uuid); + let after_len = entry.embeddings.len(); + let new_bytes = cached_embeddings_bytes(&entry.embeddings); + entry.bytes = new_bytes; + + if after_len == before_len { + return; + } + + self.total_bytes = self + .total_bytes + .saturating_sub(old_bytes.saturating_sub(new_bytes)); + self.touch(user_id); + } + + fn begin_load(&mut self, user_id: Uuid) -> CacheLoadPermit { + if let Some(load) = self.in_flight_loads.get(&user_id) { + return CacheLoadPermit::Wait(load.notify.clone()); + } + + self.in_flight_loads.insert( + user_id, + InFlightLoad { + notify: Arc::new(Notify::new()), + timeout_duplicate_started: false, + }, + ); + CacheLoadPermit::Start + } + + fn try_start_timeout_duplicate(&mut self, user_id: Uuid) -> bool { + let Some(load) = self.in_flight_loads.get_mut(&user_id) else { + return true; + }; + + if load.timeout_duplicate_started { + return false; + } + + load.timeout_duplicate_started = true; + true + } + + fn finish_load(&mut self, user_id: Uuid) { + if let Some(load) = self.in_flight_loads.remove(&user_id) { + load.notify.notify_waiters(); } } @@ -141,11 +381,126 @@ impl RagCache { .collect(); for user_id in expired { - self.evict_user(user_id); + self.evict_user_with_reason(user_id, "ttl"); } } } +fn read_env_usize(name: &str, default: usize) -> usize { + match std::env::var(name) { + Ok(value) => match value.parse::() { + Ok(parsed) if parsed > 0 => parsed, + _ => { + warn!( + target: "rag", + name, + value, + default, + "Invalid RAG usize env var; using default" + ); + default + } + }, + Err(_) => default, + } +} + +fn read_env_i64(name: &str, default: i64) -> i64 { + match std::env::var(name) { + Ok(value) => match value.parse::() { + Ok(parsed) if parsed > 0 => parsed, + _ => { + warn!( + target: "rag", + name, + value, + default, + "Invalid RAG i64 env var; using default" + ); + default + } + }, + Err(_) => default, + } +} + +#[derive(Debug, Clone, Copy)] +struct RagLimits { + scan_limit: i64, + max_insert_text_bytes: usize, + max_search_query_bytes: usize, + max_user_embeddings: i64, + max_user_stored_bytes: i64, + max_inserts_per_user_per_hour: i64, + max_project_embeddings: i64, + max_project_stored_bytes: i64, +} + +impl RagLimits { + fn from_env() -> Self { + Self { + scan_limit: read_env_i64("RAG_SCAN_LIMIT", DEFAULT_SCAN_LIMIT), + max_insert_text_bytes: read_env_usize( + "RAG_MAX_INSERT_TEXT_BYTES", + DEFAULT_MAX_INSERT_TEXT_BYTES, + ), + max_search_query_bytes: read_env_usize( + "RAG_MAX_SEARCH_QUERY_BYTES", + DEFAULT_MAX_SEARCH_QUERY_BYTES, + ), + max_user_embeddings: read_env_i64( + "RAG_MAX_USER_EMBEDDINGS", + DEFAULT_MAX_USER_EMBEDDINGS, + ), + max_user_stored_bytes: read_env_i64( + "RAG_MAX_USER_STORED_BYTES", + DEFAULT_MAX_USER_STORED_BYTES, + ), + max_inserts_per_user_per_hour: read_env_i64( + "RAG_MAX_INSERTS_PER_USER_PER_HOUR", + DEFAULT_MAX_INSERTS_PER_USER_PER_HOUR, + ), + max_project_embeddings: read_env_i64( + "RAG_MAX_PROJECT_EMBEDDINGS", + DEFAULT_MAX_PROJECT_EMBEDDINGS, + ), + max_project_stored_bytes: read_env_i64( + "RAG_MAX_PROJECT_STORED_BYTES", + DEFAULT_MAX_PROJECT_STORED_BYTES, + ), + } + } +} + +impl CachedEmbedding { + fn estimated_cache_bytes(&self) -> usize { + std::mem::size_of::() + .saturating_add(self.source_type.len()) + .saturating_add(self.vector.len().saturating_mul(std::mem::size_of::())) + } +} + +fn cached_embeddings_bytes(embeddings: &[CachedEmbedding]) -> usize { + embeddings + .iter() + .map(CachedEmbedding::estimated_cache_bytes) + .sum() +} + +fn validate_text_size(text: &str, max_bytes: usize, label: &'static str) -> Result<(), ApiError> { + if text.len() > max_bytes { + warn!( + target: "rag", + label, + bytes = text.len(), + max_bytes, + "RAG text exceeds configured byte limit" + ); + return Err(ApiError::BadRequest); + } + Ok(()) +} + pub fn serialize_f32_le(values: &[f32]) -> Vec { let mut out = Vec::with_capacity(values.len() * 4); for v in values { @@ -191,16 +546,19 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> Result { #[derive(Debug, Clone)] struct HeapItem { + id: i64, + uuid: Uuid, score: f32, token_count: i32, - content_enc: Vec, } impl Eq for HeapItem {} impl PartialEq for HeapItem { fn eq(&self, other: &Self) -> bool { - self.score.to_bits() == other.score.to_bits() && self.token_count == other.token_count + self.score.to_bits() == other.score.to_bits() + && self.token_count == other.token_count + && self.id == other.id } } @@ -209,6 +567,7 @@ impl Ord for HeapItem { self.score .total_cmp(&other.score) .then_with(|| other.token_count.cmp(&self.token_count)) + .then_with(|| other.id.cmp(&self.id)) } } @@ -253,9 +612,10 @@ fn top_k_candidates( if heap.len() < top_k { let score = cosine_similarity(query, &e.vector)?; let item = HeapItem { + id: e.id, + uuid: e.uuid, score, token_count: e.token_count, - content_enc: e.content_enc.clone(), }; heap.push(std::cmp::Reverse(item)); continue; @@ -269,9 +629,10 @@ fn top_k_candidates( if ordering == Ordering::Greater { heap.pop(); let item = HeapItem { + id: e.id, + uuid: e.uuid, score, token_count: e.token_count, - content_enc: e.content_enc.clone(), }; heap.push(std::cmp::Reverse(item)); } @@ -356,6 +717,176 @@ fn extract_tags_from_metadata(metadata: Option<&serde_json::Value>) -> Vec>, +) -> i64 { + vector_enc + .len() + .saturating_add(content_enc.len()) + .saturating_add(metadata_enc.map_or(0, Vec::len)) as i64 +} + +fn load_user_embedding_storage_stats( + conn: &mut diesel::PgConnection, + user_id: Uuid, +) -> Result { + diesel::sql_query( + r#" + SELECT + COUNT(*)::BIGINT AS row_count, + COALESCE( + SUM( + octet_length(vector_enc) + + octet_length(content_enc) + + COALESCE(octet_length(metadata_enc), 0) + ), + 0 + )::BIGINT AS stored_bytes + FROM user_embeddings + WHERE user_id = $1 + "#, + ) + .bind::(user_id) + .get_result(conn) + .map_err(|e| { + error!( + "Failed to load user embedding storage stats for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + }) +} + +fn load_project_embedding_storage_stats( + conn: &mut diesel::PgConnection, + project_id: i32, +) -> Result { + diesel::sql_query( + r#" + SELECT + COUNT(e.*)::BIGINT AS row_count, + COALESCE( + SUM( + octet_length(e.vector_enc) + + octet_length(e.content_enc) + + COALESCE(octet_length(e.metadata_enc), 0) + ), + 0 + )::BIGINT AS stored_bytes + FROM user_embeddings e + INNER JOIN users u ON u.uuid = e.user_id + WHERE u.project_id = $1 + "#, + ) + .bind::(project_id) + .get_result(conn) + .map_err(|e| { + error!( + "Failed to load project embedding storage stats for project={}: {:?}", + project_id, e + ); + ApiError::InternalServerError + }) +} + +fn ensure_storage_limits( + conn: &mut diesel::PgConnection, + user: &User, + projected_new_bytes: i64, + limits: RagLimits, +) -> Result<(), ApiError> { + use diesel::dsl::count_star; + + let recent_insert_count: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(user.uuid)) + .filter(user_embeddings::created_at.gt(Utc::now() - chrono::Duration::hours(1))) + .select(count_star()) + .first(conn) + .map_err(|e| { + error!( + "Failed to load recent embedding insert count for user={}: {:?}", + user.uuid, e + ); + ApiError::InternalServerError + })?; + if recent_insert_count >= limits.max_inserts_per_user_per_hour { + warn!( + target: "rag", + user_id = %user.uuid, + recent_insert_count, + max_inserts_per_user_per_hour = limits.max_inserts_per_user_per_hour, + "RAG per-user insert rate limit reached" + ); + return Err(ApiError::BadRequest); + } + + let user_stats = load_user_embedding_storage_stats(conn, user.uuid)?; + if user_stats.row_count.saturating_add(1) > limits.max_user_embeddings + || user_stats.stored_bytes.saturating_add(projected_new_bytes) + > limits.max_user_stored_bytes + { + warn!( + target: "rag", + user_id = %user.uuid, + rows = user_stats.row_count, + stored_bytes = user_stats.stored_bytes, + projected_new_bytes, + max_rows = limits.max_user_embeddings, + max_bytes = limits.max_user_stored_bytes, + "RAG user storage limit reached" + ); + return Err(ApiError::BadRequest); + } + + let project_stats = load_project_embedding_storage_stats(conn, user.project_id)?; + if project_stats.row_count.saturating_add(1) > limits.max_project_embeddings + || project_stats + .stored_bytes + .saturating_add(projected_new_bytes) + > limits.max_project_stored_bytes + { + warn!( + target: "rag", + project_id = user.project_id, + rows = project_stats.row_count, + stored_bytes = project_stats.stored_bytes, + projected_new_bytes, + max_rows = limits.max_project_embeddings, + max_bytes = limits.max_project_stored_bytes, + "RAG project storage limit reached" + ); + return Err(ApiError::BadRequest); + } + + Ok(()) +} + +fn cached_embedding_from_inserted( + inserted: &crate::models::user_embeddings::UserEmbedding, + vector: Vec, +) -> CachedEmbedding { + CachedEmbedding { + id: inserted.id, + uuid: inserted.uuid, + source_type: inserted.source_type.clone(), + conversation_id: inserted.conversation_id, + vector, + token_count: inserted.token_count.max(0), + created_at: inserted.created_at, + updated_at: inserted.updated_at, + } +} + pub async fn insert_archival_embedding( state: &Arc, user: &User, @@ -368,6 +899,8 @@ pub async fn insert_archival_embedding( if text.is_empty() { return Err(ApiError::BadRequest); } + let limits = RagLimits::from_env(); + validate_text_size(text, limits.max_insert_text_bytes, "insert")?; let user_id = user.uuid; let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; @@ -385,12 +918,15 @@ pub async fn insert_archival_embedding( let tags = extract_tags_from_metadata(metadata); let tags_enc = encrypt_tags_b64(user_key, &tags); + let projected_new_bytes = + encrypted_embedding_bytes(&vector_enc, &content_enc, metadata_enc.as_ref()); let mut conn = state .db .get_pool() .get() .map_err(|_| ApiError::InternalServerError)?; + ensure_storage_limits(&mut conn, user, projected_new_bytes, limits)?; let inserted = NewUserEmbedding { uuid: Uuid::new_v4(), @@ -413,7 +949,15 @@ pub async fn insert_archival_embedding( ApiError::InternalServerError })?; - state.rag_cache.lock().await.evict_user(user_id); + let cached = cached_embedding_from_inserted(&inserted, vector); + let append_result = state.rag_cache.lock().await.append(user_id, cached); + debug!( + target: "rag", + user_id = %user_id, + embedding_uuid = %inserted.uuid, + ?append_result, + "rag_cache_append_after_insert" + ); Ok(inserted) } @@ -434,18 +978,22 @@ pub async fn insert_message_embedding( if text.is_empty() { return Err(ApiError::BadRequest); } + let limits = RagLimits::from_env(); + validate_text_size(text, limits.max_insert_text_bytes, "insert")?; let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; let vector_bytes = serialize_f32_le(&vector); let vector_enc = encrypt_with_key(user_key, &vector_bytes).await; let content_enc = encrypt_with_key(user_key, text.as_bytes()).await; + let projected_new_bytes = encrypted_embedding_bytes(&vector_enc, &content_enc, None); let mut conn = state .db .get_pool() .get() .map_err(|_| ApiError::InternalServerError)?; + ensure_storage_limits(&mut conn, user, projected_new_bytes, limits)?; let inserted = NewUserEmbedding { uuid: Uuid::new_v4(), @@ -468,290 +1016,603 @@ pub async fn insert_message_embedding( ApiError::InternalServerError })?; - state.rag_cache.lock().await.evict_user(user_id); + let cached = cached_embedding_from_inserted(&inserted, vector); + let append_result = state.rag_cache.lock().await.append(user_id, cached); + debug!( + target: "rag", + user_id = %user_id, + embedding_uuid = %inserted.uuid, + ?append_result, + "rag_cache_append_after_insert" + ); Ok(inserted) } -async fn load_all_user_embeddings( +#[derive(Debug)] +struct LoadedEmbeddings { + embeddings: Arc>, + scan_limit_hit: bool, + scanned_rows: usize, + skipped_rows: usize, + db_read_bytes: usize, +} + +#[derive(Debug, Clone, Copy)] +struct LoadFilters<'a> { + source_types: Option<&'a [String]>, + conversation_id: Option, + tags_enc_filter: Option<&'a [Option]>, + begin_date: Option>, + end_date: Option>, +} + +impl LoadFilters<'_> { + fn is_cacheable_broad_load(&self) -> bool { + self.source_types.is_none() + && self.conversation_id.is_none() + && self.tags_enc_filter.is_none() + && self.begin_date.is_none() + && self.end_date.is_none() + } +} + +#[derive(Queryable)] +struct EmbeddingScanRow { + id: i64, + uuid: Uuid, + source_type: String, + conversation_id: Option, + vector_enc: Vec, + token_count: i32, + vector_dim: i32, + created_at: DateTime, + updated_at: DateTime, +} + +async fn load_user_embeddings_for_search( state: &AppState, user_id: Uuid, user_key: &SecretKey, -) -> Result>, ApiError> { + filters: LoadFilters<'_>, + scan_limit: i64, +) -> Result { let mut conn = state .db .get_pool() .get() .map_err(|_| ApiError::InternalServerError)?; - let mut last_id: i64 = 0; - let mut out: Vec = Vec::new(); + let started_at = Instant::now(); + let mut query = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::embedding_model.eq(DEFAULT_EMBEDDING_MODEL)) + .into_boxed(); - #[derive(Queryable)] - struct EmbeddingScanRow { - source_type: String, - conversation_id: Option, - vector_enc: Vec, - content_enc: Vec, - token_count: i32, - vector_dim: i32, - id: i64, + if let Some(source_types) = filters.source_types { + query = query.filter(user_embeddings::source_type.eq_any(source_types)); } - loop { - let rows: Vec = user_embeddings::table - .filter(user_embeddings::user_id.eq(user_id)) - .filter(user_embeddings::embedding_model.eq(DEFAULT_EMBEDDING_MODEL)) - .filter(user_embeddings::id.gt(last_id)) - .order(user_embeddings::id.asc()) - .select(( - user_embeddings::source_type, - user_embeddings::conversation_id, - user_embeddings::vector_enc, - user_embeddings::content_enc, - user_embeddings::token_count, - user_embeddings::vector_dim, - user_embeddings::id, - )) - .limit(DB_SCAN_BATCH_SIZE) - .load(&mut conn) - .map_err(|e| { - error!( - "Failed to load embeddings batch for user={} after id={}: {:?}", - user_id, last_id, e - ); - ApiError::InternalServerError - })?; + if let Some(conversation_id) = filters.conversation_id { + query = query.filter(user_embeddings::conversation_id.eq(Some(conversation_id))); + } - if rows.is_empty() { - break; - } + if let Some(tags_enc_filter) = filters.tags_enc_filter { + query = query.filter(user_embeddings::tags_enc.overlaps_with(tags_enc_filter.to_vec())); + } - for row in rows { - let EmbeddingScanRow { - source_type, - conversation_id, - vector_enc, - content_enc, - token_count, - vector_dim, - id, - } = row; - last_id = id; - - let vector_bytes = decrypt_with_key(user_key, &vector_enc) - .map_err(|_| ApiError::InternalServerError)?; - let vector = deserialize_f32_le(&vector_bytes)?; - - if vector.len() != vector_dim as usize { - debug!( - "Skipping embedding id={} for user={} due to dim mismatch (expected={}, got={})", - id, - user_id, - vector_dim, - vector.len() + if let Some(begin_date) = filters.begin_date { + query = query.filter(user_embeddings::created_at.ge(begin_date)); + } + + if let Some(end_date) = filters.end_date { + query = query.filter(user_embeddings::created_at.le(end_date)); + } + + let mut rows: Vec = query + .order(( + user_embeddings::created_at.desc(), + user_embeddings::id.desc(), + )) + .select(( + user_embeddings::id, + user_embeddings::uuid, + user_embeddings::source_type, + user_embeddings::conversation_id, + user_embeddings::vector_enc, + user_embeddings::token_count, + user_embeddings::vector_dim, + user_embeddings::created_at, + user_embeddings::updated_at, + )) + .limit(scan_limit.saturating_add(1)) + .load(&mut conn) + .map_err(|e| { + error!("Failed to load embeddings for user={}: {:?}", user_id, e); + ApiError::InternalServerError + })?; + + let scan_limit_hit = rows.len() as i64 > scan_limit; + if scan_limit_hit { + rows.truncate(scan_limit as usize); + } + + let db_read_bytes = rows.iter().map(|row| row.vector_enc.len()).sum(); + let scanned_rows = rows.len(); + let mut skipped_rows = 0usize; + let mut out: Vec = Vec::with_capacity(rows.len()); + let decrypt_started_at = Instant::now(); + + for row in rows { + let vector_bytes = match decrypt_with_key(user_key, &row.vector_enc) { + Ok(bytes) => bytes, + Err(_) => { + skipped_rows = skipped_rows.saturating_add(1); + warn!( + target: "rag", + user_id = %user_id, + embedding_id = row.id, + embedding_uuid = %row.uuid, + reason = "vector_decrypt", + "Skipping corrupt RAG row" ); continue; } + }; - out.push(CachedEmbedding { - source_type, - conversation_id, - vector, - content_enc, - token_count, - }); + let vector = match deserialize_f32_le(&vector_bytes) { + Ok(vector) => vector, + Err(_) => { + skipped_rows = skipped_rows.saturating_add(1); + warn!( + target: "rag", + user_id = %user_id, + embedding_id = row.id, + embedding_uuid = %row.uuid, + reason = "vector_deserialize", + "Skipping corrupt RAG row" + ); + continue; + } + }; + + if vector.len() != row.vector_dim as usize { + skipped_rows = skipped_rows.saturating_add(1); + warn!( + target: "rag", + user_id = %user_id, + embedding_id = row.id, + embedding_uuid = %row.uuid, + expected_dim = row.vector_dim, + actual_dim = vector.len(), + reason = "dimension_mismatch", + "Skipping corrupt RAG row" + ); + continue; } + + out.push(CachedEmbedding { + id: row.id, + uuid: row.uuid, + source_type: row.source_type, + conversation_id: row.conversation_id, + vector, + token_count: row.token_count.max(0), + created_at: row.created_at, + updated_at: row.updated_at, + }); } - Ok(Arc::new(out)) + info!( + target: "rag", + user_id = %user_id, + cacheable = filters.is_cacheable_broad_load(), + scanned_rows, + loaded_rows = out.len(), + skipped_rows, + db_read_bytes, + scan_limit, + scan_limit_hit, + db_and_decrypt_ms = started_at.elapsed().as_millis() as u64, + decrypt_ms = decrypt_started_at.elapsed().as_millis() as u64, + "rag_load_embeddings" + ); + + Ok(LoadedEmbeddings { + embeddings: Arc::new(out), + scan_limit_hit, + scanned_rows, + skipped_rows, + db_read_bytes, + }) +} + +#[derive(Queryable)] +struct ContentFetchRow { + id: i64, + uuid: Uuid, + content_enc: Vec, +} + +fn overfetch_limit(limit: usize) -> usize { + limit + .saturating_mul(FINAL_RESULT_OVERFETCH_MULTIPLIER) + .saturating_add(FINAL_RESULT_OVERFETCH_MIN_EXTRA) + .clamp(limit, FINAL_RESULT_OVERFETCH_MAX) +} + +async fn load_cacheable_embeddings_with_coordination( + state: &Arc, + user_id: Uuid, + user_key: &SecretKey, + limits: RagLimits, +) -> Result { + loop { + let permit = { + let mut cache = state.rag_cache.lock().await; + if let Some((embeddings, scan_limit_hit)) = cache.get(user_id) { + info!( + target: "rag", + user_id = %user_id, + cache_hit = true, + loaded_rows = embeddings.len(), + cache_bytes_per_user = cached_embeddings_bytes(&embeddings), + total_cache_bytes = cache.total_bytes, + scan_limit_hit, + "rag_cache_lookup" + ); + return Ok(LoadedEmbeddings { + embeddings, + scan_limit_hit, + scanned_rows: 0, + skipped_rows: 0, + db_read_bytes: 0, + }); + } + + info!( + target: "rag", + user_id = %user_id, + cache_hit = false, + total_cache_bytes = cache.total_bytes, + "rag_cache_lookup" + ); + cache.begin_load(user_id) + }; + + match permit { + CacheLoadPermit::Start => { + let filters = LoadFilters { + source_types: None, + conversation_id: None, + tags_enc_filter: None, + begin_date: None, + end_date: None, + }; + let loaded = load_user_embeddings_for_search( + state, + user_id, + user_key, + filters, + limits.scan_limit, + ) + .await; + + let mut cache = state.rag_cache.lock().await; + match &loaded { + Ok(loaded) => { + let cached = + cache.put(user_id, loaded.embeddings.clone(), loaded.scan_limit_hit); + info!( + target: "rag", + user_id = %user_id, + cached, + loaded_rows = loaded.embeddings.len(), + cache_bytes_per_user = cached_embeddings_bytes(&loaded.embeddings), + total_cache_bytes = cache.total_bytes, + scan_limit_hit = loaded.scan_limit_hit, + "rag_cache_store_after_load" + ); + } + Err(_) => { + warn!( + target: "rag", + user_id = %user_id, + "RAG cacheable load failed" + ); + } + } + cache.finish_load(user_id); + return loaded; + } + CacheLoadPermit::Wait(notify) => { + if timeout(CACHE_LOAD_WAIT_TIMEOUT, notify.notified()) + .await + .is_ok() + { + continue; + } + + let should_duplicate = { + let mut cache = state.rag_cache.lock().await; + cache.try_start_timeout_duplicate(user_id) + }; + + warn!( + target: "rag", + user_id = %user_id, + should_duplicate, + wait_timeout_secs = CACHE_LOAD_WAIT_TIMEOUT.as_secs(), + "RAG cache load wait timed out" + ); + + if !should_duplicate { + sleep(CACHE_LOAD_TIMEOUT_BACKOFF).await; + continue; + } + + let filters = LoadFilters { + source_types: None, + conversation_id: None, + tags_enc_filter: None, + begin_date: None, + end_date: None, + }; + let loaded = load_user_embeddings_for_search( + state, + user_id, + user_key, + filters, + limits.scan_limit, + ) + .await?; + + let mut cache = state.rag_cache.lock().await; + let cached = cache.put(user_id, loaded.embeddings.clone(), loaded.scan_limit_hit); + cache.finish_load(user_id); + info!( + target: "rag", + user_id = %user_id, + cached, + loaded_rows = loaded.embeddings.len(), + cache_bytes_per_user = cached_embeddings_bytes(&loaded.embeddings), + total_cache_bytes = cache.total_bytes, + scan_limit_hit = loaded.scan_limit_hit, + reason = "timeout_duplicate", + "rag_cache_store_after_load" + ); + return Ok(loaded); + } + } + } } -async fn load_user_embeddings_by_tags( +async fn fetch_ranked_content( state: &AppState, user_id: Uuid, user_key: &SecretKey, - source_types: Option<&[String]>, - conversation_id: Option, - tags_enc_filter: &[Option], -) -> Result>, ApiError> { + candidates: Vec, + limit: usize, +) -> Result<(Vec, usize), ApiError> { + if candidates.is_empty() { + return Ok((Vec::new(), 0)); + } + + let ids: Vec = candidates.iter().map(|candidate| candidate.id).collect(); let mut conn = state .db .get_pool() .get() .map_err(|_| ApiError::InternalServerError)?; - let tags_enc_filter = tags_enc_filter.to_vec(); - - let mut last_id: i64 = 0; - let mut out: Vec = Vec::new(); - - #[derive(Queryable)] - struct EmbeddingScanRow { - source_type: String, - conversation_id: Option, - vector_enc: Vec, - content_enc: Vec, - token_count: i32, - vector_dim: i32, - id: i64, - } - - loop { - let mut query = user_embeddings::table - .filter(user_embeddings::user_id.eq(user_id)) - .filter(user_embeddings::embedding_model.eq(DEFAULT_EMBEDDING_MODEL)) - .filter(user_embeddings::id.gt(last_id)) - .filter(user_embeddings::tags_enc.overlaps_with(tags_enc_filter.clone())) - .into_boxed(); - - if let Some(source_types) = source_types { - query = query.filter(user_embeddings::source_type.eq_any(source_types)); - } + let rows: Vec = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::id.eq_any(&ids)) + .select(( + user_embeddings::id, + user_embeddings::uuid, + user_embeddings::content_enc, + )) + .load(&mut conn) + .map_err(|e| { + error!( + "Failed to fetch RAG top-k content for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; - if let Some(conversation_id) = conversation_id { - query = query.filter(user_embeddings::conversation_id.eq(Some(conversation_id))); - } + let content_by_id: HashMap = + rows.into_iter().map(|row| (row.id, row)).collect(); + let mut results: Vec = Vec::with_capacity(limit); + let mut skipped_rows = 0usize; + + for candidate in candidates { + let Some(row) = content_by_id.get(&candidate.id) else { + skipped_rows = skipped_rows.saturating_add(1); + warn!( + target: "rag", + user_id = %user_id, + embedding_id = candidate.id, + embedding_uuid = %candidate.uuid, + reason = "content_missing", + "Skipping RAG candidate" + ); + continue; + }; - let rows: Vec = query - .order(user_embeddings::id.asc()) - .select(( - user_embeddings::source_type, - user_embeddings::conversation_id, - user_embeddings::vector_enc, - user_embeddings::content_enc, - user_embeddings::token_count, - user_embeddings::vector_dim, - user_embeddings::id, - )) - .limit(DB_SCAN_BATCH_SIZE) - .load(&mut conn) - .map_err(|e| { - error!( - "Failed to load tag-filtered embeddings batch for user={} after id={}: {:?}", - user_id, last_id, e + let plaintext = match decrypt_with_key(user_key, &row.content_enc) { + Ok(plaintext) => plaintext, + Err(_) => { + skipped_rows = skipped_rows.saturating_add(1); + warn!( + target: "rag", + user_id = %user_id, + embedding_id = candidate.id, + embedding_uuid = %row.uuid, + reason = "content_decrypt", + "Skipping corrupt RAG candidate" ); - ApiError::InternalServerError - })?; - - if rows.is_empty() { - break; - } + continue; + } + }; - for row in rows { - let EmbeddingScanRow { - source_type, - conversation_id, - vector_enc, - content_enc, - token_count, - vector_dim, - id, - } = row; - last_id = id; - - let vector_bytes = decrypt_with_key(user_key, &vector_enc) - .map_err(|_| ApiError::InternalServerError)?; - let vector = deserialize_f32_le(&vector_bytes)?; - - if vector.len() != vector_dim as usize { - debug!( - "Skipping embedding id={} for user={} due to dim mismatch (expected={}, got={})", - id, - user_id, - vector_dim, - vector.len() + let content = match String::from_utf8(plaintext) { + Ok(content) => content, + Err(_) => { + skipped_rows = skipped_rows.saturating_add(1); + warn!( + target: "rag", + user_id = %user_id, + embedding_id = candidate.id, + embedding_uuid = %row.uuid, + reason = "content_utf8", + "Skipping corrupt RAG candidate" ); continue; } + }; - out.push(CachedEmbedding { - source_type, - conversation_id, - vector, - content_enc, - token_count, - }); + results.push(RagSearchResult { + content, + score: candidate.score, + token_count: candidate.token_count, + }); + + if results.len() >= limit { + break; } } - Ok(Arc::new(out)) + Ok((results, skipped_rows)) } -#[allow(clippy::too_many_arguments)] -pub async fn search_user_embeddings( +pub async fn search_user_embeddings_with_options( state: &Arc, user: &User, auth_method: AuthMethod, user_key: &SecretKey, query: &str, - top_k: usize, - max_tokens: Option, - source_types: Option<&[String]>, - conversation_id: Option, - tags: Option<&[String]>, -) -> Result, ApiError> { - let top_k = top_k.clamp(1, 20); + options: RagSearchOptions, +) -> Result { + let started_at = Instant::now(); + let limit = options.limit.clamp(1, 20); + let candidate_limit = overfetch_limit(limit); + let limits = RagLimits::from_env(); + let query = query.trim(); + if query.is_empty() { + return Err(ApiError::BadRequest); + } + validate_text_size(query, limits.max_search_query_bytes, "search")?; let user_id = user.uuid; let (query_vec, _query_tokens) = embed_text_via_tinfoil(state, user, auth_method, query).await?; - let tags_enc_filter = tags + let tags_enc_filter = options + .filters + .tags + .as_ref() .map(|t| normalize_tags(t.iter().map(|s| s.as_str()))) .filter(|t| !t.is_empty()) .map(|t| encrypt_tags_b64(user_key, &t)); - let embeddings = if let Some(tags_enc_filter) = tags_enc_filter.as_ref() { - load_user_embeddings_by_tags( - state, - user_id, - user_key, - source_types, - conversation_id, - tags_enc_filter, - ) - .await? - } else { - let cached = { - let mut cache = state.rag_cache.lock().await; - cache.get(user_id) - }; + let load_filters = LoadFilters { + source_types: options.filters.source_types.as_deref(), + conversation_id: options.filters.conversation_id, + tags_enc_filter: tags_enc_filter.as_deref(), + begin_date: options.filters.begin_date, + end_date: options.filters.end_date, + }; - if let Some(hit) = cached { - hit - } else { - let loaded = load_all_user_embeddings(state, user_id, user_key).await?; - state.rag_cache.lock().await.put(user_id, loaded.clone()); - loaded - } + let loaded = if load_filters.is_cacheable_broad_load() { + load_cacheable_embeddings_with_coordination(state, user_id, user_key, limits).await? + } else { + load_user_embeddings_for_search(state, user_id, user_key, load_filters, limits.scan_limit) + .await? }; + let score_started_at = Instant::now(); let candidates = top_k_candidates( &query_vec, - &embeddings, - top_k, - source_types, - conversation_id, + &loaded.embeddings, + candidate_limit, + options.filters.source_types.as_deref(), + options.filters.conversation_id, )?; - let mut results: Vec = Vec::with_capacity(candidates.len()); - for c in candidates { - let plaintext = decrypt_with_key(user_key, &c.content_enc) - .map_err(|_| ApiError::InternalServerError)?; - let content = String::from_utf8(plaintext).map_err(|_| ApiError::InternalServerError)?; - results.push(RagSearchResult { - content, - score: c.score, - token_count: c.token_count, - }); - } + let score_ms = score_started_at.elapsed().as_millis() as u64; + let content_started_at = Instant::now(); + let (mut results, content_skipped_rows) = + fetch_ranked_content(state, user_id, user_key, candidates, limit).await?; + let content_ms = content_started_at.elapsed().as_millis() as u64; - if let Some(budget) = max_tokens { + if let Some(budget) = options.max_tokens { results = apply_token_budget(results, budget); } - Ok(results) + let feedback = if loaded.scan_limit_hit { + Some(format!( + "RAG search reached the internal scan limit of {} candidate rows. Older or out-of-window matches may exist; retry with begin_date/end_date for a narrower time range.", + limits.scan_limit + )) + } else { + None + }; + + let skipped_rows = loaded.skipped_rows.saturating_add(content_skipped_rows); + info!( + target: "rag", + user_id = %user_id, + returned_results = results.len(), + scanned_rows = loaded.scanned_rows, + skipped_rows, + db_read_bytes = loaded.db_read_bytes, + scan_limit = limits.scan_limit, + scan_limit_hit = loaded.scan_limit_hit, + score_ms, + content_fetch_ms = content_ms, + total_search_ms = started_at.elapsed().as_millis() as u64, + "rag_search_complete" + ); + + Ok(RagSearchOutcome { + results, + feedback, + scan_limit_hit: loaded.scan_limit_hit, + scanned_rows: loaded.scanned_rows, + skipped_rows, + }) +} + +#[allow(clippy::too_many_arguments)] +pub async fn search_user_embeddings( + state: &Arc, + user: &User, + auth_method: AuthMethod, + user_key: &SecretKey, + query: &str, + top_k: usize, + max_tokens: Option, + source_types: Option<&[String]>, + conversation_id: Option, + tags: Option<&[String]>, +) -> Result, ApiError> { + let options = RagSearchOptions { + limit: top_k, + max_tokens, + filters: RagSearchFilters { + source_types: source_types.map(|s| s.to_vec()), + conversation_id, + tags: tags.map(|t| t.to_vec()), + begin_date: None, + end_date: None, + }, + }; + let outcome = + search_user_embeddings_with_options(state, user, auth_method, user_key, query, options) + .await?; + Ok(outcome.results) } pub async fn delete_all_user_embeddings(state: &AppState, user_id: Uuid) -> Result<(), ApiError> { @@ -804,7 +1665,11 @@ pub async fn delete_user_embedding_by_uuid( return Err(ApiError::NotFound); } - state.rag_cache.lock().await.evict_user(user_id); + state + .rag_cache + .lock() + .await + .remove_embedding_by_uuid(user_id, embedding_uuid); Ok(()) } @@ -874,6 +1739,25 @@ pub async fn embeddings_status( mod tests { use super::*; + fn cached_embedding( + id: i64, + source_type: &str, + conversation_id: Option, + vector: Vec, + token_count: i32, + ) -> CachedEmbedding { + CachedEmbedding { + id, + uuid: Uuid::new_v4(), + source_type: source_type.to_string(), + conversation_id, + vector, + token_count, + created_at: Utc::now(), + updated_at: Utc::now(), + } + } + #[test] fn serialize_deserialize_roundtrip() { let v = vec![0.0f32, 1.5, -2.25, 42.0]; @@ -994,27 +1878,9 @@ mod tests { let query = vec![1.0f32, 0.0]; let embeddings = vec![ - CachedEmbedding { - source_type: SOURCE_TYPE_ARCHIVAL.to_string(), - conversation_id: None, - vector: vec![1.0, 0.0], - content_enc: b"a".to_vec(), - token_count: 5, - }, - CachedEmbedding { - source_type: SOURCE_TYPE_MESSAGE.to_string(), - conversation_id: Some(123), - vector: vec![0.0, 1.0], - content_enc: b"b".to_vec(), - token_count: 7, - }, - CachedEmbedding { - source_type: SOURCE_TYPE_MESSAGE.to_string(), - conversation_id: Some(123), - vector: vec![0.8, 0.2], - content_enc: b"c".to_vec(), - token_count: 9, - }, + cached_embedding(1, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5), + cached_embedding(2, SOURCE_TYPE_MESSAGE, Some(123), vec![0.0, 1.0], 7), + cached_embedding(3, SOURCE_TYPE_MESSAGE, Some(123), vec![0.8, 0.2], 9), ]; let items = top_k_candidates( @@ -1028,8 +1894,8 @@ mod tests { assert_eq!(items.len(), 2); assert!(items[0].score >= items[1].score); - assert_eq!(items[0].content_enc, b"c"); - assert_eq!(items[1].content_enc, b"b"); + assert_eq!(items[0].id, 3); + assert_eq!(items[1].id, 2); } #[test] @@ -1037,25 +1903,13 @@ mod tests { let query = vec![1.0f32, 0.0]; let embeddings = vec![ - CachedEmbedding { - source_type: SOURCE_TYPE_ARCHIVAL.to_string(), - conversation_id: None, - vector: vec![1.0, 0.0], - content_enc: b"a".to_vec(), - token_count: 10, - }, - CachedEmbedding { - source_type: SOURCE_TYPE_ARCHIVAL.to_string(), - conversation_id: None, - vector: vec![1.0, 0.0], - content_enc: b"b".to_vec(), - token_count: 5, - }, + cached_embedding(1, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 10), + cached_embedding(2, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5), ]; let items = top_k_candidates(&query, &embeddings, 1, None, None).unwrap(); assert_eq!(items.len(), 1); - assert_eq!(items[0].content_enc, b"b"); + assert_eq!(items[0].id, 2); } #[test] @@ -1063,20 +1917,8 @@ mod tests { let query = vec![1.0f32, 0.0]; let embeddings = vec![ - CachedEmbedding { - source_type: SOURCE_TYPE_ARCHIVAL.to_string(), - conversation_id: None, - vector: vec![1.0, 0.0], - content_enc: b"a".to_vec(), - token_count: 5, - }, - CachedEmbedding { - source_type: SOURCE_TYPE_MESSAGE.to_string(), - conversation_id: Some(123), - vector: vec![0.0, 1.0], - content_enc: b"b".to_vec(), - token_count: 7, - }, + cached_embedding(1, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5), + cached_embedding(2, SOURCE_TYPE_MESSAGE, Some(123), vec![0.0, 1.0], 7), ]; let total = embeddings.len(); @@ -1094,10 +1936,10 @@ mod tests { #[test] fn rag_cache_evict_user_removes_entry() { - let mut cache = RagCache::new(10, Duration::from_secs(60)); + let mut cache = RagCache::new(1024 * 1024, 1024 * 1024, Duration::from_secs(60)); let user_id = Uuid::new_v4(); - cache.put(user_id, Arc::new(vec![])); + cache.put(user_id, Arc::new(vec![]), false); assert!(cache.entries.contains_key(&user_id)); cache.evict_user(user_id); @@ -1107,22 +1949,24 @@ mod tests { } #[tokio::test] - async fn rag_cache_lru_eviction() { - let mut cache = RagCache::new(2, Duration::from_secs(60)); + async fn rag_cache_byte_lru_eviction() { + let row = cached_embedding(1, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5); + let row_bytes = row.estimated_cache_bytes(); + let mut cache = RagCache::new(row_bytes * 2, row_bytes * 10, Duration::from_secs(60)); - let v1 = Arc::new(vec![]); - let v2 = Arc::new(vec![]); - let v3 = Arc::new(vec![]); + let v1 = Arc::new(vec![row.clone()]); + let v2 = Arc::new(vec![row.clone()]); + let v3 = Arc::new(vec![row]); let u1 = Uuid::new_v4(); let u2 = Uuid::new_v4(); let u3 = Uuid::new_v4(); - cache.put(u1, v1); - cache.put(u2, v2); + assert!(cache.put(u1, v1, false)); + assert!(cache.put(u2, v2, false)); // touch u1 so u2 becomes LRU cache.get(u1); - cache.put(u3, v3); + assert!(cache.put(u3, v3, false)); assert!(cache.entries.contains_key(&u1)); assert!(!cache.entries.contains_key(&u2)); @@ -1131,11 +1975,58 @@ mod tests { #[tokio::test] async fn rag_cache_ttl_expiration() { - let mut cache = RagCache::new(10, Duration::from_millis(5)); + let mut cache = RagCache::new(1024 * 1024, 1024 * 1024, Duration::from_millis(5)); let user = Uuid::new_v4(); - cache.put(user, Arc::new(vec![])); + cache.put(user, Arc::new(vec![]), false); tokio::time::sleep(Duration::from_millis(10)).await; assert!(cache.get(user).is_none()); } + + #[test] + fn rag_cache_append_updates_present_entry() { + let mut cache = RagCache::new(1024 * 1024, 1024 * 1024, Duration::from_secs(60)); + let user = Uuid::new_v4(); + let first = cached_embedding(1, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5); + let second = cached_embedding(2, SOURCE_TYPE_ARCHIVAL, None, vec![0.0, 1.0], 7); + + assert!(cache.put(user, Arc::new(vec![first]), false)); + assert_eq!(cache.append(user, second), CacheAppendResult::Appended); + + let (cached, _) = cache.get(user).unwrap(); + assert_eq!(cached.len(), 2); + assert_eq!(cached[1].id, 2); + } + + #[test] + fn rag_cache_append_evicts_when_over_user_byte_cap() { + let first = cached_embedding(1, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5); + let second = cached_embedding(2, SOURCE_TYPE_ARCHIVAL, None, vec![0.0, 1.0], 7); + let row_bytes = first.estimated_cache_bytes(); + let mut cache = RagCache::new(row_bytes * 10, row_bytes + 1, Duration::from_secs(60)); + let user = Uuid::new_v4(); + + assert!(cache.put(user, Arc::new(vec![first]), false)); + assert_eq!( + cache.append(user, second), + CacheAppendResult::EvictedOverLimit + ); + assert!(cache.get(user).is_none()); + } + + #[test] + fn rag_cache_remove_embedding_by_uuid_updates_entry() { + let mut cache = RagCache::new(1024 * 1024, 1024 * 1024, Duration::from_secs(60)); + let user = Uuid::new_v4(); + let first = cached_embedding(1, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5); + let second = cached_embedding(2, SOURCE_TYPE_ARCHIVAL, None, vec![0.0, 1.0], 7); + let second_uuid = second.uuid; + + assert!(cache.put(user, Arc::new(vec![first, second]), false)); + cache.remove_embedding_by_uuid(user, second_uuid); + + let (cached, _) = cache.get(user).unwrap(); + assert_eq!(cached.len(), 1); + assert_eq!(cached[0].id, 1); + } } diff --git a/src/web/mod.rs b/src/web/mod.rs index 5e7eff65..5448a9a2 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -8,7 +8,6 @@ mod openai; pub mod openai_auth; pub mod platform; pub mod protected_routes; -pub mod rag; pub mod responses; pub use health_routes::router_with_state as health_routes_with_state; @@ -18,7 +17,6 @@ pub use openai::get_embedding_vector; pub use openai::router as openai_routes; pub use platform::router as platform_routes; pub use protected_routes::router as protected_routes; -pub use rag::router as rag_routes; pub use responses::conversation_projects_router as conversation_projects_routes; pub use responses::conversations_router as conversations_routes; pub use responses::instructions_router as instructions_routes; diff --git a/src/web/rag.rs b/src/web/rag.rs deleted file mode 100644 index 657a3e92..00000000 --- a/src/web/rag.rs +++ /dev/null @@ -1,216 +0,0 @@ -use axum::{ - extract::{Path, State}, - middleware::from_fn_with_state, - routing::{delete, get, post}, - Extension, Json, Router, -}; -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::sync::Arc; -use uuid::Uuid; - -use crate::jwt::AuthContext; -use crate::models::users::User; -use crate::rag; -use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; -use crate::web::openai_auth::AuthMethod; -use crate::web::responses::error_mapping; -use crate::{ApiError, AppMode, AppState}; - -#[derive(Debug, Clone, Deserialize)] -struct InsertEmbeddingRequest { - text: String, - metadata: Option, -} - -#[derive(Debug, Clone, Serialize)] -struct InsertEmbeddingResponse { - id: Uuid, - source_type: String, - embedding_model: String, - token_count: i32, - created_at: DateTime, -} - -#[derive(Debug, Clone, Deserialize)] -struct SearchRequest { - query: String, - top_k: Option, - max_tokens: Option, - source_types: Option>, - conversation_id: Option, -} - -#[derive(Debug, Clone, Serialize)] -struct SearchResponse { - results: Vec, -} - -pub fn router(app_state: Arc) -> Router<()> { - // Experimental endpoints: only enabled in Local/Dev. - if !matches!(app_state.app_mode, AppMode::Local | AppMode::Dev) { - return Router::new().with_state(app_state); - } - - Router::new() - .route( - "/v1/rag/embeddings", - post(insert_archival_embedding).layer(from_fn_with_state( - app_state.clone(), - decrypt_request::, - )), - ) - .route( - "/v1/rag/embeddings", - delete(delete_all_embeddings) - .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .route( - "/v1/rag/embeddings/:id", - delete(delete_embedding) - .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .route( - "/v1/rag/search", - post(search).layer(from_fn_with_state( - app_state.clone(), - decrypt_request::, - )), - ) - .route( - "/v1/rag/embeddings/status", - get(status).layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .with_state(app_state) -} - -async fn insert_archival_embedding( - State(state): State>, - Extension(session_id): Extension, - Extension(user): Extension, - Extension(auth_context): Extension, - Extension(body): Extension, -) -> Result { - if body.text.trim().is_empty() { - return Err(ApiError::BadRequest); - } - if let Some(m) = &body.metadata { - if !m.is_object() { - return Err(ApiError::BadRequest); - } - } - - let user_key = state - .get_user_key(&user, &auth_context, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let inserted = rag::insert_archival_embedding( - &state, - &user, - AuthMethod::Jwt, - &user_key, - &body.text, - body.metadata.as_ref(), - ) - .await?; - - let response = InsertEmbeddingResponse { - id: inserted.uuid, - source_type: inserted.source_type, - embedding_model: inserted.embedding_model, - token_count: inserted.token_count, - created_at: inserted.created_at, - }; - - let encrypted = encrypt_response(&state, &session_id, &response).await?; - Ok((axum::http::StatusCode::CREATED, encrypted)) -} - -async fn search( - State(state): State>, - Extension(session_id): Extension, - Extension(user): Extension, - Extension(auth_context): Extension, - Extension(body): Extension, -) -> Result>, ApiError> { - if body.query.trim().is_empty() { - return Err(ApiError::BadRequest); - } - - let top_k = body.top_k.unwrap_or(5); - if top_k == 0 || top_k > 20 { - return Err(ApiError::BadRequest); - } - - if let Some(max_tokens) = body.max_tokens { - if max_tokens <= 0 { - return Err(ApiError::BadRequest); - } - } - - if let Some(source_types) = &body.source_types { - if source_types.is_empty() { - return Err(ApiError::BadRequest); - } - } - - let conversation_internal_id = if let Some(conversation_uuid) = body.conversation_id { - let conversation = state - .db - .get_conversation_by_uuid_and_user(conversation_uuid, user.uuid) - .map_err(error_mapping::map_conversation_error)?; - Some(conversation.id) - } else { - None - }; - - let user_key = state - .get_user_key(&user, &auth_context, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let results = rag::search_user_embeddings( - &state, - &user, - AuthMethod::Jwt, - &user_key, - &body.query, - top_k, - body.max_tokens, - body.source_types.as_deref(), - conversation_internal_id, - None, - ) - .await?; - - let response = SearchResponse { results }; - encrypt_response(&state, &session_id, &response).await -} - -async fn delete_all_embeddings( - State(state): State>, - Extension(user): Extension, -) -> Result { - rag::delete_all_user_embeddings(&state, user.uuid).await?; - Ok(axum::http::StatusCode::NO_CONTENT) -} - -async fn delete_embedding( - State(state): State>, - Path(embedding_id): Path, - Extension(user): Extension, -) -> Result { - rag::delete_user_embedding_by_uuid(&state, user.uuid, embedding_id).await?; - Ok(axum::http::StatusCode::NO_CONTENT) -} - -async fn status( - State(state): State>, - Extension(session_id): Extension, - Extension(user): Extension, -) -> Result>, ApiError> { - let response = rag::embeddings_status(&state, user.uuid).await?; - encrypt_response(&state, &session_id, &response).await -} From 023bebaa24ccae6ceed5d8b932929c8a23542420 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 7 Jul 2026 04:28:46 +0000 Subject: [PATCH 3/6] test(rag): add db tamper regression coverage --- src/rag.rs | 247 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 247 insertions(+) diff --git a/src/rag.rs b/src/rag.rs index dc5a2b65..92c9a877 100644 --- a/src/rag.rs +++ b/src/rag.rs @@ -1738,6 +1738,13 @@ pub async fn embeddings_status( #[cfg(test)] mod tests { use super::*; + use crate::{ + db::setup_db, + models::{org_projects::OrgProject, schema::org_projects, users::NewUser}, + AppMode, AppState, AppStateBuilder, + }; + use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl}; + use tokio::sync::RwLock; fn cached_embedding( id: i64, @@ -2029,4 +2036,244 @@ mod tests { assert_eq!(cached.len(), 1); assert_eq!(cached[0].id, 1); } + + #[tokio::test] + #[ignore = "requires RAG_TAMPER_TEST_DATABASE_URL pointing at disposable migrated local Postgres"] + async fn db_copied_rag_row_does_not_decrypt_under_attacker_key() { + let Some(database_url) = rag_test_database_url() else { + eprintln!("skipping: RAG_TAMPER_TEST_DATABASE_URL is not set"); + return; + }; + + let app_state = build_rag_db_test_app_state(database_url).await; + let project = first_active_project(&app_state); + let victim = create_rag_test_user(&app_state, project.id, "victim"); + let attacker = create_rag_test_user(&app_state, project.id, "attacker"); + let victim_key = SecretKey::from_slice(&[1u8; 32]).unwrap(); + let attacker_key = SecretKey::from_slice(&[2u8; 32]).unwrap(); + + let victim_row = + insert_direct_archival_embedding(&app_state, victim.uuid, &victim_key, "victim secret") + .await; + copy_embedding_ciphertext_to_user(&app_state, &victim_row, attacker.uuid); + + let victim_loaded = + load_direct_test_embeddings(&app_state, victim.uuid, &victim_key, 25).await; + assert_eq!(victim_loaded.embeddings.len(), 1); + assert_eq!(victim_loaded.skipped_rows, 0); + + let attacker_loaded = + load_direct_test_embeddings(&app_state, attacker.uuid, &attacker_key, 25).await; + assert_eq!(attacker_loaded.scanned_rows, 1); + assert_eq!(attacker_loaded.embeddings.len(), 0); + assert_eq!(attacker_loaded.skipped_rows, 1); + + let _ = app_state.db.delete_user(&victim); + let _ = app_state.db.delete_user(&attacker); + } + + #[tokio::test] + #[ignore = "requires RAG_TAMPER_TEST_DATABASE_URL pointing at disposable migrated local Postgres"] + async fn db_corrupt_rag_vector_skips_row_without_failing_search_load() { + let Some(database_url) = rag_test_database_url() else { + eprintln!("skipping: RAG_TAMPER_TEST_DATABASE_URL is not set"); + return; + }; + + let app_state = build_rag_db_test_app_state(database_url).await; + let project = first_active_project(&app_state); + let user = create_rag_test_user(&app_state, project.id, "corrupt-vector"); + let key = SecretKey::from_slice(&[3u8; 32]).unwrap(); + + insert_direct_archival_embedding_with_ciphertext( + &app_state, + user.uuid, + vec![1, 2, 3], + encrypt_with_key(&key, b"still encrypted").await, + ); + + let loaded = load_direct_test_embeddings(&app_state, user.uuid, &key, 25).await; + assert_eq!(loaded.scanned_rows, 1); + assert_eq!(loaded.embeddings.len(), 0); + assert_eq!(loaded.skipped_rows, 1); + + let _ = app_state.db.delete_user(&user); + } + + #[tokio::test] + #[ignore = "requires RAG_TAMPER_TEST_DATABASE_URL pointing at disposable migrated local Postgres"] + async fn db_corrupt_top_candidate_content_is_skipped_and_next_candidate_fills() { + let Some(database_url) = rag_test_database_url() else { + eprintln!("skipping: RAG_TAMPER_TEST_DATABASE_URL is not set"); + return; + }; + + let app_state = build_rag_db_test_app_state(database_url).await; + let project = first_active_project(&app_state); + let user = create_rag_test_user(&app_state, project.id, "corrupt-content"); + let key = SecretKey::from_slice(&[4u8; 32]).unwrap(); + + let bad_content = insert_direct_archival_embedding_with_ciphertext( + &app_state, + user.uuid, + encrypt_with_key(&key, &serialize_f32_le(&[1.0, 0.0])).await, + vec![9, 9, 9], + ); + let good_content = + insert_direct_archival_embedding(&app_state, user.uuid, &key, "recoverable memory") + .await; + + let candidates = vec![ + HeapItem { + id: bad_content.id, + uuid: bad_content.uuid, + score: 1.0, + token_count: 1, + }, + HeapItem { + id: good_content.id, + uuid: good_content.uuid, + score: 0.9, + token_count: 1, + }, + ]; + + let (results, skipped_rows) = + fetch_ranked_content(&app_state, user.uuid, &key, candidates, 1) + .await + .unwrap(); + assert_eq!(skipped_rows, 1); + assert_eq!(results.len(), 1); + assert_eq!(results[0].content, "recoverable memory"); + + let _ = app_state.db.delete_user(&user); + } + + fn rag_test_database_url() -> Option { + std::env::var("RAG_TAMPER_TEST_DATABASE_URL") + .ok() + .or_else(|| std::env::var("AEAD_TAMPER_TEST_DATABASE_URL").ok()) + } + + async fn build_rag_db_test_app_state(database_url: String) -> AppState { + let db = setup_db(database_url); + AppStateBuilder::default() + .app_mode(AppMode::Local) + .db(db) + .enclave_key([42u8; 32].to_vec()) + .aws_credential_manager(Arc::new(RwLock::new(None))) + .openai_api_base("http://localhost:9".to_string()) + .tinfoil_api_base("http://localhost:9".to_string()) + .jwt_secret([24u8; 32].to_vec()) + .build() + .await + .expect("local RAG test app state should build") + } + + fn first_active_project(app_state: &AppState) -> OrgProject { + let conn = &mut app_state + .db + .get_pool() + .get() + .expect("test database connection should be available"); + + org_projects::table + .filter(org_projects::status.eq("active")) + .order(org_projects::id.asc()) + .first::(conn) + .expect("test database should contain at least one active project") + } + + fn create_rag_test_user(app_state: &AppState, project_id: i32, label: &str) -> User { + let marker = Uuid::new_v4(); + app_state + .db + .create_user(NewUser::new( + Some(format!("rag-{label}-{marker}@example.com")), + None, + project_id, + )) + .expect("RAG test user should insert") + } + + async fn insert_direct_archival_embedding( + app_state: &AppState, + user_id: Uuid, + user_key: &SecretKey, + content: &str, + ) -> crate::models::user_embeddings::UserEmbedding { + insert_direct_archival_embedding_with_ciphertext( + app_state, + user_id, + encrypt_with_key(user_key, &serialize_f32_le(&[1.0, 0.0])).await, + encrypt_with_key(user_key, content.as_bytes()).await, + ) + } + + fn insert_direct_archival_embedding_with_ciphertext( + app_state: &AppState, + user_id: Uuid, + vector_enc: Vec, + content_enc: Vec, + ) -> crate::models::user_embeddings::UserEmbedding { + let conn = &mut app_state + .db + .get_pool() + .get() + .expect("test database connection should be available"); + + NewUserEmbedding { + uuid: Uuid::new_v4(), + user_id, + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + user_message_id: None, + assistant_message_id: None, + conversation_id: None, + vector_enc, + embedding_model: DEFAULT_EMBEDDING_MODEL.to_string(), + vector_dim: 2, + content_enc, + metadata_enc: None, + tags_enc: Vec::new(), + token_count: 1, + } + .insert(conn) + .expect("direct RAG embedding should insert") + } + + fn copy_embedding_ciphertext_to_user( + app_state: &AppState, + source: &crate::models::user_embeddings::UserEmbedding, + target_user_id: Uuid, + ) -> crate::models::user_embeddings::UserEmbedding { + insert_direct_archival_embedding_with_ciphertext( + app_state, + target_user_id, + source.vector_enc.clone(), + source.content_enc.clone(), + ) + } + + async fn load_direct_test_embeddings( + app_state: &AppState, + user_id: Uuid, + user_key: &SecretKey, + scan_limit: i64, + ) -> LoadedEmbeddings { + load_user_embeddings_for_search( + app_state, + user_id, + user_key, + LoadFilters { + source_types: None, + conversation_id: None, + tags_enc_filter: None, + begin_date: None, + end_date: None, + }, + scan_limit, + ) + .await + .expect("direct RAG load should not hard-fail on row-local corruption") + } } From 632d557cc3ffdfea28beddd57a82736afa457f04 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 7 Jul 2026 04:33:15 +0000 Subject: [PATCH 4/6] fix(rag): apply review follow-ups --- .../down.sql | 6 +- .../metadata.toml | 2 + .../up.sql | 6 +- src/rag.rs | 125 ++++++++++++++---- 4 files changed, 109 insertions(+), 30 deletions(-) create mode 100644 migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/metadata.toml diff --git a/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql index c0057325..bb41b7e7 100644 --- a/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql +++ b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql @@ -1,3 +1,3 @@ -DROP INDEX IF EXISTS idx_user_embeddings_user_model_conversation_created_id; -DROP INDEX IF EXISTS idx_user_embeddings_user_model_source_created_id; -DROP INDEX IF EXISTS idx_user_embeddings_user_model_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_conversation_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_source_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_created_id; diff --git a/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/metadata.toml b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/metadata.toml new file mode 100644 index 00000000..c6ce5bdf --- /dev/null +++ b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/metadata.toml @@ -0,0 +1,2 @@ +[general] +run_in_transaction = false diff --git a/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql index b857135b..8407e441 100644 --- a/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql +++ b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql @@ -1,9 +1,9 @@ -- Hot RAG retrieval paths: active embedding model for a user, newest first. -CREATE INDEX idx_user_embeddings_user_model_created_id +CREATE INDEX CONCURRENTLY idx_user_embeddings_user_model_created_id ON user_embeddings(user_id, embedding_model, created_at DESC, id DESC); -CREATE INDEX idx_user_embeddings_user_model_source_created_id +CREATE INDEX CONCURRENTLY idx_user_embeddings_user_model_source_created_id ON user_embeddings(user_id, embedding_model, source_type, created_at DESC, id DESC); -CREATE INDEX idx_user_embeddings_user_model_conversation_created_id +CREATE INDEX CONCURRENTLY idx_user_embeddings_user_model_conversation_created_id ON user_embeddings(user_id, embedding_model, conversation_id, created_at DESC, id DESC); diff --git a/src/rag.rs b/src/rag.rs index 92c9a877..8ccb899b 100644 --- a/src/rag.rs +++ b/src/rag.rs @@ -625,7 +625,8 @@ fn top_k_candidates( if let Some(std::cmp::Reverse(min)) = heap.peek() { let ordering = score .total_cmp(&min.score) - .then_with(|| min.token_count.cmp(&e.token_count)); + .then_with(|| min.token_count.cmp(&e.token_count)) + .then_with(|| min.id.cmp(&e.id)); if ordering == Ordering::Greater { heap.pop(); let item = HeapItem { @@ -805,6 +806,48 @@ fn ensure_storage_limits( projected_new_bytes: i64, limits: RagLimits, ) -> Result<(), ApiError> { + let user_stats = ensure_pre_embedding_limits(conn, user, limits)?; + if user_stats.stored_bytes.saturating_add(projected_new_bytes) > limits.max_user_stored_bytes { + warn!( + target: "rag", + user_id = %user.uuid, + rows = user_stats.row_count, + stored_bytes = user_stats.stored_bytes, + projected_new_bytes, + max_rows = limits.max_user_embeddings, + max_bytes = limits.max_user_stored_bytes, + "RAG user storage limit reached" + ); + return Err(ApiError::BadRequest); + } + + let project_stats = load_project_embedding_storage_stats(conn, user.project_id)?; + if project_stats + .stored_bytes + .saturating_add(projected_new_bytes) + > limits.max_project_stored_bytes + { + warn!( + target: "rag", + project_id = user.project_id, + rows = project_stats.row_count, + stored_bytes = project_stats.stored_bytes, + projected_new_bytes, + max_rows = limits.max_project_embeddings, + max_bytes = limits.max_project_stored_bytes, + "RAG project storage limit reached" + ); + return Err(ApiError::BadRequest); + } + + Ok(()) +} + +fn ensure_pre_embedding_limits( + conn: &mut diesel::PgConnection, + user: &User, + limits: RagLimits, +) -> Result { use diesel::dsl::count_star; let recent_insert_count: i64 = user_embeddings::table @@ -832,15 +875,13 @@ fn ensure_storage_limits( let user_stats = load_user_embedding_storage_stats(conn, user.uuid)?; if user_stats.row_count.saturating_add(1) > limits.max_user_embeddings - || user_stats.stored_bytes.saturating_add(projected_new_bytes) - > limits.max_user_stored_bytes + || user_stats.stored_bytes >= limits.max_user_stored_bytes { warn!( target: "rag", user_id = %user.uuid, rows = user_stats.row_count, stored_bytes = user_stats.stored_bytes, - projected_new_bytes, max_rows = limits.max_user_embeddings, max_bytes = limits.max_user_stored_bytes, "RAG user storage limit reached" @@ -850,17 +891,13 @@ fn ensure_storage_limits( let project_stats = load_project_embedding_storage_stats(conn, user.project_id)?; if project_stats.row_count.saturating_add(1) > limits.max_project_embeddings - || project_stats - .stored_bytes - .saturating_add(projected_new_bytes) - > limits.max_project_stored_bytes + || project_stats.stored_bytes >= limits.max_project_stored_bytes { warn!( target: "rag", project_id = user.project_id, rows = project_stats.row_count, stored_bytes = project_stats.stored_bytes, - projected_new_bytes, max_rows = limits.max_project_embeddings, max_bytes = limits.max_project_stored_bytes, "RAG project storage limit reached" @@ -868,7 +905,7 @@ fn ensure_storage_limits( return Err(ApiError::BadRequest); } - Ok(()) + Ok(user_stats) } fn cached_embedding_from_inserted( @@ -901,6 +938,14 @@ pub async fn insert_archival_embedding( } let limits = RagLimits::from_env(); validate_text_size(text, limits.max_insert_text_bytes, "insert")?; + { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + ensure_pre_embedding_limits(&mut conn, user, limits)?; + } let user_id = user.uuid; let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; @@ -980,6 +1025,14 @@ pub async fn insert_message_embedding( } let limits = RagLimits::from_env(); validate_text_size(text, limits.max_insert_text_bytes, "insert")?; + { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + ensure_pre_embedding_limits(&mut conn, user, limits)?; + } let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; @@ -1364,23 +1417,33 @@ async fn load_cacheable_embeddings_with_coordination( filters, limits.scan_limit, ) - .await?; + .await; let mut cache = state.rag_cache.lock().await; - let cached = cache.put(user_id, loaded.embeddings.clone(), loaded.scan_limit_hit); + if let Ok(loaded) = &loaded { + let cached = + cache.put(user_id, loaded.embeddings.clone(), loaded.scan_limit_hit); + info!( + target: "rag", + user_id = %user_id, + cached, + loaded_rows = loaded.embeddings.len(), + cache_bytes_per_user = cached_embeddings_bytes(&loaded.embeddings), + total_cache_bytes = cache.total_bytes, + scan_limit_hit = loaded.scan_limit_hit, + reason = "timeout_duplicate", + "rag_cache_store_after_load" + ); + } else { + warn!( + target: "rag", + user_id = %user_id, + reason = "timeout_duplicate", + "RAG duplicate cacheable load failed" + ); + } cache.finish_load(user_id); - info!( - target: "rag", - user_id = %user_id, - cached, - loaded_rows = loaded.embeddings.len(), - cache_bytes_per_user = cached_embeddings_bytes(&loaded.embeddings), - total_cache_bytes = cache.total_bytes, - scan_limit_hit = loaded.scan_limit_hit, - reason = "timeout_duplicate", - "rag_cache_store_after_load" - ); - return Ok(loaded); + return loaded; } } } @@ -1919,6 +1982,20 @@ mod tests { assert_eq!(items[0].id, 2); } + #[test] + fn top_k_tie_break_prefers_lower_id_after_score_and_tokens() { + let query = vec![1.0f32, 0.0]; + + let embeddings = vec![ + cached_embedding(2, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5), + cached_embedding(1, SOURCE_TYPE_ARCHIVAL, None, vec![1.0, 0.0], 5), + ]; + + let items = top_k_candidates(&query, &embeddings, 1, None, None).unwrap(); + assert_eq!(items.len(), 1); + assert_eq!(items[0].id, 1); + } + #[test] fn top_k_when_k_gt_embeddings_returns_all() { let query = vec![1.0f32, 0.0]; From 37045ab92f575e1ea477208d242253c93cb781e8 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 7 Jul 2026 05:08:42 +0000 Subject: [PATCH 5/6] fix(rag): offload blocking work from async paths --- src/rag.rs | 414 +++++++++++++++++++++++++++++++++++------------------ 1 file changed, 275 insertions(+), 139 deletions(-) diff --git a/src/rag.rs b/src/rag.rs index 8ccb899b..d2d9768c 100644 --- a/src/rag.rs +++ b/src/rag.rs @@ -17,6 +17,7 @@ use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::Notify; +use tokio::task; use tokio::time::{sleep, timeout}; use tracing::{debug, error, info, warn}; use uuid::Uuid; @@ -436,6 +437,22 @@ struct RagLimits { max_project_stored_bytes: i64, } +async fn run_blocking_rag(task_name: &'static str, f: F) -> Result +where + T: Send + 'static, + F: FnOnce() -> Result + Send + 'static, +{ + task::spawn_blocking(f).await.map_err(|e| { + error!( + target: "rag", + task = task_name, + error = ?e, + "RAG blocking task failed" + ); + ApiError::InternalServerError + })? +} + impl RagLimits { fn from_env() -> Self { Self { @@ -645,6 +662,25 @@ fn top_k_candidates( Ok(out) } +async fn top_k_candidates_blocking( + query: Vec, + embeddings: Arc>, + top_k: usize, + source_types: Option>, + conversation_id: Option, +) -> Result, ApiError> { + run_blocking_rag("rag_top_k_candidates", move || { + top_k_candidates( + &query, + embeddings.as_slice(), + top_k, + source_types.as_deref(), + conversation_id, + ) + }) + .await +} + fn apply_token_budget(results: Vec, budget: i32) -> Vec { let mut total: i32 = 0; let mut limited: Vec = Vec::new(); @@ -924,6 +960,56 @@ fn cached_embedding_from_inserted( } } +async fn ensure_pre_embedding_limits_blocking( + state: &AppState, + user: &User, + limits: RagLimits, +) -> Result<(), ApiError> { + let db = state.db.clone(); + let user = user.clone(); + + run_blocking_rag("rag_pre_embedding_limits", move || { + let mut conn = db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + ensure_pre_embedding_limits(&mut conn, &user, limits)?; + Ok(()) + }) + .await +} + +async fn insert_embedding_row_with_limits( + state: &AppState, + user: &User, + new_embedding: NewUserEmbedding, + projected_new_bytes: i64, + limits: RagLimits, + task_name: &'static str, +) -> Result { + let db = state.db.clone(); + let user = user.clone(); + + run_blocking_rag(task_name, move || { + let mut conn = db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + ensure_storage_limits(&mut conn, &user, projected_new_bytes, limits)?; + + new_embedding.insert(&mut conn).map_err(|e| { + error!( + target: "rag", + task = task_name, + error = ?e, + "Failed to insert RAG embedding" + ); + ApiError::InternalServerError + }) + }) + .await +} + pub async fn insert_archival_embedding( state: &Arc, user: &User, @@ -938,14 +1024,7 @@ pub async fn insert_archival_embedding( } let limits = RagLimits::from_env(); validate_text_size(text, limits.max_insert_text_bytes, "insert")?; - { - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - ensure_pre_embedding_limits(&mut conn, user, limits)?; - } + ensure_pre_embedding_limits_blocking(state, user, limits).await?; let user_id = user.uuid; let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; @@ -966,14 +1045,7 @@ pub async fn insert_archival_embedding( let projected_new_bytes = encrypted_embedding_bytes(&vector_enc, &content_enc, metadata_enc.as_ref()); - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - ensure_storage_limits(&mut conn, user, projected_new_bytes, limits)?; - - let inserted = NewUserEmbedding { + let new_embedding = NewUserEmbedding { uuid: Uuid::new_v4(), user_id, source_type: SOURCE_TYPE_ARCHIVAL.to_string(), @@ -987,12 +1059,16 @@ pub async fn insert_archival_embedding( metadata_enc, tags_enc, token_count, - } - .insert(&mut conn) - .map_err(|e| { - error!("Failed to insert archival embedding: {:?}", e); - ApiError::InternalServerError - })?; + }; + let inserted = insert_embedding_row_with_limits( + state, + user, + new_embedding, + projected_new_bytes, + limits, + "rag_insert_archival_embedding", + ) + .await?; let cached = cached_embedding_from_inserted(&inserted, vector); let append_result = state.rag_cache.lock().await.append(user_id, cached); @@ -1025,14 +1101,7 @@ pub async fn insert_message_embedding( } let limits = RagLimits::from_env(); validate_text_size(text, limits.max_insert_text_bytes, "insert")?; - { - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - ensure_pre_embedding_limits(&mut conn, user, limits)?; - } + ensure_pre_embedding_limits_blocking(state, user, limits).await?; let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; @@ -1041,14 +1110,7 @@ pub async fn insert_message_embedding( let content_enc = encrypt_with_key(user_key, text.as_bytes()).await; let projected_new_bytes = encrypted_embedding_bytes(&vector_enc, &content_enc, None); - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - ensure_storage_limits(&mut conn, user, projected_new_bytes, limits)?; - - let inserted = NewUserEmbedding { + let new_embedding = NewUserEmbedding { uuid: Uuid::new_v4(), user_id, source_type: SOURCE_TYPE_MESSAGE.to_string(), @@ -1062,12 +1124,16 @@ pub async fn insert_message_embedding( metadata_enc: None, tags_enc: Vec::new(), token_count, - } - .insert(&mut conn) - .map_err(|e| { - error!("Failed to insert message embedding: {:?}", e); - ApiError::InternalServerError - })?; + }; + let inserted = insert_embedding_row_with_limits( + state, + user, + new_embedding, + projected_new_bytes, + limits, + "rag_insert_message_embedding", + ) + .await?; let cached = cached_embedding_from_inserted(&inserted, vector); let append_result = state.rag_cache.lock().await.append(user_id, cached); @@ -1107,6 +1173,35 @@ impl LoadFilters<'_> { && self.begin_date.is_none() && self.end_date.is_none() } + + fn to_owned_filters(self) -> OwnedLoadFilters { + OwnedLoadFilters { + source_types: self.source_types.map(|source_types| source_types.to_vec()), + conversation_id: self.conversation_id, + tags_enc_filter: self.tags_enc_filter.map(|tags| tags.to_vec()), + begin_date: self.begin_date, + end_date: self.end_date, + } + } +} + +#[derive(Debug, Clone)] +struct OwnedLoadFilters { + source_types: Option>, + conversation_id: Option, + tags_enc_filter: Option>>, + begin_date: Option>, + end_date: Option>, +} + +impl OwnedLoadFilters { + fn is_cacheable_broad_load(&self) -> bool { + self.source_types.is_none() + && self.conversation_id.is_none() + && self.tags_enc_filter.is_none() + && self.begin_date.is_none() + && self.end_date.is_none() + } } #[derive(Queryable)] @@ -1129,28 +1224,43 @@ async fn load_user_embeddings_for_search( filters: LoadFilters<'_>, scan_limit: i64, ) -> Result { - let mut conn = state - .db + let db = state.db.clone(); + let user_key = *user_key; + let filters = filters.to_owned_filters(); + + run_blocking_rag("rag_load_user_embeddings_for_search", move || { + load_user_embeddings_for_search_blocking(db, user_id, &user_key, filters, scan_limit) + }) + .await +} + +fn load_user_embeddings_for_search_blocking( + db: Arc, + user_id: Uuid, + user_key: &SecretKey, + filters: OwnedLoadFilters, + scan_limit: i64, +) -> Result { + let mut conn = db .get_pool() .get() .map_err(|_| ApiError::InternalServerError)?; - let started_at = Instant::now(); let mut query = user_embeddings::table .filter(user_embeddings::user_id.eq(user_id)) .filter(user_embeddings::embedding_model.eq(DEFAULT_EMBEDDING_MODEL)) .into_boxed(); - if let Some(source_types) = filters.source_types { - query = query.filter(user_embeddings::source_type.eq_any(source_types)); + if let Some(source_types) = filters.source_types.as_ref() { + query = query.filter(user_embeddings::source_type.eq_any(source_types.as_slice())); } if let Some(conversation_id) = filters.conversation_id { query = query.filter(user_embeddings::conversation_id.eq(Some(conversation_id))); } - if let Some(tags_enc_filter) = filters.tags_enc_filter { - query = query.filter(user_embeddings::tags_enc.overlaps_with(tags_enc_filter.to_vec())); + if let Some(tags_enc_filter) = filters.tags_enc_filter.as_ref() { + query = query.filter(user_embeddings::tags_enc.overlaps_with(tags_enc_filter.clone())); } if let Some(begin_date) = filters.begin_date { @@ -1460,9 +1570,24 @@ async fn fetch_ranked_content( return Ok((Vec::new(), 0)); } + let db = state.db.clone(); + let user_key = *user_key; + + run_blocking_rag("rag_fetch_ranked_content", move || { + fetch_ranked_content_blocking(db, user_id, &user_key, candidates, limit) + }) + .await +} + +fn fetch_ranked_content_blocking( + db: Arc, + user_id: Uuid, + user_key: &SecretKey, + candidates: Vec, + limit: usize, +) -> Result<(Vec, usize), ApiError> { let ids: Vec = candidates.iter().map(|candidate| candidate.id).collect(); - let mut conn = state - .db + let mut conn = db .get_pool() .get() .map_err(|_| ApiError::InternalServerError)?; @@ -1596,13 +1721,14 @@ pub async fn search_user_embeddings_with_options( }; let score_started_at = Instant::now(); - let candidates = top_k_candidates( - &query_vec, - &loaded.embeddings, + let candidates = top_k_candidates_blocking( + query_vec, + loaded.embeddings.clone(), candidate_limit, - options.filters.source_types.as_deref(), + options.filters.source_types.clone(), options.filters.conversation_id, - )?; + ) + .await?; let score_ms = score_started_at.elapsed().as_millis() as u64; let content_started_at = Instant::now(); @@ -1679,21 +1805,25 @@ pub async fn search_user_embeddings( } pub async fn delete_all_user_embeddings(state: &AppState, user_id: Uuid) -> Result<(), ApiError> { - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; + let db = state.db.clone(); + run_blocking_rag("rag_delete_all_user_embeddings", move || { + let mut conn = db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; - diesel::delete(user_embeddings::table.filter(user_embeddings::user_id.eq(user_id))) - .execute(&mut conn) - .map_err(|e| { - error!( - "Failed to delete all embeddings for user={}: {:?}", - user_id, e - ); - ApiError::InternalServerError - })?; + diesel::delete(user_embeddings::table.filter(user_embeddings::user_id.eq(user_id))) + .execute(&mut conn) + .map_err(|e| { + error!( + "Failed to delete all embeddings for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + Ok(()) + }) + .await?; state.rag_cache.lock().await.evict_user(user_id); Ok(()) @@ -1704,25 +1834,28 @@ pub async fn delete_user_embedding_by_uuid( user_id: Uuid, embedding_uuid: Uuid, ) -> Result<(), ApiError> { - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; + let db = state.db.clone(); + let affected = run_blocking_rag("rag_delete_user_embedding_by_uuid", move || { + let mut conn = db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; - let affected = diesel::delete( - user_embeddings::table - .filter(user_embeddings::user_id.eq(user_id)) - .filter(user_embeddings::uuid.eq(embedding_uuid)), - ) - .execute(&mut conn) - .map_err(|e| { - error!( - "Failed to delete embedding user={} uuid={}: {:?}", - user_id, embedding_uuid, e - ); - ApiError::InternalServerError - })?; + diesel::delete( + user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::uuid.eq(embedding_uuid)), + ) + .execute(&mut conn) + .map_err(|e| { + error!( + "Failed to delete embedding user={} uuid={}: {:?}", + user_id, embedding_uuid, e + ); + ApiError::InternalServerError + }) + }) + .await?; if affected == 0 { return Err(ApiError::NotFound); @@ -1740,62 +1873,65 @@ pub async fn embeddings_status( state: &AppState, user_id: Uuid, ) -> Result { - use diesel::dsl::count_star; + let db = state.db.clone(); + run_blocking_rag("rag_embeddings_status", move || { + use diesel::dsl::count_star; - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; + let mut conn = db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; - let total_embeddings: i64 = user_embeddings::table - .filter(user_embeddings::user_id.eq(user_id)) - .select(count_star()) - .first(&mut conn) - .map_err(|e| { - error!( - "Failed to count embeddings for user={} (total): {:?}", - user_id, e - ); - ApiError::InternalServerError - })?; + let total_embeddings: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .select(count_star()) + .first(&mut conn) + .map_err(|e| { + error!( + "Failed to count embeddings for user={} (total): {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; - let grouped: Vec<(String, i64)> = user_embeddings::table - .filter(user_embeddings::user_id.eq(user_id)) - .group_by(user_embeddings::embedding_model) - .select((user_embeddings::embedding_model, count_star())) - .load(&mut conn) - .map_err(|e| { - error!( - "Failed to group embeddings by model for user={}: {:?}", - user_id, e - ); - ApiError::InternalServerError - })?; + let grouped: Vec<(String, i64)> = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .group_by(user_embeddings::embedding_model) + .select((user_embeddings::embedding_model, count_star())) + .load(&mut conn) + .map_err(|e| { + error!( + "Failed to group embeddings by model for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; - let mut by_model: HashMap = HashMap::new(); - for (model, count) in grouped { - by_model.insert(model, count); - } + let mut by_model: HashMap = HashMap::new(); + for (model, count) in grouped { + by_model.insert(model, count); + } - let stale_count: i64 = user_embeddings::table - .filter(user_embeddings::user_id.eq(user_id)) - .filter(user_embeddings::embedding_model.ne(DEFAULT_EMBEDDING_MODEL)) - .select(count_star()) - .first(&mut conn) - .map_err(|e| { - error!( - "Failed to count stale embeddings for user={}: {:?}", - user_id, e - ); - ApiError::InternalServerError - })?; + let stale_count: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::embedding_model.ne(DEFAULT_EMBEDDING_MODEL)) + .select(count_star()) + .first(&mut conn) + .map_err(|e| { + error!( + "Failed to count stale embeddings for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; - Ok(RagEmbeddingsStatus { - total_embeddings, - by_model, - stale_count, + Ok(RagEmbeddingsStatus { + total_embeddings, + by_model, + stale_count, + }) }) + .await } #[cfg(test)] From 21402cf8b67c6830b86b07a658e34f7c6a943ffe Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 7 Jul 2026 05:17:36 +0000 Subject: [PATCH 6/6] fix(rag): trim hot-path duplicate work --- src/rag.rs | 44 +++++++++++++++++++++++++++----------------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/src/rag.rs b/src/rag.rs index d2d9768c..02b1a87f 100644 --- a/src/rag.rs +++ b/src/rag.rs @@ -193,15 +193,16 @@ impl RagCache { self.lru.retain(|u| *u != user_id); } - pub fn get(&mut self, user_id: Uuid) -> Option<(Arc>, bool)> { + pub fn get(&mut self, user_id: Uuid) -> Option<(Arc>, bool, usize)> { self.evict_expired(); - let (loaded_at, embeddings, scan_limit_hit) = { + let (loaded_at, embeddings, scan_limit_hit, bytes) = { let entry = self.entries.get(&user_id)?; ( entry.loaded_at, entry.embeddings.clone(), entry.scan_limit_hit, + entry.bytes, ) }; @@ -211,7 +212,7 @@ impl RagCache { } self.touch(user_id); - Some((embeddings, scan_limit_hit)) + Some((embeddings, scan_limit_hit, bytes)) } pub fn put( @@ -762,6 +763,12 @@ struct StorageStats { stored_bytes: i64, } +#[derive(Debug)] +struct StorageLimitStats { + user: StorageStats, + project: StorageStats, +} + fn encrypted_embedding_bytes( vector_enc: &[u8], content_enc: &[u8], @@ -842,13 +849,13 @@ fn ensure_storage_limits( projected_new_bytes: i64, limits: RagLimits, ) -> Result<(), ApiError> { - let user_stats = ensure_pre_embedding_limits(conn, user, limits)?; - if user_stats.stored_bytes.saturating_add(projected_new_bytes) > limits.max_user_stored_bytes { + let stats = ensure_pre_embedding_limits(conn, user, limits)?; + if stats.user.stored_bytes.saturating_add(projected_new_bytes) > limits.max_user_stored_bytes { warn!( target: "rag", user_id = %user.uuid, - rows = user_stats.row_count, - stored_bytes = user_stats.stored_bytes, + rows = stats.user.row_count, + stored_bytes = stats.user.stored_bytes, projected_new_bytes, max_rows = limits.max_user_embeddings, max_bytes = limits.max_user_stored_bytes, @@ -857,8 +864,8 @@ fn ensure_storage_limits( return Err(ApiError::BadRequest); } - let project_stats = load_project_embedding_storage_stats(conn, user.project_id)?; - if project_stats + if stats + .project .stored_bytes .saturating_add(projected_new_bytes) > limits.max_project_stored_bytes @@ -866,8 +873,8 @@ fn ensure_storage_limits( warn!( target: "rag", project_id = user.project_id, - rows = project_stats.row_count, - stored_bytes = project_stats.stored_bytes, + rows = stats.project.row_count, + stored_bytes = stats.project.stored_bytes, projected_new_bytes, max_rows = limits.max_project_embeddings, max_bytes = limits.max_project_stored_bytes, @@ -883,7 +890,7 @@ fn ensure_pre_embedding_limits( conn: &mut diesel::PgConnection, user: &User, limits: RagLimits, -) -> Result { +) -> Result { use diesel::dsl::count_star; let recent_insert_count: i64 = user_embeddings::table @@ -941,7 +948,10 @@ fn ensure_pre_embedding_limits( return Err(ApiError::BadRequest); } - Ok(user_stats) + Ok(StorageLimitStats { + user: user_stats, + project: project_stats, + }) } fn cached_embedding_from_inserted( @@ -1412,13 +1422,13 @@ async fn load_cacheable_embeddings_with_coordination( loop { let permit = { let mut cache = state.rag_cache.lock().await; - if let Some((embeddings, scan_limit_hit)) = cache.get(user_id) { + if let Some((embeddings, scan_limit_hit, cache_bytes_per_user)) = cache.get(user_id) { info!( target: "rag", user_id = %user_id, cache_hit = true, loaded_rows = embeddings.len(), - cache_bytes_per_user = cached_embeddings_bytes(&embeddings), + cache_bytes_per_user, total_cache_bytes = cache.total_bytes, scan_limit_hit, "rag_cache_lookup" @@ -2213,7 +2223,7 @@ mod tests { assert!(cache.put(user, Arc::new(vec![first]), false)); assert_eq!(cache.append(user, second), CacheAppendResult::Appended); - let (cached, _) = cache.get(user).unwrap(); + let (cached, _, _) = cache.get(user).unwrap(); assert_eq!(cached.len(), 2); assert_eq!(cached[1].id, 2); } @@ -2245,7 +2255,7 @@ mod tests { assert!(cache.put(user, Arc::new(vec![first, second]), false)); cache.remove_embedding_by_uuid(user, second_uuid); - let (cached, _) = cache.get(user).unwrap(); + let (cached, _, _) = cache.get(user).unwrap(); assert_eq!(cached.len(), 1); assert_eq!(cached[0].id, 1); }