lib/extractor/ntr_text_fragmentation/models/document.py

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

from dataclasses import dataclass

from .linker_entity import LinkerEntity, register_entity


@register_entity
@dataclass
class DocumentAsEntity(LinkerEntity):
    """
    Класс для представления документа как сущности в системе извлечения и сборки.
    """

    doc_type: str = "unknown"
    
    @classmethod
    def deserialize(cls, data: LinkerEntity) -> 'DocumentAsEntity':
        """
        Десериализует DocumentAsEntity из объекта LinkerEntity.
        
        Args:
            data: Объект LinkerEntity для преобразования в DocumentAsEntity
            
        Returns:
            Десериализованный объект DocumentAsEntity
        """
        metadata = data.metadata or {}
        
        # Получаем тип документа из метаданных или используем значение по умолчанию
        doc_type = metadata.get('_doc_type', 'unknown')
        
        # Создаем чистые метаданные без служебных полей
        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,
            type="DocumentAsEntity",
            doc_type=doc_type
        )