core/rag/src/chunking/mod.rs - main - alexefan136/flowstack

core/rag/src/chunking/mod.rs
104 строки3 KB
Alexander Efanov
Обновление репозитория
15 июл 2026, 12:19
15 июл 2026, 12:1976704c6
//! Chunking strategies для разбиения документов.
//!
//! # Доступные стратегии
//!
//! - [`FixedSizeChunker`] — простой, по количеству символов
//! - [`SentenceChunker`] — по предложениям (сохраняет семантику)
//! - [`ParagraphChunker`] — по абзацам (для структурированных документов)
//! - [`RecursiveChunker`] — рекурсивный (LangChain-style, рекомендуется по умолчанию)

pub mod fixed_size;
pub mod paragraph;
pub mod recursive;
pub mod sentence;
pub mod traits;

pub use fixed_size::FixedSizeChunker;
pub use paragraph::ParagraphChunker;
pub use recursive::RecursiveChunker;
pub use sentence::SentenceChunker;
pub use traits::{Chunker, ChunkerConfig};

use crate::domain::ChunkStrategy;
use crate::Result;

/// Factory для создания chunker по стратегии.
pub fn create_chunker(
    strategy: ChunkStrategy,
    config: ChunkerConfig,
) -> Result<Box<dyn Chunker>> {
    config.validate()?;

    let chunker: Box<dyn Chunker> = match strategy {
        ChunkStrategy::FixedSize => Box::new(FixedSizeChunker::new(config)),
        ChunkStrategy::Sentence => Box::new(SentenceChunker::new(config)),
        ChunkStrategy::Paragraph => Box::new(ParagraphChunker::new(config)),
        ChunkStrategy::Recursive => Box::new(RecursiveChunker::new(config)),
    };

    Ok(chunker)
}

/// Chunker по умолчанию (Recursive).
pub fn default_chunker() -> Box<dyn Chunker> {
    Box::new(RecursiveChunker::new(ChunkerConfig::default()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::domain::{Document, DocumentType};

    #[test]
    fn test_create_chunker_recursive() {
        let chunker = create_chunker(ChunkStrategy::Recursive, ChunkerConfig::default()).unwrap();
        assert_eq!(chunker.name(), "recursive");
    }

    #[test]
    fn test_create_chunker_fixed_size() {
        let chunker = create_chunker(ChunkStrategy::FixedSize, ChunkerConfig::default()).unwrap();
        assert_eq!(chunker.name(), "fixed_size");
    }

    #[test]
    fn test_invalid_config() {
        let config = ChunkerConfig::new(0, 0);
        let result = create_chunker(ChunkStrategy::Recursive, config);
        assert!(result.is_err());
    }

    #[test]
    fn test_overlap_greater_than_size() {
        let config = ChunkerConfig::new(100, 200);
        let result = create_chunker(ChunkStrategy::Recursive, config);
        assert!(result.is_err());
    }

    #[test]
    fn test_all_strategies_on_same_document() {
        let text = "First sentence. Second sentence.\n\nNew paragraph. More text here. And more.";
        let doc = Document::new(text.into(), "test.txt".into(), DocumentType::Txt, "ws-1".into());
        let config = ChunkerConfig::new(30, 5);

        for strategy in [
            ChunkStrategy::FixedSize,
            ChunkStrategy::Sentence,
            ChunkStrategy::Paragraph,
            ChunkStrategy::Recursive,
        ] {
            let chunker = create_chunker(strategy, config.clone()).unwrap();
            let chunks = chunker.split(&doc).unwrap();
            assert!(
                !chunks.is_empty() || text.is_empty(),
                "Strategy {:?} produced no chunks",
                strategy
            );

            // Проверяем что индексы идут подряд
            for (i, chunk) in chunks.iter().enumerate() {
                assert_eq!(chunk.index, i);
            }
        }
    }
}