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

"""Celery задачи проверки плагиата (уровни 3 и 4) и построения эмбеддингов."""

import logging
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

logger = get_task_logger(__name__)

# Пороговые значения
FAISS_SIMILARITY_THRESHOLD = 0.75  # Минимальный cosine score для подозрительного совпадения
LLM_CONFIDENCE_THRESHOLD = 0.7  # Минимальная уверенность LLM


def _get_doc_text(doc_id: int) -> str | None:
    """Получить текст документа из базы данных."""
    from app.models import Document
    with db_session() as session:
        doc = session.get(Document, doc_id)
        if doc is None:
            return None
        return doc.abstract or ""


def _get_doc_meta(doc_id: int) -> dict | None:
    """Получить метаданные документа из базы данных."""
    from app.models import Document
    with db_session() as session:
        doc = session.get(Document, doc_id)
        if doc is None:
            return None
        return {
            "id": doc.id,
            "title": doc.title,
            "url": doc.url,
            "source": doc.source,
            "authors": doc.authors,
            "year": doc.year,
        }


@celery_app.task(
    name="gpu.check_plagiarism",
    bind=True,
    max_retries=2,
    default_retry_delay=120,
)
def check_plagiarism(
    self,
    task_id: str,
    text: str,
    fragments: list[dict],
    level1_matches: list[dict] | None = None,
    level2_matches: list[dict] | None = None,
) -> dict[str, Any]:
    """
    Проверка плагиата уровни 3 (FAISS семантика) и 4 (Ollama LLM).

    Принимает результаты уровней 1 и 2 от worker-indexer и дополняет их.

    Args:
        task_id: ID задачи в PostgreSQL
        text: Полный текст документа
        fragments: Список фрагментов для проверки:
                   [{"text": str, "start": int, "end": int}, ...]
        level1_matches: Совпадения уровня 1 (Winnowing)
        level2_matches: Совпадения уровня 2 (MinHash)

    Returns:
        dict с результатами проверки
    """
    from app.models import Task

    logger.info(
        f"Проверка плагиата (ур. 3-4) для задачи {task_id!r}. "
        f"Фрагментов: {len(fragments)}"
    )

    level1_matches = level1_matches or []
    level2_matches = level2_matches or []

    # Обновить статус
    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:
        from app.faiss_manager import FAISSManager
        from app.model_manager import ModelManager
        from app.ollama_client import OllamaClient

        ollama = OllamaClient()
        semantic_matches: list[dict] = []

        for i, fragment in enumerate(fragments):
            frag_text = fragment.get("text", "")
            if len(frag_text.split()) < 10:
                # Пропустить слишком короткие фрагменты
                continue

            # Уровень 3: Семантический поиск через FAISS
            frag_vec = ModelManager.encode_single(frag_text)
            faiss_results = FAISSManager.search(frag_vec, k=10)

            for doc_id, score in faiss_results:
                if score < FAISS_SIMILARITY_THRESHOLD:
                    continue

                # Уровень 4: LLM анализ парафраза для подозрительных совпадений
                source_text = _get_doc_text(doc_id) or ""
                doc_meta = _get_doc_meta(doc_id)

                if not doc_meta:
                    continue

                llm_result = {"is_paraphrase": False, "confidence": 0.0, "reason": ""}
                if source_text:
                    llm_result = ollama.check_paraphrase(source_text, frag_text)

                if llm_result.get("is_paraphrase") and llm_result.get("confidence", 0.0) >= LLM_CONFIDENCE_THRESHOLD:
                    semantic_matches.append({
                        "fragment": frag_text[:300],
                        "position_start": fragment.get("start", 0),
                        "position_end": fragment.get("end", len(frag_text)),
                        "similarity": round(score * 100, 1),
                        "method": "semantic+llm",
                        "confidence": llm_result["confidence"],
                        "reason": llm_result.get("reason", ""),
                        "source_title": doc_meta["title"],
                        "source_url": doc_meta["url"],
                        "source_db": doc_meta["source"],
                    })

            if (i + 1) % 10 == 0:
                logger.info(f"Проверено фрагментов: {i + 1}/{len(fragments)}")

        # Объединить все совпадения и дедуплицировать по source_title
        all_matches = level1_matches + level2_matches + semantic_matches
        seen_sources: set[str] = set()
        unique_matches = []
        for m in all_matches:
            key = f"{m.get('source_title', '')}:{m.get('position_start', '')}"
            if key not in seen_sources:
                seen_sources.add(key)
                unique_matches.append(m)

        # Вычислить общий процент схожести
        total_frags = len(fragments)
        flagged_frags = len(unique_matches)
        overall_similarity = (flagged_frags / total_frags * 100) if total_frags > 0 else 0.0

        result = {
            "overall_similarity": round(overall_similarity, 2),
            "matches": unique_matches,
            "total_fragments": total_frags,
            "flagged_fragments": flagged_frags,
            "by_method": {
                "exact": len(level1_matches),
                "fuzzy": len(level2_matches),
                "semantic_llm": len(semantic_matches),
            },
        }

        # Сохранить результат
        with db_session() as session:
            task = session.get(Task, task_id)
            if task:
                task.result = result
                task.status = "done"
                task.queue_position = None
                session.commit()

        logger.info(
            f"Проверка плагиата завершена для задачи {task_id!r}. "
            f"Схожесть: {overall_similarity:.1f}%, совпадений: {flagged_frags}"
        )

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

        return result

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

        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=120)


@celery_app.task(name="gpu.embed_documents")
def embed_documents(doc_ids: list[int]) -> dict[str, Any]:
    """
    Построить эмбеддинги для документов и добавить их в FAISS индекс.

    Вызывается воркером-индексером после добавления новых документов.

    Args:
        doc_ids: Список ID документов в PostgreSQL
    """
    if not doc_ids:
        return {"status": "ok", "embedded": 0}

    from app.models import Document
    from app.faiss_manager import FAISSManager
    from app.model_manager import ModelManager
    from sqlalchemy import select

    logger.info(f"Построение эмбеддингов для {len(doc_ids)} документов...")

    with db_session() as session:
        docs = session.execute(
            select(Document).where(Document.id.in_(doc_ids))
        ).scalars().all()

        if not docs:
            logger.warning(f"Документы не найдены: {doc_ids}")
            return {"status": "not_found", "embedded": 0}

        # Формируем тексты: title + abstract
        texts = [
            f"{d.title}. {d.abstract or ''}".strip()
            for d in docs
        ]
        ids = [d.id for d in docs]

    import numpy as np
    vectors = ModelManager.encode(texts)

    FAISSManager.add_vectors(vectors, ids)
    FAISSManager.save()

    # Обновить faiss_id в PostgreSQL
    with db_session() as session:
        for doc_id in ids:
            doc = session.get(Document, doc_id)
            if doc and doc_id in FAISSManager._reverse_map:
                doc.faiss_id = FAISSManager._reverse_map[doc_id]
        session.commit()

    logger.info(f"Встроено и проиндексировано {len(ids)} документов")
    return {"status": "ok", "embedded": len(ids)}