lib/extractor/ntr_text_fragmentation/models/linker_entity.py

"""
Базовый абстрактный класс для всех сущностей с поддержкой триплетного подхода.
"""

import uuid
from abc import abstractmethod
from dataclasses import dataclass, field, fields
from uuid import UUID


@dataclass
class LinkerEntity:
    """
    Общий класс для всех сущностей в системе извлечения и сборки.
    Поддерживает триплетный подход, где каждая сущность может опционально связывать две другие сущности.
    
    Attributes:
        id (UUID): Уникальный идентификатор сущности.
        name (str): Название сущности.
        text (str): Текстое представление сущности.
        in_search_text (str | None): Текст для поиска. Если задан, используется в __str__, иначе используется обычное представление.
        metadata (dict): Метаданные сущности.
        source_id (UUID | None): Опциональный идентификатор исходной сущности. 
                             Если указан, эта сущность является связью.
        target_id (UUID | None): Опциональный идентификатор целевой сущности. 
                             Если указан, эта сущность является связью.
        number_in_relation (int | None): Используется в случае связей один-ко-многим, 
                                     указывает номер целевой сущности в списке.
        type (str): Тип сущности.
    """

    id: UUID
    name: str
    text: str
    metadata: dict  # JSON с метаданными
    in_search_text: str | None = None
    source_id: UUID | None = None
    target_id: UUID | None = None
    number_in_relation: int | None = None
    type: str = field(default_factory=lambda: "Entity")

    def __post_init__(self):
        if self.id is None:
            self.id = uuid.uuid4()
        
        # Проверяем корректность полей связи
        if (self.source_id is not None and self.target_id is None) or \
           (self.source_id is None and self.target_id is not None):
            raise ValueError("source_id и target_id должны быть либо оба указаны, либо оба None")

    def is_link(self) -> bool:
        """
        Проверяет, является ли сущность связью (имеет и source_id, и target_id).
        
        Returns:
            bool: True, если сущность является связью, иначе False
        """
        return self.source_id is not None and self.target_id is not None

    def __str__(self) -> str:
        """
        Возвращает строковое представление сущности.
        Если задан in_search_text, возвращает его, иначе возвращает стандартное представление.
        """
        if self.in_search_text is not None:
            return self.in_search_text
        return f"{self.name}: {self.text}"

    def __eq__(self, other: 'LinkerEntity') -> bool:
        """
        Сравнивает текущую сущность с другой.

        Args:
            other: Другая сущность для сравнения

        Returns:
            bool: True если сущности совпадают, иначе False
        """
        if not isinstance(other, self.__class__):
            return False

        basic_equality = (
            self.id == other.id
            and self.name == other.name
            and self.text == other.text
            and self.type == other.type
        )
        
        # Если мы имеем дело со связями, также проверяем поля связи
        if self.is_link() or other.is_link():
            return (
                basic_equality
                and self.source_id == other.source_id
                and self.target_id == other.target_id
            )
        
        return basic_equality

    def serialize(self) -> 'LinkerEntity':
        """
        Сериализует сущность в простейшую форму сущности, передавая все дополнительные поля в метаданные.
        """
        # Получаем список полей базового класса
        known_fields = {field.name for field in fields(LinkerEntity)}

        # Получаем все атрибуты текущего объекта
        dict_entity = {}
        for attr_name in dir(self):
            # Пропускаем служебные атрибуты, методы и уже известные поля
            if (
                attr_name.startswith('_')
                or attr_name in known_fields
                or callable(getattr(self, attr_name))
            ):
                continue

            # Добавляем дополнительные поля в словарь
            dict_entity[attr_name] = getattr(self, attr_name)

        # Преобразуем имена дополнительных полей, добавляя префикс "_"
        dict_entity = {f'_{name}': value for name, value in dict_entity.items()}

        # Объединяем с существующими метаданными
        dict_entity = {**dict_entity, **self.metadata}

        result_type = self.type
        if result_type == "Entity":
            result_type = self.__class__.__name__

        # Создаем базовый объект LinkerEntity с новыми метаданными
        return LinkerEntity(
            id=self.id,
            name=self.name,
            text=self.text,
            in_search_text=self.in_search_text,
            metadata=dict_entity,
            source_id=self.source_id,
            target_id=self.target_id,
            number_in_relation=self.number_in_relation,
            type=result_type,
        )

    @classmethod
    @abstractmethod
    def deserialize(cls, data: 'LinkerEntity') -> 'Self':
        """
        Десериализует сущность из простейшей формы сущности, учитывая все дополнительные поля в метаданных.
        """
        raise NotImplementedError(
            f"Метод deserialize для класса {cls.__class__.__name__} не реализован"
        )

    # Реестр для хранения всех наследников LinkerEntity
    _entity_classes = {}

    @classmethod
    def register_entity_class(cls, entity_class):
        """
        Регистрирует класс-наследник в реестре.
        
        Args:
            entity_class: Класс для регистрации
        """
        entity_type = entity_class.__name__
        cls._entity_classes[entity_type] = entity_class
        # Также регистрируем по типу, если он отличается от имени класса
        if hasattr(entity_class, 'type') and isinstance(entity_class.type, str):
            cls._entity_classes[entity_class.type] = entity_class
    
    @classmethod
    def deserialize(cls, data: 'LinkerEntity') -> 'LinkerEntity':
        """
        Десериализует сущность в нужный тип на основе поля type.
        
        Args:
            data: Сериализованная сущность типа LinkerEntity
            
        Returns:
            Десериализованная сущность правильного типа
        """
        # Получаем тип сущности
        entity_type = data.type
        
        # Проверяем реестр классов
        if entity_type in cls._entity_classes:
            try:
                return cls._entity_classes[entity_type].deserialize(data)
            except (AttributeError, NotImplementedError) as e:
                # Если метод не реализован, возвращаем исходную сущность
                return data
        
        # Если тип не найден в реестре, просто возвращаем исходную сущность
        # Больше не используем опасное сканирование sys.modules
        return data


# Декоратор для регистрации производных классов
def register_entity(cls):
    """
    Декоратор для регистрации классов-наследников LinkerEntity.
    
    Пример использования:
    
    @register_entity
    class MyEntity(LinkerEntity):
        type = "my_entity"
    
    Args:
        cls: Класс, который нужно зарегистрировать
        
    Returns:
        Исходный класс (без изменений)
    """
    # Регистрируем класс в реестр, используя его имя или указанный тип
    entity_type = getattr(cls, 'type', cls.__name__)
    LinkerEntity._entity_classes[entity_type] = cls
    return cls