lib/extractor/ntr_text_fragmentation/chunking/chunking_strategy.py

"""
Базовый класс для всех стратегий чанкинга.
"""

from abc import ABC, abstractmethod

from ntr_fileparser import ParsedDocument

from ..models import Chunk, DocumentAsEntity, LinkerEntity


class ChunkingStrategy(ABC):
    """
    Базовый абстрактный класс для всех стратегий чанкинга.
    """
    
    @abstractmethod
    def chunk(self, document: ParsedDocument, doc_entity: DocumentAsEntity | None = None) -> list[LinkerEntity]:
        """
        Разбивает документ на чанки в соответствии со стратегией.
        
        Args:
            document: ParsedDocument для извлечения текста
            doc_entity: Опциональная сущность документа для привязки чанков.
                       Если не указана, будет создана новая.
            
        Returns:
            list[LinkerEntity]: Список сущностей (документ, чанки, связи)
        """
        raise NotImplementedError("Стратегия чанкинга должна реализовать метод chunk")

    def dechunk(self, chunks: list[LinkerEntity], repository: 'EntityRepository' = None) -> str:
        """
        Собирает документ из чанков и связей.
        
        Базовая реализация сортирует чанки по chunk_index и объединяет их тексты,
        сохраняя структуру параграфов и избегая дублирования текста.
        
        Args:
            chunks: Список отфильтрованных чанков в случайном порядке
            repository: Репозиторий сущностей для получения дополнительной информации (может быть None)
            
        Returns:
            Восстановленный текст документа
        """
        import re

        # Проверяем, есть ли чанки для сборки
        if not chunks:
            return ""
            
        # Отбираем только чанки
        valid_chunks = [c for c in chunks if isinstance(c, Chunk)]
        
        # Сортируем чанки по chunk_index
        sorted_chunks = sorted(valid_chunks, key=lambda c: c.chunk_index or 0)
        
        # Собираем текст документа с учетом структуры параграфов
        result_text = ""
        
        for chunk in sorted_chunks:
            # Получаем текст чанка (предпочитаем text, а не in_search_text для избежания дублирования)
            chunk_text = chunk.text if hasattr(chunk, 'text') and chunk.text else ""
            
            # Добавляем текст чанка с сохранением структуры параграфов
            if result_text and result_text[-1] != "\n" and chunk_text and chunk_text[0] != "\n":
                result_text += " "
            result_text += chunk_text
        
        # Пост-обработка результата
        # Заменяем множественные переносы строк на одиночные
        result_text = re.sub(r'\n+', '\n', result_text)
        
        # Заменяем множественные пробелы на одиночные
        result_text = re.sub(r' +', ' ', result_text)
        
        # Убираем пробелы перед переносами строк
        result_text = re.sub(r' +\n', '\n', result_text)
        
        # Убираем пробелы после переносов строк
        result_text = re.sub(r'\n +', '\n', result_text)
        
        # Убираем лишние переносы строк в начале и конце текста
        result_text = result_text.strip()
        
        return result_text