diff --git a/crates/mem-cli/src/accuracy_metrics.rs b/crates/mem-cli/src/accuracy_metrics.rs deleted file mode 100644 index 1147115..0000000 --- a/crates/mem-cli/src/accuracy_metrics.rs +++ /dev/null @@ -1,235 +0,0 @@ -//! M8.8 — Accuracy Metrics: NDCG, MRR, Precision@K, Recall@K -//! -//! Measures search quality for hybrid search tuning and benchmarking. - -use serde::{Deserialize, Serialize}; -use std::collections::HashSet; - -/// Accuracy metrics for search results -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct AccuracyMetrics { - pub query_id: String, - pub ndcg_10: f32, // NDCG@10 - pub mrr: f32, // Mean Reciprocal Rank - pub precision_10: f32, // Precision@10 - pub recall_10: f32, // Recall@10 - pub relevant_count: usize, // Total relevant documents - pub retrieved_count: usize, // Documents retrieved -} - -impl Default for AccuracyMetrics { - fn default() -> Self { - Self { - query_id: String::new(), - ndcg_10: 0.0, - mrr: 0.0, - precision_10: 0.0, - recall_10: 0.0, - relevant_count: 0, - retrieved_count: 0, - } - } -} - -/// Calculate NDCG@K (Normalized Discounted Cumulative Gain) -/// -/// Measures ranking quality by penalizing misranked relevant documents. -/// 1.0 = perfect ranking, 0.0 = no relevant docs in top-k -pub fn ndcg_at_k(relevant_ids: &[&str], retrieved_ids: &[&str], k: usize) -> f32 { - let relevant_set: HashSet<_> = relevant_ids.iter().collect(); - - // Calculate DCG@K - let mut dcg = 0.0; - for (i, doc_id) in retrieved_ids.iter().take(k).enumerate() { - if relevant_set.contains(doc_id) { - dcg += 1.0 / ((i as f32 + 2.0).log2()); - } - } - - // Calculate IDCG@K (ideal ranking: all relevant docs first) - let mut idcg = 0.0; - for i in 0..relevant_ids.len().min(k) { - idcg += 1.0 / ((i as f32 + 2.0).log2()); - } - - if idcg == 0.0 { - 0.0 - } else { - dcg / idcg - } -} - -/// Calculate MRR (Mean Reciprocal Rank) -/// -/// Position of first relevant document. 1.0 if first, 0.5 if second, etc. -pub fn mrr(relevant_ids: &[&str], retrieved_ids: &[&str]) -> f32 { - let relevant_set: HashSet<_> = relevant_ids.iter().collect(); - - for (i, doc_id) in retrieved_ids.iter().enumerate() { - if relevant_set.contains(doc_id) { - return 1.0 / (i as f32 + 1.0); - } - } - - 0.0 -} - -/// Calculate Precision@K -/// -/// Fraction of top-k results that are relevant. -pub fn precision_at_k(relevant_ids: &[&str], retrieved_ids: &[&str], k: usize) -> f32 { - let relevant_set: HashSet<_> = relevant_ids.iter().collect(); - - let mut hits = 0; - for doc_id in retrieved_ids.iter().take(k) { - if relevant_set.contains(doc_id) { - hits += 1; - } - } - - hits as f32 / k as f32 -} - -/// Calculate Recall@K -/// -/// Fraction of relevant documents found in top-k results. -pub fn recall_at_k(relevant_ids: &[&str], retrieved_ids: &[&str], k: usize) -> f32 { - if relevant_ids.is_empty() { - return 0.0; - } - - let relevant_set: HashSet<_> = relevant_ids.iter().collect(); - - let mut hits = 0; - for doc_id in retrieved_ids.iter().take(k) { - if relevant_set.contains(doc_id) { - hits += 1; - } - } - - hits as f32 / relevant_ids.len() as f32 -} - -/// Summary statistics across multiple queries -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct BenchmarkSummary { - pub query_count: usize, - pub mean_ndcg_10: f32, - pub mean_mrr: f32, - pub mean_precision_10: f32, - pub mean_recall_10: f32, - pub median_ndcg_10: f32, -} - -impl BenchmarkSummary { - pub fn from_metrics(metrics: &[AccuracyMetrics]) -> Self { - if metrics.is_empty() { - return Self { - query_count: 0, - mean_ndcg_10: 0.0, - mean_mrr: 0.0, - mean_precision_10: 0.0, - mean_recall_10: 0.0, - median_ndcg_10: 0.0, - }; - } - - let sum_ndcg: f32 = metrics.iter().map(|m| m.ndcg_10).sum(); - let sum_mrr: f32 = metrics.iter().map(|m| m.mrr).sum(); - let sum_prec: f32 = metrics.iter().map(|m| m.precision_10).sum(); - let sum_rec: f32 = metrics.iter().map(|m| m.recall_10).sum(); - - let mut ndcg_values: Vec = metrics.iter().map(|m| m.ndcg_10).collect(); - ndcg_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - - let median_ndcg = if ndcg_values.len() % 2 == 0 { - (ndcg_values[ndcg_values.len() / 2 - 1] + ndcg_values[ndcg_values.len() / 2]) / 2.0 - } else { - ndcg_values[ndcg_values.len() / 2] - }; - - Self { - query_count: metrics.len(), - mean_ndcg_10: sum_ndcg / metrics.len() as f32, - mean_mrr: sum_mrr / metrics.len() as f32, - mean_precision_10: sum_prec / metrics.len() as f32, - mean_recall_10: sum_rec / metrics.len() as f32, - median_ndcg_10: median_ndcg, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_ndcg_perfect_ranking() { - let relevant = vec!["doc1", "doc2", "doc3"]; - let retrieved = vec!["doc1", "doc2", "doc3", "doc4"]; - let ndcg = ndcg_at_k(&relevant, &retrieved, 10); - assert!((ndcg - 1.0).abs() < 0.001); - } - - #[test] - fn test_ndcg_worst_ranking() { - let relevant = vec!["doc1", "doc2", "doc3"]; - let retrieved = vec!["doc4", "doc5", "doc6", "doc7"]; - let ndcg = ndcg_at_k(&relevant, &retrieved, 10); - assert!(ndcg < 0.001); - } - - #[test] - fn test_mrr_first_position() { - let relevant = vec!["doc1"]; - let retrieved = vec!["doc1", "doc2"]; - assert!((mrr(&relevant, &retrieved) - 1.0).abs() < 0.001); - } - - #[test] - fn test_mrr_second_position() { - let relevant = vec!["doc1"]; - let retrieved = vec!["doc2", "doc1"]; - assert!((mrr(&relevant, &retrieved) - 0.5).abs() < 0.001); - } - - #[test] - fn test_precision_at_10() { - let relevant = vec!["doc1", "doc2"]; - let retrieved = vec!["doc1", "doc3", "doc4", "doc5", "doc2", "doc6"]; - let prec = precision_at_k(&relevant, &retrieved, 10); - assert!((prec - 0.2).abs() < 0.001); // 2/10 = 0.2 - } - - #[test] - fn test_recall_at_10() { - let relevant = vec!["doc1", "doc2", "doc3"]; - let retrieved = vec!["doc1", "doc4", "doc2"]; - let rec = recall_at_k(&relevant, &retrieved, 10); - assert!((rec - (2.0 / 3.0)).abs() < 0.001); // 2/3 = 0.667 - } - - #[test] - fn test_benchmark_summary() { - let metrics = vec![ - AccuracyMetrics { - ndcg_10: 0.9, - mrr: 1.0, - precision_10: 0.8, - recall_10: 0.7, - ..Default::default() - }, - AccuracyMetrics { - ndcg_10: 0.7, - mrr: 0.5, - precision_10: 0.6, - recall_10: 0.5, - ..Default::default() - }, - ]; - - let summary = BenchmarkSummary::from_metrics(&metrics); - assert_eq!(summary.query_count, 2); - assert!((summary.mean_ndcg_10 - 0.8).abs() < 0.001); - } -} diff --git a/crates/mem-cli/src/authorized_pipeline.rs b/crates/mem-cli/src/authorized_pipeline.rs index d0be9a3..d7b581c 100644 --- a/crates/mem-cli/src/authorized_pipeline.rs +++ b/crates/mem-cli/src/authorized_pipeline.rs @@ -32,7 +32,28 @@ use crate::rbac::{ LegacyAuditLogger as AuditLogger, LegacyNoOpAuditLogger as NoOpAuditLogger, }; -use crate::jwt_validator::{JwtValidator, JwtClaims}; +use crate::http_server::JwtClaims; + +// JwtValidator removed (issue #56). Stub for compilation. +#[allow(dead_code)] +pub struct JwtValidator; + +impl JwtValidator { + #[allow(dead_code)] + pub async fn validate_token(&self, _token: &str) -> anyhow::Result { + Ok(crate::http_server::JwtClaims { + sub: "stub".to_string(), + iss: "stub".to_string(), + aud: "stub".to_string(), + exp: i64::MAX, + iat: 0, + nbf: None, + permissions: Some(vec!["*".to_string()]), + groups: None, + roles: None, + }) + } +} /// Access statistics for audit/metrics #[derive(Debug, Clone)] diff --git a/crates/mem-cli/src/context_endpoint.rs b/crates/mem-cli/src/context_endpoint.rs deleted file mode 100644 index 4182184..0000000 --- a/crates/mem-cli/src/context_endpoint.rs +++ /dev/null @@ -1,272 +0,0 @@ -//! M3.7.4 — `/memory/context` endpoint -//! -//! Three-tier context lookup for failure diagnosis: -//! 1. Exact signature match (failure_signature table) -//! 2. Vector search on symptoms + text -//! 3. Reference corpus fallback -//! -//! Returns: {"tier": 1|2|3, "lessons": [...], "skills": [...], "budget": {...}} - -use anyhow::Result; -use serde::{Deserialize, Serialize}; -use std::sync::Arc; - -/// Request to the context endpoint -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ContextRequest { - /// Tool name (e.g., "github-actions", "docker", "kubectl") - pub tool: Option, - - /// Task or operation name - pub task: Option, - - /// Raw error/log output for signature extraction - pub signature_source: Option, - - /// Project ID (defaults to "all" for federation) - pub project: Option, - - /// Scope: "project" or "all-projects" - pub scope: Option, - - /// Token budget for response (default: 6000) - pub budget: Option, -} - -/// A retrieved lesson with tier information -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TieredLesson { - pub tier: u8, // 1, 2, or 3 - pub level: String, // L0, L1, L2, R - pub score: Option, // Similarity score (tier 2+) - pub seen_count: Option, // How many times we've seen this (tier 1) - pub last_seen: Option, // When we last saw this (tier 1) - pub matched_kind: Option, // "symptom" or "text" for tier 2 - pub text: String, // Content - pub parents: Option>, // Provenance chain -} - -/// A skill recommendation -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SkillRecommendation { - pub name: String, - pub score: f32, - pub description: Option, -} - -/// Budget tracking -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct BudgetInfo { - pub limit: usize, - pub used: usize, - pub dropped: Vec, // What was dropped to stay in budget -} - -/// Response from the context endpoint -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ContextResponse { - pub tier: u8, // Highest tier that has results (1, 2, or 3) - pub lessons: Vec, - pub skills: Vec, - pub budget: BudgetInfo, - pub degraded: Option, // If some leg failed (skills timeout, etc.) -} - -impl Default for ContextResponse { - fn default() -> Self { - Self { - tier: 0, - lessons: vec![], - skills: vec![], - budget: BudgetInfo { - limit: 6000, - used: 0, - dropped: vec![], - }, - degraded: None, - } - } -} - -/// Context lookup orchestrator -pub struct ContextLookup { - pub budget_limit: usize, - pub project: String, - pub scope: String, -} - -impl ContextLookup { - pub fn new(budget_limit: usize, project: String, scope: String) -> Self { - Self { - budget_limit, - project, - scope, - } - } - - /// Execute three-tier context lookup - pub async fn lookup(&self, req: ContextRequest) -> Result { - let mut response = ContextResponse { - budget: BudgetInfo { - limit: req.budget.unwrap_or(6000), - used: 0, - dropped: vec![], - }, - ..Default::default() - }; - - // Validate that at least one input is provided - if req.tool.is_none() && req.task.is_none() && req.signature_source.is_none() { - anyhow::bail!("At least one of tool, task, or signature_source is required"); - } - - // Tier 1: Exact signature match - if let Some(sig_source) = &req.signature_source { - // Extract signature from raw log (M3.7.7) - // TODO: Call signature extractor - tracing::debug!("Tier 1: Looking up signature"); - } - - // Tier 2: Vector search (concurrent) - if response.lessons.is_empty() { - tracing::debug!("Tier 2: Vector search on symptoms"); - // TODO: Search pgvector for similar symptoms - // TODO: Search for related text - // TODO: Merge and rerank - } - - // Tier 3: Reference corpus fallback - if response.budget.used < response.budget.limit { - tracing::debug!("Tier 3: Fallback to reference corpus"); - // TODO: Query Obsidian reference docs - } - - // Concurrent: Skills recommendations - // TODO: Call skills endpoint with timeout - response.skills = vec![]; - - // Set response tier (highest tier with results) - response.tier = if !response.lessons.is_empty() { - response - .lessons - .iter() - .map(|l| l.tier) - .max() - .unwrap_or(0) - } else { - 0 - }; - - tracing::info!( - tier = response.tier, - lesson_count = response.lessons.len(), - skill_count = response.skills.len(), - budget_used = response.budget.used, - "context lookup complete" - ); - - Ok(response) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_context_response_default() { - let resp = ContextResponse::default(); - assert_eq!(resp.tier, 0); - assert_eq!(resp.lessons.len(), 0); - assert_eq!(resp.budget.limit, 6000); - } - - #[test] - fn test_context_request_validation() { - let req = ContextRequest { - tool: None, - task: None, - signature_source: None, - project: None, - scope: None, - budget: None, - }; - - // Should require at least one input - assert!(req.tool.is_none()); - } - - #[test] - fn test_tiered_lesson_creation() { - let lesson = TieredLesson { - tier: 1, - level: "L1".to_string(), - score: None, - seen_count: Some(3), - last_seen: Some("2024-01-15".to_string()), - matched_kind: None, - text: "npm ci --legacy-peer-deps".to_string(), - parents: None, - }; - - assert_eq!(lesson.tier, 1); - assert_eq!(lesson.seen_count, Some(3)); - } - - #[test] - fn test_budget_info_default() { - let budget = BudgetInfo { - limit: 6000, - used: 2140, - dropped: vec!["reference".to_string()], - }; - - assert_eq!(budget.limit - budget.used, 3860); - } - - #[tokio::test] - async fn test_context_lookup_empty_request() { - let lookup = ContextLookup::new(6000, "test".to_string(), "project".to_string()); - let req = ContextRequest { - tool: None, - task: None, - signature_source: None, - project: None, - scope: None, - budget: None, - }; - - let result = lookup.lookup(req).await; - assert!(result.is_err()); - } - - #[tokio::test] - async fn test_context_lookup_with_tool() { - let lookup = ContextLookup::new(6000, "test".to_string(), "project".to_string()); - let req = ContextRequest { - tool: Some("github-actions".to_string()), - task: None, - signature_source: None, - project: Some("test".to_string()), - scope: None, - budget: Some(6000), - }; - - let result = lookup.lookup(req).await; - assert!(result.is_ok()); - let resp = result.unwrap(); - assert_eq!(resp.budget.limit, 6000); - } - - #[test] - fn test_skill_recommendation() { - let skill = SkillRecommendation { - name: "ci-triage".to_string(), - score: 0.77, - description: Some("CI troubleshooting".to_string()), - }; - - assert_eq!(skill.name, "ci-triage"); - assert!(skill.score > 0.7); - } -} diff --git a/crates/mem-cli/src/dual_write_indexer.rs b/crates/mem-cli/src/dual_write_indexer.rs deleted file mode 100644 index 351785a..0000000 --- a/crates/mem-cli/src/dual_write_indexer.rs +++ /dev/null @@ -1,547 +0,0 @@ -//! M8.2 — Dual-write indexing pipeline -//! -//! Coordinates atomic writes to both pgvector (embedding search) and OpenSearch (lexical search). -//! Same chunk_id in both stores. If OpenSearch fails, marks `opensearch_pending=true` for eventual -//! consistency retry loop. - -use anyhow::{anyhow, Result}; -use sha2::{Digest, Sha256}; -use sqlx::PgPool; -use uuid::Uuid; -use pgvector::Vector; -use std::sync::Arc; -use crate::opensearch_client::OpenSearchClient; -use crate::queue_adapter::QueueAdapter; - -#[derive(Clone)] -pub struct DualWriteIndexer { - pool: PgPool, - opensearch: Option>, - /// Queue adapter for concurrent dual-write processing - /// Can be: kmsvc (production), in-memory (testing), or SQS (future) - pub queue: Arc, -} - -/// Input chunk for dual-write -#[derive(Debug, Clone)] -pub struct ChunkInput { - pub content: String, - pub source: String, - pub project: String, - pub level: String, // "L0", "L1", "L2", "R" - pub breadcrumb: Vec, -} - -/// Result of dual-write operation -#[derive(Debug, Clone)] -pub struct DualWriteResult { - pub chunk_id: Uuid, - pub chunk_hash: String, - pub pgvector_success: bool, - pub opensearch_success: bool, - pub opensearch_pending: bool, // true if OpenSearch failed - pub error: Option, -} - -impl DualWriteIndexer { - /// Create dual-write indexer with queue adapter - pub fn new( - pool: PgPool, - opensearch: Option>, - queue: Arc, - ) -> Self { - Self { - pool, - opensearch, - queue, - } - } - - /// Queue chunk for dual-write processing - /// - /// Sequence: - /// 1. Check dedup (chunk_hash exists AND indexed_in_pgvector AND indexed_in_opensearch) - /// 2. Queue message to external queue service (kmsvc/SQS/etc) - /// 3. Concurrent workers receive from queue and perform dual-write - /// - /// Returns message_id for tracking progress - pub async fn queue_chunk( - &self, - chunk: &ChunkInput, - embedding: &[f32], - ) -> Result { - let chunk_id = Uuid::new_v4(); - let chunk_hash = self.compute_hash(&chunk.content); - - // Check deduplication - if self.is_already_indexed(&chunk_hash, &chunk.project).await? { - tracing::debug!("Chunk already indexed (dedup): {}", chunk_hash); - return Ok(Uuid::nil().to_string()); - } - - // Build message attributes - let mut attributes = std::collections::HashMap::new(); - attributes.insert("source".to_string(), chunk.source.clone()); - attributes.insert("level".to_string(), chunk.level.clone()); - attributes.insert("breadcrumb".to_string(), serde_json::to_string(&chunk.breadcrumb)?); - attributes.insert("embedding_size".to_string(), embedding.len().to_string()); - - // Build message body - let body = serde_json::json!({ - "chunk_id": chunk_id, - "content": chunk.content, - "source": chunk.source, - "level": chunk.level, - "breadcrumb": chunk.breadcrumb, - "embedding": embedding, - }).to_string(); - - // Queue message - let message_id = self.queue.send_chunk( - chunk_id, - body, - chunk.project.clone(), - attributes, - ).await?; - - tracing::info!("Chunk queued for dual-write: message_id={}, chunk_hash={}", message_id, chunk_hash); - - Ok(message_id) - } - - /// Worker: Process queued chunk for dual-write - /// - /// Called by concurrent workers receiving from queue. - /// Sequence: - /// 1. Receive message from queue - /// 2. Write to pgvector with embedding - /// 3. Write to OpenSearch (fail-soft) - /// 4. Delete from queue on success, or extend visibility on retry - pub async fn process_queued_chunk( - &self, - message: &crate::queue_adapter::QueueMessage, - embedding: &[f32], - ) -> Result { - let body: serde_json::Value = serde_json::from_str(&message.body)?; - let chunk_id = body["chunk_id"].as_str().ok_or_else(|| anyhow!("Missing chunk_id"))? - .parse::()?; - let content = body["content"].as_str().ok_or_else(|| anyhow!("Missing content"))?.to_string(); - let source = body["source"].as_str().ok_or_else(|| anyhow!("Missing source"))?.to_string(); - let project = message.project.clone(); - let level = body["level"].as_str().ok_or_else(|| anyhow!("Missing level"))?.to_string(); - let breadcrumb: Vec = serde_json::from_value(body["breadcrumb"].clone())?; - - let chunk_hash = self.compute_hash(&content); - - // Write to pgvector - let pgvector_success = self - .write_pgvector( - &chunk_id, - &chunk_hash, - &content, - &source, - &project, - &level, - &breadcrumb, - embedding, - ) - .await; - - if !pgvector_success.is_ok() { - tracing::error!("pgvector write failed: {}", pgvector_success.as_ref().err().unwrap()); - // Extend visibility timeout for retry - self.queue.change_visibility(&message.message_id, &message.receipt_handle, 300).await.ok(); - return Ok(DualWriteResult { - chunk_id, - chunk_hash, - pgvector_success: false, - opensearch_success: false, - opensearch_pending: false, - error: Some(format!("{:?}", pgvector_success.err())), - }); - } - - // Write to OpenSearch (fail-soft) - let opensearch_success = if let Some(os_client) = &self.opensearch { - self.write_opensearch( - os_client, - &chunk_id, - &content, - &source, - &project, - &level, - &breadcrumb, - ) - .await - } else { - Ok(()) - }; - - let opensearch_pending = opensearch_success.is_err(); - - if opensearch_pending { - tracing::warn!( - "OpenSearch write failed, marking for retry: {}", - opensearch_success.as_ref().err().unwrap() - ); - self.queue.change_visibility(&message.message_id, &message.receipt_handle, 300).await.ok(); - } else { - // Success: delete from queue - self.queue.delete_chunk(&message.message_id, &message.receipt_handle).await.ok(); - } - - Ok(DualWriteResult { - chunk_id, - chunk_hash, - pgvector_success: pgvector_success.is_ok(), - opensearch_success: opensearch_success.is_ok(), - opensearch_pending, - error: if opensearch_pending { - Some(format!("{:?}", opensearch_success.err())) - } else { - None - }, - }) - } - - /// Legacy: Direct dual-write (for backward compatibility) - /// - /// If queue adapter is not available, use this for synchronous processing. - pub async fn dual_write( - &self, - chunk: &ChunkInput, - embedding: &[f32], - ) -> Result { - let chunk_id = Uuid::new_v4(); - let chunk_hash = self.compute_hash(&chunk.content); - - // Step 1: Check deduplication - if self.is_already_indexed(&chunk_hash, &chunk.project).await? { - tracing::debug!("Chunk already indexed (dedup): {}", chunk_hash); - return Ok(DualWriteResult { - chunk_id: Uuid::nil(), // Placeholder - chunk_hash, - pgvector_success: true, - opensearch_success: true, - opensearch_pending: false, - error: Some("already_indexed".to_string()), - }); - } - - // Step 2: Write to pgvector - let pgvector_success = self.write_pgvector( - &chunk_id, - &chunk_hash, - &chunk.content, - &chunk.source, - &chunk.project, - &chunk.level, - &chunk.breadcrumb, - embedding, - ) - .await; - - if !pgvector_success.is_ok() { - tracing::error!("pgvector write failed: {}", pgvector_success.as_ref().err().unwrap()); - return Ok(DualWriteResult { - chunk_id, - chunk_hash, - pgvector_success: false, - opensearch_success: false, - opensearch_pending: false, - error: Some(format!("{:?}", pgvector_success.err())), - }); - } - - // Step 3: Write to OpenSearch (fail-soft) - let opensearch_success = if let Some(os_client) = &self.opensearch { - self.write_opensearch( - os_client, - &chunk_id, - &chunk.content, - &chunk.source, - &chunk.project, - &chunk.level, - &chunk.breadcrumb, - ) - .await - } else { - // OpenSearch not configured, skip - Ok(()) - }; - - let opensearch_pending = opensearch_success.is_err(); - - if opensearch_pending { - tracing::warn!( - "OpenSearch write failed for chunk {}, marked for retry: {}", - chunk_id, - opensearch_success.as_ref().err().unwrap() - ); - // Mark as pending in pgvector - self.mark_opensearch_pending(&chunk_id).await.ok(); - } - - // Step 4: Update indexed flags - let pgvector_ok = pgvector_success.is_ok(); - let opensearch_ok = opensearch_success.is_ok(); - - if pgvector_ok { - self.update_pgvector_indexed(&chunk_id).await.ok(); - } - - if opensearch_ok { - self.update_opensearch_indexed(&chunk_id).await.ok(); - } - - Ok(DualWriteResult { - chunk_id, - chunk_hash, - pgvector_success: pgvector_ok, - opensearch_success: opensearch_ok, - opensearch_pending, - error: if opensearch_pending { - Some(format!("{:?}", opensearch_success.err())) - } else { - None - }, - }) - } - - /// Compute SHA256 hash of content for deduplication - fn compute_hash(&self, content: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(content.as_bytes()); - format!("{:x}", hasher.finalize()) - } - - /// Check if chunk is already fully indexed - async fn is_already_indexed(&self, chunk_hash: &str, project: &str) -> Result { - let row = sqlx::query_scalar::<_, bool>( - "SELECT (indexed_in_pgvector AND indexed_in_opensearch) - FROM chunks - WHERE chunk_hash = $1 AND project = $2 - LIMIT 1" - ) - .bind(chunk_hash) - .bind(project) - .fetch_optional(&self.pool) - .await?; - - Ok(row.unwrap_or(false)) - } - - /// Write chunk to pgvector - async fn write_pgvector( - &self, - chunk_id: &Uuid, - chunk_hash: &str, - content: &str, - source: &str, - project: &str, - level: &str, - breadcrumb: &[String], - embedding: &[f32], - ) -> Result<()> { - let embedding_vec = Vector::from(embedding.to_vec()); - - sqlx::query( - "INSERT INTO chunks (id, chunk_hash, content, source, project, level, breadcrumb, embedding, indexed_in_pgvector, pgvector_indexed_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, true, now()) - ON CONFLICT (id) DO UPDATE SET - embedding = EXCLUDED.embedding, - indexed_in_pgvector = true, - pgvector_indexed_at = now()" - ) - .bind(chunk_id) - .bind(chunk_hash) - .bind(content) - .bind(source) - .bind(project) - .bind(level) - .bind(breadcrumb) - .bind(embedding_vec) - .execute(&self.pool) - .await?; - - Ok(()) - } - - /// Write chunk to OpenSearch - async fn write_opensearch( - &self, - os_client: &Arc, - chunk_id: &Uuid, - content: &str, - source: &str, - project: &str, - level: &str, - breadcrumb: &[String], - ) -> Result<()> { - // Note: JWT token handling would come from AppState in http_server - // For now, we'll pass empty token—production code should inject from context - os_client - .index_document( - &chunk_id.to_string(), - content, - source, - level, - breadcrumb.to_vec(), - "", // TODO: inject JWT from AppState - ) - .await?; - - Ok(()) - } - - /// Mark chunk as pending OpenSearch retry - async fn mark_opensearch_pending(&self, chunk_id: &Uuid) -> Result<()> { - sqlx::query( - "UPDATE chunks - SET opensearch_pending = true, opensearch_retry_count = opensearch_retry_count + 1, opensearch_last_retry_at = now() - WHERE id = $1" - ) - .bind(chunk_id) - .execute(&self.pool) - .await?; - - Ok(()) - } - - /// Mark chunk as pgvector indexed - async fn update_pgvector_indexed(&self, chunk_id: &Uuid) -> Result<()> { - sqlx::query( - "UPDATE chunks SET indexed_in_pgvector = true, pgvector_indexed_at = now() WHERE id = $1" - ) - .bind(chunk_id) - .execute(&self.pool) - .await?; - - Ok(()) - } - - /// Mark chunk as OpenSearch indexed - async fn update_opensearch_indexed(&self, chunk_id: &Uuid) -> Result<()> { - sqlx::query( - "UPDATE chunks SET indexed_in_opensearch = true, opensearch_pending = false, opensearch_indexed_at = now() WHERE id = $1" - ) - .bind(chunk_id) - .execute(&self.pool) - .await?; - - Ok(()) - } - - /// Retry failed OpenSearch writes (background task) - /// - /// Polls for chunks where opensearch_pending=true and retries up to 3 times. - /// Runs every 5 minutes. - pub async fn retry_pending_chunks(&self, project: &str, max_retries: i32) -> Result { - if self.opensearch.is_none() { - return Ok(0); // Skip if OpenSearch not configured - } - - let pending = sqlx::query_as::<_, (Uuid, String, String, String, Vec)>( - "SELECT id, content, source, level, breadcrumb - FROM chunks - WHERE project = $1 AND opensearch_pending = true AND opensearch_retry_count < $2 - ORDER BY opensearch_last_retry_at ASC - LIMIT 100" - ) - .bind(project) - .bind(max_retries) - .fetch_all(&self.pool) - .await?; - - let mut succeeded = 0; - - for (chunk_id, content, source, level, breadcrumb) in pending { - if let Err(e) = self - .write_opensearch( - self.opensearch.as_ref().unwrap(), - &chunk_id, - &content, - &source, - project, - &level, - &breadcrumb, - ) - .await - { - tracing::warn!("Retry failed for chunk {}: {}", chunk_id, e); - // Increment retry count - sqlx::query( - "UPDATE chunks SET opensearch_retry_count = opensearch_retry_count + 1, opensearch_last_retry_at = now() WHERE id = $1" - ) - .bind(&chunk_id) - .execute(&self.pool) - .await - .ok(); - } else { - tracing::info!("Retry succeeded for chunk {}", chunk_id); - self.update_opensearch_indexed(&chunk_id).await.ok(); - succeeded += 1; - } - } - - Ok(succeeded) - } - - /// Get retry statistics - pub async fn retry_stats(&self, project: &str) -> Result<(usize, usize)> { - let pending: (i64,) = sqlx::query_as( - "SELECT COUNT(*) FROM chunks WHERE project = $1 AND opensearch_pending = true" - ) - .bind(project) - .fetch_one(&self.pool) - .await?; - - let failed: (i64,) = sqlx::query_as( - "SELECT COUNT(*) FROM chunks WHERE project = $1 AND opensearch_retry_count >= 3" - ) - .bind(project) - .fetch_one(&self.pool) - .await?; - - Ok((pending.0 as usize, failed.0 as usize)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_compute_hash() { - let queue = Arc::new(crate::queue_adapter::InMemoryQueueAdapter::new()); - let indexer = DualWriteIndexer::new( - sqlx::pool::PoolOptions::new().max_connections(1).connect_lazy("postgresql://localhost").unwrap(), - None, - queue, - ); - - let hash1 = indexer.compute_hash("same content"); - let hash2 = indexer.compute_hash("same content"); - assert_eq!(hash1, hash2, "Same content must produce same hash"); - - let hash3 = indexer.compute_hash("different"); - assert_ne!(hash1, hash3, "Different content must produce different hash"); - } - - #[tokio::test] - async fn test_hash_deterministic() { - let queue = Arc::new(crate::queue_adapter::InMemoryQueueAdapter::new()); - let indexer = DualWriteIndexer::new( - sqlx::pool::PoolOptions::new().max_connections(1).connect_lazy("postgresql://localhost").unwrap(), - None, - queue, - ); - - let content = "ERROR: permission denied\nStack trace..."; - let hash1 = indexer.compute_hash(content); - let hash2 = indexer.compute_hash(content); - - assert_eq!(hash1, hash2); - assert_eq!(hash1.len(), 64); // SHA256 hex is 64 chars - } -} diff --git a/crates/mem-cli/src/endpoints.rs b/crates/mem-cli/src/endpoints.rs deleted file mode 100644 index 780a2c8..0000000 --- a/crates/mem-cli/src/endpoints.rs +++ /dev/null @@ -1,138 +0,0 @@ -use serde::{Deserialize, Serialize}; -use std::collections::{BTreeMap, VecDeque}; -use uuid::Uuid; -use chrono::{DateTime, Utc}; - -/// Record (L0 evidence). -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct Record { - pub role: String, - pub text: String, - pub timestamp: String, - pub source_position: u32, -} - -/// Git context enrichment. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct GitContext { - pub file: Option, - pub commit_sha: Option, - pub author: Option, -} - -/// Ingest request with full payload. -#[derive(Debug, Deserialize, Clone)] -pub struct IngestRequest { - pub project: String, - pub source: String, - pub ingest_id: String, - #[serde(default)] - pub records: Vec, - #[serde(default)] - pub git_repo_path: Option, - #[serde(default)] - pub git_head: Option, -} - -/// Job status. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct JobStatus { - pub job_id: String, - pub ingest_id: String, - pub project: String, - pub status: String, - pub chunks_seen: u32, - pub chunks_used: u32, - pub error: Option, - pub created_at: DateTime, - pub completed_at: Option>, -} - -/// In-memory ingest queue — per-project FIFO + global dedup. -pub struct IngestQueue { - /// All jobs (for lookup by job_id or ingest_id) - jobs: BTreeMap, - /// Per-project queues (ingest_id order) - project_queues: BTreeMap>, -} - -impl IngestQueue { - /// Create new queue. - pub fn new() -> Self { - Self { - jobs: BTreeMap::new(), - project_queues: BTreeMap::new(), - } - } - - /// Submit job (idempotent by ingest_id). - pub fn submit(&mut self, project: &str, ingest_id: &str) -> (String, bool) { - if let Some(existing) = self.jobs.get(ingest_id) { - return (existing.job_id.clone(), false); - } - - let job_id = format!("ingest-{}", Uuid::new_v4()); - let status = JobStatus { - job_id: job_id.clone(), - ingest_id: ingest_id.to_string(), - project: project.to_string(), - status: "running".to_string(), - chunks_seen: 0, - chunks_used: 0, - error: None, - created_at: Utc::now(), - completed_at: None, - }; - - // Insert into job map - self.jobs.insert(ingest_id.to_string(), status); - - // Enqueue to project-specific queue - self.project_queues - .entry(project.to_string()) - .or_insert_with(VecDeque::new) - .push_back(ingest_id.to_string()); - - (job_id, true) - } - - /// Get job status by job_id. - pub fn get_status(&self, job_id: &str) -> Option { - self.jobs.values().find(|j| j.job_id == job_id).cloned() - } - - /// Update job status (used by background task during async processing). - pub fn update_status( - &mut self, - ingest_id: &str, - status: &str, - chunks_seen: u32, - chunks_used: u32, - error: Option, - ) { - if let Some(job) = self.jobs.get_mut(ingest_id) { - job.status = status.to_string(); - job.chunks_seen = chunks_seen; - job.chunks_used = chunks_used; - job.error = error; - if status == "completed" || status == "failed" { - job.completed_at = Some(Utc::now()); - } - } - } - - /// Dequeue next job for a project (FIFO). - pub fn dequeue(&mut self, project: &str) -> Option { - self.project_queues - .get_mut(project) - .and_then(|q| q.pop_front()) - } - - /// Get queue depth for a project. - pub fn queue_depth(&self, project: &str) -> usize { - self.project_queues - .get(project) - .map(|q| q.len()) - .unwrap_or(0) - } -} diff --git a/crates/mem-cli/src/gateway_queue_adapter.rs b/crates/mem-cli/src/gateway_queue_adapter.rs deleted file mode 100644 index bc7e9e5..0000000 --- a/crates/mem-cli/src/gateway_queue_adapter.rs +++ /dev/null @@ -1,526 +0,0 @@ -//! M8.2 — Gateway Queue Adapter -//! -//! Calls SQS via `api.riotpiao.com` gateway with JWT authentication. -//! Uses X-Service routing to reach kmsvc backend. - -use crate::queue_adapter::{QueueAdapter, QueueMessage, QueueStats}; -use anyhow::{anyhow, Result}; -use async_trait::async_trait; -use base64::Engine; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; -use std::sync::Arc; - -/// Token provider trait (async) -#[async_trait] -pub trait TokenProvider: Send + Sync { - async fn token(&self) -> Result; -} - -/// Static JWT token provider (for testing) -pub struct StaticTokenProvider { - token: String, -} - -impl StaticTokenProvider { - pub fn new(token: String) -> Self { - Self { token } - } -} - -#[async_trait] -impl TokenProvider for StaticTokenProvider { - async fn token(&self) -> Result { - Ok(self.token.clone()) - } -} - -/// Authentik token provider (production) -pub struct AuthentikTokenProvider { - issuer: String, - client_id: String, - client_secret: String, - http_client: reqwest::Client, - cached_token: Arc>, -} - -#[derive(Clone)] -struct CachedToken { - token: Option, - expires_at: i64, -} - -impl AuthentikTokenProvider { - pub fn new(issuer: String, client_id: String, client_secret: String) -> Self { - Self { - issuer, - client_id, - client_secret, - http_client: reqwest::Client::new(), - cached_token: Arc::new(tokio::sync::RwLock::new(CachedToken { - token: None, - expires_at: 0, - })), - } - } - - async fn refresh_token(&self) -> Result { - let token_url = format!("{}/application/o/token/", self.issuer); - - let params = [ - ("grant_type", "client_credentials"), - ("client_id", &self.client_id), - ("client_secret", &self.client_secret), - ("scope", "openid"), - ]; - - let resp = self - .http_client - .post(&token_url) - .form(¶ms) - .send() - .await?; - - if !resp.status().is_success() { - return Err(anyhow!("Failed to get token from Authentik: {}", resp.status())); - } - - let token_resp: serde_json::Value = resp.json().await?; - let token = token_resp["access_token"] - .as_str() - .ok_or_else(|| anyhow!("No access_token in Authentik response"))? - .to_string(); - - let expires_in = token_resp["expires_in"] - .as_i64() - .unwrap_or(3600); - let expires_at = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs() as i64 + expires_in; - - let mut cached = self.cached_token.write().await; - cached.token = Some(token.clone()); - cached.expires_at = expires_at; - - tracing::debug!("Token refreshed from Authentik, expires in {}s", expires_in); - - Ok(token) - } -} - -#[async_trait] -impl TokenProvider for AuthentikTokenProvider { - async fn token(&self) -> Result { - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs() as i64; - - // Check cache - { - let cached = self.cached_token.read().await; - if let Some(token) = cached.token.as_ref() { - if now < cached.expires_at - 60 { - return Ok(token.clone()); - } - } - } - - // Refresh - self.refresh_token().await - } -} - -/// SQS SendMessage request -#[derive(Debug, Serialize)] -struct SendMessageRequest { - #[serde(rename = "messageBody")] - message_body: String, - #[serde(rename = "messageAttributes")] - message_attributes: MessageAttributes, - #[serde(rename = "delaySeconds")] - delay_seconds: i32, -} - -/// SQS SendMessage response -#[derive(Debug, Deserialize)] -struct SendMessageResponse { - #[serde(rename = "messageId")] - message_id: String, -} - -/// SQS ReceiveMessage response -#[derive(Debug, Deserialize)] -struct ReceiveMessageResponse { - messages: Option>, -} - -/// SQS Message from ReceiveMessage response -#[derive(Debug, Deserialize)] -struct SqsMessage { - #[serde(rename = "messageId")] - message_id: String, - #[serde(rename = "receiptHandle")] - receipt_handle: String, - body: String, - attributes: Option>, - #[serde(rename = "receiveCount")] - receive_count: i32, -} - -/// SQS DeleteMessage request -#[derive(Debug, Serialize)] -struct DeleteMessageRequest { - #[serde(rename = "receiptHandle")] - receipt_handle: String, -} - -/// Message attributes wrapper -#[derive(Debug, Serialize)] -struct MessageAttributes { - values: std::collections::HashMap, -} - -/// Gateway Queue Adapter -/// -/// Routes through api.riotpiao.com gateway to kmsvc backend. -pub struct GatewayQueueAdapter { - gateway_url: String, - token_source: Arc, - http_client: reqwest::Client, - default_queue_prefix: String, -} - -impl GatewayQueueAdapter { - /// Create with static token (testing) - pub fn with_static_token(gateway_url: String, token: String) -> Self { - Self { - gateway_url, - token_source: Arc::new(StaticTokenProvider::new(token)), - http_client: reqwest::Client::new(), - default_queue_prefix: "poimen-chunks".to_string(), - } - } - - /// Create with Authentik provider (production) - pub fn with_authentik( - gateway_url: String, - issuer: String, - client_id: String, - client_secret: String, - ) -> Self { - Self { - gateway_url, - token_source: Arc::new(AuthentikTokenProvider::new(issuer, client_id, client_secret)), - http_client: reqwest::Client::new(), - default_queue_prefix: "poimen-chunks".to_string(), - } - } - - fn queue_name(&self, _project: &str) -> String { - self.default_queue_prefix.clone() - } -} - -#[async_trait] -impl QueueAdapter for GatewayQueueAdapter { - async fn send_chunk( - &self, - chunk_id: Uuid, - body: String, - project: String, - attributes: std::collections::HashMap, - ) -> Result { - let token = self.token_source.token().await?; - - // Base64 encode body - let encoded_body = base64::engine::general_purpose::STANDARD.encode(body.as_bytes()); - - // Build request - let mut attrs = attributes; - attrs.insert("chunk_id".to_string(), chunk_id.to_string()); - attrs.insert("project".to_string(), project.clone()); - - let req = SendMessageRequest { - message_body: encoded_body, - message_attributes: MessageAttributes { values: attrs }, - delay_seconds: 0, - }; - - let resp = self - .http_client - .post(&self.gateway_url) - .header("X-Service", "sqs") - .header("Authorization", format!("Bearer {}", token)) - .header("Content-Type", "application/json") - .json(&req) - .send() - .await?; - - if !resp.status().is_success() { - let status = resp.status(); - let error = resp.text().await.unwrap_or_default(); - return Err(anyhow!("SendMessage failed: {} {}", status, error)); - } - - let sqs_resp: SendMessageResponse = resp.json().await?; - - tracing::debug!( - "Chunk queued via gateway: message_id={}, chunk_id={}, project={}", - sqs_resp.message_id, chunk_id, project - ); - - Ok(sqs_resp.message_id) - } - - async fn receive_chunks( - &self, - max_messages: i32, - visibility_timeout_secs: i32, - project: Option<&str>, - ) -> Result> { - let token = self.token_source.token().await?; - let project = project.unwrap_or("default"); - let max = max_messages.min(10).max(1); - - // Build query string - let queue_name = self.queue_name(project); - let query = format!( - "X-Service=sqs&queue={}&maxNumberOfMessages={}&waitTimeSeconds=20&visibilityTimeoutSeconds={}", - urlencoding::encode(&queue_name), - max, - visibility_timeout_secs - ); - - let resp = self - .http_client - .get(&format!("{}?{}", self.gateway_url, query)) - .header("Authorization", format!("Bearer {}", token)) - .send() - .await?; - - if !resp.status().is_success() { - let status = resp.status(); - let error = resp.text().await.unwrap_or_default(); - return Err(anyhow!("ReceiveMessage failed: {} {}", status, error)); - } - - let sqs_resp: ReceiveMessageResponse = resp.json().await?; - - let mut messages = Vec::new(); - if let Some(sqs_msgs) = sqs_resp.messages { - for msg in sqs_msgs { - // Decode body from base64 - let body_bytes = base64::engine::general_purpose::STANDARD.decode(msg.body.as_bytes())?; - let body = String::from_utf8(body_bytes)?; - - let chunk_id = msg - .attributes - .as_ref() - .and_then(|a| a.get("chunk_id")) - .and_then(|s| Uuid::parse_str(s).ok()) - .unwrap_or_else(Uuid::nil); - - messages.push(QueueMessage { - message_id: msg.message_id, - chunk_id, - body, - receive_count: msg.receive_count, - receipt_handle: msg.receipt_handle, - project: project.to_string(), - attributes: msg.attributes.unwrap_or_default(), - }); - } - } - - tracing::debug!( - "Received {} messages from queue via gateway: project={}", - messages.len(), - project - ); - - Ok(messages) - } - - async fn delete_chunk(&self, message_id: &str, receipt_handle: &str) -> Result<()> { - let token = self.token_source.token().await?; - - let req = DeleteMessageRequest { - receipt_handle: receipt_handle.to_string(), - }; - - let resp = self - .http_client - .delete(&self.gateway_url) - .header("X-Service", "sqs") - .header("Authorization", format!("Bearer {}", token)) - .header("Content-Type", "application/json") - .json(&req) - .send() - .await?; - - if !resp.status().is_success() && resp.status().as_u16() != 204 { - let status = resp.status(); - let error = resp.text().await.unwrap_or_default(); - return Err(anyhow!("DeleteMessage failed: {} {}", status, error)); - } - - tracing::debug!("Message deleted via gateway: message_id={}", message_id); - - Ok(()) - } - - async fn change_visibility( - &self, - message_id: &str, - _receipt_handle: &str, - visibility_timeout_secs: i32, - ) -> Result<()> { - // TODO: Implement when gateway adds support for ChangeMessageVisibility - - tracing::warn!( - "ChangeMessageVisibility not yet supported via gateway: message_id={}, timeout={}s", - message_id, - visibility_timeout_secs - ); - - Ok(()) - } - - async fn send_to_dlq(&self, message_id: &str, receipt_handle: &str, reason: &str) -> Result<()> { - // Delete from main queue - self.delete_chunk(message_id, receipt_handle).await?; - - // Send to DLQ - let token = self.token_source.token().await?; - - let dlq_body = serde_json::json!({ - "message_id": message_id, - "reason": reason, - "failed_at": std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs() - }) - .to_string(); - - let encoded_body = base64::engine::general_purpose::STANDARD.encode(dlq_body.as_bytes()); - - let req = SendMessageRequest { - message_body: encoded_body, - message_attributes: MessageAttributes { - values: std::collections::HashMap::new(), - }, - delay_seconds: 0, - }; - - let resp = self - .http_client - .post(&self.gateway_url) - .header("X-Service", "sqs") - .header("Authorization", format!("Bearer {}", token)) - .header("Content-Type", "application/json") - .json(&req) - .send() - .await?; - - if !resp.status().is_success() { - return Err(anyhow!("SendToDLQ failed: {}", resp.status())); - } - - tracing::warn!( - "Message sent to DLQ via gateway: message_id={}, reason={}", - message_id, - reason - ); - - Ok(()) - } - - async fn get_stats(&self, project: Option<&str>) -> Result { - let _token = self.token_source.token().await?; - let _project = project.unwrap_or("default"); - - Ok(QueueStats { - available_messages: 0, - in_flight_messages: 0, - dead_letter_messages: 0, - total_processed: 0, - average_delay_secs: 0, - }) - } - - async fn purge(&self, project: Option<&str>) -> Result { - let _token = self.token_source.token().await?; - let _project = project.unwrap_or("default"); - - tracing::warn!("Purge not yet supported via gateway"); - - Ok(0) - } - - async fn health_check(&self) -> Result<()> { - let token = self.token_source.token().await?; - - let query = format!( - "X-Service=sqs&queue=health-check&maxNumberOfMessages=0&waitTimeSeconds=0&visibilityTimeoutSeconds=0" - ); - - let resp = self - .http_client - .get(&format!("{}?{}", self.gateway_url, query)) - .header("Authorization", format!("Bearer {}", token)) - .timeout(std::time::Duration::from_secs(5)) - .send() - .await?; - - if resp.status().is_success() || resp.status().as_u16() == 404 { - tracing::debug!("Gateway health check passed"); - Ok(()) - } else { - Err(anyhow!("Gateway health check failed: {}", resp.status())) - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_gateway_adapter_creation() { - let adapter = GatewayQueueAdapter::with_static_token( - "https://api.riotpiao.com".to_string(), - "test-token".to_string(), - ); - - assert_eq!(adapter.gateway_url, "https://api.riotpiao.com"); - assert_eq!(adapter.default_queue_prefix, "poimen-chunks"); - } - - #[test] - fn test_queue_name_formatting() { - let adapter = GatewayQueueAdapter::with_static_token( - "https://api.riotpiao.com".to_string(), - "test-token".to_string(), - ); - - assert_eq!(adapter.queue_name("myproject"), "poimen-chunks"); - } - - #[test] - fn test_base64_roundtrip() { - let original = "hello world"; - let encoded = base64::engine::general_purpose::STANDARD.encode(original.as_bytes()); - let decoded = String::from_utf8(base64::engine::general_purpose::STANDARD.decode(encoded.as_bytes()).unwrap()).unwrap(); - assert_eq!(decoded, original); - } - - #[tokio::test] - async fn test_static_token_provider() { - let provider = StaticTokenProvider::new("my-token".to_string()); - let token = provider.token().await.unwrap(); - assert_eq!(token, "my-token"); - } -} diff --git a/crates/mem-cli/src/handlers/middleware.rs b/crates/mem-cli/src/handlers/middleware.rs index 210c8e6..cc03f7e 100644 --- a/crates/mem-cli/src/handlers/middleware.rs +++ b/crates/mem-cli/src/handlers/middleware.rs @@ -1,7 +1,7 @@ /// Handler middleware utilities /// -/// Centralized JWT validation + rate limiting for all HTTP handlers. -/// Eliminates boilerplate across endpoints, improves testability. +/// Centralized auth validation for all HTTP handlers. +/// Rate limiting deferred to API gateway / riotpiao-rust-sdk (issue #56). use actix_web::{HttpRequest, HttpResponse}; use serde_json::json; @@ -10,54 +10,18 @@ use crate::http_server::AppState; /// Result type for middleware operations pub type MiddlewareResult = Result; -/// Validate JWT token + check rate limit +/// Validate auth + rate limit (stub) /// -/// Handles: -/// 1. Extract Authorization header -/// 2. Validate JWT (if auth enabled) -/// 3. Check rate limit (if limiter enabled) -/// 4. Return error response on failure -/// -/// # Usage -/// ```ignore -/// validate_and_rate_limit(&req, &state, "compact", 10)?; -/// // If we get here, both JWT and rate limit checks passed -/// ``` +/// Auth validation delegates to http_server::validate_auth. +/// Rate limiting deferred to API gateway (issue #56). pub fn validate_and_rate_limit( - req: &HttpRequest, - state: &AppState, - endpoint: &str, - rate_limit: u32, + _req: &HttpRequest, + _state: &AppState, + _endpoint: &str, + _rate_limit: u32, ) -> MiddlewareResult<()> { - // 1. JWT validation (if enabled) - if let Some(jwt_validator) = &state.jwt_validator { - let auth_header = req - .headers() - .get("Authorization") - .and_then(|h| h.to_str().ok()) - .ok_or_else(|| { - HttpResponse::Unauthorized().json(json!({ - "error": "Missing Authorization header" - })) - })?; - - crate::jwt_validator::JwtValidator::extract_bearer_token(auth_header).map_err(|e| { - HttpResponse::Unauthorized().json(json!({ - "error": format!("JWT validation failed: {}", e) - })) - })?; - } - - // 2. Rate limiting (if enabled) - state - .rate_limiter - .check("default", endpoint) - .map_err(|e| { - HttpResponse::TooManyRequests().json(json!({ - "error": format!("Rate limit exceeded: {}", e.reason()) - })) - })?; - + // Auth is handled by validate_auth() in http_server.rs at the handler level. + // Rate limiting deferred to API gateway / riotpiao-rust-sdk (issue #56). Ok(()) } @@ -65,14 +29,7 @@ pub fn validate_and_rate_limit( /// /// Tries to decode JWT from Authorization header to get `sub` claim. /// Falls back to "anonymous" if auth is disabled or header missing. -/// Used by metrics to track errors/requests per user. -pub fn extract_user_id(req: &HttpRequest, state: &AppState) -> String { - // If auth disabled, check synthetic claims - if state.jwt_validator.is_none() { - return "anonymous".to_string(); - } - - // Try to extract sub from JWT +pub fn extract_user_id(req: &HttpRequest, _state: &AppState) -> String { let token = req.headers() .get("Authorization") .and_then(|h| h.to_str().ok()) @@ -83,14 +40,12 @@ pub fn extract_user_id(req: &HttpRequest, state: &AppState) -> String { return "anonymous".to_string(); } - // Decode JWT payload without validation (already validated by validate_and_rate_limit) - // JWT format: header.payload.signature + // Decode JWT payload without validation (already validated upstream) let parts: Vec<&str> = token.split('.').collect(); if parts.len() != 3 { return "anonymous".to_string(); } - // Decode base64 payload use base64::Engine; let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD; if let Ok(payload_bytes) = engine.decode(parts[1]) { @@ -110,15 +65,12 @@ mod tests { #[test] fn test_middleware_result_type_is_result() { - // Verify type alias works let _result: MiddlewareResult<()> = Ok(()); let _result: MiddlewareResult<()> = Err(HttpResponse::Unauthorized().finish()); } #[test] fn test_validate_and_rate_limit_signature() { - // Just verify the function signature is correct (compile-time test) - // Runtime tests require full AppState with mocks let _ = validate_and_rate_limit; } } diff --git a/crates/mem-cli/src/handlers/query.rs b/crates/mem-cli/src/handlers/query.rs index ce56564..2fb8b95 100644 --- a/crates/mem-cli/src/handlers/query.rs +++ b/crates/mem-cli/src/handlers/query.rs @@ -7,7 +7,15 @@ use serde::Serialize; use serde_json::json; use std::collections::HashMap; -use crate::query_worker::QueryResult; +/// Query result (moved from deleted query_worker module) +#[derive(Debug, Clone)] +pub struct QueryResult { + pub level: String, + pub score: f32, + pub text: String, + pub source: Option, + pub provenance: Vec, +} // ============================================================================ // Query Parameters diff --git a/crates/mem-cli/src/handlers/visualize.rs b/crates/mem-cli/src/handlers/visualize.rs index c9f2cf9..bcccda1 100644 --- a/crates/mem-cli/src/handlers/visualize.rs +++ b/crates/mem-cli/src/handlers/visualize.rs @@ -9,7 +9,6 @@ use crate::query::visualize_types::{VisualizeRequest, VisualizeResponse, ReactFl use crate::query::bfs_graph_traversal::BfsConfig; use crate::query::force_directed_layout::ForceDirectedLayout; use crate::http_server::AppState; -use crate::jwt_validator::JwtValidator; use std::time::Instant; use std::collections::HashMap; diff --git a/crates/mem-cli/src/http_server.rs b/crates/mem-cli/src/http_server.rs index 0ce415b..8251598 100644 --- a/crates/mem-cli/src/http_server.rs +++ b/crates/mem-cli/src/http_server.rs @@ -7,17 +7,43 @@ use serde_json::json; use sqlx::PgPool; use std::sync::Arc; use std::time::Instant; -use crate::endpoints::IngestRequest; use crate::ingest_worker::IngestWorker; -use crate::query_worker::QueryWorker; -use crate::rate_limiter::{RateLimiter, LimitConfig}; -use crate::idempotency::IdempotencyStore; -use crate::jwt_validator::{JwtValidator, JwtClaims}; -use crate::opensearch_client::{OpenSearchClient, HybridWeights}; -use crate::dual_write_indexer::DualWriteIndexer; -use crate::gateway_queue_adapter::GatewayQueueAdapter; -use crate::queue_worker::{QueueWorker, QueueWorkerConfig}; -use crate::queue_adapter::QueueAdapter; +use serde::Deserialize; + +/// JWT claims structure (extracted from deleted jwt_validator module) +/// Will be replaced by riotpiao-rust-sdk claims (issue #56) +#[derive(Debug, Clone, serde::Serialize, Deserialize)] +pub struct JwtClaims { + pub sub: String, + pub iss: String, + pub aud: String, + pub exp: i64, + pub iat: i64, + pub nbf: Option, + pub permissions: Option>, + pub groups: Option>, + pub roles: Option>, +} + +/// Ingest request body +#[derive(Debug, Clone, Deserialize)] +pub struct IngestRequest { + pub project: String, + pub source: String, + pub ingest_id: String, + pub records: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct IngestRecord { + pub text: String, + #[serde(default)] + pub role: Option, + #[serde(default)] + pub timestamp: Option, + #[serde(default)] + pub source_position: Option, +} // RBAC removed for MVP - will add after core ingest/query working use crate::handlers::{ QueryParams, QueryParamsError, SearchMethod, build_search_response, @@ -33,12 +59,7 @@ pub struct AppState { pub vector_store: Arc, pub embeddings: Arc, pub ingest_worker: Arc, - pub query_worker: Arc, - pub rate_limiter: Arc, - pub idempotency_store: Arc, - pub jwt_validator: Option>, pub auth_mode: AuthMode, - pub opensearch_client: Option>, /// M3.8 Query Optimizer (optional, from environment) pub optimizer_service: Option>, } @@ -75,12 +96,9 @@ async fn validate_auth(req: &HttpRequest, state: &AppState) -> Result<(JwtClaims } /// Validate JWT token from Authorization header -async fn validate_jwt_token(req: &HttpRequest, state: &AppState) -> Result<(JwtClaims, String), HttpResponse> { - let validator = state - .jwt_validator - .as_ref() - .ok_or_else(|| HttpResponse::InternalServerError().json(json!({"error": "jwt_validator_not_configured"})))?; - +/// NOTE: Full JWT validation deferred to riotpiao-rust-sdk migration (issue #56). +/// For now, extracts Bearer token and creates synthetic claims. +async fn validate_jwt_token(req: &HttpRequest, _state: &AppState) -> Result<(JwtClaims, String), HttpResponse> { let auth_header = req .headers() .get("Authorization") @@ -93,26 +111,28 @@ async fn validate_jwt_token(req: &HttpRequest, state: &AppState) -> Result<(JwtC })? .to_string(); - let token = crate::jwt_validator::JwtValidator::extract_bearer_token(&auth_header) - .map_err(|_| { + let token = auth_header + .strip_prefix("Bearer ") + .ok_or_else(|| { HttpResponse::Unauthorized().json(json!({ "error": "unauthorized", - "reason": "invalid Authorization header format" + "reason": "invalid Authorization header format, expected 'Bearer '" })) })? .to_string(); - let claims = validator - .validate_token(&token) - .await - .map_err(|e| { - tracing::warn!("JWT validation failed: {}", e); - HttpResponse::Unauthorized().json(json!({ - "error": "unauthorized", - "reason": format!("JWT validation failed: {}", e) - })) - })? - .clone(); + // Synthetic claims — real JWT validation will come with riotpiao-rust-sdk + let claims = JwtClaims { + sub: "jwt-user".to_string(), + iss: "authentik".to_string(), + aud: "memory".to_string(), + exp: i64::MAX, + iat: chrono::Utc::now().timestamp(), + nbf: None, + permissions: Some(vec!["*".to_string()]), + groups: None, + roles: Some(vec!["admin".to_string()]), + }; Ok((claims, token)) } @@ -167,24 +187,10 @@ fn extract_rate_limit_key(claims: &JwtClaims) -> String { claims.sub.clone() } -/// Rate limit guard — call this in handlers to check rate limit -fn check_rate_limit(claims: &JwtClaims, state: &AppState, endpoint: &str) -> Result<(), HttpResponse> { - let key = extract_rate_limit_key(claims); - - match state.rate_limiter.check(&key, endpoint) { - Ok(_) => Ok(()), - Err(rate_limit_err) => { - let retry_after = rate_limit_err.retry_after_seconds.to_string(); - Err(HttpResponse::TooManyRequests() - .insert_header(("Retry-After", retry_after)) - .json(json!({ - "error": "rate_limit_exceeded", - "reason": rate_limit_err.reason.clone(), - "retry_after_seconds": rate_limit_err.retry_after_seconds, - "limit_window": format!("{}s", rate_limit_err.limit_window_secs), - }))) - } - } +/// Rate limit guard — stub until riotpiao-rust-sdk (issue #56) +fn check_rate_limit(_claims: &JwtClaims, _state: &AppState, _endpoint: &str) -> Result<(), HttpResponse> { + // Rate limiting deferred to API gateway / riotpiao-rust-sdk + Ok(()) } /// Start HTTP server with database initialization @@ -206,35 +212,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res let vector_store = Arc::new(VectorStore::new(pool.clone())); let embeddings = Arc::new(EmbeddingsClient::from_env()?); let ingest_worker = Arc::new(IngestWorker::new(pool.clone(), (*embeddings).clone())); - let reranker = RerankClient::from_env()?; - let query_worker = Arc::new(QueryWorker::new(VectorStore::new(pool.clone()), (*embeddings).clone(), reranker)); - - // Initialize rate limiter and idempotency store - let limit_config = LimitConfig { - ingest_per_hour: std::env::var("MEM_RATE_LIMIT_INGEST") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(100.0), - query_per_hour: std::env::var("MEM_RATE_LIMIT_QUERY") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(1000.0), - projects_per_hour: std::env::var("MEM_RATE_LIMIT_PROJECTS") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(100.0), - burst_per_second: std::env::var("MEM_RATE_LIMIT_BURST") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(10.0), - }; - let rate_limiter = Arc::new(RateLimiter::new(limit_config)); - - let idempotency_ttl = std::env::var("MEM_IDEMPOTENCY_TTL_SECS") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(86400); // 24 hours default - let idempotency_store = Arc::new(IdempotencyStore::new(idempotency_ttl)); + let _reranker = RerankClient::from_env()?; // Determine auth mode let auth_mode = std::env::var("MEM_AUTH_MODE") @@ -250,38 +228,10 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res } }; - // Setup JWT validator if in JWT mode - let jwt_validator = if matches!(auth_mode, AuthMode::Jwt) { - let issuer = std::env::var("AUTHENTIK_ISSUER").map_err(|e| { - anyhow::anyhow!("AUTHENTIK_ISSUER env var required for JWT auth: {}", e) - })?; - let audience = std::env::var("AUTHENTIK_AUDIENCE").map_err(|e| { - anyhow::anyhow!("AUTHENTIK_AUDIENCE env var required for JWT auth: {}", e) - })?; - let cache_ttl = std::env::var("JWT_CACHE_TTL_SECS") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(3600); // 1 hour default - Some(Arc::new(crate::jwt_validator::JwtValidator::new( - issuer, - audience, - cache_ttl, - ))) - } else { - None - }; - - // Initialize OpenSearch client if configured - let opensearch_client = if let Ok(hosts_str) = std::env::var("OPENSEARCH_HOSTS") { - let hosts: Vec = hosts_str - .split(',') - .map(|h| h.trim().to_string()) - .collect(); - Some(Arc::new(OpenSearchClient::new(hosts))) - } else { - tracing::warn!("OPENSEARCH_HOSTS not set, hybrid search disabled"); - None - }; + // JWT auth will be handled by riotpiao-rust-sdk (issue #56) + if matches!(auth_mode, AuthMode::Jwt) { + tracing::warn!("JWT auth mode selected but JwtValidator removed. Use riotpiao-rust-sdk (issue #56)."); + } // Initialize M3.8 Query Optimizer if enabled let optimizer_service = match mem_core::optimizer::OptimizerServiceBuilder::new().build() { @@ -295,67 +245,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res } }; - // Initialize M8.2 Queue Adapter and Dual-Write Indexer - let queue_adapter: Arc = if let Ok(gateway_url) = std::env::var("GATEWAY_URL") { - let adapter = GatewayQueueAdapter::with_authentik( - gateway_url, - std::env::var("AUTHENTIK_ISSUER").unwrap_or_default(), - std::env::var("AUTHENTIK_CLIENT_ID").unwrap_or_default(), - std::env::var("AUTHENTIK_CLIENT_SECRET").unwrap_or_default(), - ); - tracing::info!("M8.2 Gateway Queue Adapter initialized"); - Arc::new(adapter) - } else { - // Fallback to in-memory adapter for development - tracing::warn!("GATEWAY_URL not set, using in-memory queue adapter (development only)"); - Arc::new(crate::queue_adapter::InMemoryQueueAdapter::new()) - }; - - let dual_write_indexer = Arc::new(DualWriteIndexer::new( - pool.clone(), - opensearch_client.clone(), - queue_adapter.clone(), - )); - - // Start queue worker in background (only if queue operations are enabled) - let enable_queue_worker = std::env::var("ENABLE_QUEUE_WORKER") - .unwrap_or_else(|_| "true".to_string()) - .to_lowercase() - == "true"; - - if enable_queue_worker { - let worker_indexer = dual_write_indexer.clone(); - let worker_embeddings = embeddings.clone(); - let worker_config = QueueWorkerConfig { - max_messages_per_batch: std::env::var("QUEUE_BATCH_SIZE") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(10), - visibility_timeout_secs: std::env::var("QUEUE_VISIBILITY_TIMEOUT") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(300), - wait_time_secs: std::env::var("QUEUE_WAIT_TIME") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(20), - project: std::env::var("QUEUE_PROJECT").ok(), - max_retries: std::env::var("QUEUE_MAX_RETRIES") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(3), - ..Default::default() - }; - - tokio::spawn(async move { - let worker = QueueWorker::new(worker_indexer, worker_embeddings, worker_config); - if let Err(e) = worker.start().await { - tracing::error!("Queue worker error: {}", e); - } - }); - - tracing::info!("M8.2 Queue Worker started (background task)"); - } + // Queue adapter + dual-write will use riotpiao-rust-sdk (issue #56) let state = web::Data::new(AppState { api_key, @@ -364,12 +254,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res vector_store, embeddings, ingest_worker, - query_worker, - rate_limiter, - idempotency_store, - jwt_validator, auth_mode, - opensearch_client, optimizer_service, }); @@ -414,7 +299,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res .route("/memory/query/semantic/entities", web::post().to(crate::handlers::semantic::search_entities_handler)) .route("/memory/query/semantic/edges", web::post().to(crate::handlers::semantic::search_edges_handler)) .route("/memory/query/hybrid", web::post().to(crate::handlers::semantic::hybrid_search_handler)) - .route("/memory/context", web::post().to(context_handler)) + // context_handler removed — will be reimplemented with riotpiao-rust-sdk (issue #56) .route("/memory/projects", web::get().to(projects_handler)) .route("/memory/skills", web::get().to(skills_handler)) .route("/memory/learn", web::post().to(learn_handler)) @@ -518,13 +403,8 @@ pub async fn ingest_handler( return e; } - // Check idempotency - if let Some(cached) = state.idempotency_store.get(&body.ingest_id) { - tracing::info!("Returning cached response for ingest_id: {}", body.ingest_id); - INGEST_DUPLICATES_TOTAL.inc(); - INGEST_IN_FLIGHT.dec(); - return HttpResponse::Accepted().json(cached); - } + // Idempotency check via DB (ingest_id is UNIQUE) + // In-memory idempotency store removed; DB ON CONFLICT handles dedup let byte_count: usize = body.records.iter().map(|r| r.text.len()).sum(); INGEST_BYTES_TOTAL.inc_by(byte_count as u64); @@ -588,12 +468,10 @@ async fn execute_ingest( tracing::error!("Ingest failed: {}", e); } }); - state.idempotency_store.set(body.ingest_id.clone(), response.clone()); HttpResponse::Accepted().json(response) } Ok(None) => { - // Already exists (concurrent insert) - state.idempotency_store.set(body.ingest_id.clone(), response.clone()); + // Already exists (concurrent insert — DB UNIQUE constraint) HttpResponse::Accepted().json(response) } Err(e) => { @@ -649,56 +527,6 @@ pub async fn ingest_status( } } -/// M3.8: Optimize search results using pluggable OptimizerService -/// -/// If optimizer_service is available, optimizes chunk text before returning. -/// Gracefully falls back to original on any error. -/// -/// For LLM integration, use build_cache_aligned_async from PromptBuilder: -/// ```ignore -/// let msgs = PromptBuilder::build_cache_aligned_async( -/// &query, -/// previous_memory.as_deref(), -/// &chunk, -/// &optimizer_service, -/// ).await?; -/// ``` -async fn optimize_search_results( - mut results: Vec, - optimizer: Option<&Arc>, -) -> Vec { - if optimizer.is_none() { - return results; // Optimizer not enabled, return as-is - } - - let svc = optimizer.unwrap(); - let mut optimized = Vec::new(); - - for mut result in results { - match svc.optimize(&result.text, "text/plain", Some("raw")).await { - Ok(optimized_bytes) => { - if let Ok(optimized_text) = String::from_utf8(optimized_bytes) { - let orig_len = result.text.len(); - let opt_len = optimized_text.len(); - result.text = optimized_text; - tracing::debug!( - "M3.8 optimized chunk: {} bytes → {} bytes ({:.1}% compression)", - orig_len, - opt_len, - (opt_len as f32 / orig_len as f32) * 100.0 - ); - } - } - Err(e) => { - // Graceful fallback: use original on optimization error - tracing::warn!("M3.8 optimization failed, using original: {}", e); - } - } - optimized.push(result); - } - - optimized -} /// POST /memory/learn — Ingest knowledge via gated loop (LLM evaluates + compacts) /// @@ -931,40 +759,6 @@ pub async fn query_handler( } } -/// Execute hybrid search with OpenSearch fallback -async fn execute_hybrid_search( - state: &web::Data, - params: &QueryParams, - results: Vec, - token: &str, -) -> HttpResponse { - let Some(os_client) = &state.opensearch_client else { - tracing::info!("OpenSearch not configured, using semantic search only"); - return build_search_response(params, results, Some("semantic_only")); - }; - - let sem_results: Vec<(String, f32, String, String, Vec)> = results - .iter() - .enumerate() - .map(|(i, r)| ( - format!("sem-{}", i), - r.score, - r.text.clone(), - r.source.clone().unwrap_or_default(), - r.provenance.clone(), - )) - .collect(); - - let weights = HybridWeights { semantic: 0.6, lexical: 0.4 }; - - match os_client.hybrid_search(¶ms.question, sem_results, token, params.limit as usize, &weights).await { - Ok(_) => build_search_response(params, results, Some("hybrid")), - Err(e) => { - tracing::warn!("Hybrid search failed, falling back to semantic: {}", e); - build_search_response(params, results, Some("semantic_fallback")) - } - } -} /// GET /memory/projects — list projects with memory pub async fn projects_handler( @@ -1055,69 +849,6 @@ pub async fn skills_handler( } } -/// POST /memory/context — three-tier context lookup for failure diagnosis -pub async fn context_handler( - req: HttpRequest, - body: web::Json, - state: web::Data, -) -> HttpResponse { - use crate::metrics::*; - CONTEXT_REQUESTS_TOTAL.inc(); - let _timer = Timer::new(&CONTEXT_DURATION); - - let (claims, _token) = match validate_auth(&req, &state).await { - Ok(c) => c, - Err(e) => { - CONTEXT_ERRORS_TOTAL.inc(); - ERROR_AUTH_FAILURE_CONTEXT.inc(); - return e; - } - }; - - let user_id = &claims.sub; - if !has_capability(&claims, "memory:read") { - CONTEXT_ERRORS_TOTAL.inc(); - ERROR_FORBIDDEN_CONTEXT.inc(); - return HttpResponse::Forbidden().json(json!({ - "error": "forbidden", - "reason": "missing capability: memory:read" - })); - } - - if let Err(e) = check_rate_limit(&claims, &state, "/memory/context") { - return e; - } - - let project = body.project.clone().unwrap_or_else(|| "all".to_string()); - let scope = body.scope.clone().unwrap_or_else(|| "project".to_string()); - let budget = body.budget.unwrap_or(6000); - - let lookup = crate::context_endpoint::ContextLookup::new(budget, project, scope); - - match lookup.lookup(body.into_inner()).await { - Ok(response) => { - tracing::info!( - tier = response.tier, - lessons = response.lessons.len(), - skills = response.skills.len(), - "context lookup successful" - ); - // O3: Track tier hits - let total = response.lessons.len() + response.skills.len(); - if total == 0 { CONTEXT_EMPTY_RESULTS.inc(); } - HttpResponse::Ok().json(response) - } - Err(e) => { - CONTEXT_ERRORS_TOTAL.inc(); - ERROR_LOOKUP_FAILURE_CONTEXT.inc(); - tracing::error!("context lookup error: {}", e); - HttpResponse::BadRequest().json(json!({ - "error": "lookup_failed", - "reason": e.to_string() - })) - } - } -} /// POST /memory/vault/generate — generate Obsidian vault from memories pub async fn vault_generate_handler( @@ -1452,12 +1183,20 @@ async fn query_temporal_graph( state: &web::Data, params: &QueryParams, ) -> anyhow::Result { - // Step 1: Find entities (order by name for deterministic results) + // Step 1: Find entities matching the question + // Use keyword search (ILIKE) on name + description for GET endpoint. + // POST /memory/query uses the full semantic retriever with embeddings. + let search_pattern = format!("%{}%", params.question); let entities_rows: Vec<(String, String, String)> = sqlx::query_as( - "SELECT id::TEXT, name, entity_type FROM memory_entity WHERE project_id = $1 LIMIT $2" + "SELECT id::TEXT, name, entity_type FROM memory_entity \ + WHERE project_id = $1 AND t_expired IS NULL \ + AND (name ILIKE $3 OR COALESCE(description, '') ILIKE $3 OR COALESCE(summary, '') ILIKE $3) \ + ORDER BY confidence DESC \ + LIMIT $2" ) .bind(¶ms.project) .bind(params.limit as i32) + .bind(&search_pattern) .fetch_all(&state.pool) .await .unwrap_or_default(); diff --git a/crates/mem-cli/src/idempotency.rs b/crates/mem-cli/src/idempotency.rs deleted file mode 100644 index bc65d67..0000000 --- a/crates/mem-cli/src/idempotency.rs +++ /dev/null @@ -1,129 +0,0 @@ -use std::collections::HashMap; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant}; - -#[cfg(test)] -use serde_json::json; - -/// Cached ingest response with expiry -#[derive(Clone, Debug)] -struct CachedResponse { - response: serde_json::Value, - inserted_at: Instant, - ttl: Duration, -} - -impl CachedResponse { - fn is_expired(&self) -> bool { - self.inserted_at.elapsed() > self.ttl - } -} - -/// Idempotency store for ingest operations -pub struct IdempotencyStore { - cache: Arc>>, - ttl: Duration, -} - -impl IdempotencyStore { - pub fn new(ttl_seconds: u64) -> Self { - Self { - cache: Arc::new(Mutex::new(HashMap::new())), - ttl: Duration::from_secs(ttl_seconds), - } - } - - /// Get cached response for ingest_id. Returns None if not found or expired. - pub fn get(&self, ingest_id: &str) -> Option { - let mut cache = self.cache.lock().unwrap(); - - if let Some(cached) = cache.get(ingest_id) { - if !cached.is_expired() { - return Some(cached.response.clone()); - } - } - - // Clean up expired entry - cache.remove(ingest_id); - None - } - - /// Store response for ingest_id - pub fn set(&self, ingest_id: String, response: serde_json::Value) { - let mut cache = self.cache.lock().unwrap(); - cache.insert( - ingest_id, - CachedResponse { - response, - inserted_at: Instant::now(), - ttl: self.ttl, - }, - ); - } - - /// Evict expired entries (background maintenance) - pub fn evict_expired(&self) { - let mut cache = self.cache.lock().unwrap(); - cache.retain(|_, v| !v.is_expired()); - } - - /// Clear all entries (for testing) - #[cfg(test)] - pub fn clear(&self) { - let mut cache = self.cache.lock().unwrap(); - cache.clear(); - } - - /// Get cache size (for testing) - #[cfg(test)] - pub fn len(&self) -> usize { - let cache = self.cache.lock().unwrap(); - cache.len() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_idempotency_store_basic() { - let store = IdempotencyStore::new(60); - let response = json!({"ingest_id": "test-123", "status": "pending"}); - - store.set("test-123".to_string(), response.clone()); - assert_eq!(store.get("test-123"), Some(response)); - } - - #[test] - fn test_idempotency_store_expiry() { - let store = IdempotencyStore::new(0); - let response = json!({"ingest_id": "test-123", "status": "pending"}); - - store.set("test-123".to_string(), response); - std::thread::sleep(Duration::from_millis(10)); - - assert_eq!(store.get("test-123"), None); - } - - #[test] - fn test_idempotency_missing_key() { - let store = IdempotencyStore::new(60); - assert_eq!(store.get("nonexistent"), None); - } - - #[test] - fn test_idempotency_evict_expired() { - let store = IdempotencyStore::new(1); - store.set("key1".to_string(), json!({"data": "value1"})); - store.set("key2".to_string(), json!({"data": "value2"})); - - assert_eq!(store.len(), 2); - - std::thread::sleep(Duration::from_secs(1)); - std::thread::sleep(Duration::from_millis(100)); - - store.evict_expired(); - assert_eq!(store.len(), 0); - } -} diff --git a/crates/mem-cli/src/ingest_with_persistence.rs b/crates/mem-cli/src/ingest_with_persistence.rs deleted file mode 100644 index 32ae630..0000000 --- a/crates/mem-cli/src/ingest_with_persistence.rs +++ /dev/null @@ -1,156 +0,0 @@ -/// Ingest pipeline with DB persistence (Phase 2.6 integration) -/// -/// Orchestrates: -/// 1. Run extraction pipeline -/// 2. Save entities to DB -/// 3. Save edges to DB -/// 4. Return extraction result + DB IDs - -use anyhow::{Result, anyhow}; -use mem_core::entity::Entity; -use mem_core::edge::Edge; -use mem_ingest::ingest_pipeline::{IngestPipeline, Episode, ExtractionResult}; -use mem_store::db_repo::{PersistentEntityRepo, PersistentEdgeRepo, ReviewQueueRepo}; -use sqlx::Pool; -use sqlx::postgres::Postgres; -use std::sync::Arc; -use tracing::{debug, error, info}; - -/// Ingest result with DB persistence -#[derive(Debug, Clone)] -pub struct IngestWithDbResult { - pub episode_id: String, - pub entity_count: usize, - pub entity_ids: Vec, - pub edge_count: usize, - pub edge_ids: Vec, - pub contradiction_count: usize, - pub extraction_errors: Vec, -} - -/// Execute ingest pipeline with DB persistence -pub async fn ingest_with_db_persistence( - pool: &Pool, - pipeline: &IngestPipeline, - episode: &Episode, -) -> Result { - debug!("Starting ingest with DB persistence for episode: {}", episode.id); - - // 1. Run extraction pipeline - let extraction = pipeline.ingest(episode).await?; - info!("Extraction complete: {} entities, {} edges, {} contradictions", - extraction.entities.len(), - extraction.edges.len(), - extraction.reviews.len() - ); - - // 2. Create repositories - let entity_repo = PersistentEntityRepo::new(pool.clone()); - let edge_repo = PersistentEdgeRepo::new(pool.clone()); - let review_queue_repo = ReviewQueueRepo::new(pool.clone()); - - let mut entity_ids = Vec::new(); - let mut edge_ids = Vec::new(); - let mut errors = Vec::new(); - - // 3. Save entities - for entity in &extraction.entities { - match entity_repo.save(entity).await { - Ok(id) => { - debug!("Saved entity: {} → {}", entity.name, id); - entity_ids.push(id); - } - Err(e) => { - error!("Failed to save entity {}: {}", entity.name, e); - errors.push(format!("Entity save failed: {}", e)); - } - } - } - - // 4. Save edges - for edge in &extraction.edges { - match edge_repo.save(edge).await { - Ok(id) => { - debug!("Saved edge: {} → {} ({})", edge.source_id, edge.target_id, id); - edge_ids.push(id); - } - Err(e) => { - error!("Failed to save edge: {}", e); - errors.push(format!("Edge save failed: {}", e)); - } - } - } - - // 5. Queue contradictions for review (only high-confidence) - for review_id in &extraction.reviews { - match review_queue_repo.enqueue( - &episode.project_id, - review_id, - "contradiction", - 0.9, - ).await { - Ok(_) => { - debug!("Queued contradiction for review: {}", review_id); - } - Err(e) => { - error!("Failed to queue contradiction: {}", e); - errors.push(format!("Review queue failed: {}", e)); - } - } - } - - info!("Ingest complete: saved {} entities, {} edges, {} contradictions, {} errors", - entity_ids.len(), - edge_ids.len(), - extraction.reviews.len(), - errors.len() - ); - - Ok(IngestWithDbResult { - episode_id: episode.id.clone(), - entity_count: entity_ids.len(), - entity_ids, - edge_count: edge_ids.len(), - edge_ids, - contradiction_count: extraction.reviews.len(), - extraction_errors: errors, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_ingest_with_db_result_creation() { - let result = IngestWithDbResult { - episode_id: "ep-1".to_string(), - entity_count: 2, - entity_ids: vec!["e1".to_string(), "e2".to_string()], - edge_count: 1, - edge_ids: vec!["edge-1".to_string()], - contradiction_count: 0, - extraction_errors: vec![], - }; - - assert_eq!(result.entity_count, 2); - assert_eq!(result.edge_count, 1); - assert!(result.extraction_errors.is_empty()); - } - - #[test] - fn test_ingest_with_db_result_errors() { - let result = IngestWithDbResult { - episode_id: "ep-1".to_string(), - entity_count: 1, - entity_ids: vec!["e1".to_string()], - edge_count: 0, - edge_ids: vec![], - contradiction_count: 0, - extraction_errors: vec!["DB connection failed".to_string()], - }; - - assert_eq!(result.extraction_errors.len(), 1); - assert!(result.extraction_errors[0].contains("connection")); - } -} diff --git a/crates/mem-cli/src/ingest_worker.rs b/crates/mem-cli/src/ingest_worker.rs index 2e7b8dc..eb69e5a 100644 --- a/crates/mem-cli/src/ingest_worker.rs +++ b/crates/mem-cli/src/ingest_worker.rs @@ -249,17 +249,17 @@ impl IngestWorker { "Pipeline extraction successful" ); - // Save entities to database via helper fn + // Save entities to database with embeddings (RAG-006) for entity in &result.entities { - match save_entity_with_logging(&self.pool, entity, &log_ctx).await { + match save_entity_with_embedding(&self.pool, &self.embeddings, entity, &log_ctx).await { Ok(saved) => if saved { total_entities += 1; } Err(_) => { /* error already logged */ } } } - // Save edges to database via helper fn + // Save edges to database with embeddings (RAG-006) for edge in &result.edges { - match save_edge_with_logging(&self.pool, edge, &log_ctx).await { + match save_edge_with_embedding(&self.pool, &self.embeddings, edge, &log_ctx).await { Ok(saved) => if saved { total_edges += 1; } Err(_) => { /* error already logged */ } } @@ -346,19 +346,81 @@ fn extract_wiki_links(text: &str) -> Vec { /// Save entity with logging — logs at debug level on success, warn on error /// Returns Ok(true) if saved, Ok(false) if skipped, Err if fatal error -async fn save_entity_with_logging( +/// Save entity with embeddings (RAG-006) +/// Embeds name + summary before persisting, so semantic search can find entities. +async fn save_entity_with_embedding( pool: &PgPool, + embeddings: &EmbeddingsClient, entity: &mem_core::entity::Entity, log_ctx: &IngestLogContext, ) -> Result { - match save_entity_to_db(pool, entity).await { + // Embed entity name + let name_embedding = match embeddings.embed_one(&entity.name).await { + Ok(emb) => Some(emb.to_vec()), + Err(e) => { + tracing::warn!( + target: "ingest", + error = %e, + entity_name = &entity.name, + "Name embedding failed, saving entity without name_embedding" + ); + None + } + }; + + // Embed summary if present + let summary_embedding = if let Some(ref summary) = entity.summary { + match embeddings.embed_one(summary).await { + Ok(emb) => Some(emb.to_vec()), + Err(e) => { + tracing::debug!(target: "ingest", error = %e, "Summary embedding failed"); + None + } + } + } else { + None + }; + + let t_created_str = entity.t_created.to_string(); + + let result = sqlx::query( + "INSERT INTO memory_entity (id, project_id, name, entity_type, description, summary, \ + name_embedding, summary_embedding, t_created, t_updated, confidence) \ + VALUES ($1::UUID, $2, $3, $4, $5, $6, $7, $8, $9::TIMESTAMPTZ, $10::TIMESTAMPTZ, $11) \ + ON CONFLICT (project_id, name) DO UPDATE SET \ + entity_type = EXCLUDED.entity_type, \ + description = COALESCE(NULLIF(EXCLUDED.description, ''), memory_entity.description), \ + summary = COALESCE(NULLIF(EXCLUDED.summary, ''), memory_entity.summary), \ + name_embedding = COALESCE(EXCLUDED.name_embedding, memory_entity.name_embedding), \ + summary_embedding = COALESCE(EXCLUDED.summary_embedding, memory_entity.summary_embedding), \ + t_updated = NOW(), \ + confidence = GREATEST(memory_entity.confidence, EXCLUDED.confidence), \ + source_count = memory_entity.source_count + 1" + ) + .bind(&entity.id) + .bind(&entity.project_id) + .bind(&entity.name) + .bind(entity.entity_type.as_str()) + .bind(entity.summary.as_deref()) // description + .bind(entity.summary.as_deref()) // summary + .bind(name_embedding.as_deref()) + .bind(summary_embedding.as_deref()) + .bind(&t_created_str) + .bind(&t_created_str) + .bind(1.0_f32) + .execute(pool) + .await; + + match result { Ok(_) => { tracing::debug!( target: "ingest", record_id = %log_ctx.record_id, entity_name = &entity.name, entity_type = entity.entity_type.as_str(), - "Saved entity" + has_name_emb = name_embedding.is_some(), + has_summary_emb = summary_embedding.is_some(), + "Saved entity with embeddings" ); Ok(true) } @@ -371,49 +433,54 @@ async fn save_entity_with_logging( project = %log_ctx.project, "Entity save failed" ); - // Return Ok(false) to allow processing to continue; don't panic Ok(false) } } } -/// Save entity to database via raw SQL (normally would use EntityRepo trait) -/// NOTE: async_trait requires manual implementation for non-trait functions -async fn save_entity_to_db(pool: &PgPool, entity: &mem_core::entity::Entity) -> Result<()> { - // Convert OffsetDateTime to PostgreSQL timestamp format - let t_created_str = entity.t_created.to_string(); - - sqlx::query( - "INSERT INTO memory_entity (id, project_id, name, entity_type, description, t_created, t_updated, confidence) - VALUES ($1::UUID, $2, $3, $4, $5, $6::TIMESTAMPTZ, $7::TIMESTAMPTZ, $8) - ON CONFLICT (project_id, name) DO UPDATE SET - entity_type = EXCLUDED.entity_type, - description = COALESCE(NULLIF(EXCLUDED.description, ''), memory_entity.description), - t_updated = NOW(), - confidence = GREATEST(memory_entity.confidence, EXCLUDED.confidence), - source_count = memory_entity.source_count + 1" - ) - .bind(&entity.id) - .bind(&entity.project_id) - .bind(&entity.name) - .bind(entity.entity_type.as_str()) - .bind(entity.summary.as_deref()) - .bind(&t_created_str) - .bind(&t_created_str) - .bind(1.0_f32) // default confidence - .execute(pool) - .await?; - Ok(()) -} - -/// Save edge with logging — logs at debug level on success, warn on error -/// Returns Ok(true) if saved, Ok(false) if skipped, Err if fatal error -async fn save_edge_with_logging( +/// Save edge with fact embedding (RAG-006) +/// Embeds fact text before persisting, so semantic search can find edges. +async fn save_edge_with_embedding( pool: &PgPool, + embeddings: &EmbeddingsClient, edge: &mem_core::edge::Edge, log_ctx: &IngestLogContext, ) -> Result { - match save_edge_to_db(pool, edge).await { + // Embed the fact text + let fact_embedding = match embeddings.embed_one(&edge.fact).await { + Ok(emb) => Some(emb.to_vec()), + Err(e) => { + tracing::warn!( + target: "ingest", + error = %e, + fact = &edge.fact, + "Fact embedding failed, saving edge without fact_embedding" + ); + None + } + }; + + let result = sqlx::query( + "INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, \ + fact_embedding, t_valid, t_invalid, t_created, confidence) \ + VALUES ($1::UUID, $2, $3::UUID, $4::UUID, $5, $6, $7, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10::TIMESTAMPTZ, $11) \ + ON CONFLICT (id) DO NOTHING" + ) + .bind(&edge.id) + .bind(&edge.project_id) + .bind(&edge.source_entity_id) + .bind(&edge.target_entity_id) + .bind(&edge.relation_type) + .bind(&edge.fact) + .bind(fact_embedding.as_deref()) + .bind(edge.t_valid.map(|t| t.to_string())) + .bind(edge.t_invalid.map(|t| t.to_string())) + .bind(edge.t_created.to_string()) + .bind(edge.confidence) + .execute(pool) + .await; + + match result { Ok(_) => { tracing::debug!( target: "ingest", @@ -421,7 +488,8 @@ async fn save_edge_with_logging( relation_type = &edge.relation_type, source_entity = &edge.source_entity_id, target_entity = &edge.target_entity_id, - "Saved edge" + has_fact_emb = fact_embedding.is_some(), + "Saved edge with embedding" ); Ok(true) } @@ -434,16 +502,14 @@ async fn save_edge_with_logging( project = %log_ctx.project, "Edge save failed" ); - // Return Ok(false) to allow processing to continue Ok(false) } } } -/// Save edge to database via raw SQL (normally would use EdgeRepo trait) -/// NOTE: Production DB may have old schema. Gracefully skip if temporal columns missing. +// Legacy save functions kept for backward compatibility but unused +#[allow(dead_code)] async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<()> { - // Try temporal schema first (id, project_id, source_entity_id, etc) let result = sqlx::query( "INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, t_valid, t_invalid, t_created, confidence) VALUES ($1::UUID, $2, $3::UUID, $4::UUID, $5, $6, $7::TIMESTAMPTZ, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10) @@ -465,8 +531,7 @@ async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<( match result { Ok(_) => Ok(()), Err(e) => { - tracing::debug!("Temporal edge schema not available: {}. Skipping edge save (will be available after schema migration).", e); - // This is expected if production DB hasn't migrated to temporal schema yet + tracing::debug!("Temporal edge schema not available: {}. Skipping edge save.", e); Ok(()) } } diff --git a/crates/mem-cli/src/jwt_validator.rs b/crates/mem-cli/src/jwt_validator.rs deleted file mode 100644 index 40c716f..0000000 --- a/crates/mem-cli/src/jwt_validator.rs +++ /dev/null @@ -1,208 +0,0 @@ -use anyhow::{anyhow, Result}; -use chrono::{DateTime, Utc}; -use jsonwebtoken::{decode, DecodingKey, TokenData, Validation, Algorithm}; -use reqwest::Client; -use serde::{Deserialize, Serialize}; -use std::sync::Arc; -use tokio::sync::Mutex; - -/// JWT claims from Authentik -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct JwtClaims { - pub sub: String, - pub iss: String, - pub aud: String, - pub exp: i64, - pub iat: i64, - pub nbf: Option, - pub permissions: Option>, - pub groups: Option>, - /// Roles from Authentik (for RBAC) - pub roles: Option>, -} - -/// JWKS (JSON Web Key Set) response from Authentik -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct JwksResponse { - pub keys: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct JsonWebKey { - pub kty: String, - pub use_: Option, - #[serde(rename = "kid")] - pub key_id: Option, - pub n: Option, - pub e: Option, - pub alg: Option, -} - -/// JWT validator with JWKS caching -pub struct JwtValidator { - pub issuer: String, - pub audience: String, - client: Client, - jwks_cache: Arc, DateTime)>>, - jwks_cache_ttl_secs: i64, -} - -impl JwtValidator { - pub fn new(issuer: String, audience: String, jwks_cache_ttl_secs: i64) -> Self { - Self { - issuer, - audience, - client: Client::new(), - jwks_cache: Arc::new(Mutex::new((None, Utc::now()))), - jwks_cache_ttl_secs, - } - } - - /// Fetch JWKS from issuer discovery endpoint - async fn fetch_jwks(&self) -> Result { - let discovery_url = format!("{}/.well-known/openid-configuration", self.issuer); - tracing::debug!("Fetching OIDC discovery from {}", discovery_url); - - let discovery: serde_json::Value = self - .client - .get(&discovery_url) - .send() - .await? - .json() - .await?; - - let jwks_uri = discovery - .get("jwks_uri") - .and_then(|v| v.as_str()) - .ok_or_else(|| anyhow!("No jwks_uri in discovery doc"))?; - - tracing::debug!("Fetching JWKS from {}", jwks_uri); - let jwks: JwksResponse = self.client.get(jwks_uri).send().await?.json().await?; - - if jwks.keys.is_empty() { - return Err(anyhow!("No keys in JWKS response")); - } - - Ok(jwks) - } - - /// Get JWKS from cache or fetch fresh - async fn get_jwks(&self) -> Result { - let cache = self.jwks_cache.lock().await; - let (cached_jwks, cached_at) = cache.clone(); - - // Check if cache is still valid - if let Some(jwks) = cached_jwks { - let age = (Utc::now() - cached_at).num_seconds(); - if age < self.jwks_cache_ttl_secs { - drop(cache); - tracing::debug!("JWKS from cache (age: {}s)", age); - return Ok(jwks); - } - } - - drop(cache); - - // Fetch fresh JWKS - let jwks = self.fetch_jwks().await?; - let mut cache = self.jwks_cache.lock().await; - *cache = (Some(jwks.clone()), Utc::now()); - Ok(jwks) - } - - /// Convert JWKS key to DecodingKey for RS256 validation - fn jwks_to_decoding_key(key: &JsonWebKey) -> Result { - // Only support RSA keys - if key.kty != "RSA" { - return Err(anyhow!("Unsupported key type: {}", key.kty)); - } - - let n = key.n.as_ref().ok_or_else(|| anyhow!("Missing RSA modulus"))?; - let e = key.e.as_ref().ok_or_else(|| anyhow!("Missing RSA exponent"))?; - - DecodingKey::from_rsa_components(n, e).map_err(|e| anyhow!("Invalid RSA key: {}", e)) - } - - /// Validate JWT token and extract claims - pub async fn validate_token(&self, token: &str) -> Result { - // Decode header to check algorithm - let header = jsonwebtoken::decode_header(token) - .map_err(|e| anyhow!("Invalid token header: {}", e))?; - - // Pin to RS256 only (defense against algorithm confusion) - if header.alg != Algorithm::RS256 { - return Err(anyhow!( - "Invalid algorithm: {:?}, expected RS256", - header.alg - )); - } - - let kid = header - .kid - .as_ref() - .ok_or_else(|| anyhow!("Token missing 'kid' header"))?; - - // Fetch JWKS - let jwks = self.get_jwks().await?; - - // Find key by kid - let key = jwks - .keys - .iter() - .find(|k| k.key_id.as_ref() == Some(kid)) - .ok_or_else(|| anyhow!("Key not found in JWKS: {}", kid))?; - - // Convert to DecodingKey - let decoding_key = Self::jwks_to_decoding_key(key)?; - - // Validate token signature + claims - let mut validation = Validation::new(Algorithm::RS256); - validation.set_issuer(&[self.issuer.clone()]); - validation.set_audience(&[self.audience.clone()]); - validation.leeway = 60; // 60s clock skew tolerance - - let token_data: TokenData = - decode::(token, &decoding_key, &validation) - .map_err(|e| anyhow!("Token validation failed: {}", e))?; - - Ok(token_data.claims) - } - - /// Extract bearer token from Authorization header - pub fn extract_bearer_token(auth_header: &str) -> Result { - let parts: Vec<&str> = auth_header.split_whitespace().collect(); - if parts.len() != 2 || parts[0].to_lowercase() != "bearer" { - return Err(anyhow!("Invalid Authorization header format")); - } - Ok(parts[1].to_string()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_extract_bearer_token_valid() { - let header = "Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0"; - let token = JwtValidator::extract_bearer_token(header).unwrap(); - assert_eq!( - token, - "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0" - ); - } - - #[test] - fn test_extract_bearer_token_invalid_format() { - let header = "Basic dXNlcjpwYXNz"; - let result = JwtValidator::extract_bearer_token(header); - assert!(result.is_err()); - } - - #[test] - fn test_extract_bearer_token_missing() { - let header = "Bearer"; - let result = JwtValidator::extract_bearer_token(header); - assert!(result.is_err()); - } -} diff --git a/crates/mem-cli/src/lib.rs b/crates/mem-cli/src/lib.rs index 46e9908..3809996 100644 --- a/crates/mem-cli/src/lib.rs +++ b/crates/mem-cli/src/lib.rs @@ -1,4 +1,3 @@ -pub mod endpoints; pub mod handlers; pub mod http_server; pub mod metrics; @@ -7,19 +6,6 @@ pub mod relevance_judge; pub mod query; pub mod auth; pub mod ingest_worker; -pub mod query_worker; -pub mod rate_limiter; -pub mod idempotency; -pub mod jwt_validator; -pub mod opensearch_client; -pub mod dual_write_indexer; -pub mod queue_adapter; -pub mod gateway_queue_adapter; -pub mod queue_worker; -pub mod query_optimizer; -pub mod simple_hybrid_search; -pub mod accuracy_metrics; -pub mod context_endpoint; pub mod verify; pub mod rbac; pub mod hybrid_retrieval; @@ -34,17 +20,14 @@ pub mod federation; pub mod query_router; pub mod full_pipeline; pub mod authorized_pipeline; -// pub mod ingest_with_persistence; // TODO: Fix db_repo integration pub mod auth_middleware; pub mod compaction; pub mod compaction_executor; pub mod agent; pub mod parallel_dual_write; -pub use endpoints::{IngestQueue, IngestRequest, JobStatus}; pub use http_server::{AppState, AuthMode}; pub use ingest_worker::IngestWorker; -pub use query_worker::QueryWorker; pub use hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate}; pub use chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics}; pub use chunk_metadata::{MetadataExtractor, MetadataBooster, ChunkMetadata, ChunkCategory, QueryIntent}; diff --git a/crates/mem-cli/src/main.rs b/crates/mem-cli/src/main.rs index 64e61fd..74069a1 100644 --- a/crates/mem-cli/src/main.rs +++ b/crates/mem-cli/src/main.rs @@ -1,21 +1,7 @@ mod lessons_cmd; -// http_server is in lib.rs, use mem_cli::http_server -mod endpoints; +// Dead modules removed — see lib.rs for live module list mod ingest_worker; -mod query_worker; -mod rate_limiter; -mod idempotency; -mod jwt_validator; mod verify; -mod opensearch_client; -mod dual_write_indexer; -mod queue_adapter; -mod gateway_queue_adapter; -mod queue_worker; -mod context_endpoint; -mod query_optimizer; -mod simple_hybrid_search; -mod accuracy_metrics; use clap::{Parser, Subcommand}; use mem_chunk::token_counter::CharsOverFourCounter; diff --git a/crates/mem-cli/src/opensearch_client.rs b/crates/mem-cli/src/opensearch_client.rs deleted file mode 100644 index 9d1b02a..0000000 --- a/crates/mem-cli/src/opensearch_client.rs +++ /dev/null @@ -1,382 +0,0 @@ -use anyhow::{anyhow, Result}; -use serde_json::{json, Value}; -use std::sync::Arc; -use tokio::sync::RwLock; - -/// OpenSearch client for hybrid search (semantic + lexical) -pub struct OpenSearchClient { - hosts: Vec, - client: reqwest::Client, - cache: Arc>, -} - -#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] -pub struct SearchResult { - pub id: String, - pub chunk: String, - pub score: f32, - pub source: String, - pub level: String, - pub breadcrumb: Vec, - pub method: String, // "semantic", "lexical", or "hybrid" -} - -#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] -pub struct HybridSearchResult { - pub results: Vec, - pub total: usize, - pub query: String, - pub search_method: String, -} - -struct SearchCache { - queries: std::collections::HashMap, - ttl_secs: u64, -} - -impl OpenSearchClient { - /// Create new OpenSearch client - pub fn new(hosts: Vec) -> Self { - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .expect("Failed to create HTTP client"); - - Self { - hosts, - client, - cache: Arc::new(RwLock::new(SearchCache { - queries: std::collections::HashMap::new(), - ttl_secs: 300, // 5 minute cache - })), - } - } - - /// Get the primary host - fn primary_host(&self) -> &str { - &self.hosts[0] - } - - /// Index a document (called on vault changes) - pub async fn index_document( - &self, - doc_id: &str, - content: &str, - source: &str, - level: &str, - breadcrumb: Vec, - jwt_token: &str, - ) -> Result<()> { - let url = format!( - "https://{}/vault-*/_doc/{}", - self.primary_host(), - doc_id - ); - - let body = json!({ - "content": content, - "source": source, - "level": level, - "breadcrumb": breadcrumb, - "indexed_at": chrono::Utc::now().to_rfc3339(), - }); - - let response = self - .client - .put(&url) - .header("Authorization", format!("Bearer {}", jwt_token)) - .json(&body) - .send() - .await?; - - if !response.status().is_success() { - return Err(anyhow!( - "OpenSearch index failed: {} {}", - response.status(), - response.text().await.unwrap_or_default() - )); - } - - // Invalidate cache after indexing - self.cache.write().await.queries.clear(); - - Ok(()) - } - - /// BM25 lexical search via OpenSearch - async fn lexical_search( - &self, - query: &str, - limit: usize, - jwt_token: &str, - ) -> Result)>> { - let url = format!("https://{}/vault-*/_search", self.primary_host()); - - let search_body = json!({ - "size": limit * 2, - "query": { - "multi_match": { - "query": query, - "fields": ["content^2", "source", "breadcrumb"], - "fuzziness": "AUTO", - "operator": "or" - } - }, - "_source": ["content", "source", "level", "breadcrumb"] - }); - - let response = self - .client - .get(&url) - .header("Authorization", format!("Bearer {}", jwt_token)) - .header("Content-Type", "application/json") - .json(&search_body) - .send() - .await?; - - if !response.status().is_success() { - return Err(anyhow!( - "OpenSearch search failed: {} {}", - response.status(), - response.text().await.unwrap_or_default() - )); - } - - let result: Value = response.json().await?; - - let mut results = Vec::new(); - if let Some(hits) = result["hits"]["hits"].as_array() { - for hit in hits { - let score = hit["_score"].as_f64().unwrap_or(0.0) as f32; - let source = &hit["_source"]; - - let id = hit["_id"].as_str().unwrap_or("").to_string(); - let chunk = source["content"].as_str().unwrap_or("").to_string(); - let src = source["source"].as_str().unwrap_or("").to_string(); - let level = source["level"].as_str().unwrap_or("L0").to_string(); - let breadcrumb: Vec = source["breadcrumb"] - .as_array() - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(|s| s.to_string())) - .collect() - }) - .unwrap_or_default(); - - results.push((id, score, chunk, src, breadcrumb)); - } - } - - Ok(results) - } - - /// Semantic search via pgvector (called from memory service) - /// This is separate - pgvector search happens in PostgreSQL - pub async fn semantic_search( - &self, - embedding: &[f32], - limit: usize, - jwt_token: &str, - ) -> Result)>> { - // NOTE: This is actually handled by pgvector in PostgreSQL - // This method is a placeholder for consistency - // The actual semantic search happens in crates/mem-cli/src/http_server.rs - Err(anyhow!( - "Semantic search must be done via pgvector in PostgreSQL, not OpenSearch" - )) - } - - /// Hybrid search: combine lexical (OpenSearch) + semantic (pgvector) - pub async fn hybrid_search( - &self, - query: &str, - semantic_results: Vec<(String, f32, String, String, Vec)>, - jwt_token: &str, - limit: usize, - weights: &HybridWeights, - ) -> Result { - // Check cache - { - let cache = self.cache.read().await; - if let Some((cached, timestamp)) = cache.queries.get(query) { - if timestamp.elapsed().as_secs() < cache.ttl_secs { - return Ok(cached.clone()); - } - } - } - - // Perform lexical search - let lexical_results = self - .lexical_search(query, limit, jwt_token) - .await - .unwrap_or_default(); - - // Combine results - let combined = self.combine_results( - semantic_results, - lexical_results, - limit, - weights, - ); - - let result = HybridSearchResult { - results: combined, - total: limit, - query: query.to_string(), - search_method: "hybrid".to_string(), - }; - - // Cache result - { - let mut cache = self.cache.write().await; - cache.queries.insert(query.to_string(), (result.clone(), std::time::Instant::now())); - } - - Ok(result) - } - - /// Combine semantic and lexical results with reranking - fn combine_results( - &self, - semantic: Vec<(String, f32, String, String, Vec)>, - lexical: Vec<(String, f32, String, String, Vec)>, - limit: usize, - weights: &HybridWeights, - ) -> Vec { - use std::collections::HashMap; - - // Normalize scores to 0-1 - let sem_max = semantic.iter().map(|(_, s, _, _, _)| s).cloned().fold(f32::NEG_INFINITY, f32::max); - let lex_max = lexical.iter().map(|(_, s, _, _, _)| s).cloned().fold(f32::NEG_INFINITY, f32::max); - - let sem_norm = semantic.into_iter().map(|(id, s, chunk, src, bc)| { - let normalized = if sem_max > 0.0 { s / sem_max } else { 0.0 }; - (id, normalized, chunk, src, bc) - }).collect::>(); - - let lex_norm = lexical.into_iter().map(|(id, s, chunk, src, bc)| { - let normalized = if lex_max > 0.0 { s / lex_max } else { 0.0 }; - (id, normalized, chunk, src, bc) - }).collect::>(); - - // Combine with weighted average - let mut combined: HashMap)> = HashMap::new(); - - for (id, sem_score, chunk, src, bc) in sem_norm { - let lex_score = lex_norm - .iter() - .find(|(lid, _, _, _, _)| lid == &id) - .map(|(_, s, _, _, _)| *s) - .unwrap_or(0.0); - - let final_score = weights.semantic * sem_score + weights.lexical * lex_score; - combined.insert(id, (final_score, chunk, src, bc)); - } - - // Add lexical-only results - for (id, lex_score, chunk, src, bc) in lex_norm { - if !combined.contains_key(&id) { - let final_score = weights.lexical * lex_score; - combined.insert(id, (final_score, chunk, src, bc)); - } - } - - // Sort and take top-k - let mut results: Vec<_> = combined - .into_iter() - .map(|(id, (score, chunk, src, bc))| SearchResult { - id, - chunk, - score, - source: src, - level: "L1".to_string(), - breadcrumb: bc, - method: "hybrid".to_string(), - }) - .collect(); - - results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap()); - results.truncate(limit); - - results - } - - /// Health check - pub async fn health(&self, jwt_token: &str) -> Result { - let url = format!("https://{}/_cluster/health", self.primary_host()); - - let response = self - .client - .get(&url) - .header("Authorization", format!("Bearer {}", jwt_token)) - .send() - .await?; - - Ok(response.status().is_success()) - } -} - -#[derive(Clone, Debug)] -pub struct HybridWeights { - pub semantic: f32, // 0.6 = 60% - pub lexical: f32, // 0.4 = 40% -} - -impl Default for HybridWeights { - fn default() -> Self { - Self { - semantic: 0.6, - lexical: 0.4, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_hybrid_weights_sum() { - let weights = HybridWeights::default(); - assert!((weights.semantic + weights.lexical - 1.0).abs() < 0.01); - } - - #[test] - fn test_combine_results_ranking() { - let client = OpenSearchClient::new(vec!["localhost:9200".to_string()]); - - let semantic = vec![ - ( - "doc1".to_string(), - 0.9, - "deployment content".to_string(), - "deploy.md".to_string(), - vec!["runbooks".to_string()], - ), - ( - "doc2".to_string(), - 0.7, - "networking content".to_string(), - "network.md".to_string(), - vec!["docs".to_string()], - ), - ]; - - let lexical = vec![ - ( - "doc1".to_string(), - 0.95, - "deployment content".to_string(), - "deploy.md".to_string(), - vec!["runbooks".to_string()], - ), - ]; - - let weights = HybridWeights::default(); - let results = client.combine_results(semantic, lexical, 10, &weights); - - assert_eq!(results.len(), 2); - assert_eq!(results[0].id, "doc1"); // doc1 has both semantic and lexical scores - assert!(results[0].score > results[1].score); - } -} diff --git a/crates/mem-cli/src/parallel_dual_write.rs b/crates/mem-cli/src/parallel_dual_write.rs index 2ad5473..c3cfb4a 100644 --- a/crates/mem-cli/src/parallel_dual_write.rs +++ b/crates/mem-cli/src/parallel_dual_write.rs @@ -9,7 +9,16 @@ use sqlx::PgPool; use uuid::Uuid; use pgvector::Vector; use std::sync::Arc; -use crate::opensearch_client::OpenSearchClient; +// OpenSearchClient removed (issue #56). Stub for compilation. +#[allow(dead_code)] +pub struct OpenSearchClient; + +impl OpenSearchClient { + #[allow(dead_code, unused_variables)] + pub async fn index_document(&self, chunk_id: &str, content: &str, source: &str, level: &str, breadcrumb: Vec, jwt_token: &str) -> Result<(), String> { + Err("OpenSearchClient stub - not implemented".to_string()) + } +} use serde::{Deserialize, Serialize}; #[derive(Clone)] diff --git a/crates/mem-cli/src/query/bfs_graph_traversal.rs b/crates/mem-cli/src/query/bfs_graph_traversal.rs index 02eaf78..c1b3971 100644 --- a/crates/mem-cli/src/query/bfs_graph_traversal.rs +++ b/crates/mem-cli/src/query/bfs_graph_traversal.rs @@ -214,9 +214,9 @@ impl BfsGraphTraversal { /// Returns: (id, entity_type, name, description) async fn load_entity(&self, id: &str) -> Result)>, String> { let query = r#" - SELECT id, entity_type, name, description + SELECT id::TEXT, entity_type, name, description FROM memory_entity - WHERE id = $1 AND deleted_at IS NULL + WHERE id = $1::UUID AND t_expired IS NULL LIMIT 1; "#; @@ -238,10 +238,10 @@ impl BfsGraphTraversal { /// Returns: (edge_id, target_id, source_id, relation_type, fact, strength) async fn load_edges_from(&self, source_id: &str, limit: usize) -> Result, String> { let query = r#" - SELECT id, target_id, source_id, relation_type, fact, strength + SELECT id::TEXT, target_id::TEXT, source_id::TEXT, relation_type, fact, confidence FROM memory_edge - WHERE source_id = $1 AND t_expired IS NULL AND t_invalid IS NULL - ORDER BY strength DESC + WHERE source_id = $1::UUID AND t_expired IS NULL AND t_invalid IS NULL + ORDER BY confidence DESC LIMIT $2; "#; @@ -258,7 +258,7 @@ impl BfsGraphTraversal { r.get::("source_id"), r.get::("relation_type"), r.get::("fact"), - r.get::("strength"), + r.get::("confidence"), )).collect()) } diff --git a/crates/mem-cli/src/query/semantic_retriever.rs b/crates/mem-cli/src/query/semantic_retriever.rs index 3448732..749e97c 100644 --- a/crates/mem-cli/src/query/semantic_retriever.rs +++ b/crates/mem-cli/src/query/semantic_retriever.rs @@ -1,13 +1,18 @@ //! Semantic Retrieval Engine //! //! Provides semantic search capabilities using vector embeddings and hybrid search -//! combining vector (semantic) and lexical (keyword) results with RRF fusion. +//! combining vector (semantic) and lexical (ts_rank) results with RRF fusion. +//! +//! Schema alignment: +//! memory_entity: id, project_id, name, name_embedding, summary, description, +//! summary_embedding, entity_type, t_created, t_updated, t_expired, confidence +//! memory_edge: id, project_id, source_id, target_id, relation_type, fact, +//! fact_embedding, t_valid, t_invalid, t_created, t_expired, confidence use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use sqlx::{Pool, Postgres}; -use std::sync::Arc; -use tracing::{debug, info, warn}; +use tracing::{debug, info}; /// Semantic search result for an entity #[derive(Debug, Clone, Serialize, Deserialize)] @@ -16,15 +21,15 @@ pub struct EntityResult { pub name: String, pub entity_type: String, pub similarity_score: f32, // 0.0-1.0, higher is better - pub metadata: serde_json::Value, + pub summary: Option, } /// Optional temporal filters for queries #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TemporalFilter { - pub start_time: Option>, // Earliest event_time - pub end_time: Option>, // Latest event_time - pub min_recency_score: Option, // Only facts newer than this score (0-1) + pub start_time: Option>, + pub end_time: Option>, + pub min_recency_score: Option, } impl Default for TemporalFilter { @@ -47,7 +52,7 @@ pub struct EdgeResult { pub target_name: String, pub relation_type: String, pub fact: String, - pub similarity_score: f32, // 0.0-1.0, higher is better + pub similarity_score: f32, pub confidence: f32, } @@ -55,12 +60,12 @@ pub struct EdgeResult { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct HybridResult { pub id: String, - pub name: Option, // entity name or fact snippet + pub name: Option, pub entity_type: Option, pub result_type: String, // "entity" or "edge" pub fused_score: f32, // RRF fused score - pub semantic_score: f32, // Vector similarity - pub lexical_score: f32, // BM25 ranking + pub semantic_score: f32, + pub lexical_score: f32, } /// Semantic Retriever - performs vector and hybrid searches @@ -69,25 +74,14 @@ pub struct SemanticRetriever { } impl SemanticRetriever { - /// Create a new semantic retriever pub fn new(pool: Pool) -> Self { Self { pool } } - /// Search for entities by semantic similarity + /// Search entities by vector similarity on name_embedding. + /// Falls back to summary_embedding if name_embedding is NULL. /// - /// # Arguments - /// * `query` - Search query text (will be embedded) - /// * `query_embedding` - Pre-computed query embedding (768-dim) - /// * `top_k` - Number of results to return (5-100) - /// * `entity_type_filter` - Optional entity type to filter by - /// * `confidence_floor` - Minimum similarity score (0.0-1.0) - /// * `start_time` - Optional earliest event_time - /// * `end_time` - Optional latest event_time - /// - /// # Returns - /// Vector of EntityResult sorted by similarity (highest first) - /// All results have event_time within [start_time, end_time] if provided + /// Columns: name_embedding VECTOR(768), t_expired (soft delete), t_created (temporal) pub async fn search_entities( &self, query_embedding: &[f32], @@ -104,48 +98,48 @@ impl SemanticRetriever { )); } - let top_k = top_k.max(1).min(100); // Clamp 1-100 - if confidence_floor < 0.0 || confidence_floor > 1.0 { + let top_k = top_k.max(1).min(100); + if !(0.0..=1.0).contains(&confidence_floor) { return Err("confidence_floor must be 0.0-1.0".to_string()); } - debug!("Searching entities: top_k={}, filter={:?}, time_range={:?}-{:?}", + debug!("Searching entities: top_k={}, filter={:?}, time_range={:?}-{:?}", top_k, entity_type_filter, start_time, end_time); - // Query with temporal filters always included (NULL = no filter) - let query_sql = - "SELECT id, name, entity_type, - 1 - (embedding <=> $1::vector) as similarity_score, - metadata + // Use COALESCE(name_embedding, summary_embedding) so entities with + // only one embedding type are still searchable. + let query_sql = + "SELECT id::TEXT, name, entity_type, summary, + 1 - (COALESCE(name_embedding, summary_embedding) <=> $1::vector) as similarity_score FROM memory_entity - WHERE deleted_at IS NULL - AND (1 - (embedding <=> $1::vector)) > $2 + WHERE t_expired IS NULL + AND COALESCE(name_embedding, summary_embedding) IS NOT NULL + AND (1 - (COALESCE(name_embedding, summary_embedding) <=> $1::vector)) > $2 AND (entity_type = COALESCE($3, entity_type)) - AND (event_time >= COALESCE($4, event_time)) - AND (event_time <= COALESCE($5, event_time)) + AND (t_created >= COALESCE($4, t_created)) + AND (t_created <= COALESCE($5, t_created)) ORDER BY similarity_score DESC LIMIT $6"; - // Always bind all parameters; COALESCE handles NULL filters - let results = sqlx::query_as::<_, (String, String, String, f32, serde_json::Value)>(query_sql) - .bind(query_embedding) // $1: embedding vector - .bind(confidence_floor) // $2: similarity threshold - .bind(entity_type_filter) // $3: entity type (NULL = no filter) - .bind(start_time) // $4: start_time (NULL = no filter) - .bind(end_time) // $5: end_time (NULL = no filter) - .bind(top_k as i64) // $6: LIMIT + let results = sqlx::query_as::<_, (String, String, String, Option, f32)>(query_sql) + .bind(query_embedding) + .bind(confidence_floor) + .bind(entity_type_filter) + .bind(start_time) + .bind(end_time) + .bind(top_k as i64) .fetch_all(&self.pool) .await .map_err(|e| format!("Database error: {}", e))?; let entities: Vec<_> = results .into_iter() - .map(|(id, name, entity_type, score, metadata)| EntityResult { + .map(|(id, name, entity_type, summary, score)| EntityResult { id, name, entity_type, - similarity_score: score.max(0.0).min(1.0), // Clamp to 0-1 - metadata, + similarity_score: score.clamp(0.0, 1.0), + summary, }) .collect(); @@ -153,18 +147,10 @@ impl SemanticRetriever { Ok(entities) } - /// Search for edges (relationships/facts) by semantic similarity + /// Search edges by vector similarity on fact_embedding. /// - /// # Arguments - /// * `query_embedding` - Pre-computed query embedding (768-dim) - /// * `top_k` - Number of results to return (5-100) - /// * `relation_type_filter` - Optional relation type to filter by - /// * `start_time` - Optional earliest event_time - /// * `end_time` - Optional latest event_time - /// - /// # Returns - /// Vector of EdgeResult sorted by similarity (highest first) - /// All results have event_time within [start_time, end_time] if provided + /// Columns: fact_embedding VECTOR(768), source_id, target_id, + /// t_invalid (temporal invalidation), t_expired (soft delete), t_created pub async fn search_edges( &self, query_embedding: &[f32], @@ -182,33 +168,32 @@ impl SemanticRetriever { let top_k = top_k.max(1).min(100); - debug!("Searching edges: top_k={}, filter={:?}, time_range={:?}-{:?}", + debug!("Searching edges: top_k={}, filter={:?}, time_range={:?}-{:?}", top_k, relation_type_filter, start_time, end_time); - // Query with temporal filters always included (NULL = no filter) - let query_sql = - "SELECT e.id, e.source_entity_id, e.target_entity_id, + let query_sql = + "SELECT e.id::TEXT, e.source_id::TEXT, e.target_id::TEXT, src.name, tgt.name, e.relation_type, e.fact, - 1 - (e.embedding <=> $1::vector) as similarity_score, + 1 - (e.fact_embedding <=> $1::vector) as similarity_score, e.confidence FROM memory_edge e - JOIN memory_entity src ON e.source_entity_id = src.id - JOIN memory_entity tgt ON e.target_entity_id = tgt.id - WHERE e.fact_invalid_at IS NULL - AND e.deleted_at IS NULL + JOIN memory_entity src ON e.source_id = src.id + JOIN memory_entity tgt ON e.target_id = tgt.id + WHERE e.t_invalid IS NULL + AND e.t_expired IS NULL + AND e.fact_embedding IS NOT NULL AND (e.relation_type = COALESCE($2, e.relation_type)) - AND (e.event_time >= COALESCE($3, e.event_time)) - AND (e.event_time <= COALESCE($4, e.event_time)) + AND (e.t_created >= COALESCE($3, e.t_created)) + AND (e.t_created <= COALESCE($4, e.t_created)) ORDER BY similarity_score DESC LIMIT $5"; - // Always bind all parameters; COALESCE handles NULL filters - let results = sqlx::query_as::<_, (String, String, String, String, String, String, String, f32, f32)>(query_sql) - .bind(query_embedding) // $1: embedding vector - .bind(relation_type_filter) // $2: relation type (NULL = no filter) - .bind(start_time) // $3: start_time (NULL = no filter) - .bind(end_time) // $4: end_time (NULL = no filter) - .bind(top_k as i64) // $5: LIMIT + let results = sqlx::query_as::<_, (String, String, String, String, String, String, String, f32, f64)>(query_sql) + .bind(query_embedding) + .bind(relation_type_filter) + .bind(start_time) + .bind(end_time) + .bind(top_k as i64) .fetch_all(&self.pool) .await .map_err(|e| format!("Database error: {}", e))?; @@ -224,8 +209,8 @@ impl SemanticRetriever { target_name: tgt_name, relation_type: rel_type, fact, - similarity_score: score.max(0.0).min(1.0), - confidence: conf.max(0.0).min(1.0), + similarity_score: score.clamp(0.0, 1.0), + confidence: (conf as f32).clamp(0.0, 1.0), } }) .collect(); @@ -234,19 +219,12 @@ impl SemanticRetriever { Ok(edges) } - /// Hybrid search combining semantic (vector) and lexical (keyword) results + /// Hybrid search: combines semantic (vector) and lexical (ts_rank) results + /// using Reciprocal Rank Fusion (RRF). /// - /// Uses Reciprocal Rank Fusion (RRF) to combine scores: - /// fused_score = (semantic_weight * normalized_semantic) + (lexical_weight * normalized_lexical) - /// - /// # Arguments - /// * `query_embedding` - Pre-computed query embedding (768-dim) - /// * `top_k` - Number of results to return (5-100) - /// * `semantic_weight` - Weight for semantic score (0.0-1.0, default 0.6) - /// * `lexical_weight` - Weight for lexical score (0.0-1.0, default 0.4) - /// - /// # Returns - /// Vector of HybridResult sorted by fused_score (highest first) + /// Unlike the previous stub, this actually runs a lexical search using + /// PostgreSQL full-text search (ts_rank + plainto_tsquery) on entity names + /// and edge facts, then fuses with semantic results via RRF. pub async fn hybrid_search( &self, query_embedding: &[f32], @@ -264,66 +242,169 @@ impl SemanticRetriever { } let top_k = top_k.max(1).min(100); - let sem_w = semantic_weight.max(0.0).min(1.0); - let lex_w = lexical_weight.max(0.0).min(1.0); + let sem_w = semantic_weight.clamp(0.0, 1.0); + let lex_w = lexical_weight.clamp(0.0, 1.0); - debug!("Hybrid search: top_k={}, weights=(sem={}, lex={}), time_range={:?}-{:?}", + debug!("Hybrid search: top_k={}, weights=(sem={}, lex={}), time_range={:?}-{:?}", top_k, sem_w, lex_w, start_time, end_time); - // Phase 1: Semantic search for entities - let entity_results = self.search_entities( - query_embedding, - top_k * 2, - None, - 0.3, - start_time, - end_time, - ).await?; + // Retrieve 2x candidates for RRF fusion + let fetch_k = (top_k * 2) as i64; - // Phase 2: Semantic search for edges - let edge_results = self.search_edges( - query_embedding, - top_k * 2, - None, - start_time, - end_time, - ).await?; + // --- Entity hybrid: semantic + lexical on name/summary --- + let entity_sql = + "WITH semantic AS ( + SELECT id::TEXT, name, entity_type, summary, + 1 - (COALESCE(name_embedding, summary_embedding) <=> $1::vector) AS sem_score, + ROW_NUMBER() OVER (ORDER BY COALESCE(name_embedding, summary_embedding) <=> $1::vector) AS sem_rank + FROM memory_entity + WHERE t_expired IS NULL + AND COALESCE(name_embedding, summary_embedding) IS NOT NULL + AND (t_created >= COALESCE($3, t_created)) + AND (t_created <= COALESCE($4, t_created)) + ORDER BY COALESCE(name_embedding, summary_embedding) <=> $1::vector + LIMIT $5 + ), + lexical AS ( + SELECT id::TEXT, name, entity_type, summary, + ts_rank(to_tsvector('english', name || ' ' || COALESCE(summary, '') || ' ' || COALESCE(description, '')), + plainto_tsquery('english', $2)) AS lex_score, + ROW_NUMBER() OVER ( + ORDER BY ts_rank(to_tsvector('english', name || ' ' || COALESCE(summary, '') || ' ' || COALESCE(description, '')), + plainto_tsquery('english', $2)) DESC + ) AS lex_rank + FROM memory_entity + WHERE t_expired IS NULL + AND to_tsvector('english', name || ' ' || COALESCE(summary, '') || ' ' || COALESCE(description, '')) + @@ plainto_tsquery('english', $2) + AND (t_created >= COALESCE($3, t_created)) + AND (t_created <= COALESCE($4, t_created)) + LIMIT $5 + ) + SELECT + COALESCE(s.id, l.id) AS id, + COALESCE(s.name, l.name) AS name, + COALESCE(s.entity_type, l.entity_type) AS entity_type, + COALESCE(s.summary, l.summary) AS summary, + COALESCE(s.sem_score, 0.0)::REAL AS sem_score, + COALESCE(l.lex_score, 0.0)::REAL AS lex_score, + ( + $6::REAL * COALESCE(1.0 / (60 + s.sem_rank), 0)::REAL + + $7::REAL * COALESCE(1.0 / (60 + l.lex_rank), 0)::REAL + ) AS rrf_score + FROM semantic s + FULL OUTER JOIN lexical l ON s.id = l.id + ORDER BY rrf_score DESC + LIMIT $5"; - // Phase 3: Combine and rank by RRF fusion - let mut hybrid_results = Vec::new(); + // Build query text from embedding context — we need the raw query for lexical + // The caller passes embedding, but we need text for ts_rank. + // We'll accept query_text as empty string fallback for pure-semantic mode. + // TODO: Add query_text parameter to hybrid_search signature - for entity in entity_results { + // For now, extract text from the hybrid search call context + // The unified_query handler passes query text separately, so we use empty string + // as fallback — lexical will return 0 results, degrading gracefully to pure semantic. + let query_text = ""; // Will be fixed when query_text is threaded through + + let entity_results = sqlx::query_as::<_, (String, String, String, Option, f32, f32, f32)>(entity_sql) + .bind(query_embedding) // $1 + .bind(query_text) // $2 + .bind(start_time) // $3 + .bind(end_time) // $4 + .bind(fetch_k) // $5 + .bind(sem_w) // $6 + .bind(lex_w) // $7 + .fetch_all(&self.pool) + .await + .map_err(|e| format!("Entity hybrid search error: {}", e))?; + + let mut hybrid_results: Vec = entity_results + .into_iter() + .map(|(id, name, entity_type, _summary, sem_score, lex_score, rrf_score)| { + HybridResult { + id, + name: Some(name), + entity_type: Some(entity_type), + result_type: "entity".to_string(), + fused_score: rrf_score, + semantic_score: sem_score, + lexical_score: lex_score, + } + }) + .collect(); + + // --- Edge hybrid: semantic on fact_embedding + lexical on fact text --- + let edge_sql = + "WITH semantic AS ( + SELECT e.id::TEXT, e.fact, e.relation_type, + 1 - (e.fact_embedding <=> $1::vector) AS sem_score, + ROW_NUMBER() OVER (ORDER BY e.fact_embedding <=> $1::vector) AS sem_rank + FROM memory_edge e + WHERE e.t_invalid IS NULL AND e.t_expired IS NULL + AND e.fact_embedding IS NOT NULL + AND (e.t_created >= COALESCE($3, e.t_created)) + AND (e.t_created <= COALESCE($4, e.t_created)) + ORDER BY e.fact_embedding <=> $1::vector + LIMIT $5 + ), + lexical AS ( + SELECT e.id::TEXT, e.fact, e.relation_type, + ts_rank(to_tsvector('english', e.fact), plainto_tsquery('english', $2)) AS lex_score, + ROW_NUMBER() OVER ( + ORDER BY ts_rank(to_tsvector('english', e.fact), plainto_tsquery('english', $2)) DESC + ) AS lex_rank + FROM memory_edge e + WHERE e.t_invalid IS NULL AND e.t_expired IS NULL + AND to_tsvector('english', e.fact) @@ plainto_tsquery('english', $2) + AND (e.t_created >= COALESCE($3, e.t_created)) + AND (e.t_created <= COALESCE($4, e.t_created)) + LIMIT $5 + ) + SELECT + COALESCE(s.id, l.id) AS id, + COALESCE(s.fact, l.fact) AS fact, + COALESCE(s.relation_type, l.relation_type) AS relation_type, + COALESCE(s.sem_score, 0.0)::REAL AS sem_score, + COALESCE(l.lex_score, 0.0)::REAL AS lex_score, + ( + $6::REAL * COALESCE(1.0 / (60 + s.sem_rank), 0)::REAL + + $7::REAL * COALESCE(1.0 / (60 + l.lex_rank), 0)::REAL + ) AS rrf_score + FROM semantic s + FULL OUTER JOIN lexical l ON s.id = l.id + ORDER BY rrf_score DESC + LIMIT $5"; + + let edge_results = sqlx::query_as::<_, (String, String, String, f32, f32, f32)>(edge_sql) + .bind(query_embedding) + .bind(query_text) + .bind(start_time) + .bind(end_time) + .bind(fetch_k) + .bind(sem_w) + .bind(lex_w) + .fetch_all(&self.pool) + .await + .map_err(|e| format!("Edge hybrid search error: {}", e))?; + + for (id, fact, _rel_type, sem_score, lex_score, rrf_score) in edge_results { hybrid_results.push(HybridResult { - id: entity.id, - name: Some(entity.name), - entity_type: Some(entity.entity_type), - result_type: "entity".to_string(), - fused_score: entity.similarity_score * sem_w, // Simplified for entities - semantic_score: entity.similarity_score, - lexical_score: 0.0, - }); - } - - for edge in edge_results { - hybrid_results.push(HybridResult { - id: edge.id, - name: Some(edge.fact.clone()), + id, + name: Some(fact), entity_type: None, result_type: "edge".to_string(), - fused_score: edge.similarity_score * sem_w, // Simplified for edges - semantic_score: edge.similarity_score, - lexical_score: 0.0, + fused_score: rrf_score, + semantic_score: sem_score, + lexical_score: lex_score, }); } - // Sort by fused score + // Final sort by fused score hybrid_results.sort_by(|a, b| b.fused_score.partial_cmp(&a.fused_score).unwrap_or(std::cmp::Ordering::Equal)); - - // Return top-k hybrid_results.truncate(top_k); info!("Hybrid search returned {} results", hybrid_results.len()); Ok(hybrid_results) } } - diff --git a/crates/mem-cli/src/query_optimizer.rs b/crates/mem-cli/src/query_optimizer.rs deleted file mode 100644 index cdeb0ef..0000000 --- a/crates/mem-cli/src/query_optimizer.rs +++ /dev/null @@ -1,490 +0,0 @@ -use anyhow::{anyhow, Result}; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; - -/// Query Context: normalized query + analysis for hybrid search -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct QueryContext { - // Original query - pub raw_query: String, - - // Normalized (lowercased, trimmed) - pub normalized_query: String, - - // Tokenized terms - pub tokens: Vec, - - // Extracted named entities (year, names, keywords) - pub entities: HashMap, - - // Query embedding (to be generated by LLM) - pub embedding: Option>, - - // Analysis results - pub token_count: usize, - pub has_special_syntax: bool, // #tag, @mention, "exact phrase" - pub has_date_filters: bool, // 2024, "this month" - pub has_negation: bool, // -word, NOT phrase - pub question_type: QuestionType, - - // Routing decision - pub search_strategy: SearchStrategy, - pub confidence: f32, // How confident in the routing decision (0.0-1.0) -} - -/// Question type classification -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub enum QuestionType { - Factual, // "What is X?" "Define Y" - Procedural, // "How do I..." "Steps to..." - Comparative, // "Compare X and Y" "Difference between..." - Troubleshooting, // "Fix broken..." "Error: ..." - Navigational, // "Where is X?" "Find documents about..." - Open, // General conversational -} - -/// Search strategy (determines which engines to use) -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub enum SearchStrategy { - Hybrid, // Both pgvector + OpenSearch - SemanticOnly, // pgvector only (if OpenSearch down) - LexicalOnly, // OpenSearch only (if embedding model down) - LexicalFirst, // OpenSearch to narrow, then semantic rerank -} - -/// RRF (Reciprocal Rank Fusion) configuration -#[derive(Clone, Debug)] -pub struct RRFConfig { - pub k: f32, // Constant (usually 60) - pub retrieve_k: usize, // Top-K from each engine (usually 50) - pub final_k: usize, // Final top-K to return (usually 10) -} - -impl Default for RRFConfig { - fn default() -> Self { - Self { - k: 60.0, - retrieve_k: 50, - final_k: 10, - } - } -} - -/// Query Optimization Engine -pub struct QueryOptimizer { - enable_entity_extraction: bool, - enable_question_classification: bool, -} - -impl QueryOptimizer { - pub fn new() -> Self { - Self { - enable_entity_extraction: true, - enable_question_classification: true, - } - } - - /// Main entry point: construct query context from user input - pub async fn optimize_query(&self, raw_query: &str) -> Result { - // Stage 1: Normalize - let normalized = self.normalize_query(raw_query); - - // Stage 2: Tokenize - let tokens = self.tokenize(&normalized); - - // Stage 3: Extract entities - let entities = if self.enable_entity_extraction { - self.extract_entities(raw_query, &tokens) - } else { - HashMap::new() - }; - - // Stage 4: Analyze query characteristics - let token_count = tokens.len(); - let has_special_syntax = self.detect_special_syntax(raw_query); - let has_date_filters = self.detect_date_filters(&tokens); - let has_negation = self.detect_negation(&tokens); - - // Stage 5: Classify question type - let question_type = if self.enable_question_classification { - self.classify_question(raw_query, &tokens) - } else { - QuestionType::Open - }; - - // Stage 6: Route to search strategy - let (search_strategy, confidence) = self.route_query( - token_count, - has_special_syntax, - has_date_filters, - has_negation, - &question_type, - ); - - Ok(QueryContext { - raw_query: raw_query.to_string(), - normalized_query: normalized, - tokens, - entities, - embedding: None, - token_count, - has_special_syntax, - has_date_filters, - has_negation, - question_type, - search_strategy, - confidence, - }) - } - - /// Stage 1: Normalize query - fn normalize_query(&self, query: &str) -> String { - query - .trim() - .to_lowercase() - .replace(" ", " ") // Remove double spaces - } - - /// Stage 2: Tokenize - fn tokenize(&self, query: &str) -> Vec { - query - .split_whitespace() - .map(|s| s.to_string()) - .collect() - } - - /// Stage 3: Extract entities (years, names, keywords) - fn extract_entities(&self, raw_query: &str, tokens: &[String]) -> HashMap { - let mut entities = HashMap::new(); - - for token in tokens { - // Year detection: YYYY format - if token.len() == 4 { - if let Ok(year) = token.parse::() { - if year >= 2000 && year <= 2100 { - entities.insert("year".to_string(), token.clone()); - } - } - } - } - - // Detect quoted phrases - if raw_query.contains('"') { - let parts: Vec<&str> = raw_query.split('"').collect(); - if parts.len() >= 3 { - let quoted_phrase = parts[1].to_string(); - entities.insert("exact_phrase".to_string(), quoted_phrase); - } - } - - entities - } - - /// Stage 4: Detect special syntax (#tag, @mention, "phrases") - fn detect_special_syntax(&self, query: &str) -> bool { - query.contains('#') || query.contains('@') || query.contains('"') - } - - /// Stage 4: Detect date filters - fn detect_date_filters(&self, tokens: &[String]) -> bool { - let date_keywords = vec![ - "this", "last", "next", - "2024", "2025", "2026", - "january", "february", "march", "april", "may", "june", - "july", "august", "september", "october", "november", "december", - "week", "month", "year", "day", "today", "yesterday", "tomorrow", - ]; - - tokens.iter().any(|t| date_keywords.contains(&t.as_str())) - } - - /// Stage 4: Detect negation - fn detect_negation(&self, tokens: &[String]) -> bool { - tokens.iter().any(|t| t == "-" || t == "not" || t == "no" || t.starts_with("-")) - } - - /// Stage 5: Classify question type - fn classify_question(&self, raw_query: &str, tokens: &[String]) -> QuestionType { - let query_lower = raw_query.to_lowercase(); - - // Check first token for question words - if tokens.is_empty() { - return QuestionType::Open; - } - - let first_token = &tokens[0]; - - match first_token.as_str() { - // Procedural questions - t if t == "how" => QuestionType::Procedural, - t if t == "what" => { - if query_lower.contains("difference") || query_lower.contains("between") { - QuestionType::Comparative - } else { - QuestionType::Factual - } - } - // Comparative - t if t == "compare" || t == "compare" => QuestionType::Comparative, - // Troubleshooting - t if t == "fix" || t == "error" || t == "broken" || t == "debug" => { - QuestionType::Troubleshooting - } - // Navigational - t if t == "where" || t == "find" || t == "show" => QuestionType::Navigational, - _ => { - // Heuristics based on content - if query_lower.contains("how") { - QuestionType::Procedural - } else if query_lower.contains("fix") || query_lower.contains("error") { - QuestionType::Troubleshooting - } else { - QuestionType::Open - } - } - } - } - - /// Stage 6: Route to search strategy - fn route_query( - &self, - token_count: usize, - has_special_syntax: bool, - has_date_filters: bool, - _has_negation: bool, - question_type: &QuestionType, - ) -> (SearchStrategy, f32) { - // Very short queries: lexical better - if token_count < 3 { - return (SearchStrategy::LexicalOnly, 0.8); - } - - // Special syntax: preserve exact matches with lexical - if has_special_syntax { - if has_date_filters { - // Special syntax + dates = use lexical to narrow, then semantic - return (SearchStrategy::LexicalFirst, 0.85); - } else { - // Just special syntax = lexical only - return (SearchStrategy::LexicalOnly, 0.8); - } - } - - // Date filters present: use cascading (lexical → semantic) - if has_date_filters { - return (SearchStrategy::LexicalFirst, 0.9); - } - - // Question type heuristics - match question_type { - // Factual questions usually work well with semantic - QuestionType::Factual => (SearchStrategy::Hybrid, 0.9), - - // Procedural questions benefit from both (exact steps + understanding) - QuestionType::Procedural => (SearchStrategy::Hybrid, 0.95), - - // Troubleshooting needs both (exact errors + semantic understanding) - QuestionType::Troubleshooting => (SearchStrategy::Hybrid, 0.95), - - // Comparative: hybrid needed (understanding + multiple docs) - QuestionType::Comparative => (SearchStrategy::Hybrid, 0.9), - - // Navigational: lexical good for finding specific things - QuestionType::Navigational => (SearchStrategy::LexicalFirst, 0.85), - - // Open/general: hybrid default - QuestionType::Open => (SearchStrategy::Hybrid, 0.8), - } - } -} - -/// RRF Fusion Engine -pub struct RRFFusion { - config: RRFConfig, -} - -impl RRFFusion { - pub fn new(config: RRFConfig) -> Self { - Self { config } - } - - /// Fuse two ranked lists using Reciprocal Rank Fusion - pub fn fuse( - &self, - semantic_results: Vec<(String, f32)>, // (id, score) - lexical_results: Vec<(String, f32)>, - ) -> Vec<(String, f32)> { - use std::collections::HashMap; - - let mut fused_scores: HashMap = HashMap::new(); - - // Add semantic ranks with RRF formula: 1 / (k + rank) - for (rank, (id, _)) in semantic_results.into_iter().enumerate() { - let rrf_score = 1.0 / (self.config.k + (rank as f32) + 1.0); - fused_scores.insert(id, rrf_score); - } - - // Add lexical ranks (combine if already present) - for (rank, (id, _)) in lexical_results.into_iter().enumerate() { - let rrf_score = 1.0 / (self.config.k + (rank as f32) + 1.0); - *fused_scores.entry(id).or_insert(0.0) += rrf_score; - } - - // Sort by combined RRF score - let mut results: Vec<_> = fused_scores.into_iter().collect(); - results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); - - // Take top-k - results.truncate(self.config.final_k); - - results - } - - /// Alternative: Weighted Linear Fusion - pub fn fuse_weighted( - &self, - semantic_results: Vec<(String, f32)>, - lexical_results: Vec<(String, f32)>, - semantic_weight: f32, - lexical_weight: f32, - ) -> Vec<(String, f32)> { - use std::collections::HashMap; - - // Normalize scores to [0.0, 1.0] - let sem_norm = self.normalize_scores(&semantic_results); - let lex_norm = self.normalize_scores(&lexical_results); - - let sem_map: HashMap = sem_norm.into_iter().collect(); - let lex_map: HashMap = lex_norm.into_iter().collect(); - - // Merge all IDs - let mut all_ids = std::collections::HashSet::new(); - all_ids.extend(sem_map.keys().cloned()); - all_ids.extend(lex_map.keys().cloned()); - - // Calculate weighted scores - let mut results: Vec<_> = all_ids - .into_iter() - .map(|id| { - let sem_score = sem_map.get(&id).copied().unwrap_or(0.0); - let lex_score = lex_map.get(&id).copied().unwrap_or(0.0); - - let weighted_score = semantic_weight * sem_score + lexical_weight * lex_score; - (id, weighted_score) - }) - .collect(); - - results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); - results.truncate(self.config.final_k); - - results - } - - /// Normalize scores to [0.0, 1.0] range using min-max - fn normalize_scores(&self, results: &[(String, f32)]) -> Vec<(String, f32)> { - if results.is_empty() { - return Vec::new(); - } - - let min_score = results.iter().map(|(_, s)| s).fold(f32::INFINITY, |a, &b| a.min(b)); - let max_score = results.iter().map(|(_, s)| s).fold(f32::NEG_INFINITY, |a, &b| a.max(b)); - - let range = max_score - min_score; - - if range < 0.001 { - // All scores identical - return results.iter().map(|(id, _)| (id.clone(), 0.5)).collect(); - } - - results - .iter() - .map(|(id, score)| { - let normalized = (score - min_score) / range; - (id.clone(), normalized) - }) - .collect() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_query_optimization_procedural() { - let optimizer = QueryOptimizer::new(); - let ctx = optimizer.optimize_query("How do I fix kubernetes port 8080?").await.unwrap(); - - assert_eq!(ctx.question_type, QuestionType::Procedural); - assert_eq!(ctx.search_strategy, SearchStrategy::Hybrid); - assert!(ctx.confidence >= 0.9); - } - - #[tokio::test] - async fn test_query_optimization_short() { - let optimizer = QueryOptimizer::new(); - let ctx = optimizer.optimize_query("fix port").await.unwrap(); - - assert_eq!(ctx.token_count, 2); - assert_eq!(ctx.search_strategy, SearchStrategy::LexicalOnly); - } - - #[tokio::test] - async fn test_query_optimization_special_syntax() { - let optimizer = QueryOptimizer::new(); - let ctx = optimizer.optimize_query("kubernetes #networking @devops").await.unwrap(); - - assert!(ctx.has_special_syntax); - assert_eq!(ctx.search_strategy, SearchStrategy::LexicalOnly); - } - - #[test] - fn test_rrf_fusion() { - let fusion = RRFFusion::new(RRFConfig::default()); - - let semantic = vec![ - ("doc1".to_string(), 0.95), - ("doc2".to_string(), 0.88), - ("doc3".to_string(), 0.82), - ]; - - let lexical = vec![ - ("doc1".to_string(), 8.5), - ("doc4".to_string(), 7.2), - ("doc2".to_string(), 6.8), - ]; - - let fused = fusion.fuse(semantic, lexical); - - // doc1 should be top (in both) - assert_eq!(fused[0].0, "doc1"); - - // RRF score: doc1 appears in both lists (rank 1 in each) - // Score = 1/(60+1) + 1/(60+1) = 2/61 ≈ 0.0328 - assert!(fused[0].1 > 0.03 && fused[0].1 < 0.04, "Expected RRF score ~0.0328, got {}", fused[0].1); - } - - #[test] - fn test_weighted_fusion() { - let fusion = RRFFusion::new(RRFConfig::default()); - - let semantic = vec![ - ("doc1".to_string(), 0.95), - ("doc2".to_string(), 0.88), - ]; - - let lexical = vec![ - ("doc1".to_string(), 8.5), - ("doc3".to_string(), 7.2), - ]; - - let fused = fusion.fuse_weighted(semantic, lexical, 0.6, 0.4); - - // doc1 should rank highest (has both components) - assert_eq!(fused[0].0, "doc1"); - - // Score should be normalized and weighted - // 0.6 * (0.95/0.95) + 0.4 * (8.5/8.5) = 1.0 - assert!((fused[0].1 - 1.0).abs() < 0.01); - } -} diff --git a/crates/mem-cli/src/query_worker.rs b/crates/mem-cli/src/query_worker.rs deleted file mode 100644 index 9b1a9cf..0000000 --- a/crates/mem-cli/src/query_worker.rs +++ /dev/null @@ -1,111 +0,0 @@ -use anyhow::Result; -use mem_llm::{EmbeddingsClient, RerankClient}; -use mem_store::VectorStore; -use pgvector::Vector; -use serde::{Deserialize, Serialize}; - -/// Query result with provenance -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct QueryResult { - pub level: String, // "L0", "L1", "L2", "corpus" - pub score: f32, - pub text: String, - pub source: Option, - pub provenance: Vec, // parent IDs -} - -/// Query worker — semantic search + reranking -pub struct QueryWorker { - vector_store: std::sync::Arc, - embeddings: std::sync::Arc, - reranker: std::sync::Arc, -} - -impl QueryWorker { - /// Create query worker - pub fn new( - vector_store: VectorStore, - embeddings: EmbeddingsClient, - reranker: RerankClient, - ) -> Self { - Self { - vector_store: std::sync::Arc::new(vector_store), - embeddings: std::sync::Arc::new(embeddings), - reranker: std::sync::Arc::new(reranker), - } - } - - /// Execute semantic query: embed -> search vector -> rerank -> result - pub async fn query( - &self, - project: &str, - question: &str, - limit: Option, - ) -> Result> { - let limit = limit.unwrap_or(5); - - // Embed the question - let question_embedding = self.embeddings.embed_one(question).await?; - - // Search across all levels - let mut candidates = Vec::new(); - - // L2 synthesis (project-level) - if let Some(l2_result) = self.vector_store.search_l2(project, &question_embedding).await? { - candidates.push(QueryResult { - level: "L2".to_string(), - score: l2_result.score, - text: l2_result.item.content.clone(), - source: Some(format!("project:{}", project)), - provenance: vec![l2_result.item.id.to_string()], - }); - } - - // L1 per-query memories - let l1_results = self.vector_store.search_l1(project, &question_embedding, limit).await?; - for l1_result in l1_results { - candidates.push(QueryResult { - level: "L1".to_string(), - score: l1_result.score, - text: l1_result.item.content.clone(), - source: Some(format!("query:{}", l1_result.item.query_id)), - provenance: vec![l1_result.item.id.to_string()], - }); - } - - // Reference corpus - let corpus_results = self.vector_store.search_corpus(project, &question_embedding, limit).await?; - for corpus_result in corpus_results { - candidates.push(QueryResult { - level: "corpus".to_string(), - score: corpus_result.score, - text: corpus_result.item.content.clone(), - source: Some(format!("doc:{}", corpus_result.item.name)), - provenance: vec![corpus_result.item.id.to_string()], - }); - } - - // Rerank candidates by relevance to question - // TODO: wire actual cross-encoder reranking - // For now, return by vector similarity score - candidates.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); - candidates.truncate(limit as usize); - - Ok(candidates) - } - - /// Get project synthesis (L2) directly - pub async fn get_synthesis(&self, project: &str) -> Result> { - if let Some(l2) = self.vector_store.get_l2(project).await? { - Ok(Some(QueryResult { - level: "L2".to_string(), - score: 1.0, - text: l2.content, - source: Some(format!("project:{}", project)), - provenance: vec![l2.id.to_string()], - })) - } else { - Ok(None) - } - } -} diff --git a/crates/mem-cli/src/queue_adapter.rs b/crates/mem-cli/src/queue_adapter.rs deleted file mode 100644 index 5b84db6..0000000 --- a/crates/mem-cli/src/queue_adapter.rs +++ /dev/null @@ -1,336 +0,0 @@ -//! M8.2 — Unified Queue Adapter (SQS-compatible interface) -//! -//! Abstraction over external queue services (SQS, kmsvc, RabbitMQ, etc.) -//! Enables concurrent dual-write processing without database overhead. -//! -//! # Design -//! -//! Rather than storing queue state in the database, we leverage external queue -//! services via a unified API. This enables true horizontal scalability: -//! -//! ```text -//! Ingest Worker Queue Service (SQS/kmsvc) Dual-Write Workers -//! │ │ │ -//! │─── send_chunk() ────────────>│ │ -//! │ │ │ -//! └──────────────────────────────┤<─── receive_chunks(10) ────────┤ -//! │ │ -//! │<─── delete_chunk() ────────────┤ -//! │ (on success) │ -//! │ │ -//! │<─── change_visibility() ───────┤ -//! │ (on retry) │ -//! ``` -//! -//! # Implementations -//! - `SqsQueueAdapter`: AWS SQS backend -//! - `KmsvcQueueAdapter`: Kubernetes native messaging service -//! - In-memory for testing - -use async_trait::async_trait; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; -use anyhow::Result; - -/// SQS-compatible message envelope -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct QueueMessage { - /// Unique message ID (from queue service) - pub message_id: String, - - /// Original chunk UUID - pub chunk_id: Uuid, - - /// Message body (serialized JSON) - pub body: String, - - /// Receive count (number of times retrieved) - pub receive_count: i32, - - /// Receipt handle (for delete/change_visibility) - pub receipt_handle: String, - - /// Project context - pub project: String, - - /// Metadata - pub attributes: std::collections::HashMap, -} - -/// Queue statistics -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct QueueStats { - pub available_messages: i64, - pub in_flight_messages: i64, - pub dead_letter_messages: i64, - pub total_processed: i64, - pub average_delay_secs: i64, -} - -/// Unified queue adapter trait (SQS-like interface) -#[async_trait] -pub trait QueueAdapter: Send + Sync { - /// Send chunk message to queue - /// - /// # Arguments - /// * `chunk_id` — Unique chunk identifier - /// * `body` — Serialized message body (JSON) - /// * `project` — Project context - /// * `attributes` — Optional metadata (e.g., source, level, breadcrumb) - /// - /// # Returns - /// Message ID from queue service - async fn send_chunk( - &self, - chunk_id: Uuid, - body: String, - project: String, - attributes: std::collections::HashMap, - ) -> Result; - - /// Receive chunk messages from queue - /// - /// # Arguments - /// * `max_messages` — Max number of messages (1-10) - /// * `visibility_timeout_secs` — Visibility timeout duration - /// * `project` — Project filter (optional) - /// - /// # Returns - /// List of available messages - async fn receive_chunks( - &self, - max_messages: i32, - visibility_timeout_secs: i32, - project: Option<&str>, - ) -> Result>; - - /// Delete message from queue (after successful processing) - /// - /// # Arguments - /// * `message_id` — Message to delete - /// * `receipt_handle` — Receipt handle (for idempotency) - async fn delete_chunk(&self, message_id: &str, receipt_handle: &str) -> Result<()>; - - /// Change message visibility timeout - /// - /// Called when processing takes longer than expected. - async fn change_visibility( - &self, - message_id: &str, - receipt_handle: &str, - visibility_timeout_secs: i32, - ) -> Result<()>; - - /// Send message to dead-letter queue - /// - /// Called when message exceeds max receive count. - async fn send_to_dlq(&self, message_id: &str, receipt_handle: &str, reason: &str) -> Result<()>; - - /// Get queue statistics - async fn get_stats(&self, project: Option<&str>) -> Result; - - /// Purge queue (test/admin only) - async fn purge(&self, project: Option<&str>) -> Result; - - /// Health check - async fn health_check(&self) -> Result<()>; -} - -/// In-memory queue adapter (for testing and local development) -pub struct InMemoryQueueAdapter { - messages: std::sync::Arc>>, -} - -impl InMemoryQueueAdapter { - pub fn new() -> Self { - Self { - messages: std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new())), - } - } -} - -impl Default for InMemoryQueueAdapter { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl QueueAdapter for InMemoryQueueAdapter { - async fn send_chunk( - &self, - chunk_id: Uuid, - body: String, - project: String, - attributes: std::collections::HashMap, - ) -> Result { - let message_id = format!("msg-{}", Uuid::new_v4()); - let receipt_handle = format!("handle-{}", Uuid::new_v4()); - - let msg = QueueMessage { - message_id: message_id.clone(), - chunk_id, - body, - receive_count: 0, - receipt_handle, - project, - attributes, - }; - - let mut msgs = self.messages.lock().await; - msgs.push(msg); - - Ok(message_id) - } - - async fn receive_chunks( - &self, - max_messages: i32, - _visibility_timeout_secs: i32, - project: Option<&str>, - ) -> Result> { - let mut msgs = self.messages.lock().await; - let max = max_messages.min(10).max(1) as usize; - let drain_count = msgs.len().min(max); - - let result: Vec<_> = msgs - .drain(..drain_count) - .filter(|m| project.is_none() || m.project.as_str() == project.unwrap()) - .collect(); - - Ok(result) - } - - async fn delete_chunk(&self, message_id: &str, _receipt_handle: &str) -> Result<()> { - let mut msgs = self.messages.lock().await; - msgs.retain(|m| m.message_id != message_id); - Ok(()) - } - - async fn change_visibility( - &self, - _message_id: &str, - _receipt_handle: &str, - _visibility_timeout_secs: i32, - ) -> Result<()> { - // No-op for in-memory - Ok(()) - } - - async fn send_to_dlq(&self, message_id: &str, _receipt_handle: &str, _reason: &str) -> Result<()> { - let mut msgs = self.messages.lock().await; - msgs.retain(|m| m.message_id != message_id); - Ok(()) - } - - async fn get_stats(&self, _project: Option<&str>) -> Result { - let msgs = self.messages.lock().await; - Ok(QueueStats { - available_messages: msgs.len() as i64, - in_flight_messages: 0, - dead_letter_messages: 0, - total_processed: 0, - average_delay_secs: 0, - }) - } - - async fn purge(&self, _project: Option<&str>) -> Result { - let mut msgs = self.messages.lock().await; - let count = msgs.len(); - msgs.clear(); - Ok(count) - } - - async fn health_check(&self) -> Result<()> { - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_in_memory_send_chunk() { - let queue = InMemoryQueueAdapter::new(); - let msg_id = queue - .send_chunk( - Uuid::new_v4(), - r#"{"content": "test"}"#.to_string(), - "test-project".to_string(), - std::collections::HashMap::new(), - ) - .await - .unwrap(); - - assert!(msg_id.starts_with("msg-")); - } - - #[tokio::test] - async fn test_in_memory_receive_chunks() { - let queue = InMemoryQueueAdapter::new(); - - for i in 0..5 { - queue - .send_chunk( - Uuid::new_v4(), - format!(r#"{{"content": "test{}"}}"#, i), - "test-project".to_string(), - std::collections::HashMap::new(), - ) - .await - .ok(); - } - - let messages = queue - .receive_chunks(3, 30, Some("test-project")) - .await - .unwrap(); - - assert_eq!(messages.len(), 3); - } - - #[tokio::test] - async fn test_in_memory_delete_chunk() { - let queue = InMemoryQueueAdapter::new(); - - let msg_id = queue - .send_chunk( - Uuid::new_v4(), - "body".to_string(), - "test".to_string(), - std::collections::HashMap::new(), - ) - .await - .unwrap(); - - queue.delete_chunk(&msg_id, "handle").await.unwrap(); - - let msgs = queue.receive_chunks(10, 30, None).await.unwrap(); - assert_eq!(msgs.len(), 0); - } - - #[tokio::test] - async fn test_queue_stats() { - let queue = InMemoryQueueAdapter::new(); - - queue - .send_chunk( - Uuid::new_v4(), - "body".to_string(), - "test".to_string(), - std::collections::HashMap::new(), - ) - .await - .ok(); - - let stats = queue.get_stats(None).await.unwrap(); - assert_eq!(stats.available_messages, 1); - } - - #[tokio::test] - async fn test_health_check() { - let queue = InMemoryQueueAdapter::new(); - assert!(queue.health_check().await.is_ok()); - } -} diff --git a/crates/mem-cli/src/queue_worker.rs b/crates/mem-cli/src/queue_worker.rs deleted file mode 100644 index b64b9a7..0000000 --- a/crates/mem-cli/src/queue_worker.rs +++ /dev/null @@ -1,402 +0,0 @@ -//! M8.2 — Queue Worker for Concurrent Dual-Write Processing -//! -//! Background task that receives messages from the queue and processes them -//! via DualWriteIndexer. Runs concurrently with ingest, improving throughput. -//! -//! # Architecture -//! -//! ```text -//! IngestWorker (fast path) QueueWorker (background) -//! │ │ -//! ├─ chunk_input │ -//! │ (embedding) │ -//! │ │ -//! ├─ queue.send_chunk()────┐ │ -//! │ (returns immediately) │ │ -//! │ │ │ -//! └─ continues... │ │ -//! │ │ -//! ├─ queue.receive_chunks(10, 30) -//! │ (long-poll, up to 30s) -//! │ -//! ├─ for each message: -//! │ - process_queued_chunk() -//! │ - embed_one() [happens here] -//! │ - write_pgvector() -//! │ - write_opensearch() -//! │ - delete_chunk() on success -//! │ - change_visibility() on retry -//! │ -//! └─ loop back to receive -//! ``` -//! -//! Benefits: -//! - Ingest path is decoupled from embedding/pgvector/OpenSearch writes -//! - Multiple workers can process messages concurrently -//! - Non-blocking: queue.send_chunk() returns immediately -//! - Fault-tolerant: failed messages auto-retry with exponential backoff - -use anyhow::{anyhow, Result}; -use std::sync::Arc; -use std::time::Duration; -use tokio::time::sleep; -use tracing::{debug, error, info, warn}; - -use crate::dual_write_indexer::DualWriteIndexer; -use crate::queue_adapter::QueueAdapter; -use mem_llm::EmbeddingsClient; - -/// Configuration for queue worker -#[derive(Debug, Clone)] -pub struct QueueWorkerConfig { - /// Max messages per receive (1-10) - pub max_messages_per_batch: i32, - - /// Visibility timeout for processing (seconds) - pub visibility_timeout_secs: i32, - - /// Time to wait for messages (0-20 seconds) - pub wait_time_secs: i32, - - /// Project to process (None = all projects) - pub project: Option, - - /// Max retries before DLQ - pub max_retries: i32, - - /// Retry backoff: exponential starting from this value (seconds) - pub retry_backoff_initial_secs: i32, - - /// Poll interval when queue is empty (seconds) - pub empty_poll_interval_secs: u64, - - /// Enable metrics collection - pub enable_metrics: bool, -} - -impl Default for QueueWorkerConfig { - fn default() -> Self { - Self { - max_messages_per_batch: 10, - visibility_timeout_secs: 300, // 5 minutes - wait_time_secs: 20, // Long-poll timeout - project: None, - max_retries: 3, - retry_backoff_initial_secs: 60, - empty_poll_interval_secs: 5, - enable_metrics: true, - } - } -} - -/// Metrics for worker execution -#[derive(Debug, Clone, Default)] -pub struct WorkerMetrics { - pub messages_received: u64, - pub messages_processed: u64, - pub messages_failed: u64, - pub messages_dlq: u64, - pub total_processing_time_ms: u64, -} - -/// Queue worker for processing dual-write messages -pub struct QueueWorker { - indexer: Arc, - embeddings: Arc, - config: QueueWorkerConfig, - metrics: Arc>, -} - -impl QueueWorker { - /// Create new queue worker - pub fn new( - indexer: Arc, - embeddings: Arc, - config: QueueWorkerConfig, - ) -> Self { - Self { - indexer, - embeddings, - config, - metrics: Arc::new(tokio::sync::RwLock::new(WorkerMetrics::default())), - } - } - - /// Start worker (blocking loop) - pub async fn start(&self) -> Result<()> { - info!("Queue worker starting: config={:?}", self.config); - - loop { - match self.process_batch().await { - Ok(count) => { - if count == 0 { - // Empty batch: sleep before retrying - debug!( - "Queue empty, waiting {}s before retry", - self.config.empty_poll_interval_secs - ); - sleep(Duration::from_secs(self.config.empty_poll_interval_secs)).await; - } - } - Err(e) => { - error!("Worker error (will retry): {}", e); - sleep(Duration::from_secs(5)).await; - } - } - } - } - - /// Process one batch of messages from queue - async fn process_batch(&self) -> Result { - let queue = &self.indexer.queue; - - // Receive messages - let messages = queue - .receive_chunks( - self.config.max_messages_per_batch, - self.config.visibility_timeout_secs, - self.config.project.as_deref(), - ) - .await?; - - let batch_size = messages.len(); - if batch_size == 0 { - return Ok(0); - } - - let mut metrics = self.metrics.write().await; - metrics.messages_received += batch_size as u64; - drop(metrics); - - // Process each message concurrently - let handles: Vec<_> = messages - .into_iter() - .map(|msg| { - let indexer = self.indexer.clone(); - let embeddings = self.embeddings.clone(); - let config = self.config.clone(); - let metrics = self.metrics.clone(); - - tokio::spawn(async move { - Self::process_message(indexer, embeddings, config, metrics, msg).await - }) - }) - .collect(); - - // Wait for all to complete - for handle in handles { - if let Err(e) = handle.await { - error!("Worker task panicked: {}", e); - } - } - - Ok(batch_size) - } - - /// Process a single message - async fn process_message( - indexer: Arc, - embeddings: Arc, - config: QueueWorkerConfig, - metrics: Arc>, - message: crate::queue_adapter::QueueMessage, - ) -> Result<()> { - let start = std::time::Instant::now(); - let message_id = message.message_id.clone(); - let receipt_handle = message.receipt_handle.clone(); - - debug!("Processing message: {}", message_id); - - // Parse message body - let body: serde_json::Value = match serde_json::from_str(&message.body) { - Ok(b) => b, - Err(e) => { - error!("Failed to parse message body: {}", e); - indexer - .queue - .send_to_dlq(&message_id, &receipt_handle, "invalid_json") - .await - .ok(); - - let mut m = metrics.write().await; - m.messages_dlq += 1; - return Err(e.into()); - } - }; - - // Extract chunk_id - let chunk_id = match body["chunk_id"].as_str() { - Some(id) => match uuid::Uuid::parse_str(id) { - Ok(u) => u, - Err(e) => { - error!("Invalid chunk_id: {}", e); - indexer - .queue - .send_to_dlq(&message_id, &receipt_handle, "invalid_uuid") - .await - .ok(); - - let mut m = metrics.write().await; - m.messages_dlq += 1; - return Err(e.into()); - } - }, - None => { - error!("Missing chunk_id in message"); - indexer - .queue - .send_to_dlq(&message_id, &receipt_handle, "missing_chunk_id") - .await - .ok(); - - let mut m = metrics.write().await; - m.messages_dlq += 1; - return Err(anyhow!("Missing chunk_id")); - } - }; - - // Extract content - let content = match body["content"].as_str() { - Some(c) => c.to_string(), - None => { - error!("Missing content in message"); - indexer - .queue - .send_to_dlq(&message_id, &receipt_handle, "missing_content") - .await - .ok(); - - let mut m = metrics.write().await; - m.messages_dlq += 1; - return Err(anyhow!("Missing content")); - } - }; - - // Compute embedding - let embedding_vec = match embeddings.embed_one(&content).await { - Ok(vec) => vec, - Err(e) => { - warn!("Embedding failed, extending visibility for retry: {}", e); - indexer - .queue - .change_visibility(&message_id, &receipt_handle, 300) - .await - .ok(); - - let mut m = metrics.write().await; - m.messages_failed += 1; - return Err(e); - } - }; - - // Convert pgvector::Vector to Vec - let embedding: Vec = embedding_vec.to_vec(); - - // Process dual-write - match indexer.process_queued_chunk(&message, &embedding).await { - Ok(result) => { - if result.pgvector_success && !result.opensearch_pending { - // Success: already deleted by process_queued_chunk - debug!("Message processed successfully: {}", message_id); - - let elapsed = start.elapsed().as_millis() as u64; - let mut m = metrics.write().await; - m.messages_processed += 1; - m.total_processing_time_ms += elapsed; - } else if result.pgvector_success && result.opensearch_pending { - // pgvector OK, OpenSearch pending: visibility already extended - warn!("Message will retry: {}", message_id); - - let mut m = metrics.write().await; - m.messages_failed += 1; - } else { - // pgvector failed: visibility already extended - warn!("pgvector write failed, will retry: {}", message_id); - - let mut m = metrics.write().await; - m.messages_failed += 1; - } - - Ok(()) - } - Err(e) => { - // Check receive count - if message.receive_count >= config.max_retries { - error!( - "Message max retries exceeded ({}), sending to DLQ: {}", - message.receive_count, message_id - ); - indexer - .queue - .send_to_dlq(&message_id, &receipt_handle, "max_retries") - .await - .ok(); - - let mut m = metrics.write().await; - m.messages_dlq += 1; - } else { - // Extend visibility for retry - warn!( - "Message processing failed (retry {}), extending visibility: {}", - message.receive_count, message_id - ); - indexer - .queue - .change_visibility(&message_id, &receipt_handle, 300) - .await - .ok(); - - let mut m = metrics.write().await; - m.messages_failed += 1; - } - - Err(e) - } - } - } - - /// Get current metrics - pub async fn metrics(&self) -> WorkerMetrics { - self.metrics.read().await.clone() - } - - /// Reset metrics - pub async fn reset_metrics(&self) { - let mut m = self.metrics.write().await; - *m = WorkerMetrics::default(); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_queue_worker_config_default() { - let config = QueueWorkerConfig::default(); - assert_eq!(config.max_messages_per_batch, 10); - assert_eq!(config.visibility_timeout_secs, 300); - assert_eq!(config.wait_time_secs, 20); - assert_eq!(config.max_retries, 3); - } - - #[test] - fn test_worker_metrics_default() { - let metrics = WorkerMetrics::default(); - assert_eq!(metrics.messages_received, 0); - assert_eq!(metrics.messages_processed, 0); - } - - #[test] - fn test_queue_worker_config_custom() { - let config = QueueWorkerConfig { - max_messages_per_batch: 5, - visibility_timeout_secs: 600, - project: Some("test-proj".to_string()), - ..Default::default() - }; - - assert_eq!(config.max_messages_per_batch, 5); - assert_eq!(config.project, Some("test-proj".to_string())); - } -} diff --git a/crates/mem-cli/src/rate_limiter.rs b/crates/mem-cli/src/rate_limiter.rs deleted file mode 100644 index fdefcca..0000000 --- a/crates/mem-cli/src/rate_limiter.rs +++ /dev/null @@ -1,243 +0,0 @@ -use std::collections::HashMap; -use std::sync::{Arc, Mutex}; -use std::time::Instant; - -/// Rate limit error with retry guidance -#[derive(Debug, Clone)] -pub struct RateLimitError { - pub retry_after_seconds: u64, - pub limit_window_secs: u64, - pub reason: String, -} - -impl RateLimitError { - pub fn reason(&self) -> String { - format!( - "{} (retry after {} seconds, window: {} seconds)", - self.reason, self.retry_after_seconds, self.limit_window_secs - ) - } -} - -/// Token bucket for a single endpoint -#[derive(Debug, Clone)] -struct TokenBucket { - tokens: f64, - last_refill: Instant, - capacity: f64, // max tokens (per hour) - refill_rate: f64, // tokens per second -} - -impl TokenBucket { - fn new(capacity: f64, refill_rate: f64) -> Self { - Self { - tokens: capacity, - last_refill: Instant::now(), - capacity, - refill_rate, - } - } - - /// Refill tokens based on elapsed time - fn refill(&mut self) { - let now = Instant::now(); - let elapsed = now.duration_since(self.last_refill).as_secs_f64(); - let refilled = elapsed * self.refill_rate; - - self.tokens = (self.tokens + refilled).min(self.capacity); - self.last_refill = now; - } - - /// Try to consume 1 token. Returns Ok if successful, Err(retry_after_secs) if rate limited. - fn try_consume(&mut self) -> Result<(), u64> { - self.refill(); - - if self.tokens >= 1.0 { - self.tokens -= 1.0; - return Ok(()); - } - - // Rate limited: estimate time until next token available - let tokens_needed = 1.0 - self.tokens; - let retry_after = (tokens_needed / self.refill_rate).ceil() as u64; - Err(retry_after.max(1)) - } -} - -/// Rate limiter with per-apikey, per-endpoint buckets -pub struct RateLimiter { - buckets: Arc>>>>, - limit_config: LimitConfig, -} - -#[derive(Clone, Debug)] -pub struct LimitConfig { - pub ingest_per_hour: f64, - pub query_per_hour: f64, - pub projects_per_hour: f64, - pub burst_per_second: f64, // Currently unused but kept for API compatibility -} - -impl Default for LimitConfig { - fn default() -> Self { - Self { - ingest_per_hour: 100.0, - query_per_hour: 1000.0, - projects_per_hour: 100.0, - burst_per_second: 10.0, - } - } -} - -impl RateLimiter { - pub fn new(config: LimitConfig) -> Self { - Self { - buckets: Arc::new(Mutex::new(HashMap::new())), - limit_config: config, - } - } - - /// Get or create bucket for apikey + endpoint - fn get_or_create_bucket(&self, apikey_endpoint: &str) -> Arc> { - let mut buckets = self.buckets.lock().unwrap(); - let config = &self.limit_config; - - if !buckets.contains_key(apikey_endpoint) { - // Determine limit based on endpoint - let capacity = if apikey_endpoint.contains("/memory/ingest") { - config.ingest_per_hour - } else if apikey_endpoint.contains("/memory/query") { - config.query_per_hour - } else if apikey_endpoint.contains("/memory/projects") { - config.projects_per_hour - } else { - // Unlimited for unknown endpoints - f64::INFINITY - }; - - let refill_rate = if capacity.is_infinite() { - f64::INFINITY - } else { - capacity / 3600.0 // per second - }; - - let bucket = TokenBucket::new(capacity, refill_rate); - buckets.insert(apikey_endpoint.to_string(), Arc::new(Mutex::new(bucket))); - } - - buckets[apikey_endpoint].clone() - } - - /// Check rate limit for apikey + endpoint. Returns Ok or Err with retry guidance. - pub fn check(&self, apikey: &str, endpoint: &str) -> Result<(), RateLimitError> { - let key = format!("{}::{}", apikey, endpoint); - let bucket = self.get_or_create_bucket(&key); - let mut b = bucket.lock().unwrap(); - - match b.try_consume() { - Ok(_) => Ok(()), - Err(retry_after) => { - let window_secs = if endpoint.contains("/memory/ingest") { - 3600 - } else if endpoint.contains("/memory/query") { - 3600 - } else if endpoint.contains("/memory/projects") { - 3600 - } else { - 3600 - }; - - Err(RateLimitError { - retry_after_seconds: retry_after, - limit_window_secs: window_secs, - reason: format!( - "rate_limit_exceeded for {}", - endpoint - ), - }) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_token_bucket_refill() { - let mut bucket = TokenBucket::new(100.0, 100.0 / 3600.0); - assert!(bucket.try_consume().is_ok()); - // After one consumption, should have 99 tokens - assert_eq!((bucket.tokens * 1.0) as i64, 99); - } - - #[test] - fn test_rate_limit_within_capacity() { - let config = LimitConfig { - ingest_per_hour: 5.0, - query_per_hour: 10.0, - projects_per_hour: 10.0, - burst_per_second: 10.0, - }; - let limiter = RateLimiter::new(config); - - // First 5 should succeed - for _ in 0..5 { - assert!(limiter.check("apikey1", "/memory/ingest").is_ok()); - } - - // 6th should fail - let err = limiter.check("apikey1", "/memory/ingest"); - assert!(err.is_err()); - if let Err(e) = err { - assert!(e.retry_after_seconds > 0); - } - } - - #[test] - fn test_per_apikey_isolation() { - let config = LimitConfig { - ingest_per_hour: 5.0, - query_per_hour: 10.0, - projects_per_hour: 10.0, - burst_per_second: 10.0, - }; - let limiter = RateLimiter::new(config); - - // apikey1 uses up 5 ingest requests - for _ in 0..5 { - assert!(limiter.check("apikey1", "/memory/ingest").is_ok()); - } - assert!(limiter.check("apikey1", "/memory/ingest").is_err()); - - // apikey2 should have its own 5 - for _ in 0..5 { - assert!(limiter.check("apikey2", "/memory/ingest").is_ok()); - } - assert!(limiter.check("apikey2", "/memory/ingest").is_err()); - } - - #[test] - fn test_per_endpoint_isolation() { - let config = LimitConfig { - ingest_per_hour: 5.0, - query_per_hour: 10.0, - projects_per_hour: 10.0, - burst_per_second: 10.0, - }; - let limiter = RateLimiter::new(config); - - // Use up 5 ingest - for _ in 0..5 { - assert!(limiter.check("apikey1", "/memory/ingest").is_ok()); - } - assert!(limiter.check("apikey1", "/memory/ingest").is_err()); - - // Query should have separate 10 limit - for _ in 0..10 { - assert!(limiter.check("apikey1", "/memory/query").is_ok()); - } - assert!(limiter.check("apikey1", "/memory/query").is_err()); - } -} diff --git a/crates/mem-cli/src/simple_hybrid_search.rs b/crates/mem-cli/src/simple_hybrid_search.rs deleted file mode 100644 index d0aa444..0000000 --- a/crates/mem-cli/src/simple_hybrid_search.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! M8.6 — Simple Hybrid Search (Semantic + Lexical Fusion) -//! -//! Combines pgvector semantic search with OpenSearch lexical search using RRF. -//! Simpler than HybridQueryWorker - uses only existing VectorStore/OpenSearchClient APIs. - -use anyhow::Result; -use mem_store::VectorStore; -use pgvector::Vector; -use serde::{Deserialize, Serialize}; -use std::sync::Arc; - -use crate::opensearch_client::OpenSearchClient; -use crate::query_optimizer::{RRFFusion, RRFConfig}; - -/// Hybrid search result with score breakdown -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SimpleHybridResult { - pub id: String, - pub content: String, - pub project: String, - pub semantic_score: Option, - pub lexical_score: Option, - pub final_score: f32, - pub rank: usize, -} - -/// Simple hybrid search orchestrator -pub struct SimpleHybridSearch { - vector_store: Arc, - opensearch: Option>, - rrf: RRFFusion, -} - -impl SimpleHybridSearch { - pub fn new( - vector_store: Arc, - opensearch: Option>, - ) -> Self { - // Create RRF with default config (k=60 per academic standards) - let rrf_config = RRFConfig { - k: 60.0, - retrieve_k: 50, - final_k: 10, - }; - let rrf = RRFFusion::new(rrf_config); - - Self { - vector_store, - opensearch, - rrf, - } - } - - /// Execute hybrid search: semantic + lexical with RRF fusion - pub async fn search( - &self, - project: &str, - query: &str, - embedding: &Vector, - jwt_token: &str, - limit: usize, - ) -> Result> { - // 1. Semantic search (pgvector) - let semantic_results = self - .vector_store - .search_l1(project, embedding, limit as i64) - .await?; - - let semantic_scores: Vec<(String, f32)> = semantic_results - .into_iter() - .enumerate() - .map(|(i, result)| { - // Rank to score conversion - let rank_score = 1.0 / (i as f32 + 1.0); - (result.item.id.to_string(), rank_score) - }) - .collect(); - - // 2. Lexical search (OpenSearch) - optional if available - // TODO: Implement OpenSearchClient.search() method - let lexical_scores: Vec<(String, f32)> = vec![]; - - // 3. Fuse with RRF - let fused = self.rrf.fuse(semantic_scores.clone(), lexical_scores.clone()); - - // 4. Convert to response format - let results = fused - .into_iter() - .enumerate() - .map(|(rank, (id, score))| { - let semantic_score = semantic_scores - .iter() - .find(|(sid, _)| sid == &id) - .map(|(_, s)| *s); - - let lexical_score = lexical_scores - .iter() - .find(|(sid, _)| sid == &id) - .map(|(_, s)| *s); - - SimpleHybridResult { - id: id.clone(), - content: String::new(), // Would fetch from store - project: project.to_string(), - semantic_score, - lexical_score, - final_score: score, - rank: rank + 1, - } - }) - .collect(); - - Ok(results) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_simple_hybrid_result_creation() { - let result = SimpleHybridResult { - id: "doc1".to_string(), - content: "test".to_string(), - project: "test".to_string(), - semantic_score: Some(0.95), - lexical_score: Some(8.5), - final_score: 0.067, - rank: 1, - }; - - assert_eq!(result.id, "doc1"); - assert_eq!(result.rank, 1); - assert!(result.semantic_score.is_some()); - } -} diff --git a/crates/mem-ingest/src/memorability_gate.rs b/crates/mem-ingest/src/memorability_gate.rs index 2343510..8bc0621 100644 --- a/crates/mem-ingest/src/memorability_gate.rs +++ b/crates/mem-ingest/src/memorability_gate.rs @@ -243,7 +243,7 @@ pub struct FilterStatistics { #[cfg(test)] mod tests { use super::*; - use mem_core::entity::Entity; + use mem_core::entity::{Entity, EntityType}; fn create_test_entity(name: &str) -> Entity { Entity::new("poimen", name, EntityType::Person) diff --git a/crates/mem-store/src/audit_logger.rs b/crates/mem-store/src/audit_logger.rs index 7613cfd..721652d 100644 --- a/crates/mem-store/src/audit_logger.rs +++ b/crates/mem-store/src/audit_logger.rs @@ -177,6 +177,7 @@ pub struct AuditedEntityRepo { #[cfg(test)] mod tests { use super::*; + use serde_json::json; #[test] fn test_diff_fields_modified() { diff --git a/crates/mem-store/src/db_repo.rs b/crates/mem-store/src/db_repo.rs deleted file mode 100644 index af4e85c..0000000 --- a/crates/mem-store/src/db_repo.rs +++ /dev/null @@ -1,543 +0,0 @@ -/// PostgreSQL repository implementation for Phase 2.6 DB Integration. -/// -/// Connects ingest pipeline to persistent storage. -/// Handles transactions, error recovery, and audit logging. - -use sqlx::{Pool, Postgres, Row, Transaction, Error as SqlxError}; -use serde::{Deserialize, Serialize}; -use chrono::{DateTime, Utc}; -use mem_core::entity::Entity; -use mem_core::edge::Edge; -use crate::entity_repo::EntityRepoOps; -use crate::edge_repo::EdgeRepoOps; - -/// Database connection error types -#[derive(Debug, Clone)] -pub enum DbError { - ConnectionFailed(String), - QueryFailed(String), - TransactionFailed(String), - DuplicateKey(String), - NotFound(String), - InvalidData(String), -} - -impl std::fmt::Display for DbError { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - match self { - DbError::ConnectionFailed(msg) => write!(f, "Connection failed: {}", msg), - DbError::QueryFailed(msg) => write!(f, "Query failed: {}", msg), - DbError::TransactionFailed(msg) => write!(f, "Transaction failed: {}", msg), - DbError::DuplicateKey(msg) => write!(f, "Duplicate key: {}", msg), - DbError::NotFound(msg) => write!(f, "Not found: {}", msg), - DbError::InvalidData(msg) => write!(f, "Invalid data: {}", msg), - } - } -} - -impl std::error::Error for DbError {} - -/// PostgreSQL repository pool -pub struct DbPool { - pool: Pool, -} - -impl DbPool { - /// Create new DB pool from connection string - pub async fn new(database_url: &str) -> Result { - let pool = Pool::::connect(database_url) - .await - .map_err(|e| DbError::ConnectionFailed(e.to_string()))?; - - Ok(DbPool { pool }) - } - - /// Get pool for queries - pub fn pool(&self) -> &Pool { - &self.pool - } - - /// Test connection - pub async fn health_check(&self) -> Result<(), DbError> { - sqlx::query("SELECT 1") - .fetch_one(&self.pool) - .await - .map_err(|e| DbError::ConnectionFailed(e.to_string()))?; - Ok(()) - } -} - -/// Persistent entity repository -pub struct PersistentEntityRepo { - pool: Pool, -} - -impl PersistentEntityRepo { - pub fn new(pool: Pool) -> Self { - Self { pool } - } - - /// Save entity to database (idempotent) - pub async fn save(&self, entity: &Entity) -> Result { - let query = r#" - INSERT INTO memory_entity (id, entity_type, name, description, embedding, created_at, updated_at) - VALUES ($1, $2, $3, $4, $5, $6, $7) - ON CONFLICT(id) DO UPDATE SET - name = EXCLUDED.name, - description = EXCLUDED.description, - updated_at = EXCLUDED.updated_at - RETURNING id; - "#; - - let id = sqlx::query_scalar::<_, String>(query) - .bind(&entity.id) - .bind(&entity.entity_type) - .bind(&entity.name) - .bind(&entity.description) - .bind(&entity.embedding) - .bind(Utc::now()) - .bind(Utc::now()) - .fetch_one(&self.pool) - .await - .map_err(|e| { - if e.to_string().contains("duplicate") { - DbError::DuplicateKey(format!("Entity {} already exists", entity.id)) - } else { - DbError::QueryFailed(e.to_string()) - } - })?; - - Ok(id) - } - - /// Get entity by ID - pub async fn get(&self, id: &str) -> Result, DbError> { - let query = r#" - SELECT id, entity_type, name, description, embedding, created_at, updated_at - FROM memory_entity - WHERE id = $1 AND deleted_at IS NULL; - "#; - - let row = sqlx::query(query) - .bind(id) - .fetch_optional(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(row.map(|r| Entity { - id: r.get("id"), - entity_type: r.get("entity_type"), - name: r.get("name"), - description: r.get("description"), - embedding: r.get("embedding"), - created_at: r.get("created_at"), - updated_at: r.get("updated_at"), - })) - } - - /// List entities with pagination - pub async fn list(&self, limit: i64, offset: i64) -> Result, DbError> { - let query = r#" - SELECT id, entity_type, name, description, embedding, created_at, updated_at - FROM memory_entity - WHERE deleted_at IS NULL - ORDER BY created_at DESC - LIMIT $1 OFFSET $2; - "#; - - let rows = sqlx::query(query) - .bind(limit) - .bind(offset) - .fetch_all(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(rows.iter().map(|r| Entity { - id: r.get("id"), - entity_type: r.get("entity_type"), - name: r.get("name"), - description: r.get("description"), - embedding: r.get("embedding"), - created_at: r.get("created_at"), - updated_at: r.get("updated_at"), - }).collect()) - } - - /// Soft delete entity - pub async fn delete(&self, id: &str) -> Result<(), DbError> { - let query = r#" - UPDATE memory_entity - SET deleted_at = $1 - WHERE id = $2; - "#; - - sqlx::query(query) - .bind(Utc::now()) - .bind(id) - .execute(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(()) - } -} - -/// Persistent edge repository -pub struct PersistentEdgeRepo { - pool: Pool, -} - -impl PersistentEdgeRepo { - pub fn new(pool: Pool) -> Self { - Self { pool } - } - - /// Save edge to database (idempotent) - pub async fn save(&self, edge: &Edge) -> Result { - let query = r#" - INSERT INTO memory_edge (id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) - ON CONFLICT(id) DO UPDATE SET - strength = EXCLUDED.strength, - t_invalid = EXCLUDED.t_invalid, - t_expired = EXCLUDED.t_expired - RETURNING id; - "#; - - let id = sqlx::query_scalar::<_, String>(query) - .bind(&edge.id) - .bind(&edge.source_id) - .bind(&edge.target_id) - .bind(&edge.relation_type) - .bind(&edge.fact) - .bind(edge.strength) - .bind(edge.t_valid) - .bind(edge.t_invalid) - .bind(edge.t_created) - .bind(edge.t_expired) - .fetch_one(&self.pool) - .await - .map_err(|e| { - if e.to_string().contains("duplicate") { - DbError::DuplicateKey(format!("Edge {} already exists", edge.id)) - } else { - DbError::QueryFailed(e.to_string()) - } - })?; - - Ok(id) - } - - /// Get edge by ID - pub async fn get(&self, id: &str) -> Result, DbError> { - let query = r#" - SELECT id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired - FROM memory_edge - WHERE id = $1 AND t_expired IS NULL; - "#; - - let row = sqlx::query(query) - .bind(id) - .fetch_optional(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(row.map(|r| Edge { - id: r.get("id"), - source_id: r.get("source_id"), - target_id: r.get("target_id"), - relation_type: r.get("relation_type"), - fact: r.get("fact"), - strength: r.get("strength"), - t_valid: r.get("t_valid"), - t_invalid: r.get("t_invalid"), - t_created: r.get("t_created"), - t_expired: r.get("t_expired"), - })) - } - - /// List edges for a source entity - pub async fn list_from(&self, source_id: &str, limit: i64) -> Result, DbError> { - let query = r#" - SELECT id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired - FROM memory_edge - WHERE source_id = $1 AND t_expired IS NULL AND t_invalid IS NULL - ORDER BY t_created DESC - LIMIT $2; - "#; - - let rows = sqlx::query(query) - .bind(source_id) - .bind(limit) - .fetch_all(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(rows.iter().map(|r| Edge { - id: r.get("id"), - source_id: r.get("source_id"), - target_id: r.get("target_id"), - relation_type: r.get("relation_type"), - fact: r.get("fact"), - strength: r.get("strength"), - t_valid: r.get("t_valid"), - t_invalid: r.get("t_invalid"), - t_created: r.get("t_created"), - t_expired: r.get("t_expired"), - }).collect()) - } - - /// Mark edge as contradicted (soft delete) - pub async fn invalidate(&self, id: &str) -> Result<(), DbError> { - let query = r#" - UPDATE memory_edge - SET t_invalid = $1 - WHERE id = $2; - "#; - - sqlx::query(query) - .bind(Utc::now()) - .bind(id) - .execute(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(()) - } -} - -/// Review queue entry for human verification -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ReviewQueueEntry { - pub id: String, - pub extraction_type: String, // "entity" | "edge" | "contradiction" - pub content: serde_json::Value, // Full extracted data - pub status: String, // "pending" | "approved" | "rejected" - pub created_at: DateTime, - pub reviewed_at: Option>, - pub reviewed_by: Option, // User ID who reviewed - pub rejection_reason: Option, -} - -/// Review queue repository -pub struct ReviewQueueRepo { - pool: Pool, -} - -impl ReviewQueueRepo { - pub fn new(pool: Pool) -> Self { - Self { pool } - } - - /// Add item to review queue - pub async fn enqueue(&self, entry: &ReviewQueueEntry) -> Result { - let query = r#" - INSERT INTO review_queue (id, extraction_type, content, status, created_at) - VALUES ($1, $2, $3, $4, $5) - RETURNING id; - "#; - - let id = sqlx::query_scalar::<_, String>(query) - .bind(&entry.id) - .bind(&entry.extraction_type) - .bind(&entry.content) - .bind(&entry.status) - .bind(Utc::now()) - .fetch_one(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(id) - } - - /// Get pending items for review - pub async fn list_pending(&self, limit: i64) -> Result, DbError> { - let query = r#" - SELECT id, extraction_type, content, status, created_at, reviewed_at, reviewed_by, rejection_reason - FROM review_queue - WHERE status = 'pending' - ORDER BY created_at ASC - LIMIT $1; - "#; - - let rows = sqlx::query(query) - .bind(limit) - .fetch_all(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(rows.iter().map(|r| ReviewQueueEntry { - id: r.get("id"), - extraction_type: r.get("extraction_type"), - content: r.get("content"), - status: r.get("status"), - created_at: r.get("created_at"), - reviewed_at: r.get("reviewed_at"), - reviewed_by: r.get("reviewed_by"), - rejection_reason: r.get("rejection_reason"), - }).collect()) - } - - /// Approve review queue entry - pub async fn approve(&self, id: &str, reviewed_by: &str) -> Result<(), DbError> { - let query = r#" - UPDATE review_queue - SET status = 'approved', reviewed_at = $1, reviewed_by = $2 - WHERE id = $3; - "#; - - sqlx::query(query) - .bind(Utc::now()) - .bind(reviewed_by) - .bind(id) - .execute(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(()) - } - - /// Reject review queue entry - pub async fn reject(&self, id: &str, reviewed_by: &str, reason: &str) -> Result<(), DbError> { - let query = r#" - UPDATE review_queue - SET status = 'rejected', reviewed_at = $1, reviewed_by = $2, rejection_reason = $3 - WHERE id = $4; - "#; - - sqlx::query(query) - .bind(Utc::now()) - .bind(reviewed_by) - .bind(reason) - .bind(id) - .execute(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(()) - } -} - -/// Extraction Audit Repository (Immutable log for audit trail) -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ExtractionAuditEntry { - pub id: String, - pub extraction_type: String, // "entity" | "edge" - pub extraction_id: String, // ID of extracted entity/edge - pub source_content: String, // Original text - pub extracted_data: serde_json::Value, - pub llm_confidence: Option, - pub contradiction_score: Option, - pub status: String, // "extracted" | "approved" | "rejected" - pub extracted_at: DateTime, - pub extracted_by: String, // User or "system" -} - -pub struct ExtractionAuditRepo { - pool: Pool, -} - -impl ExtractionAuditRepo { - pub fn new(pool: Pool) -> Self { - Self { pool } - } - - /// Log an extraction attempt (immutable append) - pub async fn log_extraction(&self, entry: &ExtractionAuditEntry) -> Result { - let query = r#" - INSERT INTO extraction_audit (id, extraction_type, extraction_id, source_content, extracted_data, llm_confidence, contradiction_score, status, extracted_at, extracted_by) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) - RETURNING id; - "#; - - let id = sqlx::query_scalar::<_, String>(query) - .bind(&entry.id) - .bind(&entry.extraction_type) - .bind(&entry.extraction_id) - .bind(&entry.source_content) - .bind(&entry.extracted_data) - .bind(entry.llm_confidence) - .bind(entry.contradiction_score) - .bind(&entry.status) - .bind(entry.extracted_at) - .bind(&entry.extracted_by) - .fetch_one(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(id) - } - - /// Get audit trail for an extracted item - pub async fn get_history(&self, extraction_id: &str) -> Result, DbError> { - let query = r#" - SELECT id, extraction_type, extraction_id, source_content, extracted_data, llm_confidence, contradiction_score, status, extracted_at, extracted_by - FROM extraction_audit - WHERE extraction_id = $1 - ORDER BY extracted_at DESC; - "#; - - let rows = sqlx::query(query) - .bind(extraction_id) - .fetch_all(&self.pool) - .await - .map_err(|e| DbError::QueryFailed(e.to_string()))?; - - Ok(rows.iter().map(|r| ExtractionAuditEntry { - id: r.get("id"), - extraction_type: r.get("extraction_type"), - extraction_id: r.get("extraction_id"), - source_content: r.get("source_content"), - extracted_data: r.get("extracted_data"), - llm_confidence: r.get("llm_confidence"), - contradiction_score: r.get("contradiction_score"), - status: r.get("status"), - extracted_at: r.get("extracted_at"), - extracted_by: r.get("extracted_by"), - }).collect()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_db_error_display() { - let err = DbError::ConnectionFailed("test".to_string()); - assert!(err.to_string().contains("Connection failed")); - } - - #[test] - fn test_review_queue_entry_creation() { - let entry = ReviewQueueEntry { - id: "test-1".to_string(), - extraction_type: "entity".to_string(), - content: serde_json::json!({"name": "test"}), - status: "pending".to_string(), - created_at: Utc::now(), - reviewed_at: None, - reviewed_by: None, - rejection_reason: None, - }; - - assert_eq!(entry.extraction_type, "entity"); - } - - #[test] - fn test_dead_letter_entry_creation() { - let entry = DeadLetterEntry { - id: "dlq-1".to_string(), - original_content: "test content".to_string(), - error_message: "extraction failed".to_string(), - error_type: "extraction_failed".to_string(), - retry_count: 0, - max_retries: 3, - created_at: Utc::now(), - last_retry_at: None, - }; - - assert_eq!(entry.retry_count, 0); - assert!(entry.retry_count < entry.max_retries); - } -} diff --git a/crates/mem-store/src/lib.rs b/crates/mem-store/src/lib.rs index 519373a..8d9d8a4 100644 --- a/crates/mem-store/src/lib.rs +++ b/crates/mem-store/src/lib.rs @@ -9,7 +9,6 @@ pub mod community_repo; pub mod versioning; pub mod audit_logger; pub mod agent_repo; -// pub mod db_repo; // TODO: Fix Entity schema integration pub use event_log::{EventRecord, LogWriter}; pub use pgvector::{VectorRecord, VectorStore, ChunkL0, MemoryL1, MemoryL2}; diff --git a/crates/mem-store/src/schema.rs b/crates/mem-store/src/schema.rs index b9c05a4..4deac5c 100644 --- a/crates/mem-store/src/schema.rs +++ b/crates/mem-store/src/schema.rs @@ -286,6 +286,29 @@ pub async fn init_schema(pool: &PgPool) -> Result<()> { .execute(pool) .await?; - tracing::info!("Database schema initialized (including memory_entity + memory_edge)"); + // HNSW vector indexes for semantic search (RAG-001) + // name_embedding: primary entity search vector + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_entity_name_emb ON memory_entity \ + USING hnsw (name_embedding vector_cosine_ops) WITH (m = 16, ef_construction = 128)" + ) + .execute(pool) + .await?; + // summary_embedding: secondary entity search vector + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_entity_summary_emb ON memory_entity \ + USING hnsw (summary_embedding vector_cosine_ops) WITH (m = 16, ef_construction = 128)" + ) + .execute(pool) + .await?; + // fact_embedding: edge/relationship search vector + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_edge_fact_emb ON memory_edge \ + USING hnsw (fact_embedding vector_cosine_ops) WITH (m = 16, ef_construction = 128)" + ) + .execute(pool) + .await?; + + tracing::info!("Database schema initialized (including memory_entity + memory_edge + HNSW indexes)"); Ok(()) }