405dea3fa3
- /rag/rag_orchestrator.py - /tests/test_hierarchical_summarize.py - /template_bot_profile/bot.conf.sample
95 lines
3.8 KiB
Python
95 lines
3.8 KiB
Python
# -*- 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 |