xz-rag 0.1.1

Multi-channel Retrieval-Augmented Generation engine
Documentation
/// Chunking strategy implementations.
pub mod chunker;

use std::sync::Arc;

use serde::{Deserialize, Serialize};

use crate::error::RagError;
use crate::indexing::chunker::ChunkStrategy;
use crate::types::chunk::{Chunk, ChunkMetadata};

/// Document to be indexed.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexDocument {
    /// Unique document identifier.
    pub id: String,
    /// Document text content.
    pub content: String,
    /// Optional document title.
    pub title: Option<String>,
    /// Metadata attached to the document.
    pub metadata: ChunkMetadata,
}

/// Document indexing pipeline.
pub struct DocumentIndexer {
    chunk_strategy: Arc<dyn ChunkStrategy>,
    namespace: String,
}

impl DocumentIndexer {
    /// Create a new `DocumentIndexer` with the given chunk strategy and namespace.
    pub fn new(chunk_strategy: Arc<dyn ChunkStrategy>, namespace: impl Into<String>) -> Self {
        Self { chunk_strategy, namespace: namespace.into() }
    }

    /// Index a single document into chunks.
    pub fn index_document(&self, doc: IndexDocument) -> Result<Vec<Chunk>, RagError> {
        let texts = self.chunk_strategy.chunk(&doc.content);
        let mut chunks = Vec::new();

        for (i, text) in texts.into_iter().enumerate() {
            let chunk_id = format!("{}-{}", doc.id, i);
            let mut chunk = Chunk::new(chunk_id, doc.id.clone(), text, i as u32);
            chunk.metadata = doc.metadata.clone();
            chunk.metadata.document_title = doc.title.clone();
            chunk.metadata.namespace = Some(self.namespace.clone());
            chunks.push(chunk);
        }

        Ok(chunks)
    }
}