lib/extractor/ntr_text_fragmentation/models/linker_entity.py

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

import logging
import uuid
from dataclasses import dataclass, field, fields
from uuid import UUID

logger = logging.getLogger(__name__)


@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 = field(default_factory=uuid.uuid4)
    name: str = field(default="")
    text: str = field(default="")
    metadata: dict = field(default_factory=dict)
    in_search_text: str | None = None
    source_id: UUID | None = None
    target_id: UUID | None = None
    number_in_relation: int | None = None
    groupper: str | None = None
    type: str | None = None

    @property
    def owner_id(self) -> UUID | None:
        """
        Возвращает идентификатор владельца сущности.
        """
        if self.is_link():
            return None
        return self.target_id

    @owner_id.setter
    def owner_id(self, value: UUID | None):
        """
        Устанавливает идентификатор владельца сущности.
        """
        if self.is_link():
            raise ValueError("Связь не может иметь владельца")
        self.target_id = value

    def __post_init__(self):
        if self.id is None:
            self.id = uuid.uuid4()

        if self.type is None:
            self.type = self.__class__.__name__

    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':
        """
        Сериализует сущность в базовый класс `LinkerEntity`, сохраняя все дополнительные поля в метаданные.
        """
        base_fields = {f.name for f in fields(LinkerEntity)}
        current_fields = {f.name for f in fields(self.__class__)}
        extra_field_names = current_fields - base_fields

        # Собираем только дополнительные поля, определенные в подклассе
        extra_fields_dict = {name: getattr(self, name) for name in extra_field_names}

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

        # Объединяем с существующими метаданными (если они были установлены вручную)
        final_metadata = {**prefixed_extra_fields, **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=final_metadata,  # Используем собранные метаданные
            source_id=self.source_id,
            target_id=self.target_id,
            number_in_relation=self.number_in_relation,
            groupper=self.groupper,
            type=result_type,
        )

    def deserialize(self) -> 'LinkerEntity':
        """
        Десериализует сущность в нужный тип на основе поля type.
        """
        return self._deserialize(self)

    # Реестр для хранения всех наследников 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_to_me(data)
            except Exception as e:
                logger.error(f"Ошибка при вызове _deserialize_to_me для {entity_type}: {e}", exc_info=True)
                return data

        return data

    @classmethod
    def _deserialize_to_me(cls, data: 'LinkerEntity') -> 'LinkerEntity':
        """
        Десериализует сущность в нужный тип на основе поля type.
        """
        return cls(
            id=data.id,
            name=data.name,
            text=data.text,
            in_search_text=data.in_search_text,
            metadata=data.metadata,
            source_id=data.source_id,
            target_id=data.target_id,
            number_in_relation=data.number_in_relation,
            type=data.type,
            groupper=data.groupper,
        )


# Алиасы для удобства
Link = LinkerEntity
Entity = LinkerEntity


# Декоратор для регистрации производных классов
def register_entity(cls):
    """
    Декоратор для регистрации классов-наследников LinkerEntity.

    Пример использования:

    @register_entity
    class MyEntity(LinkerEntity):
        type = "my_entity"

    Args:
        cls: Класс, который нужно зарегистрировать

    Returns:
        Исходный класс (без изменений)
    """
    # Регистрируем класс в реестр, используя его имя или указанный тип
    entity_type = cls.__name__
    LinkerEntity._entity_classes[entity_type] = cls

    if hasattr(cls, 'type') and isinstance(cls.type, str):
        LinkerEntity._entity_classes[cls.type] = cls

    return cls