476 lines
16 KiB
Rust
476 lines
16 KiB
Rust
|
//! 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 = 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 = 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)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_entity_result_creation() {
|
||
|
|
let result = EntityResult {
|
||
|
|
id: "e1".to_string(),
|
||
|
|
name: "Test".to_string(),
|
||
|
|
entity_type: "concept".to_string(),
|
||
|
|
similarity_score: 0.95,
|
||
|
|
metadata: serde_json::json!({"key": "value"}),
|
||
|
|
};
|
||
|
|
assert_eq!(result.id, "e1");
|
||
|
|
assert_eq!(result.similarity_score, 0.95);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_edge_result_creation() {
|
||
|
|
let result = EdgeResult {
|
||
|
|
id: "e1".to_string(),
|
||
|
|
source_entity_id: "src".to_string(),
|
||
|
|
target_entity_id: "tgt".to_string(),
|
||
|
|
source_name: "A".to_string(),
|
||
|
|
target_name: "B".to_string(),
|
||
|
|
relation_type: "related_to".to_string(),
|
||
|
|
fact: "A is related to B".to_string(),
|
||
|
|
similarity_score: 0.88,
|
||
|
|
confidence: 0.90,
|
||
|
|
};
|
||
|
|
assert_eq!(result.similarity_score, 0.88);
|
||
|
|
assert_eq!(result.confidence, 0.90);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_hybrid_result_creation() {
|
||
|
|
let result = HybridResult {
|
||
|
|
id: "h1".to_string(),
|
||
|
|
name: Some("Test".to_string()),
|
||
|
|
entity_type: Some("concept".to_string()),
|
||
|
|
result_type: "entity".to_string(),
|
||
|
|
fused_score: 0.85,
|
||
|
|
semantic_score: 0.90,
|
||
|
|
lexical_score: 0.75,
|
||
|
|
};
|
||
|
|
assert!(result.fused_score >= 0.0 && result.fused_score <= 1.0);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_embedding_dimension_validation() {
|
||
|
|
let invalid_embedding = vec![0.5; 512]; // Wrong size
|
||
|
|
assert_eq!(invalid_embedding.len(), 512);
|
||
|
|
assert_ne!(invalid_embedding.len(), 768);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_confidence_floor_bounds() {
|
||
|
|
let floor = 0.5;
|
||
|
|
assert!(floor >= 0.0 && floor <= 1.0);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_top_k_bounds() {
|
||
|
|
let top_k = 50;
|
||
|
|
let clamped = top_k.max(1).min(100);
|
||
|
|
assert_eq!(clamped, 50);
|
||
|
|
|
||
|
|
let too_small = 0;
|
||
|
|
assert_eq!(too_small.max(1).min(100), 1);
|
||
|
|
|
||
|
|
let too_large = 500;
|
||
|
|
assert_eq!(too_large.max(1).min(100), 100);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_weight_normalization() {
|
||
|
|
let sem_w = 0.6;
|
||
|
|
let lex_w = 0.4;
|
||
|
|
let normalized_sem = sem_w.max(0.0).min(1.0);
|
||
|
|
let normalized_lex = lex_w.max(0.0).min(1.0);
|
||
|
|
assert_eq!(normalized_sem, 0.6);
|
||
|
|
assert_eq!(normalized_lex, 0.4);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_score_clamping() {
|
||
|
|
let scores = vec![0.5, 1.0, 1.5, -0.1, 0.999];
|
||
|
|
for score in scores {
|
||
|
|
let clamped = score.max(0.0).min(1.0);
|
||
|
|
assert!(clamped >= 0.0 && clamped <= 1.0);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_hybrid_result_type_values() {
|
||
|
|
let entity_result = HybridResult {
|
||
|
|
id: "e1".to_string(),
|
||
|
|
name: Some("Entity".to_string()),
|
||
|
|
entity_type: Some("concept".to_string()),
|
||
|
|
result_type: "entity".to_string(),
|
||
|
|
fused_score: 0.9,
|
||
|
|
semantic_score: 0.92,
|
||
|
|
lexical_score: 0.85,
|
||
|
|
};
|
||
|
|
assert_eq!(entity_result.result_type, "entity");
|
||
|
|
|
||
|
|
let edge_result = HybridResult {
|
||
|
|
id: "edge1".to_string(),
|
||
|
|
name: Some("fact".to_string()),
|
||
|
|
entity_type: None,
|
||
|
|
result_type: "edge".to_string(),
|
||
|
|
fused_score: 0.85,
|
||
|
|
semantic_score: 0.87,
|
||
|
|
lexical_score: 0.80,
|
||
|
|
};
|
||
|
|
assert_eq!(edge_result.result_type, "edge");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_sorting_by_score() {
|
||
|
|
let mut results = vec![
|
||
|
|
HybridResult {
|
||
|
|
id: "1".to_string(),
|
||
|
|
name: None,
|
||
|
|
entity_type: None,
|
||
|
|
result_type: "entity".to_string(),
|
||
|
|
fused_score: 0.5,
|
||
|
|
semantic_score: 0.5,
|
||
|
|
lexical_score: 0.5,
|
||
|
|
},
|
||
|
|
HybridResult {
|
||
|
|
id: "2".to_string(),
|
||
|
|
name: None,
|
||
|
|
entity_type: None,
|
||
|
|
result_type: "entity".to_string(),
|
||
|
|
fused_score: 0.9,
|
||
|
|
semantic_score: 0.9,
|
||
|
|
lexical_score: 0.9,
|
||
|
|
},
|
||
|
|
];
|
||
|
|
|
||
|
|
results.sort_by(|a, b| b.fused_score.partial_cmp(&a.fused_score).unwrap_or(std::cmp::Ordering::Equal));
|
||
|
|
assert_eq!(results[0].id, "2");
|
||
|
|
assert_eq!(results[1].id, "1");
|
||
|
|
}
|
||
|
|
}
|