Add new file

rag/history_manager.py

@@ -0,0 +1,206 @@
+# -*- coding: utf-8 -*-
+"""
+Менеджер истории диалогов.
+Отвечает за:
+- Получение истории из PostgreSQL (через PostgresService).
+- Сжатие длинной истории с помощью иерархического резюмирования.
+- Форматирование истории в список сообщений для передачи в LLM.
+- Сохранение новых сообщений в историю.
+
+Используется в RAGOrchestrator для управления контекстом диалога.
+"""
+
+import logging
+from typing import List, Dict, Optional, Any
+
+from core.services.postgres_service import PostgresService
+from core.services.giga_client import GigaClient
+from core.utils.text_utils import count_tokens
+from core.functions.hierarchical_summarize import hierarchical_summarize
+
+logger = logging.getLogger(__name__)
+
+
+class HistoryManager:
+    """
+    Управляет историей диалога для конкретного пользователя/комнаты.
+    """
+
+    def __init__(
+        self,
+        db: PostgresService,
+        giga: GigaClient,
+        config,
+        default_prompts: Optional[Dict[str, str]] = None
+    ):
+        """
+        Инициализация менеджера истории.
+
+        Аргументы:
+            db: сервис PostgreSQL (для получения и сохранения истории)
+            giga: клиент GigaChat (для сжатия истории)
+            config: объект конфигурации (BotConfig)
+            default_prompts: словарь промптов по умолчанию (нужен для hierarchical_summary)
+        """
+        self.db = db
+        self.giga = giga
+        self.config = config
+        self.default_prompts = default_prompts or {}
+
+    # ------------------------------------------------------------------
+    # Получение и форматирование истории
+    # ------------------------------------------------------------------
+
+    async def get_history(
+        self,
+        user_jid: str,
+        room_jid: Optional[str] = None,
+        limit: int = 100
+    ) -> List[Dict[str, str]]:
+        """
+        Получает историю диалога из БД.
+
+        Аргументы:
+            user_jid: JID пользователя (без ресурса)
+            room_jid: JID комнаты (None для личного чата)
+            limit: максимальное количество сообщений
+
+        Возвращает:
+            List[Dict[str, str]]: список сообщений с ключами 'role' и 'content'
+        """
+        raw_history = await self.db.get_history(user_jid, room_jid, limit=limit)
+        # Преобразуем в формат, ожидаемый LLM
+        return [
+            {"role": rec['role'], "content": rec['content']}
+            for rec in raw_history
+        ]
+
+    async def compress_history_if_needed(
+        self,
+        history: List[Dict[str, str]],
+        max_tokens: int,
+        prompt_template: Optional[str] = None
+    ) -> List[Dict[str, str]]:
+        """
+        Сжимает историю, если она превышает лимит токенов.
+        Использует иерархическое резюмирование для сжатия всей истории в одно системное сообщение.
+
+        Аргументы:
+            history: список сообщений (role, content)
+            max_tokens: максимальное допустимое количество токенов для истории
+            prompt_template: шаблон промпта для резюмирования (если None, используется стандартный)
+
+        Возвращает:
+            List[Dict[str, str]]: либо исходная история (если укладывается в лимит),
+                                  либо сжатая история в виде одного системного сообщения.
+        """
+        if not history:
+            return []
+
+        # Подсчитываем общее количество токенов в истории
+        history_text = "\n".join([f"{msg['role']}: {msg['content']}" for msg in history])
+        total_tokens = count_tokens(history_text)
+
+        if total_tokens <= max_tokens:
+            logger.debug(f"История укладывается в {max_tokens} токенов (фактически {total_tokens})")
+            return history
+
+        logger.info(
+            f"История слишком длинная ({total_tokens} токенов), "
+            f"превышает лимит {max_tokens}, применяем сжатие"
+        )
+
+        # Загружаем промпт для иерархического резюмирования
+        if prompt_template is None:
+            prompt_template = self.default_prompts.get('hierarchical_summary', '')
+            if not prompt_template:
+                # Пытаемся загрузить из файла, если он есть
+                try:
+                    prompt_path = self.config.prompts_dir / 'hierarchical_summary.txt'
+                    with open(prompt_path, 'r', encoding='utf-8') as f:
+                        prompt_template = f.read()
+                except Exception as e:
+                    logger.warning(f"Не удалось загрузить промпт hierarchical_summary: {e}")
+                    # Стандартный промпт
+                    prompt_template = (
+                        "Кратко изложи суть диалога, сохранив ключевые факты, вопросы и ответы.\n"
+                        "Объём: не более 30% от исходного, но не менее 2 предложений.\n"
+                        "Стиль: нейтральный, без оценок.\n\n"
+                        "ТЕКСТ ДИАЛОГА:\n{text}\n\nКРАТКОЕ ИЗЛОЖЕНИЕ:"
+                    )
+
+        # Параметры сжатия из конфига
+        summarization_config = getattr(self.config, 'summarization', {})
+        target_tokens = summarization_config.get('target_tokens_after_summary', max_tokens)
+        chunk_size = summarization_config.get('chunk_size_tokens', 500)
+        max_depth = summarization_config.get('max_depth', 2)
+        temperature = getattr(self.config, 'ai_temperature', 0.1)
+
+        try:
+            compressed_text = await hierarchical_summarize(
+                text=history_text,
+                giga=self.giga,
+                prompt_template=prompt_template,
+                target_tokens=target_tokens,
+                chunk_size_tokens=chunk_size,
+                max_depth=max_depth,
+                temperature=temperature
+            )
+            logger.info(
+                f"История сжата с {total_tokens} до {count_tokens(compressed_text)} токенов"
+            )
+            # Возвращаем сжатую историю как одно системное сообщение
+            return [
+                {"role": "system", "content": f"Сжатая история диалога:\n{compressed_text}"}
+            ]
+        except Exception as e:
+            logger.error(f"Ошибка при сжатии истории: {e}, выполняем простую обрезку")
+            # Fallback: простая обрезка по токенам
+            truncated_history = []
+            total = 0
+            for msg in reversed(history):
+                tokens = count_tokens(msg['content'])
+                if total + tokens <= max_tokens:
+                    truncated_history.append(msg)
+                    total += tokens
+                else:
+                    break
+            truncated_history.reverse()
+            return truncated_history
+
+    async def save_message(
+        self,
+        user_jid: str,
+        role: str,
+        content: str,
+        room_jid: Optional[str] = None
+    ) -> None:
+        """
+        Сохраняет одно сообщение в историю диалога.
+
+        Аргументы:
+            user_jid: JID пользователя
+            role: 'user' или 'assistant'
+            content: текст сообщения
+            room_jid: JID комнаты (None для личного чата)
+        """
+        await self.db.add_history(user_jid, role, content, room_jid)
+        logger.debug(f"Сохранено сообщение {role} для {user_jid} (room={room_jid})")
+
+    async def clear_history(
+        self,
+        user_jid: str,
+        room_jid: Optional[str] = None
+    ) -> None:
+        """
+        Очищает историю диалога (личную или комнатную).
+
+        Аргументы:
+            user_jid: JID пользователя
+            room_jid: JID комнаты (None для личной истории)
+        """
+        if room_jid:
+            await self.db.clear_room_history(room_jid)
+        else:
+            await self.db.clear_user_history(user_jid)
+        logger.info(f"История очищена для {user_jid} (room={room_jid})")