"""Celery задача семантического поиска источников.""" import hashlib import json 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, 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)