diff --git a/core/functions/rerank_context.py b/core/functions/rerank_context.py new file mode 100644 index 0000000..10c5439 --- /dev/null +++ b/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 \ No newline at end of file