lib/extractor/ntr_text_fragmentation/models/document.py

"""
Класс для представления документа как сущности.
"""

from dataclasses import dataclass, field

from .linker_entity import Entity, register_entity


@register_entity
@dataclass
class DocumentAsEntity(Entity):
    """
    Класс для представления документа как сущности в системе извлечения и сборки.
    Содержит ссылки на классы стратегии чанкинга и обработчика таблиц,
    использовавшихся при деструктуризации.
    """

    doc_type: str = "unknown"

    chunking_strategy_ref: str | None = None

    type: str = field(default="DocumentAsEntity")

    @classmethod
    def _deserialize_to_me(cls, data: Entity) -> 'DocumentAsEntity':
        """
        Десериализует DocumentAsEntity из объекта LinkerEntity.

        Args:
            data: Объект LinkerEntity для преобразования в DocumentAsEntity

        Returns:
            Десериализованный объект DocumentAsEntity
        """
        if not isinstance(data, Entity):
            raise TypeError(f"Ожидался LinkerEntity, получен {type(data)}")

        metadata = data.metadata or {}

        # Получаем поля из атрибутов или метаданных
        doc_type = getattr(data, 'doc_type', metadata.get('_doc_type', 'unknown'))
        strategy_ref = getattr(
            data, 'chunking_strategy_ref', metadata.get('_chunking_strategy_ref', None)
        )

        # Создаем чистые метаданные без служебных полей
        clean_metadata = {k: v for k, v in metadata.items() if not k.startswith('_')}

        return cls(
            id=data.id,
            name=data.name,
            text=data.text,
            in_search_text=data.in_search_text,
            metadata=clean_metadata,
            source_id=data.source_id,
            target_id=data.target_id,
            number_in_relation=data.number_in_relation,
            groupper=data.groupper,
            type=cls.__name__,
            doc_type=doc_type,
            chunking_strategy_ref=strategy_ref,
        )