oxirs_chat/rag/

mod.rs

//! RAG (Retrieval-Augmented Generation) System for OxiRS Chat
//!
//! Implements multi-stage retrieval with semantic search, graph traversal,
//! and intelligent context assembly for knowledge graph exploration.
//!
//! This module is organized into specialized sub-modules for different aspects
//! of the RAG system:
//!
//! - **quantum_rag**: Quantum-inspired retrieval optimization
//! - **consciousness**: Consciousness-aware processing with memory traces
//! - **vector_search**: Vector-based semantic search and document management
//! - **embedding_providers**: Enhanced embedding models and multiple providers
//! - **graph_traversal**: Knowledge graph exploration and entity expansion
//! - **entity_extraction**: LLM-powered entity and relationship extraction
//! - **query_processing**: Query constraint processing and analysis utilities
//!
//! # Examples
//!
//! ```rust,no_run
//! use oxirs_chat::rag::{RagEngine, RagConfig};
//! use oxirs_core::ConcreteStore;
//! use std::sync::Arc;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let config = RagConfig::default();
//! let store = Arc::new(ConcreteStore::new()?);
//! let mut rag_engine = RagEngine::new(config, store as Arc<dyn oxirs_core::Store>);
//! rag_engine.initialize().await?;
//! # Ok(())
//! # }
//! ```

pub mod advanced_reasoning;
pub mod advanced_retrieval; // NEW: Advanced retrieval strategies (DPR, ColBERT, BM25+, LTR)
pub mod consciousness;
pub mod consciousness_types;
pub mod context;
pub mod embedding;
pub mod embedding_providers;
pub mod entity_extraction;
pub mod graph_traversal;
pub mod knowledge_extraction;
pub mod quantum;
pub mod quantum_rag;
pub mod query_processing;
pub mod ranking; // NEW: Result ranking system
pub mod retrieval;
pub mod types;
pub mod vector_search;

// Re-export main types for convenience
pub use advanced_reasoning::{
    AdvancedReasoningEngine, ReasoningChain, ReasoningConfig, ReasoningQuality, ReasoningResult,
    ReasoningType, UncertaintyFactor,
};
pub use consciousness::{
    AdvancedConsciousInsight, AdvancedConsciousResponse, AdvancedConsciousnessMetadata,
    AdvancedInsightType, ConsciousInsight, ConsciousnessConfig, ConsciousnessIntegration,
    ConsciousnessModel, EmotionalState, InsightType, MemoryTrace,
};
pub use consciousness_types::*;
pub use embedding_providers::{
    EmbeddingConfig, EmbeddingProviderType, EnhancedEmbeddingModel, SimpleEmbeddingModel,
};
pub use entity_extraction::{EntityExtractor, LLMEntityExtraction};
pub use graph_traversal::{EntityType, ExtractedEntity, ExtractedRelationship, GraphTraversal};
pub use knowledge_extraction::{
    EntityType as KnowledgeEntityType, ExtractedEntity as KnowledgeExtractedEntity,
    ExtractedKnowledge, ExtractedRelationship as KnowledgeExtractedRelationship,
    KnowledgeExtractionConfig, KnowledgeExtractionEngine, RelationshipType,
};
pub use quantum_rag::{QuantumRetrievalState, QuantumSearchResult, RagDocument};
pub use query_processing::{ConstraintType, QueryConstraint, QueryIntent, QueryProcessor};
pub use vector_search::{EnhancedVectorIndex, RagIndex, SearchDocument};

// Additional imports from submodules
pub use context::*;
pub use embedding::*;
pub use retrieval::*;
pub use types::*;

use anyhow::{anyhow, Result};
use chrono::{DateTime, Utc};
use oxirs_core::{
    model::{triple::Triple, NamedNode, Object, Subject},
    Store,
};
use oxirs_embed::{
    EmbeddingModel, ModelConfig, ModelStats, TrainingStats, Triple as EmbedTriple,
    Vector as EmbedVector,
};
use oxirs_vec::VectorIndex;
use oxirs_vec::{
    index::{
        AdvancedVectorIndex, DistanceMetric, IndexConfig, IndexType,
        SearchResult as VecSearchResult,
    },
    similarity,
};
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
    time::Duration,
};
use tracing::{debug, info, warn};
use uuid::Uuid;

/// RAG search result with metadata
#[derive(Debug, Clone)]
pub struct RagSearchResult {
    pub triple: Triple,
    pub score: f32,
    pub search_type: SearchType,
}

/// Search type for categorizing results
#[derive(Debug, Clone, PartialEq)]
pub enum SearchType {
    SemanticSimilarity,
    GraphTraversal,
    KeywordMatch,
    EntityExpansion,
}

/// RAG configuration
#[derive(Debug, Clone)]
pub struct RagConfig {
    pub retrieval: RetrievalConfig,
    pub quantum: QuantumConfig,
    pub consciousness: ConsciousnessConfig,
    pub embedding: EmbeddingConfig,
    pub graph: GraphConfig,
    pub max_context_length: usize,
    pub context_overlap: usize,
}

impl Default for RagConfig {
    fn default() -> Self {
        Self {
            retrieval: RetrievalConfig::default(),
            quantum: QuantumConfig::default(),
            consciousness: ConsciousnessConfig::default(),
            embedding: EmbeddingConfig::default(),
            graph: GraphConfig::default(),
            max_context_length: 4096,
            context_overlap: 200,
        }
    }
}

/// Retrieval configuration
#[derive(Debug, Clone)]
pub struct RetrievalConfig {
    pub max_results: usize,
    pub similarity_threshold: f32,
    pub graph_traversal_depth: usize,
    pub enable_entity_expansion: bool,
    pub enable_quantum_enhancement: bool,
    pub enable_consciousness_integration: bool,
}

impl Default for RetrievalConfig {
    fn default() -> Self {
        Self {
            max_results: 10,
            similarity_threshold: 0.7,
            graph_traversal_depth: 2,
            enable_entity_expansion: true,
            enable_quantum_enhancement: false,
            enable_consciousness_integration: false,
        }
    }
}

/// Quantum configuration
#[derive(Debug, Clone)]
pub struct QuantumConfig {
    pub enabled: bool,
    pub superposition_threshold: f64,
    pub entanglement_factor: f64,
    pub coherence_time: Duration,
}

impl Default for QuantumConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            superposition_threshold: 0.5,
            entanglement_factor: 0.8,
            coherence_time: Duration::from_secs(30),
        }
    }
}

/// Graph configuration
#[derive(Debug, Clone)]
pub struct GraphConfig {
    pub max_traversal_depth: usize,
    pub entity_expansion_limit: usize,
    pub relationship_weights: HashMap<String, f32>,
}

impl Default for GraphConfig {
    fn default() -> Self {
        Self {
            max_traversal_depth: 3,
            entity_expansion_limit: 50,
            relationship_weights: HashMap::new(),
        }
    }
}

/// Context assembled from various retrieval stages
#[derive(Debug, Clone)]
pub struct AssembledContext {
    pub retrieved_triples: Option<Vec<Triple>>,
    pub semantic_results: Vec<RagSearchResult>,
    pub graph_results: Vec<RagSearchResult>,
    pub quantum_results: Option<Vec<quantum_rag::QuantumSearchResult>>,
    pub consciousness_insights: Option<Vec<consciousness::ConsciousInsight>>,
    pub extracted_entities: Vec<graph_traversal::ExtractedEntity>,
    pub extracted_relationships: Vec<graph_traversal::ExtractedRelationship>,
    pub query_constraints: Vec<query_processing::QueryConstraint>,
    pub reasoning_results: Option<advanced_reasoning::ReasoningResult>,
    pub extracted_knowledge: Option<knowledge_extraction::ExtractedKnowledge>,
    pub context_score: f32,
    pub assembly_time: Duration,
}

impl AssembledContext {
    pub fn new() -> Self {
        Self {
            retrieved_triples: None,
            semantic_results: Vec::new(),
            graph_results: Vec::new(),
            quantum_results: None,
            consciousness_insights: None,
            extracted_entities: Vec::new(),
            extracted_relationships: Vec::new(),
            query_constraints: Vec::new(),
            reasoning_results: None,
            extracted_knowledge: None,
            context_score: 0.0,
            assembly_time: Duration::from_secs(0),
        }
    }
}

impl Default for AssembledContext {
    fn default() -> Self {
        Self::new()
    }
}

/// Main RAG engine that coordinates all components
pub struct RagEngine {
    pub config: RagConfig,
    pub store: Arc<dyn Store>,
    pub vector_index: Option<RagIndex>,
    pub embedding_model: Option<EnhancedEmbeddingModel>,
    pub quantum_state: Option<quantum_rag::QuantumRetrievalState>,
    pub consciousness: Option<consciousness::ConsciousnessIntegration>,
    pub graph_traversal: graph_traversal::GraphTraversal,
    pub entity_extractor: entity_extraction::EntityExtractor,
    pub query_processor: query_processing::QueryProcessor,
    pub reasoning_engine: Option<advanced_reasoning::AdvancedReasoningEngine>,
    pub knowledge_extractor: Option<knowledge_extraction::KnowledgeExtractionEngine>,
}

impl RagEngine {
    /// Create a new RAG engine with the given configuration
    pub fn new(config: RagConfig, store: Arc<dyn Store>) -> Self {
        Self {
            config: config.clone(),
            store: store.clone(),
            vector_index: None,
            embedding_model: None,
            quantum_state: None,
            consciousness: None,
            graph_traversal: graph_traversal::GraphTraversal::new(store.clone()),
            entity_extractor: entity_extraction::EntityExtractor::new(),
            query_processor: query_processing::QueryProcessor::new(),
            reasoning_engine: None,
            knowledge_extractor: None,
        }
    }

    /// Create a new RAG engine with vector index configuration
    pub async fn with_vector_index(
        config: RagConfig,
        store: Arc<dyn Store>,
        _vector_dimensions: usize,
    ) -> Result<Self> {
        let mut engine = Self::new(config, store);
        engine.initialize().await?;
        Ok(engine)
    }

    /// Initialize the RAG engine with all components
    pub async fn initialize(&mut self) -> Result<()> {
        // Initialize embedding model
        self.embedding_model = Some(EnhancedEmbeddingModel::new(self.config.embedding.clone())?);

        // Initialize vector index
        self.vector_index = Some(RagIndex::new().await?);

        // Initialize optional components
        if self.config.quantum.enabled {
            self.quantum_state = Some(quantum_rag::QuantumRetrievalState::new(0.5));
        }

        if self.config.consciousness.enabled {
            self.consciousness = Some(consciousness::ConsciousnessIntegration::new(
                self.config.consciousness.clone(),
            ));
        }

        // Initialize advanced reasoning engine
        let reasoning_config = advanced_reasoning::ReasoningConfig::default();
        self.reasoning_engine = Some(advanced_reasoning::AdvancedReasoningEngine::new(
            reasoning_config,
        ));

        // Initialize knowledge extraction engine
        let extraction_config = knowledge_extraction::KnowledgeExtractionConfig::default();
        self.knowledge_extractor = Some(knowledge_extraction::KnowledgeExtractionEngine::new(
            extraction_config,
        )?);

        info!("RAG engine initialized successfully with Version 1.2 features");
        Ok(())
    }

    /// Perform comprehensive retrieval for a query
    pub async fn retrieve(&mut self, query: &str) -> Result<AssembledContext> {
        let start_time = std::time::Instant::now();
        let mut context = AssembledContext::new();

        // Extract entities and constraints
        let (entities, relationships) = self
            .entity_extractor
            .extract_entities_and_relationships(query)
            .await?;
        let constraints = self
            .query_processor
            .extract_constraints(query, &entities)
            .await?;

        context.extracted_entities = entities;
        context.extracted_relationships = relationships;
        context.query_constraints = constraints;

        // Semantic search
        if let Some(ref mut vector_index) = self.vector_index {
            let semantic_docs = vector_index
                .search(query, self.config.retrieval.max_results)
                .await?;
            context.semantic_results = semantic_docs
                .into_iter()
                .map(|doc| RagSearchResult {
                    triple: doc.document,
                    score: doc.score,
                    search_type: SearchType::SemanticSimilarity,
                })
                .collect();
        }

        // Graph traversal
        let graph_results = self
            .graph_traversal
            .perform_graph_search(
                query,
                &context.extracted_entities,
                self.config.retrieval.graph_traversal_depth,
            )
            .await?;
        context.graph_results = graph_results;

        // Quantum enhancement (if enabled)
        if let Some(ref quantum_state) = self.quantum_state {
            if self.config.retrieval.enable_quantum_enhancement {
                let quantum_docs: Vec<RagDocument> = context
                    .semantic_results
                    .iter()
                    .map(|result| RagDocument {
                        id: uuid::Uuid::new_v4().to_string(),
                        content: result.triple.object().to_string(),
                        triple: Some(result.triple.clone()),
                        metadata: HashMap::new(),
                        embedding: None,
                    })
                    .collect();

                context.quantum_results = Some(quantum_state.superposition_search(&quantum_docs)?);
            }
        }

        // Consciousness integration (if enabled)
        if let Some(ref mut consciousness) = self.consciousness {
            if self.config.retrieval.enable_consciousness_integration {
                context.consciousness_insights = Some(
                    consciousness
                        .process_query_with_consciousness(query, &context)
                        .await?,
                );
            }
        }

        // Advanced reasoning (if enabled)
        if let Some(ref mut reasoning_engine) = self.reasoning_engine {
            debug!("Applying advanced reasoning to assembled context");
            match reasoning_engine.reason(query, &context).await {
                Ok(reasoning_result) => {
                    context.reasoning_results = Some(reasoning_result);
                    debug!("Advanced reasoning completed successfully");
                }
                Err(e) => {
                    warn!("Advanced reasoning failed: {}", e);
                    // Continue without reasoning results
                }
            }
        }

        // Knowledge extraction (if enabled)
        if let Some(ref mut knowledge_extractor) = self.knowledge_extractor {
            debug!("Extracting structured knowledge from query and context");
            // Create text for knowledge extraction from query and semantic results
            let mut extraction_text = query.to_string();
            for result in &context.semantic_results {
                extraction_text.push_str(&format!(" {}", result.triple.object()));
            }

            match knowledge_extractor
                .extract_knowledge(&extraction_text)
                .await
            {
                Ok(extracted_knowledge) => {
                    context.extracted_knowledge = Some(extracted_knowledge);
                    debug!("Knowledge extraction completed successfully");
                }
                Err(e) => {
                    warn!("Knowledge extraction failed: {}", e);
                    // Continue without extracted knowledge
                }
            }
        }

        // Calculate context score
        context.context_score = self.calculate_context_score(&context);
        context.assembly_time = start_time.elapsed();

        Ok(context)
    }

    /// Calculate overall context quality score
    fn calculate_context_score(&self, context: &AssembledContext) -> f32 {
        let mut score = 0.0;
        let mut components = 0;

        // Semantic results score
        if !context.semantic_results.is_empty() {
            let avg_semantic_score = context
                .semantic_results
                .iter()
                .map(|r| r.score)
                .sum::<f32>()
                / context.semantic_results.len() as f32;
            score += avg_semantic_score;
            components += 1;
        }

        // Graph results score
        if !context.graph_results.is_empty() {
            let avg_graph_score = context.graph_results.iter().map(|r| r.score).sum::<f32>()
                / context.graph_results.len() as f32;
            score += avg_graph_score;
            components += 1;
        }

        // Entity extraction score
        if !context.extracted_entities.is_empty() {
            let avg_entity_confidence = context
                .extracted_entities
                .iter()
                .map(|e| e.confidence)
                .sum::<f32>()
                / context.extracted_entities.len() as f32;
            score += avg_entity_confidence;
            components += 1;
        }

        // Advanced reasoning score
        if let Some(ref reasoning_results) = context.reasoning_results {
            let reasoning_score = reasoning_results.reasoning_quality.overall_quality as f32;
            score += reasoning_score;
            components += 1;
        }

        // Knowledge extraction score
        if let Some(ref extracted_knowledge) = context.extracted_knowledge {
            let knowledge_score = extracted_knowledge.confidence_score as f32;
            score += knowledge_score;
            components += 1;
        }

        // Quantum results score (if available)
        if let Some(ref quantum_results) = context.quantum_results {
            if !quantum_results.is_empty() {
                let avg_quantum_score = quantum_results
                    .iter()
                    .map(|r| r.quantum_probability as f32)
                    .sum::<f32>()
                    / quantum_results.len() as f32;
                score += avg_quantum_score;
                components += 1;
            }
        }

        // Consciousness insights score (if available)
        if let Some(ref consciousness_insights) = context.consciousness_insights {
            if !consciousness_insights.is_empty() {
                let avg_consciousness_score = consciousness_insights
                    .iter()
                    .map(|insight| insight.confidence)
                    .sum::<f64>()
                    / consciousness_insights.len() as f64;
                score += avg_consciousness_score as f32;
                components += 1;
            }
        }

        if components > 0 {
            score / components as f32
        } else {
            0.0
        }
    }

    /// Get current configuration
    pub fn config(&self) -> &RagConfig {
        &self.config
    }

    /// Update configuration
    pub fn set_config(&mut self, config: RagConfig) {
        self.config = config;
    }
}

// Type aliases for compatibility with lib.rs
pub type RAGSystem = RagEngine;
pub type RAGConfig = RagConfig;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rag_config_default() {
        let config = RagConfig::default();
        assert_eq!(config.retrieval.max_results, 10);
        assert_eq!(config.retrieval.similarity_threshold, 0.7);
        assert_eq!(config.retrieval.graph_traversal_depth, 2);
        assert!(config.retrieval.enable_entity_expansion);
    }

    #[test]
    fn test_assembled_context_creation() {
        let context = AssembledContext::new();
        assert!(context.semantic_results.is_empty());
        assert!(context.graph_results.is_empty());
        assert!(context.extracted_entities.is_empty());
        assert_eq!(context.context_score, 0.0);
    }
}