xz-rag 0.1.1

use serde::{Deserialize, Serialize};

use super::retrieval::{RetrieveRequest, RetrieveResult};

// === RAG Request ===

/// End-to-end RAG request including retrieval, context, and generation config.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RagRequest {
    /// User query string.
    pub query: String,
    /// Optional system prompt override.
    pub system_prompt: Option<String>,
    /// Chat history for conversational context.
    pub history: Vec<ChatMessage>,
    /// Retrieval configuration.
    pub retrieve_config: RetrieveRequest,
    /// Generation parameters.
    pub generation: RagGenerationConfig,
    /// Optional context assembly configuration.
    pub context_config: Option<ContextConfig>,
    /// Optional prompt template name.
    pub prompt_template: Option<String>,
    /// Additional request options.
    pub options: RequestOptions,
}

impl RagRequest {
    /// Create a new builder for `RagRequest`.
    pub fn builder(query: impl Into<String>) -> RagRequestBuilder {
        RagRequestBuilder {
            query: query.into(),
            system_prompt: None,
            history: vec![],
            retrieve_config: RetrieveRequest::builder("").build(),
            generation: RagGenerationConfig::default(),
            context_config: None,
            prompt_template: None,
            options: RequestOptions::default(),
        }
    }
}

/// Builder for `RagRequest`.
pub struct RagRequestBuilder {
    query: String,
    system_prompt: Option<String>,
    history: Vec<ChatMessage>,
    retrieve_config: RetrieveRequest,
    generation: RagGenerationConfig,
    context_config: Option<ContextConfig>,
    prompt_template: Option<String>,
    options: RequestOptions,
}

impl RagRequestBuilder {
    /// Set the system prompt.
    pub fn system_prompt(mut self, sp: impl Into<String>) -> Self {
        self.system_prompt = Some(sp.into());
        self
    }

    /// Set the chat history.
    pub fn history(mut self, h: Vec<ChatMessage>) -> Self {
        self.history = h;
        self
    }

    /// Set the retrieval configuration.
    pub fn retrieve_config(mut self, rc: RetrieveRequest) -> Self {
        self.retrieve_config = rc;
        self
    }

    /// Set the generation parameters.
    pub fn generation(mut self, g: RagGenerationConfig) -> Self {
        self.generation = g;
        self
    }

    /// Set the context assembly configuration.
    pub fn context_config(mut self, cc: ContextConfig) -> Self {
        self.context_config = Some(cc);
        self
    }

    /// Set the prompt template name.
    pub fn prompt_template(mut self, pt: impl Into<String>) -> Self {
        self.prompt_template = Some(pt.into());
        self
    }

    /// Set additional request options.
    pub fn options(mut self, opts: RequestOptions) -> Self {
        self.options = opts;
        self
    }

    /// Build the `RagRequest`.
    pub fn build(self) -> RagRequest {
        RagRequest {
            query: self.query,
            system_prompt: self.system_prompt,
            history: self.history,
            retrieve_config: self.retrieve_config,
            generation: self.generation,
            context_config: self.context_config,
            prompt_template: self.prompt_template,
            options: self.options,
        }
    }
}

// === Chat Message ===

/// A single chat message in the conversation history.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    /// Role of the message sender.
    pub role: ChatRole,
    /// Content of the message.
    pub content: String,
}

/// Role of a chat message participant.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ChatRole {
    /// System instruction message.
    System,
    /// User message.
    User,
    /// Assistant response.
    Assistant,
}

// === Request Options ===

/// Additional options for a RAG request.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct RequestOptions {
    /// Optional namespace for scoping the retrieval.
    pub namespace: Option<String>,
    /// Optional timeout in milliseconds.
    pub timeout_ms: Option<u64>,
    /// Optional retry count for transient failures.
    pub retry_count: Option<u32>,
    /// Whether to stream the response.
    pub stream: bool,
}

// === RAG Generation Config ===

/// Generation parameters for the LLM.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RagGenerationConfig {
    /// Maximum context tokens for generation.
    pub max_context_tokens: usize,
    /// Model identifier (e.g. "gpt-4o", "claude-3").
    pub model: Option<String>,
    /// Temperature for sampling (0.0-1.0).
    pub temperature: Option<f32>,
    /// Maximum output tokens.
    pub max_output_tokens: Option<usize>,
    /// Whether to stream the generation.
    pub stream: bool,
}

impl Default for RagGenerationConfig {
    fn default() -> Self {
        Self {
            max_context_tokens: 4096,
            model: None,
            temperature: Some(0.7),
            max_output_tokens: Some(1024),
            stream: false,
        }
    }
}

// === Context Config ===

/// Configuration for context assembly from retrieved chunks.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextConfig {
    /// 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,
    /// Minimum overlap between chunks in the context.
    pub min_chunk_overlap: usize,
    /// Citation format to use.
    pub citation_format: CitationFormat,
}

impl Default for ContextConfig {
    fn default() -> Self {
        Self {
            max_context_tokens: 4096,
            system_prompt_reserve: 256,
            query_reserve: 128,
            output_reserve: 512,
            min_chunk_overlap: 50,
            citation_format: CitationFormat::Numeric,
        }
    }
}

/// Citation format style.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CitationFormat {
    /// Numeric citations like `[1]`, `[2]`, etc.
    Numeric,
    /// Use chunk IDs as citation markers.
    ChunkId,
    /// Use source document names as citation markers.
    SourceName,
}

// === RAG Response ===

/// Full RAG response including answer, citations, and usage statistics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RagResponse {
    /// Generated answer text.
    pub answer: String,
    /// List of citations referenced in the answer.
    pub citations: Vec<Citation>,
    /// Token usage statistics.
    pub usage: RagTokenUsage,
    /// Retrieval statistics from the pipeline.
    pub retrieve_stats: RetrieveResult,
    /// Total latency in milliseconds.
    pub total_latency_ms: u64,
    /// Model used for generation.
    pub model: Option<String>,
}

// === Citation ===

/// A single citation referencing a retrieved chunk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Citation {
    /// Citation index number.
    pub index: usize,
    /// ID of the referenced chunk.
    pub chunk_id: String,
    /// Content of the referenced chunk.
    pub content: String,
    /// Optional title of the source document.
    pub document_title: Option<String>,
    /// Relevance score of the chunk.
    pub score: f32,
    /// Channel that produced this chunk.
    pub channel: String,
}

// === Token Usage ===

/// Token usage statistics for a RAG operation.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct RagTokenUsage {
    /// Tokens consumed by the context.
    pub context_tokens: usize,
    /// Tokens consumed by the prompt.
    pub prompt_tokens: usize,
    /// Tokens consumed by the completion.
    pub completion_tokens: usize,
    /// Total tokens consumed.
    pub total_tokens: usize,
    /// Number of chunks included in the context.
    pub chunks_used: usize,
    /// Number of chunks dropped due to budget limits.
    pub chunks_dropped: usize,
}

// === Prompt Template ===

/// Prompt template for formatting the LLM prompt.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptTemplate {
    /// Template name.
    pub name: String,
    /// System prompt content.
    pub system: String,
    /// User message template with `{query}` and `{context}` placeholders.
    pub user_template: String,
    /// Prefix prepended to the context block.
    pub context_prefix: String,
    /// Format string for each chunk in the context.
    pub chunk_format: String,
    /// Suffix appended to the context block.
    pub context_suffix: String,
    /// Citation instruction appended to the prompt.
    pub citation_instruction: String,
}

impl PromptTemplate {
    /// Create a default QA prompt template.
    pub fn default_qa() -> Self {
        Self {
            name: "default_qa".into(),
            system: "You are a helpful assistant. Answer the user's question based on the provided context.".into(),
            user_template: "Question: {query}\n\nContext:\n{context}\n\nAnswer:".into(),
            context_prefix: "Relevant information:\n".into(),
            chunk_format: "[{index}] {content}\n".into(),
            context_suffix: "\n".into(),
            citation_instruction: "Please cite sources using [N] notation when referencing the context.".into(),
        }
    }

    /// Render the user template with the given query and context.
    pub fn render(&self, query: &str, context: &str) -> String {
        self.user_template.replace("{query}", query).replace("{context}", context)
    }
}

// === Streaming Event ===

/// Events emitted during a streaming RAG operation.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum RagStreamEvent {
    /// Retrieval phase has started.
    RetrievalStarted {
        /// Number of channels being queried.
        channel_count: usize,
    },
    /// A single channel has completed retrieval.
    ChannelDone {
        /// Channel identifier.
        channel: String,
        /// Number of hits from this channel.
        hits: usize,
        /// Latency for this channel in milliseconds.
        latency_ms: u64,
    },
    /// Generation phase has started with the assembled context.
    GenerationStarted {
        /// Number of chunks included in the context.
        context_chunks: usize,
        /// Tokens consumed by the context.
        context_tokens: usize,
    },
    /// A delta of generated content.
    ContentDelta {
        /// Text delta chunk.
        delta: String,
    },
    /// A citation has been identified in the generated content.
    Citation {
        /// Referenced chunk ID.
        chunk_id: String,
        /// Citation index.
        index: usize,
    },
    /// Generation is complete.
    Done {
        /// Total latency in milliseconds.
        total_latency_ms: u64,
        /// Full list of citations.
        citations: Vec<Citation>,
        /// Token usage statistics.
        usage: RagTokenUsage,
    },
}

// === Built Context ===

/// Assembled context from retrieved chunks, ready for generation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BuiltContext {
    /// Formatted context text.
    pub context_text: String,
    /// Citations extracted from the context.
    pub citations: Vec<Citation>,
    /// Number of chunks included in the context.
    pub chunks_used: usize,
    /// Number of chunks dropped due to budget limits.
    pub chunks_dropped: usize,
    /// Tokens consumed by the context.
    pub tokens_used: usize,
}