From fecb58ce859a496cc63abaf75b1561c35a68b6b0 Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 13:49:26 +0000 Subject: [PATCH] Add new file --- rag/history_manager.py | 206 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 206 insertions(+) create mode 100644 rag/history_manager.py diff --git a/rag/history_manager.py b/rag/history_manager.py new file mode 100644 index 0000000..c74e901 --- /dev/null +++ b/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})") \ No newline at end of file