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

core/rag/src/chunking/recursive.rs
301 строка10 KB
Alexander Efanov
upd fix
04 авг 2026, 14:09
04 авг 2026, 14:09f63ab86
//! Recursive chunker (как в LangChain RecursiveCharacterTextSplitter).
//!
//! Пытается разбить текст по иерархии разделителей:
//! "\n\n" → "\n" → ". " → " " → ""
//!
//! Это лучшая стратегия по умолчанию для большинства случаев.

use super::traits::{Chunker, ChunkerConfig};
use crate::domain::{Chunk, Document};
use crate::Result;

/// Разделители по умолчанию (от самых сильных к самым слабым).
const DEFAULT_SEPARATORS: &[&str] = &["\n\n", "\n", ". ", " ", ""];

pub struct RecursiveChunker {
    config: ChunkerConfig,
    separators: Vec<String>,
}

impl RecursiveChunker {
    pub fn new(config: ChunkerConfig) -> Self {
        Self {
            config,
            separators: DEFAULT_SEPARATORS.iter().map(|s| s.to_string()).collect(),
        }
    }

    pub fn with_separators(config: ChunkerConfig, separators: Vec<String>) -> Self {
        Self { config, separators }
    }

    /// Рекурсивно разбить текст используя иерархию разделителей.
    ///
    /// Возвращает Vec<(текст_чанка, символьная_позиция_в_исходнике)>
    fn split_text(
        &self,
        text: &str,
        separators: &[String],
        base_offset: usize,
    ) -> Vec<(String, usize)> {
        // Если текст помещается в один чанк
        if text.chars().count() <= self.config.chunk_size {
            return if text.trim().is_empty() {
                vec![]
            } else {
                vec![(text.to_string(), base_offset)]
            };
        }

        // Найти подходящий разделитель
        let mut separator_idx = separators.len() - 1;
        let mut separator = String::new();

        for (i, sep) in separators.iter().enumerate() {
            if sep.is_empty() || text.contains(sep.as_str()) {
                separator_idx = i;
                separator = sep.clone();
                break;
            }
        }

        // Разбить текст по найденному разделителю
        let splits: Vec<String> = if separator.is_empty() {
            // Для пустого разделителя разбиваем по символам (без leak)
            text.chars().map(|c| c.to_string()).collect()
        } else {
            text.split(&separator).map(|s| s.to_string()).collect()
        };

        let mut final_chunks: Vec<(String, usize)> = Vec::new();
        let mut good_splits: Vec<(String, usize)> = Vec::new();
        let mut current_length: usize = 0;

        // Определить merge_separator
        let merge_separator = if separator_idx >= separators.len().saturating_sub(2) {
            String::new()
        } else {
            separator.clone()
        };
        let merge_sep_len = merge_separator.chars().count();

        let mut char_offset: usize = base_offset;

        for s in &splits {
            let s_chars = s.chars().count();

            if s_chars < self.config.chunk_size {
                let added_len = if good_splits.is_empty() {
                    s_chars
                } else {
                    s_chars + merge_sep_len
                };

                // Если добавление этого фрагмента превысит лимит, сохраним текущие
                if current_length + added_len > self.config.chunk_size && !good_splits.is_empty() {
                    let merged: String = good_splits
                        .iter()
                        .map(|(text, _)| text.as_str())
                        .collect::<Vec<_>>()
                        .join(&merge_separator);

                    if !merged.trim().is_empty() {
                        let first_offset = good_splits[0].1;
                        final_chunks.push((merged, first_offset));
                    }
                    good_splits.clear();
                    current_length = 0;
                }

                good_splits.push((s.clone(), char_offset));
                current_length += added_len;
            } else {
                // Фрагмент слишком большой, сохраним текущие good_splits
                if !good_splits.is_empty() {
                    let merged: String = good_splits
                        .iter()
                        .map(|(text, _)| text.as_str())
                        .collect::<Vec<_>>()
                        .join(&merge_separator);

                    if !merged.trim().is_empty() {
                        let first_offset = good_splits[0].1;
                        final_chunks.push((merged, first_offset));
                    }
                    good_splits.clear();
                    current_length = 0;
                }

                // Рекурсивно разбить большой фрагмент
                if separator_idx < separators.len() - 1 {
                    let sub_separators = &separators[separator_idx + 1..];
                    let sub_chunks = self.split_text(s, sub_separators, char_offset);
                    final_chunks.extend(sub_chunks);
                } else if !s.trim().is_empty() {
                    final_chunks.push((s.clone(), char_offset));
                }
            }

            // Обновить offset для следующего фрагмента
            char_offset += s_chars + separator.chars().count();
        }

        // Сохранить оставшиеся good_splits
        if !good_splits.is_empty() {
            let merged: String = good_splits
                .iter()
                .map(|(text, _)| text.as_str())
                .collect::<Vec<_>>()
                .join(&merge_separator);

            if !merged.trim().is_empty() {
                let first_offset = good_splits[0].1;
                final_chunks.push((merged, first_offset));
            }
        }

        final_chunks
    }
}

impl Chunker for RecursiveChunker {
    fn split(&self, document: &Document) -> Result<Vec<Chunk>> {
        self.config.validate()?;

        let text = &document.content;
        if text.is_empty() {
            return Ok(vec![]);
        }

        let parts = self.split_text(text, &self.separators, 0);

        let mut chunks = Vec::with_capacity(parts.len());

        for (index, (content, start_char)) in parts.iter().enumerate() {
            if content.trim().is_empty() {
                continue;
            }

            chunks.push(Chunk::from_document(
                document,
                content.clone(),
                index,
                *start_char,
            ));
        }

        Ok(chunks)
    }

    fn config(&self) -> &ChunkerConfig {
        &self.config
    }

    fn name(&self) -> &'static str {
        "recursive"
    }
}

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

    #[test]
    fn test_recursive_split() {
        let text = "Para 1. Sentence 1. Sentence 2.\n\nPara 2. Sentence 3.";
        let doc = Document::new(
            text.into(),
            "test.txt".into(),
            DocumentType::Txt,
            "ws-1".into(),
        );
        let chunker = RecursiveChunker::new(ChunkerConfig::new(30, 5));
        let chunks = chunker.split(&doc).unwrap();
        assert!(chunks.len() >= 2);
    }

    #[test]
    fn test_small_document() {
        let text = "Short text";
        let doc = Document::new(
            text.into(),
            "test.txt".into(),
            DocumentType::Txt,
            "ws-1".into(),
        );
        let chunker = RecursiveChunker::new(ChunkerConfig::new(1000, 200));
        let chunks = chunker.split(&doc).unwrap();
        assert_eq!(chunks.len(), 1);
        assert_eq!(chunks[0].content, "Short text");
    }

    #[test]
    fn test_empty_document() {
        let doc = Document::new(
            "".into(),
            "test.txt".into(),
            DocumentType::Txt,
            "ws-1".into(),
        );
        let chunker = RecursiveChunker::new(ChunkerConfig::new(100, 20));
        let chunks = chunker.split(&doc).unwrap();
        assert_eq!(chunks.len(), 0);
    }

    #[test]
    fn test_unicode_text() {
        let text = "Липидный комплекс (липидограмма) Холестерин липопротеинов низкой плотности (ЛПНП) (венозная кровь) A09.05.028 (Приказ МЗ РФ № 804н) 2.92 мм";
        let doc = Document::new(
            text.into(),
            "test.txt".into(),
            DocumentType::Txt,
            "ws-1".into(),
        );
        let chunker = RecursiveChunker::new(ChunkerConfig::new(50, 10));
        let chunks = chunker.split(&doc).unwrap();

        for chunk in &chunks {
            assert!(!chunk.content.is_empty());
            // ✅ FIX: start_char, не start_pos
            assert!(chunk.start_char <= text.len());
        }
    }

    #[test]
    fn test_mixed_unicode_ascii() {
        let text = "Hello мир! This is a test с русскими словами. Another sentence here.";
        let doc = Document::new(
            text.into(),
            "test.txt".into(),
            DocumentType::Txt,
            "ws-1".into(),
        );
        let chunker = RecursiveChunker::new(ChunkerConfig::new(30, 5));
        let chunks = chunker.split(&doc).unwrap();

        assert!(chunks.len() >= 2);
        for chunk in &chunks {
            assert!(!chunk.content.is_empty());
            assert!(chunk.start_char <= text.len());
        }
    }

    #[test]
    fn test_long_document() {
        let text = "Sentence one. Sentence two. Sentence three.\n\n".repeat(100);
        let doc = Document::new(
            text.into(),
            "test.txt".into(),
            DocumentType::Txt,
            "ws-1".into(),
        );
        let chunker = RecursiveChunker::new(ChunkerConfig::new(100, 20));
        let chunks = chunker.split(&doc).unwrap();

        assert!(chunks.len() > 10);
        for chunk in &chunks {
            assert!(chunk.content.chars().count() <= 100);
        }
    }
}