fckbot/rag/functions/rerank_context.py

# -*- 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