lens_core/semantic/
mod.rs

1//! # Semantic Search Module
2//!
3//! Phase 3: Advanced semantic search with performance constraints
4//! - 2048-token encoder for long code context
5//! - Hard negatives from SymbolGraph relationships  
6//! - Learned reranking with isotonic regression
7//! - Optional cross-encoder for precision boost
8//! - Calibration preservation with ECE ≤ 0.005 drift
9
10pub mod encoder;
11pub mod hard_negatives; 
12pub mod rerank;
13pub mod cross_encoder;
14pub mod calibration;
15pub mod isotonic_calibration;
16pub mod sla_bounded_evaluation;
17pub mod pipeline;
18pub mod validation;
19pub mod ltr_trainer;
20
21// New Rust-based semantic processing modules  
22pub mod embedding;
23pub mod query_classifier;
24pub mod intent_router;
25pub mod conformal_router;
26pub mod benchmarks;
27pub mod integration;
28pub mod examples;
29
30// Re-export main types for easier usage
31pub use pipeline::{SemanticPipeline, SemanticSearchRequest, SemanticSearchResponse, initialize_semantic_pipeline};
32pub use encoder::{SemanticEncoder, CodeEmbedding};
33pub use rerank::{LearnedReranker, SearchResult, RerankedResult};
34pub use cross_encoder::{CrossEncoder, QueryAnalysis};
35pub use calibration::{CalibrationSystem, CalibrationMetrics};
36pub use isotonic_calibration::{IsotonicCalibrationSystem, IsotonicCalibrationConfig, CalibrationTrainingSample};
37pub use sla_bounded_evaluation::{SLABoundedEvaluator, SLAEvaluationConfig, SLABoundedEvaluationResult};
38pub use hard_negatives::{HardNegativesGenerator, ContrastivePair};
39pub use validation::{validate_phase3_implementation, ValidationResults};
40pub use ltr_trainer::{LTRTrainer, LTRConfig, LTRObjective, BoundedLTRModel, TrainingReport};
41
42// Re-export new Rust-based semantic processing types
43pub use embedding::{SemanticEncoder as RustSemanticEncoder, CodeEmbedding as RustCodeEmbedding, EmbeddingConfig};
44pub use query_classifier::{QueryClassifier, QueryClassification, QueryIntent as RustQueryIntent, ClassifierConfig};
45pub use intent_router::{IntentRouter, IntentRouterConfig};
46pub use conformal_router::{ConformalRouter, ConformalRouterConfig, UpshiftType};
47pub use integration::{
48    SemanticSearchIntegration, 
49    SemanticIntegrationConfig, 
50    SemanticSearchRequest as IntegratedSemanticSearchRequest,
51    SemanticSearchResponse as IntegratedSemanticSearchResponse,
52    SearchEngineSemanticExt,
53    SemanticHealthStatus,
54    IntegrationMetrics,
55};
56
57use anyhow::Result;
58use serde::{Deserialize, Serialize};
59use std::time::Duration;
60
61/// Configuration for semantic search components
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct SemanticConfig {
64    /// 2048-token encoder settings
65    pub encoder: EncoderConfig,
66    /// Learned reranking configuration  
67    pub rerank: RerankConfig,
68    /// Cross-encoder for precision boost
69    pub cross_encoder: CrossEncoderConfig,
70    /// Calibration preservation settings
71    pub calibration: CalibrationConfig,
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct EncoderConfig {
76    /// Model architecture (CodeT5/UniXcoder-class)
77    pub model_type: String,
78    /// Maximum token context window
79    pub max_tokens: usize,
80    /// Embedding dimension
81    pub embedding_dim: usize,
82    /// Batch size for inference
83    pub batch_size: usize,
84    /// Device for inference (cpu/cuda/mps)
85    pub device: String,
86}
87
88#[derive(Debug, Clone, Serialize, Deserialize)]  
89pub struct RerankConfig {
90    /// Top-K results to rerank
91    pub top_k: usize,
92    /// Use isotonic regression for score calibration
93    pub use_isotonic: bool,
94    /// Learning rate for isotonic fitting
95    pub learning_rate: f32,
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct CrossEncoderConfig {
100    /// Enable cross-encoder for precision boost  
101    pub enabled: bool,
102    /// Maximum inference time budget (≤50ms p95)
103    pub max_inference_ms: u64,
104    /// Query complexity threshold for activation
105    pub complexity_threshold: f32,
106    /// Top-K candidates for cross-encoding
107    pub top_k: usize,
108}
109
110#[derive(Debug, Clone, Serialize, Deserialize)]
111pub struct CalibrationConfig {
112    /// Maximum ECE drift allowed (≤0.005)
113    pub max_ece_drift: f32,
114    /// Cap dense features in log-odds space
115    pub log_odds_cap: f32,
116    /// Temperature scaling factor
117    pub temperature: f32,
118}
119
120impl Default for SemanticConfig {
121    fn default() -> Self {
122        Self {
123            encoder: EncoderConfig {
124                model_type: "codet5-base".to_string(),
125                max_tokens: 2048,
126                embedding_dim: 768,
127                batch_size: 16,
128                device: "cpu".to_string(),
129            },
130            rerank: RerankConfig {
131                top_k: 100,
132                use_isotonic: true,
133                learning_rate: 0.01,
134            },
135            cross_encoder: CrossEncoderConfig {
136                enabled: false, // Start disabled for performance
137                max_inference_ms: 50,
138                complexity_threshold: 0.7,
139                top_k: 10,
140            },
141            calibration: CalibrationConfig {
142                max_ece_drift: 0.005,
143                log_odds_cap: 5.0,
144                temperature: 1.0,
145            },
146        }
147    }
148}
149
150/// Performance constraints for semantic search
151pub const SEMANTIC_INFERENCE_TARGET_MS: u64 = 50;
152pub const SEMANTIC_P95_TARGET_MS: u64 = 50;
153pub const COIR_NDCG_TARGET: f32 = 0.52;
154pub const NL_IMPROVEMENT_TARGET_PP: f32 = 4.0; // 4-6pp target
155
156/// Semantic search metrics for performance tracking
157#[derive(Debug, Clone, Serialize, Deserialize)]
158pub struct SemanticMetrics {
159    /// Inference latency statistics
160    pub latency_p50_ms: f64,
161    pub latency_p95_ms: f64,
162    pub latency_p99_ms: f64,
163    
164    /// Quality metrics
165    pub ndcg_at_10: f32,
166    pub success_at_10: f32,
167    pub recall_at_50: f32,
168    
169    /// Calibration metrics
170    pub expected_calibration_error: f32,
171    pub ece_drift_from_baseline: f32,
172    
173    /// Natural language query performance  
174    pub nl_slice_improvement_pp: f32,
175    
176    /// Performance gates
177    pub meets_coir_target: bool,
178    pub meets_latency_target: bool,
179    pub meets_ece_target: bool,
180}
181
182impl SemanticMetrics {
183    /// Check if all performance gates are met
184    pub fn passes_gates(&self) -> bool {
185        self.meets_coir_target && self.meets_latency_target && self.meets_ece_target
186    }
187    
188    /// Validate against Phase 3 success criteria
189    pub fn validate_phase3_gates(&self) -> Result<()> {
190        if self.ndcg_at_10 < COIR_NDCG_TARGET {
191            anyhow::bail!(
192                "CoIR nDCG@10 {} < target {}", 
193                self.ndcg_at_10, 
194                COIR_NDCG_TARGET
195            );
196        }
197        
198        if self.latency_p95_ms > SEMANTIC_P95_TARGET_MS as f64 {
199            anyhow::bail!(
200                "Semantic p95 latency {}ms > target {}ms",
201                self.latency_p95_ms,
202                SEMANTIC_P95_TARGET_MS
203            );
204        }
205        
206        if self.ece_drift_from_baseline > 0.005 {
207            anyhow::bail!(
208                "ECE drift {} > target 0.005",
209                self.ece_drift_from_baseline
210            );
211        }
212        
213        if self.nl_slice_improvement_pp < NL_IMPROVEMENT_TARGET_PP {
214            anyhow::bail!(
215                "NL improvement {}pp < target {}pp",
216                self.nl_slice_improvement_pp,
217                NL_IMPROVEMENT_TARGET_PP
218            );
219        }
220        
221        Ok(())
222    }
223}
224
225/// Initialize new Rust-based semantic integration system
226pub async fn initialize_semantic_integration(config: &SemanticConfig) -> Result<SemanticSearchIntegration> {
227    tracing::info!("Initializing Rust-based semantic integration system");
228    
229    let integration_config = SemanticIntegrationConfig {
230        enabled: true,
231        nl_upshift_threshold: 0.7,
232        max_processing_time_ms: 100, // Stay within search SLA budget
233        enable_conformal_routing: true,
234        fallback_on_error: true,
235        enable_result_caching: config.cross_encoder.enabled,
236        similarity_threshold: 0.5,
237    };
238    
239    let integration = SemanticSearchIntegration::new(integration_config).await?;
240    
241    tracing::info!("✅ Rust-based semantic integration system initialized successfully");
242    Ok(integration)
243}
244
245/// Initialize semantic search module (original TypeScript-compatible function)
246pub async fn initialize_semantic(config: &SemanticConfig) -> Result<()> {
247    tracing::info!("Initializing semantic search module");
248    tracing::info!("Encoder: {} with {} tokens", config.encoder.model_type, config.encoder.max_tokens);
249    tracing::info!("Rerank: top-{} with isotonic={}", config.rerank.top_k, config.rerank.use_isotonic);
250    tracing::info!("Cross-encoder: enabled={}", config.cross_encoder.enabled);
251    tracing::info!("Performance targets: CoIR nDCG@10 ≥ {}, p95 ≤ {}ms", 
252                   COIR_NDCG_TARGET, SEMANTIC_P95_TARGET_MS);
253    
254    // Initialize encoder
255    encoder::initialize_encoder(&config.encoder).await?;
256    
257    // Initialize reranker
258    let rerank_config = rerank::RerankConfig {
259        top_k: config.rerank.top_k,
260        use_isotonic: config.rerank.use_isotonic,
261        learning_rate: config.rerank.learning_rate,
262        l2_regularization: 0.01, // Default value since not in semantic config
263        min_training_samples: 100, // Default value
264        combination_strategy: rerank::CombinationStrategy::LearnedWeights,
265    };
266    rerank::initialize_reranker(&rerank_config).await?;
267    
268    // Initialize cross-encoder if enabled
269    if config.cross_encoder.enabled {
270        let cross_encoder_config = cross_encoder::CrossEncoderConfig {
271            enabled: config.cross_encoder.enabled,
272            max_inference_ms: config.cross_encoder.max_inference_ms,
273            complexity_threshold: config.cross_encoder.complexity_threshold,
274            top_k: config.cross_encoder.top_k,
275            model_type: "bert-base-uncased".to_string(), // Default model type
276            max_batch_size: 16, // Default value
277            budget_strategy: cross_encoder::BudgetStrategy::FixedPerQuery,
278        };
279        cross_encoder::initialize_cross_encoder(&cross_encoder_config).await?;
280    }
281    
282    // Initialize calibration system
283    let calibration_config = calibration::CalibrationConfig {
284        max_ece_drift: config.calibration.max_ece_drift,
285        log_odds_cap: config.calibration.log_odds_cap,
286        temperature: config.calibration.temperature,
287        min_samples_for_calibration: 1000, // Default value
288        measurement_window_size: 10000,    // Default value
289        auto_temperature_adjustment: true,  // Enable automatic adjustment
290    };
291    calibration::initialize_calibration(&calibration_config).await?;
292    
293    tracing::info!("Semantic search module initialized successfully");
294    Ok(())
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300
301    #[test]
302    fn test_semantic_config_default() {
303        let config = SemanticConfig::default();
304        assert_eq!(config.encoder.max_tokens, 2048);
305        assert_eq!(config.encoder.model_type, "codet5-base");
306        assert!(config.rerank.use_isotonic);
307        assert!(!config.cross_encoder.enabled); // Start disabled
308    }
309
310    #[test] 
311    fn test_semantic_metrics_validation() {
312        let mut metrics = SemanticMetrics {
313            latency_p50_ms: 25.0,
314            latency_p95_ms: 45.0, // Within 50ms target
315            latency_p99_ms: 75.0,
316            ndcg_at_10: 0.53, // Above 0.52 target
317            success_at_10: 0.65,
318            recall_at_50: 0.85,
319            expected_calibration_error: 0.012,
320            ece_drift_from_baseline: 0.003, // Within 0.005 target
321            nl_slice_improvement_pp: 5.2, // Above 4pp target
322            meets_coir_target: true,
323            meets_latency_target: true, 
324            meets_ece_target: true,
325        };
326        
327        // Should pass all gates
328        assert!(metrics.validate_phase3_gates().is_ok());
329        
330        // Test failure cases
331        metrics.ndcg_at_10 = 0.49; // Below target
332        assert!(metrics.validate_phase3_gates().is_err());
333        
334        metrics.ndcg_at_10 = 0.53;
335        metrics.latency_p95_ms = 65.0; // Above target  
336        assert!(metrics.validate_phase3_gates().is_err());
337    }
338}
lens_core/semantic/mod.rs

lens_core/semantic/
mod.rs
