xz-rag 0.1.1

use serde::{Deserialize, Serialize};

/// RAG engine information metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RagEngineInfo {
    /// Engine name.
    pub name: String,
    /// Engine version.
    pub version: String,
    /// List of channel types this engine supports.
    pub supported_channels: Vec<String>,
    /// Whether streaming generation is supported.
    pub supports_streaming: bool,
    /// Whether reranking is enabled.
    pub reranking_enabled: bool,
    /// Maximum context window size in tokens.
    pub max_context_window: usize,
}

/// RAG engine configuration loaded from YAML/JSON.
///
/// TODO (@future): Wire RagConfig into DefaultRagEngineBuilder::from_config()
/// Currently engine builder does not accept this config.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RagConfig {
    /// Engine-level settings.
    pub engine: EngineSection,
    /// Channel configuration.
    pub channels: ChannelSection,
    /// Fusion algorithm settings.
    pub fusion: FusionSection,
    /// Reranking configuration.
    pub reranking: RerankingSection,
    /// Context window configuration.
    pub context: ContextSection,
    /// Chunking configuration.
    pub chunking: ChunkingSection,
    /// Query preprocessing configuration.
    pub query_preprocessing: QueryPreprocessingSection,
    /// Prompt template definitions.
    pub prompt_templates: Vec<TemplateSection>,
    /// Cache configuration.
    pub cache: CacheSection,
}

/// Engine-level configuration section.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineSection {
    /// Engine name identifier.
    pub name: String,
    /// Optional default namespace.
    pub namespace: Option<String>,
}

impl Default for EngineSection {
    fn default() -> Self {
        Self { name: "default".into(), namespace: None }
    }
}

/// Channel configuration section.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChannelSection {
    /// Semantic search channel config.
    pub semantic: Option<ChannelDef>,
    /// BM25 full-text search channel config.
    pub bm25: Option<ChannelDef>,
    /// Metadata search channel config.
    pub metadata: Option<ChannelDef>,
}

impl Default for ChannelSection {
    fn default() -> Self {
        Self {
            semantic: Some(ChannelDef { weight: 0.5, top_k: 10, min_score: Some(0.1) }),
            bm25: None,
            metadata: Some(ChannelDef { weight: 0.2, top_k: 5, min_score: None }),
        }
    }
}

/// Per-channel weight and threshold definitions.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChannelDef {
    /// Channel weight for RRF fusion.
    pub weight: f32,
    /// Maximum number of results from this channel.
    pub top_k: usize,
    /// Minimum score threshold (results below are discarded).
    pub min_score: Option<f32>,
}

/// Fusion algorithm configuration section.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FusionSection {
    /// Fusion algorithm name (e.g. "rrf").
    pub algorithm: String,
    /// RRF k constant for score smoothing.
    pub rrf_k: usize,
    /// Whether to normalize scores before fusion.
    pub normalize_scores: bool,
}

impl Default for FusionSection {
    fn default() -> Self {
        Self { algorithm: "rrf".into(), rrf_k: 60, normalize_scores: true }
    }
}

/// Reranking configuration section.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RerankingSection {
    /// Whether reranking is enabled.
    pub enabled: bool,
    /// Reranker provider name (e.g. "cohere", "jina").
    pub reranker: Option<String>,
}

/// Context window configuration section.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextSection {
    /// Maximum context window size in tokens.
    pub max_context_tokens: usize,
    /// Tokens reserved for the system prompt.
    pub system_prompt_reserve: usize,
    /// Tokens reserved for the user query.
    pub query_reserve: usize,
    /// Tokens reserved for the model output.
    pub output_reserve: usize,
    /// Citation format ("numeric", "chunk_id", or "source_name").
    pub citation_format: String,
    /// Overlap between consecutive chunks in the context.
    pub chunk_overlap: usize,
    /// Separator string between chunks.
    pub separator: String,
}

impl Default for ContextSection {
    fn default() -> Self {
        Self {
            max_context_tokens: 4096,
            system_prompt_reserve: 256,
            query_reserve: 128,
            output_reserve: 512,
            citation_format: "numeric".into(),
            chunk_overlap: 50,
            separator: "\n---\n".into(),
        }
    }
}

/// Chunking configuration section.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChunkingSection {
    /// Default chunking strategy name.
    pub default_strategy: String,
    /// Recursive chunker configuration.
    pub recursive: Option<ChunkerDef>,
    /// Fixed-size chunker configuration.
    pub fixed: Option<FixedChunkerDef>,
}

impl Default for ChunkingSection {
    fn default() -> Self {
        Self {
            default_strategy: "recursive".into(),
            recursive: Some(ChunkerDef {
                chunk_size: 512,
                overlap: 50,
                separators: vec!["\n\n".into(), "\n".into(), ". ".into(), " ".into()],
            }),
            fixed: Some(FixedChunkerDef { chunk_size: 512, overlap: 50 }),
        }
    }
}

/// Recursive chunker configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChunkerDef {
    /// Target chunk size in characters.
    pub chunk_size: usize,
    /// Overlap between consecutive chunks.
    pub overlap: usize,
    /// Ordered list of separators to try splitting on.
    pub separators: Vec<String>,
}

/// Fixed-size chunker configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FixedChunkerDef {
    /// Target chunk size in characters.
    pub chunk_size: usize,
    /// Overlap between consecutive chunks.
    pub overlap: usize,
}

/// Query preprocessing configuration section.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct QueryPreprocessingSection {
    /// HYDE expansion configuration.
    pub hyde: Option<HydeDef>,
    /// Query expansion configuration.
    pub query_expansion: Option<QueryExpansionDef>,
}

/// HYDE (Hypothetical Document Embeddings) configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HydeDef {
    /// Prompt template for generating hypothetical documents.
    pub prompt_template: String,
}

/// Query expansion configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryExpansionDef {
    /// Number of query variations to generate.
    pub count: usize,
}

/// Prompt template definition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TemplateSection {
    /// Template name.
    pub name: String,
    /// System prompt content.
    pub system: String,
    /// User message template with placeholders.
    pub user_template: String,
}

/// Cache configuration section.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheSection {
    /// Whether caching is enabled.
    pub enabled: bool,
    /// Time-to-live for cached entries in seconds.
    pub ttl_seconds: u64,
    /// Maximum number of entries in the cache.
    pub max_entries: usize,
}

impl Default for CacheSection {
    fn default() -> Self {
        Self { enabled: false, ttl_seconds: 3600, max_entries: 1000 }
    }
}