fckbot/rag/functions/expand_query.py

# -*- coding: utf-8 -*-
"""
Модуль расширения запроса ключевыми словами.
Добавляет синонимы, аббревиатуры и связанные термины для увеличения релевантности.

АДАПТАЦИЯ: теперь функция принимает промпт как аргумент.
"""

import logging
import time
from core.services.giga_client import GigaClient

logger = logging.getLogger(__name__)

_expand_cache = {}
_CACHE_TTL = 300


async def expand_query(
    giga: GigaClient,
    query: str,
    prompt_text: str,
    bot_config,
) -> str:
    """
    Расширяет исходный запрос пользователя набором ключевых слов.

    Аргументы:
        giga (GigaClient): клиент GigaChat
        query (str): исходный запрос
        prompt_text (str): содержимое промпта расширения
        bot_config (BotConfig): объект конфигурации

    Возвращает:
        str: расширенный запрос или исходный, если расширение не удалось
    """
    if not prompt_text:
        logger.debug("Промпт расширения пуст, возвращаем исходный запрос")
        return query

    cache_key = query.strip().lower()
    now = time.time()
    if cache_key in _expand_cache:
        cached_time, cached_result = _expand_cache[cache_key]
        if now - cached_time < _CACHE_TTL:
            logger.debug(f"Кэш расширения: '{cache_key}' -> '{cached_result}'")
            return cached_result
        else:
            del _expand_cache[cache_key]

    expand_cfg = getattr(bot_config, 'expand', {})
    temperature = expand_cfg.get('temperature', 0.1)
    if isinstance(temperature, dict):
        temperature = 0.1

    prompt = prompt_text.format(query=query)

    try:
        response = await giga.chat(
            history=[],
            query=prompt,
            system_prompt=None,
            file_id=None,
            temperature=temperature
        )
        expanded = response.strip()
        if expanded and expanded != query:
            logger.debug(f"Расширенный запрос: {expanded}")
            _expand_cache[cache_key] = (now, expanded)
            return expanded
        else:
            _expand_cache[cache_key] = (now, query)
            return query
    except Exception as e:
        logger.error(f"Ошибка расширения запроса: {e}", exc_info=True)
        return query