206 lines
9.3 KiB
Python
206 lines
9.3 KiB
Python
# -*- 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})") |