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..bb41b7e7 --- /dev/null +++ b/migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql @@ -0,0 +1,3 @@ +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 new file mode 100644 index 00000000..8407e441 --- /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 CONCURRENTLY idx_user_embeddings_user_model_created_id + ON user_embeddings(user_id, embedding_model, created_at DESC, id DESC); + +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 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/main.rs b/src/main.rs index 86cc5ff6..baaa0206 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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())), @@ -1706,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(); @@ -2096,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() { 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..02b1a87f --- /dev/null +++ b/src/rag.rs @@ -0,0 +1,2502 @@ +#![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; +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 tokio::sync::Notify; +use tokio::task; +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 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 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"; +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 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, + pub by_model: HashMap, + pub stale_count: i64, +} + +#[derive(Debug, Clone)] +pub struct CachedEmbedding { + pub id: i64, + pub uuid: Uuid, + pub source_type: String, + pub conversation_id: Option, + pub vector: 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_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::from_env() + } +} + +impl RagCache { + 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_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.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<(Arc>, bool, usize)> { + self.evict_expired(); + + 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, + ) + }; + + if loaded_at.elapsed() > self.ttl { + self.evict_user_with_reason(user_id, "ttl"); + return None; + } + + self.touch(user_id); + Some((embeddings, scan_limit_hit, bytes)) + } + + 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 + } + + 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(); + } + } + + 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_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, +} + +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 { + 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 { + 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 { + id: i64, + uuid: Uuid, + score: f32, + token_count: i32, +} + +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.id == other.id + } +} + +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)) + .then_with(|| other.id.cmp(&self.id)) + } +} + +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 { + id: e.id, + uuid: e.uuid, + score, + token_count: e.token_count, + }; + 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)) + .then_with(|| min.id.cmp(&e.id)); + if ordering == Ordering::Greater { + heap.pop(); + let item = HeapItem { + id: e.id, + uuid: e.uuid, + score, + token_count: e.token_count, + }; + 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) +} + +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(); + 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(), + } +} + +#[derive(Debug, QueryableByName)] +struct StorageStats { + #[diesel(sql_type = diesel::sql_types::BigInt)] + row_count: i64, + #[diesel(sql_type = diesel::sql_types::BigInt)] + stored_bytes: i64, +} + +#[derive(Debug)] +struct StorageLimitStats { + user: StorageStats, + project: StorageStats, +} + +fn encrypted_embedding_bytes( + vector_enc: &[u8], + content_enc: &[u8], + metadata_enc: Option<&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> { + 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 = 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, + "RAG user storage limit reached" + ); + return Err(ApiError::BadRequest); + } + + if stats + .project + .stored_bytes + .saturating_add(projected_new_bytes) + > limits.max_project_stored_bytes + { + warn!( + target: "rag", + project_id = user.project_id, + 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, + "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 + .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 >= limits.max_user_stored_bytes + { + warn!( + target: "rag", + user_id = %user.uuid, + rows = user_stats.row_count, + stored_bytes = user_stats.stored_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 >= limits.max_project_stored_bytes + { + warn!( + target: "rag", + project_id = user.project_id, + rows = project_stats.row_count, + stored_bytes = project_stats.stored_bytes, + max_rows = limits.max_project_embeddings, + max_bytes = limits.max_project_stored_bytes, + "RAG project storage limit reached" + ); + return Err(ApiError::BadRequest); + } + + Ok(StorageLimitStats { + user: user_stats, + project: project_stats, + }) +} + +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, + } +} + +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, + 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 limits = RagLimits::from_env(); + validate_text_size(text, limits.max_insert_text_bytes, "insert")?; + 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?; + + 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 projected_new_bytes = + encrypted_embedding_bytes(&vector_enc, &content_enc, metadata_enc.as_ref()); + + let new_embedding = 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, + }; + 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); + debug!( + target: "rag", + user_id = %user_id, + embedding_uuid = %inserted.uuid, + ?append_result, + "rag_cache_append_after_insert" + ); + 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 limits = RagLimits::from_env(); + validate_text_size(text, limits.max_insert_text_bytes, "insert")?; + ensure_pre_embedding_limits_blocking(state, user, limits).await?; + + 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 new_embedding = 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, + }; + 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); + debug!( + target: "rag", + user_id = %user_id, + embedding_uuid = %inserted.uuid, + ?append_result, + "rag_cache_append_after_insert" + ); + Ok(inserted) +} + +#[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() + } + + 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)] +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, + filters: LoadFilters<'_>, + scan_limit: i64, +) -> Result { + 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.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.as_ref() { + query = query.filter(user_embeddings::tags_enc.overlaps_with(tags_enc_filter.clone())); + } + + 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; + } + }; + + 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, + }); + } + + 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_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, + 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; + 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); + return loaded; + } + } + } +} + +async fn fetch_ranked_content( + state: &AppState, + user_id: Uuid, + user_key: &SecretKey, + candidates: Vec, + limit: usize, +) -> Result<(Vec, usize), ApiError> { + if candidates.is_empty() { + 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 = db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + 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 + })?; + + 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 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" + ); + continue; + } + }; + + 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; + } + }; + + results.push(RagSearchResult { + content, + score: candidate.score, + token_count: candidate.token_count, + }); + + if results.len() >= limit { + break; + } + } + + Ok((results, skipped_rows)) +} + +pub async fn search_user_embeddings_with_options( + state: &Arc, + user: &User, + auth_method: AuthMethod, + user_key: &SecretKey, + query: &str, + 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 = 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 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, + }; + + 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_blocking( + query_vec, + loaded.embeddings.clone(), + candidate_limit, + 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(); + 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) = options.max_tokens { + results = apply_token_budget(results, budget); + } + + 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> { + 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 + })?; + Ok(()) + }) + .await?; + + 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 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)?; + + 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); + } + + state + .rag_cache + .lock() + .await + .remove_embedding_by_uuid(user_id, embedding_uuid); + Ok(()) +} + +pub async fn embeddings_status( + state: &AppState, + user_id: Uuid, +) -> Result { + let db = state.db.clone(); + run_blocking_rag("rag_embeddings_status", move || { + use diesel::dsl::count_star; + + 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 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, + }) + }) + .await +} + +#[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, + 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]; + 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![ + 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( + &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].id, 3); + assert_eq!(items[1].id, 2); + } + + #[test] + fn top_k_tie_break_prefers_fewer_tokens() { + let query = vec![1.0f32, 0.0]; + + let embeddings = vec![ + 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].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]; + + let embeddings = vec![ + 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(); + 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(1024 * 1024, 1024 * 1024, Duration::from_secs(60)); + let user_id = Uuid::new_v4(); + + cache.put(user_id, Arc::new(vec![]), false); + 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_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![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(); + + assert!(cache.put(u1, v1, false)); + assert!(cache.put(u2, v2, false)); + // touch u1 so u2 becomes LRU + cache.get(u1); + assert!(cache.put(u3, v3, false)); + + 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(1024 * 1024, 1024 * 1024, Duration::from_millis(5)); + let user = Uuid::new_v4(); + + 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); + } + + #[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") + } +} diff --git a/src/web/mod.rs b/src/web/mod.rs index dd9a5cc0..5448a9a2 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -13,6 +13,7 @@ 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; 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 }