From 405dea3fa3628bead1521c58f6203fc5ae10e60f Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 12:51:14 +0000 Subject: [PATCH] Update 3 files - /rag/rag_orchestrator.py - /tests/test_hierarchical_summarize.py - /template_bot_profile/bot.conf.sample --- rag/rag_orchestrator.py | 76 +++++++++++++++------- template_bot_profile/bot.conf.sample | 1 + tests/test_hierarchical_summarize.py | 95 ++++++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 24 deletions(-) create mode 100644 tests/test_hierarchical_summarize.py diff --git a/rag/rag_orchestrator.py b/rag/rag_orchestrator.py index 8522495..cdf8fc0 100644 --- a/rag/rag_orchestrator.py +++ b/rag/rag_orchestrator.py @@ -33,6 +33,7 @@ from core.functions.check_consistency import check_consistency from core.functions.critique_answer import critique_answer from core.functions.rerank_context import rerank_context from core.functions.check_spelling import check_spelling +from core.functions.hierarchical_summarize import hierarchical_summarize from core.utils.text_utils import count_tokens logger = logging.getLogger(__name__) @@ -208,31 +209,58 @@ class RAGOrchestrator: f"промпт: {token_info['prompt_tokens']} токенов" ) - # ----- 4. Обрезка истории по токенам (с учётом лимита) ----- - # Идём с конца (свежие сообщения) к началу, чтобы сохранить самые последние - truncated_history = [] - total_history_tokens = 0 - history_limit = min(available_for_history_and_context // 2, 2000) # выделяем половину для истории + # ----- 4. Обрезка истории (с использованием иерархического резюмирования) ----- + max_history_tokens = min(available_for_history_and_context // 2, 2000) - for record in reversed(history): - tokens = count_tokens(record['content']) - if total_history_tokens + tokens <= history_limit: - truncated_history.append(record) - total_history_tokens += tokens - else: - break - truncated_history.reverse() # восстанавливаем хронологический порядок - - # Преобразуем в формат, ожидаемый GigaChat - formatted_history = [ - {"role": rec['role'], "content": rec['content']} - for rec in truncated_history - ] - - logger.debug( - f"История обрезана: {len(truncated_history)} сообщений, " - f"{total_history_tokens} токенов (лимит {history_limit})" - ) + # Если история слишком длинная, сжимаем её иерархически + history_text = "\n".join([f"{rec['role']}: {rec['content']}" for rec in history]) + if count_tokens(history_text) > max_history_tokens: + logger.info(f"История слишком длинная ({count_tokens(history_text)} токенов), применяем иерархическое резюмирование") + summary_prompt = prompts.get('hierarchical_summary', '') + if not summary_prompt: + # Загружаем из файла, если не передано + try: + with open(self.config.prompts_dir / 'hierarchical_summary.txt', 'r', encoding='utf-8') as f: + summary_prompt = f.read() + except Exception: + summary_prompt = "Кратко изложи суть диалога, сохранив ключевые факты и вопросы.\n\nТЕКСТ:\n{text}" + try: + compressed_history = await hierarchical_summarize( + text=history_text, + giga=self.giga, + prompt_template=summary_prompt, + target_tokens=max_history_tokens, + chunk_size_tokens=500, + max_depth=2, + temperature=getattr(self.config, 'ai_temperature', 0.1) + ) + # Заменяем историю одним сжатым сообщением + formatted_history = [{"role": "system", "content": "Сжатая история диалога:\n" + compressed_history}] + logger.info(f"История сжата с {count_tokens(history_text)} до {count_tokens(compressed_history)} токенов") + except Exception as e: + logger.error(f"Ошибка при иерархическом резюмировании истории: {e}") + # Fallback: простая обрезка + truncated_history = [] + total_tokens = 0 + for record in reversed(history): + tokens = count_tokens(record['content']) + if total_tokens + tokens <= max_history_tokens: + truncated_history.append(record) + total_tokens += tokens + else: + break + truncated_history.reverse() + formatted_history = [ + {"role": rec['role'], "content": rec['content']} + for rec in truncated_history + ] + else: + # Если история укладывается, просто форматируем + truncated_history = history[-10:] # ограничим последними 10 сообщениями для простоты, но можно взять больше + formatted_history = [ + {"role": rec['role'], "content": rec['content']} + for rec in truncated_history + ] # ----- 5. Классификация намерений (если не переопределена) ----- intent = intent_override diff --git a/template_bot_profile/bot.conf.sample b/template_bot_profile/bot.conf.sample index 93c70d0..84f5e5b 100644 --- a/template_bot_profile/bot.conf.sample +++ b/template_bot_profile/bot.conf.sample @@ -23,6 +23,7 @@ prompts: expand: "prompts/expand.txt" metrics: "prompts/metrics_extract.txt" summary: "prompts/smart_summary.txt" + hierarchical_summary: "prompts/hierarchical_summary.txt" consistency: "prompts/consistency_check.txt" critique: "prompts/self_critique.txt" spellcheck: "prompts/spellcheck.txt" diff --git a/tests/test_hierarchical_summarize.py b/tests/test_hierarchical_summarize.py new file mode 100644 index 0000000..8749a1f --- /dev/null +++ b/tests/test_hierarchical_summarize.py @@ -0,0 +1,95 @@ +# -*- coding: utf-8 -*- +""" +Тесты для модуля иерархического резюмирования. +""" + +import pytest +from core.functions.hierarchical_summarize import ( + hierarchical_summarize, + _split_into_sentence_blocks, + _truncate_by_tokens +) +from core.utils.text_utils import count_tokens + + +@pytest.mark.asyncio +async def test_hierarchical_summarize_basic(mock_giga): + """ + Тест: базовое резюмирование короткого текста. + """ + text = "Это тестовый текст для проверки иерархического резюмирования. Он содержит несколько предложений. Все они должны быть сжаты." + mock_giga.chat.return_value = "Краткое резюме текста." + + result = await hierarchical_summarize( + text=text, + giga=mock_giga, + prompt_template="Резюмируй: {text}", + target_tokens=50, + chunk_size_tokens=100, + max_depth=2 + ) + assert result == "Краткое резюме текста." + mock_giga.chat.assert_called_once() + + +@pytest.mark.asyncio +async def test_hierarchical_summarize_no_change(mock_giga): + """ + Тест: если текст уже короткий, возвращается без изменений. + """ + text = "Короткий текст." + result = await hierarchical_summarize( + text=text, + giga=mock_giga, + prompt_template="Резюмируй: {text}", + target_tokens=100, + chunk_size_tokens=100 + ) + assert result == text + mock_giga.chat.assert_not_called() + + +@pytest.mark.asyncio +async def test_hierarchical_summarize_recursive(mock_giga): + """ + Тест: рекурсивное резюмирование, когда одно резюме всё ещё слишком длинное. + """ + long_text = "A " * 1000 # 1000 слов "A" + mock_giga.chat.side_effect = [ + "Краткое резюме части 1. Краткое резюме части 2.", # первое резюмирование + "Итоговое резюме." # второе (рекурсивное) + ] + + result = await hierarchical_summarize( + text=long_text, + giga=mock_giga, + prompt_template="Резюмируй: {text}", + target_tokens=50, + chunk_size_tokens=200, + max_depth=3 + ) + assert result == "Итоговое резюме." + assert mock_giga.chat.call_count == 2 # два вызова: для блоков и для объединённого + + +def test_split_into_sentence_blocks(): + """ + Тест разбиения на блоки по предложениям. + """ + text = "Первое предложение. Второе предложение! Третье предложение? Четвёртое." + blocks = _split_into_sentence_blocks(text, 20) # маленький размер, чтобы разбить + # Ожидаем, что блоков будет несколько (или все предложения в одном, если они короткие) + assert len(blocks) >= 1 + # Проверяем, что каждый блок содержит целые предложения + for block in blocks: + assert block.endswith(('.', '!', '?')) or count_tokens(block) < 20 + + +def test_truncate_by_tokens(): + """ + Тест обрезки по токенам. + """ + text = "Один два три четыре пять шесть семь восемь девять десять." + truncated = _truncate_by_tokens(text, 10) # должно обрезать до ~30 символов + # Проверяем, что длина в токенах стала меньше или равна целевой + assert count_tokens(truncated) <= 10 \ No newline at end of file