//! Semantic Retrieval Engine //! //! Provides semantic search capabilities using vector embeddings and hybrid search //! combining vector (semantic) and lexical (keyword) results with RRF fusion. use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use sqlx::{Pool, Postgres}; use std::sync::Arc; use tracing::{debug, info, warn}; /// Semantic search result for an entity #[derive(Debug, Clone, Serialize, Deserialize)] pub struct EntityResult { pub id: String, pub name: String, pub entity_type: String, pub similarity_score: f32, // 0.0-1.0, higher is better pub metadata: serde_json::Value, } /// 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) } impl Default for TemporalFilter { fn default() -> Self { Self { start_time: None, end_time: None, min_recency_score: None, } } } /// Semantic search result for an edge (relationship) #[derive(Debug, Clone, Serialize, Deserialize)] pub struct EdgeResult { pub id: String, pub source_entity_id: String, pub target_entity_id: String, pub source_name: String, pub target_name: String, pub relation_type: String, pub fact: String, pub similarity_score: f32, // 0.0-1.0, higher is better pub confidence: f32, } /// Hybrid search result combining semantic and lexical scores #[derive(Debug, Clone, Serialize, Deserialize)] pub struct HybridResult { pub id: String, pub name: Option, // entity name or fact snippet 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 } /// Semantic Retriever - performs vector and hybrid searches pub struct SemanticRetriever { pub pool: Pool, } impl SemanticRetriever { /// Create a new semantic retriever pub fn new(pool: Pool) -> Self { Self { pool } } /// Search for entities by semantic similarity /// /// # 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 pub async fn search_entities( &self, query_embedding: &[f32], top_k: usize, entity_type_filter: Option<&str>, confidence_floor: f32, start_time: Option>, end_time: Option>, ) -> Result, String> { if query_embedding.len() != 768 { return Err(format!( "Invalid embedding dimension: expected 768, got {}", query_embedding.len() )); } let top_k = top_k.max(1).min(100); // Clamp 1-100 if confidence_floor < 0.0 || confidence_floor > 1.0 { return Err("confidence_floor must be 0.0-1.0".to_string()); } 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 FROM memory_entity WHERE deleted_at IS NULL AND (1 - (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)) 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 .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 { id, name, entity_type, similarity_score: score.max(0.0).min(1.0), // Clamp to 0-1 metadata, }) .collect(); info!("Found {} entities", entities.len()); Ok(entities) } /// Search for edges (relationships/facts) by semantic similarity /// /// # 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 pub async fn search_edges( &self, query_embedding: &[f32], top_k: usize, relation_type_filter: Option<&str>, start_time: Option>, end_time: Option>, ) -> Result, String> { if query_embedding.len() != 768 { return Err(format!( "Invalid embedding dimension: expected 768, got {}", query_embedding.len() )); } let top_k = top_k.max(1).min(100); 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, src.name, tgt.name, e.relation_type, e.fact, 1 - (e.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 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)) 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 .fetch_all(&self.pool) .await .map_err(|e| format!("Database error: {}", e))?; let edges: Vec<_> = results .into_iter() .map(|(id, src_id, tgt_id, src_name, tgt_name, rel_type, fact, score, conf)| { EdgeResult { id, source_entity_id: src_id, target_entity_id: tgt_id, source_name: src_name, 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), } }) .collect(); info!("Found {} edges", edges.len()); Ok(edges) } /// Hybrid search combining semantic (vector) and lexical (keyword) results /// /// 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) pub async fn hybrid_search( &self, query_embedding: &[f32], top_k: usize, semantic_weight: f32, lexical_weight: f32, start_time: Option>, end_time: Option>, ) -> Result, String> { if query_embedding.len() != 768 { return Err(format!( "Invalid embedding dimension: expected 768, got {}", query_embedding.len() )); } 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); 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?; // Phase 2: Semantic search for edges let edge_results = self.search_edges( query_embedding, top_k * 2, None, start_time, end_time, ).await?; // Phase 3: Combine and rank by RRF fusion let mut hybrid_results = Vec::new(); for entity in entity_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()), 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, }); } // 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) } }