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

"""Elasticsearch клиент для полнотекстового поиска (BM25)."""

import logging
from typing import Any

from elasticsearch import Elasticsearch, exceptions as es_exceptions

from app.config import settings

logger = logging.getLogger(__name__)

_es_client: Elasticsearch | None = None

INDEX_NAME = "documents"


def get_es_client() -> Elasticsearch:
    """Получить или создать синглтон ES клиент."""
    global _es_client
    if _es_client is None:
        _es_client = Elasticsearch(
            settings.ELASTICSEARCH_URL,
            retry_on_timeout=True,
            max_retries=3,
            request_timeout=30,
        )
    return _es_client


def search_fulltext(
    query: str,
    lang: str | None = None,
    year_from: int | None = None,
    year_to: int | None = None,
    category: str | None = None,
    size: int = 50,
) -> list[dict[str, Any]]:
    """
    Полнотекстовый BM25 поиск в Elasticsearch.

    Args:
        query: Поисковый запрос
        lang: Фильтр языка (ru, en, None = все)
        year_from: Фильтр года публикации (от)
        year_to: Фильтр года публикации (до)
        category: Фильтр тематической категории
        size: Максимальное количество результатов

    Returns:
        Список словарей с полями doc_id и score
    """
    es = get_es_client()

    # Составить bool запрос
    must_clauses: list[dict] = [
        {
            "multi_match": {
                "query": query,
                "fields": ["title^3", "abstract^2", "authors"],
                "type": "best_fields",
                "operator": "or",
                "minimum_should_match": "30%",
            }
        }
    ]

    filter_clauses: list[dict] = []

    if lang:
        filter_clauses.append({"term": {"lang": lang}})

    if year_from or year_to:
        range_filter: dict = {"range": {"year": {}}}
        if year_from:
            range_filter["range"]["year"]["gte"] = year_from
        if year_to:
            range_filter["range"]["year"]["lte"] = year_to
        filter_clauses.append(range_filter)

    body: dict[str, Any] = {
        "query": {
            "bool": {
                "must": must_clauses,
                "filter": filter_clauses,
            }
        },
        "size": size,
        "_source": ["doc_id"],
    }

    try:
        response = es.search(index=INDEX_NAME, body=body)
        hits = response["hits"]["hits"]
        return [
            {
                "doc_id": hit["_source"]["doc_id"],
                "es_score": hit["_score"],
                "es_id": hit["_id"],
            }
            for hit in hits
        ]
    except es_exceptions.ConnectionError as e:
        logger.error(f"Elasticsearch недоступен: {e}")
        return []
    except es_exceptions.NotFoundError:
        logger.warning(f"Индекс {INDEX_NAME!r} не найден в Elasticsearch")
        return []
    except Exception as e:
        logger.error(f"Ошибка поиска в Elasticsearch: {e}")
        return []


def index_document(doc_id: int, doc_data: dict[str, Any]) -> bool:
    """
    Индексировать документ в Elasticsearch.

    Args:
        doc_id: ID документа в PostgreSQL
        doc_data: Словарь с метаданными документа

    Returns:
        True при успехе, False при ошибке
    """
    es = get_es_client()

    doc = {
        "doc_id": doc_id,
        "source": doc_data.get("source"),
        "title": doc_data.get("title", ""),
        "abstract": doc_data.get("abstract", ""),
        "authors": " ".join(
            f"{a.get('last_name', '')} {a.get('first_name', '')}"
            for a in doc_data.get("authors", [])
        ),
        "year": doc_data.get("year"),
        "lang": doc_data.get("lang"),
        "journal": doc_data.get("journal"),
        "doi": doc_data.get("doi"),
        "url": doc_data.get("url"),
    }

    try:
        es.index(
            index=INDEX_NAME,
            id=str(doc_id),
            document=doc,
        )
        return True
    except Exception as e:
        logger.error(f"Ошибка индексации документа {doc_id} в ES: {e}")
        return False