#![ allow( dead_code, missing_debug_implementations, missing_docs ) ]
mod private
{
use serde::{ Deserialize, Serialize };
use std::collections::{ HashMap, BTreeMap };
use std::sync::{ Arc, RwLock };
use core::sync::atomic::{ AtomicU64, Ordering };
use core::time::Duration;
use std::time::{ SystemTime, Instant };
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct SearchResult
{
pub id : String,
pub score : f32,
pub content : String,
pub metadata : Option< HashMap< String, String > >,
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct SemanticRetrievalConfig
{
pub endpoint : Option< String >,
pub timeout_seconds : u64,
pub max_results : usize,
}
impl Default for SemanticRetrievalConfig
{
fn default() -> Self
{
Self {
endpoint : None,
timeout_seconds : 30,
max_results : 10,
}
}
}
pub trait VectorIndex : Send + Sync
{
fn add_vector( &mut self, id : &str, vector : &[ f32 ], metadata : Option< HashMap< String, String > > ) -> Result< (), crate::error::Error >;
fn search( &self, query_vector : &[ f32 ], limit : usize, threshold : f32 ) -> Result< Vec< VectorSearchResult >, crate::error::Error >;
fn remove_vector( &mut self, id : &str ) -> Result< bool, crate::error::Error >;
fn get_stats( &self ) -> IndexStats;
fn optimize( &mut self ) -> Result< (), crate::error::Error >;
}
pub trait CacheStrategy< K, V >: Send + Sync
{
fn put( &mut self, key : K, value : V );
fn get( &self, key : &K ) -> Option< V >;
fn remove( &mut self, key : &K ) -> Option< V >;
fn clear( &mut self );
fn stats( &self ) -> CacheStats;
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct VectorSearchResult
{
pub id : String,
pub score : f32,
pub distance : f32,
pub metadata : Option< HashMap< String, String > >,
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct IndexStats
{
pub vector_count : u64,
pub memory_usage_bytes : u64,
pub avg_search_time_ms : f64,
pub total_searches : u64,
pub build_time_ms : u64,
pub last_optimized_at : Option< SystemTime >,
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct CacheStats
{
pub hits : u64,
pub misses : u64,
pub size : usize,
pub capacity : usize,
pub hit_ratio : f64,
pub memory_usage_bytes : u64,
}
#[ derive( Debug ) ]
pub struct FlatVectorIndex
{
vectors : HashMap< String, ( Vec< f32 >, Option< HashMap< String, String > > ) >,
stats : Arc< RwLock< IndexStats > >,
dimensions : usize,
}
impl FlatVectorIndex
{
#[ inline ]
#[ must_use ]
pub fn new( dimensions : usize ) -> Self
{
Self {
vectors : HashMap::new(),
stats : Arc::new( RwLock::new( IndexStats {
vector_count : 0,
memory_usage_bytes : 0,
avg_search_time_ms : 0.0,
total_searches : 0,
build_time_ms : 0,
last_optimized_at : None,
} ) ),
dimensions,
}
}
#[ inline ]
fn cosine_similarity( &self, a : &[ f32 ], b : &[ f32 ] ) -> f32
{
if a.len() != b.len() || a.is_empty()
{
return 0.0;
}
let dot_product : f32 = a.iter().zip( b.iter() ).map( | ( x, y ) | x * y ).sum();
let magnitude_a : f32 = a.iter().map( | x | x * x ).sum::< f32 >().sqrt();
let magnitude_b : f32 = b.iter().map( | x | x * x ).sum::< f32 >().sqrt();
if magnitude_a == 0.0 || magnitude_b == 0.0
{
0.0
} else {
dot_product / ( magnitude_a * magnitude_b )
}
}
}
impl VectorIndex for FlatVectorIndex
{
#[ inline ]
fn add_vector( &mut self, id : &str, vector : &[ f32 ], metadata : Option< HashMap< String, String > > ) -> Result< (), crate::error::Error >
{
if vector.len() != self.dimensions
{
return Err( crate::error::Error::InvalidArgument( format!( "Vector dimension mismatch : expected {}, got {}", self.dimensions, vector.len() ) ) );
}
self.vectors.insert( id.to_string(), ( vector.to_vec(), metadata ) );
if let Ok( mut stats ) = self.stats.write()
{
stats.vector_count = self.vectors.len() as u64;
stats.memory_usage_bytes = self.vectors.len() as u64 * ( self.dimensions as u64 * 4 + 64 ); }
Ok( () )
}
#[ inline ]
fn search( &self, query_vector : &[ f32 ], limit : usize, threshold : f32 ) -> Result< Vec< VectorSearchResult >, crate::error::Error >
{
let start_time = Instant::now();
if query_vector.len() != self.dimensions
{
return Err( crate::error::Error::InvalidArgument( format!( "Query vector dimension mismatch : expected {}, got {}", self.dimensions, query_vector.len() ) ) );
}
let mut results : Vec< VectorSearchResult > = Vec::new();
for ( id, ( vector, metadata ) ) in &self.vectors
{
let similarity = self.cosine_similarity( query_vector, vector );
if similarity >= threshold
{
results.push( VectorSearchResult {
id : id.clone(),
score : similarity,
distance : 1.0 - similarity, metadata : metadata.clone(),
} );
}
}
results.sort_by( | a, b | b.score.partial_cmp( &a.score ).unwrap_or( std::cmp::Ordering::Equal ) );
results.truncate( limit );
let search_time_ms = start_time.elapsed().as_millis() as f64;
if let Ok( mut stats ) = self.stats.write()
{
stats.total_searches += 1;
stats.avg_search_time_ms = ( stats.avg_search_time_ms * ( stats.total_searches - 1 ) as f64 + search_time_ms ) / stats.total_searches as f64;
}
Ok( results )
}
#[ inline ]
fn remove_vector( &mut self, id : &str ) -> Result< bool, crate::error::Error >
{
let removed = self.vectors.remove( id ).is_some();
if removed
{
if let Ok( mut stats ) = self.stats.write()
{
stats.vector_count = self.vectors.len() as u64;
stats.memory_usage_bytes = self.vectors.len() as u64 * ( self.dimensions as u64 * 4 + 64 );
}
}
Ok( removed )
}
#[ inline ]
fn get_stats( &self ) -> IndexStats
{
self.stats.read().unwrap_or_else( | poisoned | {
poisoned.into_inner()
} ).clone()
}
#[ inline ]
fn optimize( &mut self ) -> Result< (), crate::error::Error >
{
let start_time = Instant::now();
let optimized_vectors : HashMap< String, ( Vec< f32 >, Option< HashMap< String, String > > ) > =
self.vectors.iter().map( | ( k, v ) | ( k.clone(), v.clone() ) ).collect();
self.vectors = optimized_vectors;
if let Ok( mut stats ) = self.stats.write()
{
stats.last_optimized_at = Some( SystemTime::now() );
stats.build_time_ms = start_time.elapsed().as_millis() as u64;
}
Ok( () )
}
}
#[ derive( Debug ) ]
pub struct AdaptiveLruCache< K, V >
where
K: Clone + Eq + std::hash::Hash + Send + Sync,
V: Clone + Send + Sync,
{
cache : HashMap< K, CacheEntry< V > >,
access_order : BTreeMap< u64, K >,
access_counter : AtomicU64,
capacity : usize,
stats : Arc< RwLock< CacheStats > >,
ttl : Option< Duration >,
}
#[ derive( Debug, Clone ) ]
struct CacheEntry< V >
{
value : V,
last_accessed : SystemTime,
access_count : u64,
created_at : SystemTime,
}
impl< K, V > AdaptiveLruCache< K, V >
where
K: Clone + Eq + std::hash::Hash + Send + Sync,
V: Clone + Send + Sync,
{
pub fn new( capacity : usize ) -> Self
{
Self {
cache : HashMap::new(),
access_order : BTreeMap::new(),
access_counter : AtomicU64::new( 0 ),
capacity,
stats : Arc::new( RwLock::new( CacheStats {
hits : 0,
misses : 0,
size : 0,
capacity,
hit_ratio : 0.0,
memory_usage_bytes : 0,
} ) ),
ttl : None,
}
}
pub fn with_ttl( capacity : usize, ttl : Duration ) -> Self
{
let mut cache = Self::new( capacity );
cache.ttl = Some( ttl );
cache
}
fn is_expired( &self, entry : &CacheEntry< V > ) -> bool
{
if let Some( ttl ) = self.ttl
{
if let Ok( elapsed ) = entry.created_at.elapsed()
{
return elapsed > ttl;
}
}
false
}
fn evict_lru( &mut self )
{
while self.cache.len() >= self.capacity
{
if let Some( ( _, oldest_key ) ) = self.access_order.first_key_value()
{
let oldest_key = oldest_key.clone();
self.access_order.remove( &self.access_counter.load( Ordering::Relaxed ) );
self.cache.remove( &oldest_key );
} else {
break;
}
}
}
fn update_stats( &self, hit : bool )
{
if let Ok( mut stats ) = self.stats.write()
{
if hit
{
stats.hits += 1;
} else {
stats.misses += 1;
}
let total = stats.hits + stats.misses;
stats.hit_ratio = if total > 0 { stats.hits as f64 / total as f64 } else { 0.0 };
stats.size = self.cache.len();
stats.memory_usage_bytes = self.cache.len() as u64 * 256; }
}
}
impl< K, V > CacheStrategy< K, V > for AdaptiveLruCache< K, V >
where
K: Clone + Eq + std::hash::Hash + Send + Sync,
V: Clone + Send + Sync,
{
fn put( &mut self, key : K, value : V )
{
let now = SystemTime::now();
if let Some( ttl ) = self.ttl
{
let ttl_duration = ttl;
self.cache.retain( | _, entry | {
if let Ok( elapsed ) = entry.created_at.elapsed()
{
elapsed <= ttl_duration
} else {
true }
} );
}
self.evict_lru();
let access_id = self.access_counter.fetch_add( 1, Ordering::Relaxed );
let entry = CacheEntry {
value,
last_accessed : now,
access_count : 1,
created_at : now,
};
self.cache.insert( key.clone(), entry );
self.access_order.insert( access_id, key );
self.update_stats( false );
}
fn get( &self, key : &K ) -> Option< V >
{
if let Some( entry ) = self.cache.get( key )
{
if self.is_expired( entry )
{
self.update_stats( false );
return None;
}
self.update_stats( true );
Some( entry.value.clone() )
} else {
self.update_stats( false );
None
}
}
fn remove( &mut self, key : &K ) -> Option< V >
{
if let Some( entry ) = self.cache.remove( key )
{
self.access_order.retain( | _, k | k != key );
Some( entry.value )
} else {
None
}
}
fn clear( &mut self )
{
self.cache.clear();
self.access_order.clear();
self.access_counter.store( 0, Ordering::Relaxed );
if let Ok( mut stats ) = self.stats.write()
{
stats.size = 0;
stats.memory_usage_bytes = 0;
}
}
fn stats( &self ) -> CacheStats
{
self.stats.read().unwrap_or_else( | poisoned | poisoned.into_inner() ).clone()
}
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct OptimizedRetrievalConfig
{
pub base : SemanticRetrievalConfig,
pub index_type : OptimizedIndexType,
pub cache_config : CacheConfig,
pub search_config : SearchOptimizationConfig,
pub monitoring_config : MonitoringConfig,
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub enum OptimizedIndexType
{
OptimizedFlat
{
dimensions : usize
},
HNSW
{
dimensions : usize,
max_connections : u32,
ef_construction : u32
},
LSH
{
dimensions : usize,
hash_tables : u32,
hash_functions : u32
},
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq, Default ) ]
pub struct CacheConfig
{
pub capacity : usize,
pub ttl_seconds : Option< u64 >,
pub adaptive_sizing : bool,
pub warming_strategy : CacheWarmingStrategy,
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq, Default ) ]
pub enum CacheWarmingStrategy
{
#[ default ]
None,
CommonQueries,
AccessPatterns,
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct SearchOptimizationConfig
{
pub enable_query_preprocessing : bool,
pub enable_reranking : bool,
pub enable_parallel_search : bool,
pub search_timeout_ms : u64,
pub enable_query_expansion : bool,
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct MonitoringConfig
{
pub enable_metrics : bool,
pub metrics_interval_seconds : u64,
pub enable_detailed_timing : bool,
pub max_metrics_history : usize,
}
impl Default for OptimizedRetrievalConfig
{
fn default() -> Self
{
Self {
base : SemanticRetrievalConfig::default(),
index_type : OptimizedIndexType::OptimizedFlat { dimensions : 1536 },
cache_config : CacheConfig {
capacity : 10000,
ttl_seconds : Some( 3600 ), adaptive_sizing : true,
warming_strategy : CacheWarmingStrategy::CommonQueries,
},
search_config : SearchOptimizationConfig {
enable_query_preprocessing : true,
enable_reranking : true,
enable_parallel_search : true,
search_timeout_ms : 5000, enable_query_expansion : false,
},
monitoring_config : MonitoringConfig {
enable_metrics : true,
metrics_interval_seconds : 60,
enable_detailed_timing : true,
max_metrics_history : 1000,
},
}
}
}
pub struct OptimizedSemanticRetrievalApi< 'a >
{
client : &'a crate::client::Client,
index : Arc< RwLock< dyn VectorIndex > >,
cache : Arc< RwLock< dyn CacheStrategy< String, Vec< SearchResult > > > >,
embedding_cache : Arc< RwLock< dyn CacheStrategy< String, Vec< f32 > > > >,
config : OptimizedRetrievalConfig,
metrics : Arc< RwLock< PerformanceMetrics > >,
}
#[ derive( Debug, Clone, Serialize, Deserialize, PartialEq ) ]
pub struct PerformanceMetrics
{
pub total_searches : u64,
pub avg_search_latency_ms : f64,
pub p95_search_latency_ms : f64,
pub total_indexing_ops : u64,
pub avg_indexing_latency_ms : f64,
pub search_cache_hit_ratio : f64,
pub embedding_cache_hit_ratio : f64,
pub memory_usage_bytes : u64,
pub last_updated : SystemTime,
}
impl Default for PerformanceMetrics
{
fn default() -> Self
{
Self {
total_searches : 0,
avg_search_latency_ms : 0.0,
p95_search_latency_ms : 0.0,
total_indexing_ops : 0,
avg_indexing_latency_ms : 0.0,
search_cache_hit_ratio : 0.0,
embedding_cache_hit_ratio : 0.0,
memory_usage_bytes : 0,
last_updated : SystemTime::now(),
}
}
}
impl< 'a > OptimizedSemanticRetrievalApi< 'a >
{
pub fn new( client : &'a crate::client::Client ) -> Self
{
Self::with_config( client, OptimizedRetrievalConfig::default() )
}
pub fn with_config( client : &'a crate::client::Client, config : OptimizedRetrievalConfig ) -> Self
{
let dimensions = match &config.index_type
{
OptimizedIndexType::OptimizedFlat { dimensions } => *dimensions,
OptimizedIndexType::HNSW { dimensions, .. } => *dimensions,
OptimizedIndexType::LSH { dimensions, .. } => *dimensions,
};
let index : Arc< RwLock< dyn VectorIndex > > = Arc::new( RwLock::new( FlatVectorIndex::new( dimensions ) ) );
let cache_capacity = config.cache_config.capacity;
let cache_ttl = config.cache_config.ttl_seconds.map( Duration::from_secs );
let search_cache : Arc< RwLock< dyn CacheStrategy< String, Vec< SearchResult > > > > =
if let Some( ttl ) = cache_ttl
{
Arc::new( RwLock::new( AdaptiveLruCache::with_ttl( cache_capacity, ttl ) ) )
} else {
Arc::new( RwLock::new( AdaptiveLruCache::new( cache_capacity ) ) )
};
let embedding_cache : Arc< RwLock< dyn CacheStrategy< String, Vec< f32 > > > > =
if let Some( ttl ) = cache_ttl
{
Arc::new( RwLock::new( AdaptiveLruCache::with_ttl( cache_capacity / 2, ttl ) ) )
} else {
Arc::new( RwLock::new( AdaptiveLruCache::new( cache_capacity / 2 ) ) )
};
Self {
client,
index,
cache : search_cache,
embedding_cache,
config,
metrics : Arc::new( RwLock::new( PerformanceMetrics::default() ) ),
}
}
pub fn get_metrics( &self ) -> PerformanceMetrics
{
self.metrics.read().unwrap_or_else( | poisoned | poisoned.into_inner() ).clone()
}
pub async fn optimize_index( &self ) -> Result< (), crate::error::Error >
{
if let Ok( mut index ) = self.index.write()
{
index.optimize()?;
}
Ok( () )
}
pub fn clear_caches( &self )
{
if let Ok( mut cache ) = self.cache.write()
{
cache.clear();
}
if let Ok( mut embedding_cache ) = self.embedding_cache.write()
{
embedding_cache.clear();
}
}
pub fn get_index_stats( &self ) -> Option< IndexStats >
{
if let Ok( index ) = self.index.read()
{
Some( index.get_stats() )
} else {
None
}
}
pub fn get_cache_stats( &self ) -> ( CacheStats, CacheStats )
{
let search_cache_stats = if let Ok( cache ) = self.cache.read()
{
cache.stats()
} else {
CacheStats {
hits : 0, misses : 0, size : 0, capacity : 0, hit_ratio : 0.0, memory_usage_bytes : 0
}
};
let embedding_cache_stats = if let Ok( cache ) = self.embedding_cache.read()
{
cache.stats()
} else {
CacheStats {
hits : 0, misses : 0, size : 0, capacity : 0, hit_ratio : 0.0, memory_usage_bytes : 0
}
};
( search_cache_stats, embedding_cache_stats )
}
}
impl< 'a > std::fmt::Debug for OptimizedSemanticRetrievalApi< 'a >
{
fn fmt( &self, f : &mut std::fmt::Formatter< '_ > ) -> std::fmt::Result
{
f.debug_struct( "OptimizedSemanticRetrievalApi" )
.field( "config", &self.config )
.finish_non_exhaustive()
}
}
}
pub use private::
{
SearchResult, SemanticRetrievalConfig,
VectorIndex, CacheStrategy, VectorSearchResult, IndexStats, CacheStats,
FlatVectorIndex, AdaptiveLruCache, OptimizedRetrievalConfig,
OptimizedIndexType, CacheConfig, CacheWarmingStrategy,
SearchOptimizationConfig, MonitoringConfig, OptimizedSemanticRetrievalApi,
PerformanceMetrics
};