anti-plagiarism/services/worker-gpu/app/tasks/search.py

"""Celery задача семантического поиска источников."""

import hashlib
import json
from typing import Any

from celery.utils.log import get_task_logger
from sqlalchemy import select

from app.celery_app import celery_app
from app.db import db_session, get_redis

logger = get_task_logger(__name__)

# Время жизни кэша поиска (1 час)
SEARCH_CACHE_TTL = 3600

# Минимальный cosine score для включения в результаты
FAISS_MIN_SCORE = 0.3

# Вес BM25 и семантического скора при ранжировании
WEIGHT_SEMANTIC = 0.6
WEIGHT_BM25 = 0.4


def _get_cache_key(query: str, lang: str | None, year_from: int | None, year_to: int | None) -> str:
    """Получить ключ кэша для поискового запроса."""
    key_str = f"{query}:{lang}:{year_from}:{year_to}"
    return f"search:{hashlib.md5(key_str.encode()).hexdigest()}"


def _merge_and_rank(
    faiss_results: list[tuple[int, float]],
    es_results: list[dict],
) -> list[dict[str, Any]]:
    """
    Объединить и ранжировать результаты FAISS и Elasticsearch.

    Нормализует оба скора в диапазон [0, 1] и взвешенно суммирует.

    Args:
        faiss_results: [(doc_id, cosine_score), ...] из FAISS
        es_results: [{"doc_id": int, "es_score": float}, ...] из ES

    Returns:
        Список словарей с doc_id и combined_score, отсортированный по убыванию
    """
    scores: dict[int, dict[str, float]] = {}

    # Нормализация FAISS скоров
    if faiss_results:
        max_faiss = max(s for _, s in faiss_results) or 1.0
        for doc_id, score in faiss_results:
            scores[doc_id] = {"faiss": score / max_faiss, "es": 0.0}

    # Нормализация BM25 скоров
    if es_results:
        max_es = max(r["es_score"] for r in es_results) or 1.0
        for r in es_results:
            doc_id = r["doc_id"]
            norm_score = r["es_score"] / max_es
            if doc_id in scores:
                scores[doc_id]["es"] = norm_score
            else:
                scores[doc_id] = {"faiss": 0.0, "es": norm_score}

    # Комбинированный скор
    ranked = []
    for doc_id, s in scores.items():
        combined = WEIGHT_SEMANTIC * s["faiss"] + WEIGHT_BM25 * s["es"]
        ranked.append({
            "doc_id": doc_id,
            "combined_score": round(combined, 4),
            "faiss_score": round(s["faiss"], 4),
            "es_score": round(s["es"], 4),
        })

    ranked.sort(key=lambda x: x["combined_score"], reverse=True)
    return ranked


def _enrich_from_db(ranked: list[dict], session) -> list[dict[str, Any]]:
    """
    Обогатить результаты метаданными из PostgreSQL и добавить ГОСТ-цитату.

    Args:
        ranked: Список с doc_id и скорами
        session: SQLAlchemy сессия

    Returns:
        Список источников с полными метаданными
    """
    from app.models import Document

    if not ranked:
        return []

    doc_ids = [r["doc_id"] for r in ranked]
    docs = session.execute(
        select(Document).where(Document.id.in_(doc_ids))
    ).scalars().all()

    doc_map = {d.id: d for d in docs}

    results = []
    for r in ranked:
        doc = doc_map.get(r["doc_id"])
        if not doc:
            continue

        # Простая ГОСТ-цитата (полная форматирует worker-gost)
        gost_citation = _format_gost_simple(doc)

        results.append({
            "id": doc.id,
            "title": doc.title,
            "authors": doc.authors,
            "year": doc.year,
            "journal": doc.journal,
            "abstract": doc.abstract[:300] if doc.abstract else None,
            "url": doc.url,
            "doi": doc.doi,
            "relevance_score": r["combined_score"],
            "gost_citation": gost_citation,
            "source_db": doc.source,
        })

    return results


def _format_gost_simple(doc) -> str:
    """Простое ГОСТ-форматирование (без полной логики worker-gost)."""
    authors = doc.authors or []
    author_str = ""
    if authors:
        first = authors[0]
        last_name = first.get("last_name", "")
        initials = first.get("initials", "")
        if len(authors) == 1:
            author_str = f"{last_name} {initials}"
        elif len(authors) <= 3:
            parts = [f"{a.get('last_name', '')} {a.get('initials', '')}" for a in authors]
            author_str = ", ".join(parts)
        else:
            author_str = f"{last_name} {initials} и др."

    title = doc.title or ""
    year = str(doc.year) if doc.year else "б. г."
    journal = doc.journal or ""
    pages = doc.pages or ""

    if journal:
        citation = f"{author_str}. {title} // {journal}. — {year}."
        if pages:
            citation += f" — С. {pages}."
    else:
        citation = f"{author_str}. {title}. — {year}."

    return citation.strip()


@celery_app.task(
    name="gpu.search_semantic",
    bind=True,
    max_retries=3,
    default_retry_delay=60,
)
def search_semantic(
    self,
    task_id: str,
    query: str,
    lang: str | None = None,
    year_from: int | None = None,
    year_to: int | None = None,
) -> dict[str, Any]:
    """
    Семантический поиск источников.

    Алгоритм:
    1. Нормализация запроса
    2. Проверка Redis кэша
    3. Параллельно: FAISS GPU cosine + Elasticsearch BM25
    4. Merge + дедупликация + ранжирование
    5. Обогащение метаданными из PostgreSQL
    6. Сохранение результата в task.result
    7. Диспатч notify.send_task_done

    Args:
        task_id: ID задачи в PostgreSQL
        query: Поисковый запрос
        lang: Фильтр языка
        year_from: Фильтр года (от)
        year_to: Фильтр года (до)
    """
    from app.models import Task

    logger.info(f"Начало поиска для задачи {task_id!r}, запрос: {query[:50]!r}")

    # Обновить статус на processing
    with db_session() as session:
        task = session.get(Task, task_id)
        if task is None:
            logger.error(f"Задача {task_id!r} не найдена в БД")
            return {}
        task.status = "processing"
        session.commit()

    try:
        redis = get_redis()
        cache_key = _get_cache_key(query, lang, year_from, year_to)

        # Проверить кэш
        cached = redis.get(cache_key)
        if cached:
            logger.info(f"Результат поиска получен из кэша: {cache_key}")
            results = json.loads(cached)
        else:
            # Нормализация запроса
            query_normalized = query.strip()

            # Закодировать запрос в вектор
            from app.model_manager import ModelManager
            query_vec = ModelManager.encode_single(query_normalized)

            # FAISS семантический поиск
            from app.faiss_manager import FAISSManager
            faiss_results = FAISSManager.search(query_vec, k=50)
            logger.info(f"FAISS: найдено {len(faiss_results)} результатов")

            # Elasticsearch BM25 поиск
            from app.es_client import search_fulltext
            es_results = search_fulltext(query_normalized, lang=lang, year_from=year_from, year_to=year_to)
            logger.info(f"ES BM25: найдено {len(es_results)} результатов")

            # Объединить и ранжировать
            ranked = _merge_and_rank(faiss_results, es_results)

            # Обогатить метаданными
            with db_session() as session:
                results = _enrich_from_db(ranked[:30], session)  # Топ-30 результатов

            # Кэшировать на 1 час
            redis.setex(cache_key, SEARCH_CACHE_TTL, json.dumps(results, ensure_ascii=False, default=str))

        # Сохранить результат в задачу
        with db_session() as session:
            task = session.get(Task, task_id)
            if task:
                task.result = {"sources": results, "total": len(results), "query": query}
                task.status = "done"
                task.queue_position = None
                session.commit()

        logger.info(f"Задача {task_id!r} выполнена. Найдено {len(results)} источников.")

        # Уведомить пользователя
        celery_app.send_task(
            "notify.send_task_done",
            args=[task_id],
            queue="queue.notify",
        )

        return {"task_id": task_id, "total": len(results)}

    except Exception as exc:
        logger.error(f"Ошибка при выполнении поиска для задачи {task_id!r}: {exc}", exc_info=True)

        # Обновить статус задачи на failed
        try:
            with db_session() as session:
                task = session.get(Task, task_id)
                if task:
                    task.status = "failed"
                    task.error = str(exc)
                    session.commit()
        except Exception as db_exc:
            logger.error(f"Не удалось обновить статус задачи: {db_exc}")

        # Повторить попытку
        raise self.retry(exc=exc, countdown=60) from exc