poimen-memory/crates/mem-cli/src/query/semantic_retriever.rs

//! 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<DateTime<Utc>>,     // Earliest event_time
    pub end_time: Option<DateTime<Utc>>,       // Latest event_time
    pub min_recency_score: Option<f32>,        // 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<String>,  // entity name or fact snippet
    pub entity_type: Option<String>,
    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<Postgres>,
}

impl SemanticRetriever {
    /// Create a new semantic retriever
    pub fn new(pool: Pool<Postgres>) -> 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<DateTime<Utc>>,
        end_time: Option<DateTime<Utc>>,
    ) -> Result<Vec<EntityResult>, 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<DateTime<Utc>>,
        end_time: Option<DateTime<Utc>>,
    ) -> Result<Vec<EdgeResult>, 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<DateTime<Utc>>,
        end_time: Option<DateTime<Utc>>,
    ) -> Result<Vec<HybridResult>, 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)
    }
}