Add new file

core/functions/rerank_context.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+"""
+Модуль переранжирования контекста: отбирает только те фрагменты,
+которые действительно релевантны вопросу.
+Использует локальный кросс-энкодер.
+
+АДАПТАЦИЯ: теперь функция принимает bot (для доступа к конфигу) и использует RerankerService.
+"""
+
+import logging
+from core.services.reranker_service import RerankerService
+
+logger = logging.getLogger(__name__)
+
+_reranker = None
+
+
+def get_reranker(bot_config) -> RerankerService:
+    """Возвращает глобальный экземпляр RerankerService (ленивая инициализация)."""
+    global _reranker
+    if _reranker is None:
+        model_name = getattr(bot_config, 'rerank_model', None)
+        device = getattr(bot_config, 'rerank_device', None)
+        _reranker = RerankerService(model_name=model_name, device=device)
+        logger.info(f"RerankerService инициализирован (модель: {model_name or 'default'})")
+    return _reranker
+
+
+async def rerank_context(
+    bot,
+    query: str,
+    context: str,
+    prompt_text: str = None,  # сохранён для обратной совместимости, но не используется
+    bot_config=None,
+) -> str:
+    """
+    Переранжирует контекст, оставляя только релевантные фрагменты.
+    Использует кросс-энкодер вместо LLM.
+
+    Аргументы:
+        bot: экземпляр Bot (для доступа к конфигу) — может быть None, если передан bot_config
+        query: исходный запрос пользователя
+        context: исходный контекст (абзацы, разделённые двойными переносами строки "\n\n")
+        prompt_text: не используется
+        bot_config: объект конфигурации (если bot не передан)
+
+    Возвращает:
+        str: отфильтрованный контекст
+    """
+    if not context:
+        return context
+
+    # Получаем конфиг либо из bot, либо из bot_config
+    if bot is not None and hasattr(bot, 'config'):
+        config = bot.config
+    elif bot_config is not None:
+        config = bot_config
+    else:
+        logger.error("Ни bot, ни bot_config не переданы в rerank_context")
+        return context
+
+    rerank_cfg = getattr(config, 'rerank', {})
+    min_length = rerank_cfg.get('min_length', 5000)
+
+    # Если контекст короткий, переранжирование не требуется
+    if len(context) < min_length:
+        logger.debug(f"Контекст слишком короткий ({len(context)} < {min_length}), пропускаем переранжирование")
+        return context
+
+    # Получаем сервис переранжирования
+    reranker = get_reranker(config)
+
+    try:
+        top_k = rerank_cfg.get('top_k', None)
+        ranked_context = reranker.rerank_context(query, context, top_k=top_k)
+        logger.debug(f"Переранжирование выполнено: исходный контекст {len(context)} символов")
+        return ranked_context
+    except Exception as e:
+        logger.error(f"Ошибка при переранжировании: {e}. Возвращаем исходный контекст.")
+        return context