anti-plagiarism/services/worker-gpu/app/faiss_manager.py

"""Singleton менеджер FAISS индекса.

Используется плоский индекс IndexFlatIP, обёрнутый в IndexIDMap2, что даёт:
- отсутствие этапа обучения (в отличие от IVFFlat) — индекс работоспособен сразу,
  начиная с первого вектора;
- хранение doc_id прямо внутри индекса (add_with_ids) — не нужен отдельный
  маппинг на диске, а поиск сразу возвращает doc_id из PostgreSQL;
- корректную персистентность через faiss.write_index / read_index.

Векторы модели нормализованы (normalize_embeddings=True), поэтому inner product
эквивалентен косинусной близости. Для масштаба проекта (сотни тысяч — единицы
миллионов документов) полный перебор по FlatIP по скорости приемлем.
"""

import logging
import os
from pathlib import Path

import numpy as np

from app.config import settings

logger = logging.getLogger(__name__)


class FAISSManager:
    """Singleton для управления FAISS индексом (IndexIDMap2 поверх IndexFlatIP)."""

    _index = None
    # doc_id -> faiss_id. Для IDMap2 faiss_id == doc_id, но маппинг сохраняем
    # для совместимости с вызывающим кодом (plagiarism.embed_documents).
    _reverse_map: dict[int, int] = {}
    _use_gpu: bool = False

    @classmethod
    def _new_index(cls):
        """Создать новый пустой индекс нужного типа."""
        import faiss

        base = faiss.IndexFlatIP(settings.EMBED_DIM)
        return faiss.IndexIDMap2(base)

    @classmethod
    def load_or_create(cls) -> None:
        """Загрузить индекс с диска или создать новый.

        Старый несовместимый индекс (например, IVFFlat от предыдущей версии,
        который не хранит id_map) безопасно пересоздаётся — полезных векторов в
        нём всё равно не было.
        """
        import faiss

        index_path = settings.FAISS_INDEX_PATH

        if os.path.exists(index_path):
            try:
                loaded = faiss.read_index(index_path)
                if hasattr(loaded, "id_map"):
                    cls._index = loaded
                    cls._rebuild_reverse_map()
                    logger.info(
                        f"Загрузка FAISS индекса из {index_path} "
                        f"({cls._index.ntotal} векторов)"
                    )
                else:
                    logger.warning(
                        "На диске несовместимый FAISS индекс (без id_map) — "
                        "пересоздаём как IndexIDMap2(IndexFlatIP)"
                    )
                    cls._index = cls._new_index()
                    cls._reverse_map = {}
            except Exception as e:
                logger.warning(f"Не удалось загрузить FAISS индекс ({e}) — создаём новый")
                cls._index = cls._new_index()
                cls._reverse_map = {}
        else:
            logger.info("Создание нового FAISS индекса IndexIDMap2(IndexFlatIP)...")
            cls._index = cls._new_index()
            cls._reverse_map = {}

        cls._use_gpu = False  # FlatIP на CPU достаточно быстр для целевого масштаба

    @classmethod
    def _ensure(cls) -> None:
        """Ленивая инициализация индекса при первом обращении."""
        if cls._index is None:
            cls.load_or_create()

    @classmethod
    def _rebuild_reverse_map(cls) -> None:
        """Восстановить _reverse_map из id, хранящихся внутри загруженного индекса."""
        import faiss

        try:
            ids = faiss.vector_to_array(cls._index.id_map)
            cls._reverse_map = {int(i): int(i) for i in ids}
        except Exception as e:
            logger.warning(f"Не удалось восстановить reverse_map из индекса: {e}")
            cls._reverse_map = {}

    @classmethod
    def search(cls, query_vector: np.ndarray, k: int = 20) -> list[tuple[int, float]]:
        """Поиск k ближайших векторов.

        Args:
            query_vector: Нормализованный вектор запроса, форма (768,)
            k: Количество результатов

        Returns:
            Список кортежей (doc_id, cosine_score), отсортированных по убыванию score
        """
        cls._ensure()

        if cls._index.ntotal == 0:
            return []

        try:
            query = query_vector.reshape(1, -1).astype(np.float32)
            distances, ids = cls._index.search(query, min(k, cls._index.ntotal))

            results = []
            for idx, dist in zip(ids[0], distances[0]):
                if idx == -1:
                    continue
                # Для IDMap2 idx — это уже doc_id из PostgreSQL
                results.append((int(idx), float(dist)))

            return results

        except Exception as e:
            logger.error(f"Ошибка поиска FAISS: {e}")
            return []

    @classmethod
    def add_vectors(cls, vectors: np.ndarray, doc_ids: list[int]) -> None:
        """Добавить (или обновить) векторы в индекс.

        Идемпотентно по doc_id: при повторном эмбеддинге старый вектор документа
        удаляется перед добавлением нового, чтобы не плодить дубли.

        Args:
            vectors: numpy массив формы (N, 768)
            doc_ids: Список doc_id из PostgreSQL
        """
        cls._ensure()

        if len(doc_ids) == 0:
            return

        vectors = np.asarray(vectors, dtype=np.float32)
        ids = np.asarray(doc_ids, dtype=np.int64)

        # Удалить существующие id, чтобы повторный эмбеддинг не создавал дубли
        try:
            cls._index.remove_ids(ids)
        except Exception:
            pass

        cls._index.add_with_ids(vectors, ids)
        for doc_id in doc_ids:
            cls._reverse_map[int(doc_id)] = int(doc_id)

        logger.info(f"Добавлено {len(doc_ids)} векторов в FAISS. Всего: {cls._index.ntotal}")

    @classmethod
    def save(cls) -> None:
        """Сохранить индекс на диск."""
        cls._ensure()

        import faiss

        index_path = Path(settings.FAISS_INDEX_PATH)
        index_path.parent.mkdir(parents=True, exist_ok=True)

        faiss.write_index(cls._index, str(index_path))
        logger.info(f"FAISS индекс сохранён: {index_path} ({cls._index.ntotal} векторов)")

    @classmethod
    def total_vectors(cls) -> int:
        """Количество векторов в индексе."""
        if cls._index is None:
            return 0
        return cls._index.ntotal