oxirs_chat/rag/
mod.rs

1//! RAG (Retrieval-Augmented Generation) System for OxiRS Chat
2//!
3//! Implements multi-stage retrieval with semantic search, graph traversal,
4//! and intelligent context assembly for knowledge graph exploration.
5//!
6//! This module is organized into specialized sub-modules for different aspects
7//! of the RAG system:
8//!
9//! - **quantum_rag**: Quantum-inspired retrieval optimization
10//! - **consciousness**: Consciousness-aware processing with memory traces
11//! - **vector_search**: Vector-based semantic search and document management
12//! - **embedding_providers**: Enhanced embedding models and multiple providers
13//! - **graph_traversal**: Knowledge graph exploration and entity expansion
14//! - **entity_extraction**: LLM-powered entity and relationship extraction
15//! - **query_processing**: Query constraint processing and analysis utilities
16//!
17//! # Examples
18//!
19//! ```rust,no_run
20//! use oxirs_chat::rag::{RagEngine, RagConfig};
21//! use oxirs_core::ConcreteStore;
22//! use std::sync::Arc;
23//!
24//! # async fn example() -> anyhow::Result<()> {
25//! let config = RagConfig::default();
26//! let store = Arc::new(ConcreteStore::new()?);
27//! let mut rag_engine = RagEngine::new(config, store as Arc<dyn oxirs_core::Store>);
28//! rag_engine.initialize().await?;
29//! # Ok(())
30//! # }
31//! ```
32
33pub mod advanced_reasoning;
34pub mod advanced_retrieval; // NEW: Advanced retrieval strategies (DPR, ColBERT, BM25+, LTR)
35pub mod consciousness;
36pub mod consciousness_types;
37pub mod context;
38pub mod embedding;
39pub mod embedding_providers;
40pub mod entity_extraction;
41pub mod graph_traversal;
42pub mod knowledge_extraction;
43pub mod pipeline; // NEW: Hybrid BM25+vector RAG pipeline
44pub mod quantum;
45pub mod quantum_rag;
46pub mod query_processing;
47pub mod ranking; // NEW: Result ranking system
48pub mod retrieval;
49pub mod types;
50pub mod vector_search;
51
52// Re-export main types for convenience
53pub use advanced_reasoning::{
54    AdvancedReasoningEngine, ReasoningChain, ReasoningConfig, ReasoningQuality, ReasoningResult,
55    ReasoningType, UncertaintyFactor,
56};
57pub use consciousness::{
58    AdvancedConsciousInsight, AdvancedConsciousResponse, AdvancedConsciousnessMetadata,
59    AdvancedInsightType, ConsciousInsight, ConsciousnessConfig, ConsciousnessIntegration,
60    ConsciousnessModel, EmotionalState, InsightType, MemoryTrace,
61};
62pub use consciousness_types::*;
63pub use embedding_providers::{
64    EmbeddingConfig, EmbeddingProviderType, EnhancedEmbeddingModel, SimpleEmbeddingModel,
65};
66pub use entity_extraction::{EntityExtractor, LLMEntityExtraction};
67pub use graph_traversal::{EntityType, ExtractedEntity, ExtractedRelationship, GraphTraversal};
68pub use knowledge_extraction::{
69    EntityType as KnowledgeEntityType, ExtractedEntity as KnowledgeExtractedEntity,
70    ExtractedKnowledge, ExtractedRelationship as KnowledgeExtractedRelationship,
71    KnowledgeExtractionConfig, KnowledgeExtractionEngine, RelationshipType,
72};
73pub use quantum_rag::{QuantumRetrievalState, QuantumSearchResult, RagDocument};
74pub use query_processing::{ConstraintType, QueryConstraint, QueryIntent, QueryProcessor};
75pub use vector_search::{EnhancedVectorIndex, RagIndex, SearchDocument};
76
77// Additional imports from submodules
78pub use context::*;
79pub use embedding::*;
80pub use retrieval::*;
81pub use types::*;
82
83use anyhow::{anyhow, Result};
84use chrono::{DateTime, Utc};
85use oxirs_core::{
86    model::{triple::Triple, NamedNode, Object, Subject},
87    Store,
88};
89use oxirs_embed::{
90    EmbeddingModel, ModelConfig, ModelStats, TrainingStats, Triple as EmbedTriple,
91    Vector as EmbedVector,
92};
93use oxirs_vec::VectorIndex;
94use oxirs_vec::{
95    index::{
96        AdvancedVectorIndex, DistanceMetric, IndexConfig, IndexType,
97        SearchResult as VecSearchResult,
98    },
99    similarity,
100};
101use regex::Regex;
102use serde::{Deserialize, Serialize};
103use std::{
104    collections::{HashMap, HashSet},
105    sync::Arc,
106    time::Duration,
107};
108use tracing::{debug, info, warn};
109use uuid::Uuid;
110
111/// RAG search result with metadata
112#[derive(Debug, Clone)]
113pub struct RagSearchResult {
114    pub triple: Triple,
115    pub score: f32,
116    pub search_type: SearchType,
117}
118
119/// Search type for categorizing results
120#[derive(Debug, Clone, PartialEq)]
121pub enum SearchType {
122    SemanticSimilarity,
123    GraphTraversal,
124    KeywordMatch,
125    EntityExpansion,
126}
127
128/// RAG configuration
129#[derive(Debug, Clone)]
130pub struct RagConfig {
131    pub retrieval: RetrievalConfig,
132    pub quantum: QuantumConfig,
133    pub consciousness: ConsciousnessConfig,
134    pub embedding: EmbeddingConfig,
135    pub graph: GraphConfig,
136    pub max_context_length: usize,
137    pub context_overlap: usize,
138}
139
140impl Default for RagConfig {
141    fn default() -> Self {
142        Self {
143            retrieval: RetrievalConfig::default(),
144            quantum: QuantumConfig::default(),
145            consciousness: ConsciousnessConfig::default(),
146            embedding: EmbeddingConfig::default(),
147            graph: GraphConfig::default(),
148            max_context_length: 4096,
149            context_overlap: 200,
150        }
151    }
152}
153
154/// Retrieval configuration
155#[derive(Debug, Clone)]
156pub struct RetrievalConfig {
157    pub max_results: usize,
158    pub similarity_threshold: f32,
159    pub graph_traversal_depth: usize,
160    pub enable_entity_expansion: bool,
161    pub enable_quantum_enhancement: bool,
162    pub enable_consciousness_integration: bool,
163}
164
165impl Default for RetrievalConfig {
166    fn default() -> Self {
167        Self {
168            max_results: 10,
169            similarity_threshold: 0.7,
170            graph_traversal_depth: 2,
171            enable_entity_expansion: true,
172            enable_quantum_enhancement: false,
173            enable_consciousness_integration: false,
174        }
175    }
176}
177
178/// Quantum configuration
179#[derive(Debug, Clone)]
180pub struct QuantumConfig {
181    pub enabled: bool,
182    pub superposition_threshold: f64,
183    pub entanglement_factor: f64,
184    pub coherence_time: Duration,
185}
186
187impl Default for QuantumConfig {
188    fn default() -> Self {
189        Self {
190            enabled: false,
191            superposition_threshold: 0.5,
192            entanglement_factor: 0.8,
193            coherence_time: Duration::from_secs(30),
194        }
195    }
196}
197
198/// Graph configuration
199#[derive(Debug, Clone)]
200pub struct GraphConfig {
201    pub max_traversal_depth: usize,
202    pub entity_expansion_limit: usize,
203    pub relationship_weights: HashMap<String, f32>,
204}
205
206impl Default for GraphConfig {
207    fn default() -> Self {
208        Self {
209            max_traversal_depth: 3,
210            entity_expansion_limit: 50,
211            relationship_weights: HashMap::new(),
212        }
213    }
214}
215
216/// Context assembled from various retrieval stages
217#[derive(Debug, Clone)]
218pub struct AssembledContext {
219    pub retrieved_triples: Option<Vec<Triple>>,
220    pub semantic_results: Vec<RagSearchResult>,
221    pub graph_results: Vec<RagSearchResult>,
222    pub quantum_results: Option<Vec<quantum_rag::QuantumSearchResult>>,
223    pub consciousness_insights: Option<Vec<consciousness::ConsciousInsight>>,
224    pub extracted_entities: Vec<graph_traversal::ExtractedEntity>,
225    pub extracted_relationships: Vec<graph_traversal::ExtractedRelationship>,
226    pub query_constraints: Vec<query_processing::QueryConstraint>,
227    pub reasoning_results: Option<advanced_reasoning::ReasoningResult>,
228    pub extracted_knowledge: Option<knowledge_extraction::ExtractedKnowledge>,
229    pub context_score: f32,
230    pub assembly_time: Duration,
231}
232
233impl AssembledContext {
234    pub fn new() -> Self {
235        Self {
236            retrieved_triples: None,
237            semantic_results: Vec::new(),
238            graph_results: Vec::new(),
239            quantum_results: None,
240            consciousness_insights: None,
241            extracted_entities: Vec::new(),
242            extracted_relationships: Vec::new(),
243            query_constraints: Vec::new(),
244            reasoning_results: None,
245            extracted_knowledge: None,
246            context_score: 0.0,
247            assembly_time: Duration::from_secs(0),
248        }
249    }
250}
251
252impl Default for AssembledContext {
253    fn default() -> Self {
254        Self::new()
255    }
256}
257
258/// Main RAG engine that coordinates all components
259pub struct RagEngine {
260    pub config: RagConfig,
261    pub store: Arc<dyn Store>,
262    pub vector_index: Option<RagIndex>,
263    pub embedding_model: Option<EnhancedEmbeddingModel>,
264    pub quantum_state: Option<quantum_rag::QuantumRetrievalState>,
265    pub consciousness: Option<consciousness::ConsciousnessIntegration>,
266    pub graph_traversal: graph_traversal::GraphTraversal,
267    pub entity_extractor: entity_extraction::EntityExtractor,
268    pub query_processor: query_processing::QueryProcessor,
269    pub reasoning_engine: Option<advanced_reasoning::AdvancedReasoningEngine>,
270    pub knowledge_extractor: Option<knowledge_extraction::KnowledgeExtractionEngine>,
271}
272
273impl RagEngine {
274    /// Create a new RAG engine with the given configuration
275    pub fn new(config: RagConfig, store: Arc<dyn Store>) -> Self {
276        Self {
277            config: config.clone(),
278            store: store.clone(),
279            vector_index: None,
280            embedding_model: None,
281            quantum_state: None,
282            consciousness: None,
283            graph_traversal: graph_traversal::GraphTraversal::new(store.clone()),
284            entity_extractor: entity_extraction::EntityExtractor::new(),
285            query_processor: query_processing::QueryProcessor::new(),
286            reasoning_engine: None,
287            knowledge_extractor: None,
288        }
289    }
290
291    /// Create a new RAG engine with vector index configuration
292    pub async fn with_vector_index(
293        config: RagConfig,
294        store: Arc<dyn Store>,
295        _vector_dimensions: usize,
296    ) -> Result<Self> {
297        let mut engine = Self::new(config, store);
298        engine.initialize().await?;
299        Ok(engine)
300    }
301
302    /// Initialize the RAG engine with all components
303    pub async fn initialize(&mut self) -> Result<()> {
304        // Initialize embedding model
305        self.embedding_model = Some(EnhancedEmbeddingModel::new(self.config.embedding.clone())?);
306
307        // Initialize vector index
308        self.vector_index = Some(RagIndex::new().await?);
309
310        // Initialize optional components
311        if self.config.quantum.enabled {
312            self.quantum_state = Some(quantum_rag::QuantumRetrievalState::new(0.5));
313        }
314
315        if self.config.consciousness.enabled {
316            self.consciousness = Some(consciousness::ConsciousnessIntegration::new(
317                self.config.consciousness.clone(),
318            ));
319        }
320
321        // Initialize advanced reasoning engine
322        let reasoning_config = advanced_reasoning::ReasoningConfig::default();
323        self.reasoning_engine = Some(advanced_reasoning::AdvancedReasoningEngine::new(
324            reasoning_config,
325        ));
326
327        // Initialize knowledge extraction engine
328        let extraction_config = knowledge_extraction::KnowledgeExtractionConfig::default();
329        self.knowledge_extractor = Some(knowledge_extraction::KnowledgeExtractionEngine::new(
330            extraction_config,
331        )?);
332
333        info!("RAG engine initialized successfully with Version 1.2 features");
334        Ok(())
335    }
336
337    /// Perform comprehensive retrieval for a query
338    pub async fn retrieve(&mut self, query: &str) -> Result<AssembledContext> {
339        let start_time = std::time::Instant::now();
340        let mut context = AssembledContext::new();
341
342        // Extract entities and constraints
343        let (entities, relationships) = self
344            .entity_extractor
345            .extract_entities_and_relationships(query)
346            .await?;
347        let constraints = self
348            .query_processor
349            .extract_constraints(query, &entities)
350            .await?;
351
352        context.extracted_entities = entities;
353        context.extracted_relationships = relationships;
354        context.query_constraints = constraints;
355
356        // Semantic search
357        if let Some(ref mut vector_index) = self.vector_index {
358            let semantic_docs = vector_index
359                .search(query, self.config.retrieval.max_results)
360                .await?;
361            context.semantic_results = semantic_docs
362                .into_iter()
363                .map(|doc| RagSearchResult {
364                    triple: doc.document,
365                    score: doc.score,
366                    search_type: SearchType::SemanticSimilarity,
367                })
368                .collect();
369        }
370
371        // Graph traversal
372        let graph_results = self
373            .graph_traversal
374            .perform_graph_search(
375                query,
376                &context.extracted_entities,
377                self.config.retrieval.graph_traversal_depth,
378            )
379            .await?;
380        context.graph_results = graph_results;
381
382        // Quantum enhancement (if enabled)
383        if let Some(ref quantum_state) = self.quantum_state {
384            if self.config.retrieval.enable_quantum_enhancement {
385                let quantum_docs: Vec<RagDocument> = context
386                    .semantic_results
387                    .iter()
388                    .map(|result| RagDocument {
389                        id: uuid::Uuid::new_v4().to_string(),
390                        content: result.triple.object().to_string(),
391                        triple: Some(result.triple.clone()),
392                        metadata: HashMap::new(),
393                        embedding: None,
394                    })
395                    .collect();
396
397                context.quantum_results = Some(quantum_state.superposition_search(&quantum_docs)?);
398            }
399        }
400
401        // Consciousness integration (if enabled)
402        if let Some(ref mut consciousness) = self.consciousness {
403            if self.config.retrieval.enable_consciousness_integration {
404                context.consciousness_insights = Some(
405                    consciousness
406                        .process_query_with_consciousness(query, &context)
407                        .await?,
408                );
409            }
410        }
411
412        // Advanced reasoning (if enabled)
413        if let Some(ref mut reasoning_engine) = self.reasoning_engine {
414            debug!("Applying advanced reasoning to assembled context");
415            match reasoning_engine.reason(query, &context).await {
416                Ok(reasoning_result) => {
417                    context.reasoning_results = Some(reasoning_result);
418                    debug!("Advanced reasoning completed successfully");
419                }
420                Err(e) => {
421                    warn!("Advanced reasoning failed: {}", e);
422                    // Continue without reasoning results
423                }
424            }
425        }
426
427        // Knowledge extraction (if enabled)
428        if let Some(ref mut knowledge_extractor) = self.knowledge_extractor {
429            debug!("Extracting structured knowledge from query and context");
430            // Create text for knowledge extraction from query and semantic results
431            let mut extraction_text = query.to_string();
432            for result in &context.semantic_results {
433                extraction_text.push_str(&format!(" {}", result.triple.object()));
434            }
435
436            match knowledge_extractor
437                .extract_knowledge(&extraction_text)
438                .await
439            {
440                Ok(extracted_knowledge) => {
441                    context.extracted_knowledge = Some(extracted_knowledge);
442                    debug!("Knowledge extraction completed successfully");
443                }
444                Err(e) => {
445                    warn!("Knowledge extraction failed: {}", e);
446                    // Continue without extracted knowledge
447                }
448            }
449        }
450
451        // Calculate context score
452        context.context_score = self.calculate_context_score(&context);
453        context.assembly_time = start_time.elapsed();
454
455        Ok(context)
456    }
457
458    /// Calculate overall context quality score
459    fn calculate_context_score(&self, context: &AssembledContext) -> f32 {
460        let mut score = 0.0;
461        let mut components = 0;
462
463        // Semantic results score
464        if !context.semantic_results.is_empty() {
465            let avg_semantic_score = context
466                .semantic_results
467                .iter()
468                .map(|r| r.score)
469                .sum::<f32>()
470                / context.semantic_results.len() as f32;
471            score += avg_semantic_score;
472            components += 1;
473        }
474
475        // Graph results score
476        if !context.graph_results.is_empty() {
477            let avg_graph_score = context.graph_results.iter().map(|r| r.score).sum::<f32>()
478                / context.graph_results.len() as f32;
479            score += avg_graph_score;
480            components += 1;
481        }
482
483        // Entity extraction score
484        if !context.extracted_entities.is_empty() {
485            let avg_entity_confidence = context
486                .extracted_entities
487                .iter()
488                .map(|e| e.confidence)
489                .sum::<f32>()
490                / context.extracted_entities.len() as f32;
491            score += avg_entity_confidence;
492            components += 1;
493        }
494
495        // Advanced reasoning score
496        if let Some(ref reasoning_results) = context.reasoning_results {
497            let reasoning_score = reasoning_results.reasoning_quality.overall_quality as f32;
498            score += reasoning_score;
499            components += 1;
500        }
501
502        // Knowledge extraction score
503        if let Some(ref extracted_knowledge) = context.extracted_knowledge {
504            let knowledge_score = extracted_knowledge.confidence_score as f32;
505            score += knowledge_score;
506            components += 1;
507        }
508
509        // Quantum results score (if available)
510        if let Some(ref quantum_results) = context.quantum_results {
511            if !quantum_results.is_empty() {
512                let avg_quantum_score = quantum_results
513                    .iter()
514                    .map(|r| r.quantum_probability as f32)
515                    .sum::<f32>()
516                    / quantum_results.len() as f32;
517                score += avg_quantum_score;
518                components += 1;
519            }
520        }
521
522        // Consciousness insights score (if available)
523        if let Some(ref consciousness_insights) = context.consciousness_insights {
524            if !consciousness_insights.is_empty() {
525                let avg_consciousness_score = consciousness_insights
526                    .iter()
527                    .map(|insight| insight.confidence)
528                    .sum::<f64>()
529                    / consciousness_insights.len() as f64;
530                score += avg_consciousness_score as f32;
531                components += 1;
532            }
533        }
534
535        if components > 0 {
536            score / components as f32
537        } else {
538            0.0
539        }
540    }
541
542    /// Get current configuration
543    pub fn config(&self) -> &RagConfig {
544        &self.config
545    }
546
547    /// Update configuration
548    pub fn set_config(&mut self, config: RagConfig) {
549        self.config = config;
550    }
551}
552
553// Type aliases for compatibility with lib.rs
554pub type RAGSystem = RagEngine;
555pub type RAGConfig = RagConfig;
556
557#[cfg(test)]
558mod tests {
559    use super::*;
560
561    #[test]
562    fn test_rag_config_default() {
563        let config = RagConfig::default();
564        assert_eq!(config.retrieval.max_results, 10);
565        assert_eq!(config.retrieval.similarity_threshold, 0.7);
566        assert_eq!(config.retrieval.graph_traversal_depth, 2);
567        assert!(config.retrieval.enable_entity_expansion);
568    }
569
570    #[test]
571    fn test_assembled_context_creation() {
572        let context = AssembledContext::new();
573        assert!(context.semantic_results.is_empty());
574        assert!(context.graph_results.is_empty());
575        assert!(context.extracted_entities.is_empty());
576        assert_eq!(context.context_score, 0.0);
577    }
578}
oxirs_chat/rag/mod.rs

oxirs_chat/rag/
mod.rs
