244 lines
12 KiB
Python
244 lines
12 KiB
Python
"""Тесты плагина `QwenLocal` (map-reduce) с мок-LLM-клиентом (без сети/токенизатора).
|
||
|
||
`QwenLocal._get_client()` подменяется фейковым клиентом с методом `complete`,
|
||
поэтому чанкер и подстановка промптов проверяются изолированно от
|
||
`OpenAICompatClient`/`httpx`. Подсчёт токенов подменяется простым фейком
|
||
(число слов), чтобы детерминированно управлять группировкой в reduce.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import pytest
|
||
|
||
from core.plugins.qwen_local import QwenLocal
|
||
|
||
PROMPTS_DIR = Path(__file__).parent.parent.parent / "workers" / "summarizer" / "prompts"
|
||
|
||
|
||
class _FakeLlmClient:
|
||
"""Фейковый LLM-клиент: запоминает все промпты и возвращает по очереди ответы."""
|
||
|
||
def __init__(self, responses: list[str] | None = None) -> None:
|
||
self.prompts: list[str] = []
|
||
self.max_tokens_seen: list[int | None] = []
|
||
self._responses = responses
|
||
self.closed = False
|
||
|
||
def complete(self, prompt: str, *, max_tokens: int | None = None) -> str:
|
||
self.prompts.append(prompt)
|
||
self.max_tokens_seen.append(max_tokens)
|
||
if self._responses is not None:
|
||
return self._responses[len(self.prompts) - 1]
|
||
return f"резюме №{len(self.prompts)}"
|
||
|
||
def close(self) -> None:
|
||
"""Заглушка `OpenAICompatClient.close()` — фиксирует факт закрытия."""
|
||
self.closed = True
|
||
|
||
|
||
def _make_plugin(fake_client: _FakeLlmClient, count_tokens: Any = None, **kwargs: Any) -> QwenLocal:
|
||
plugin = QwenLocal(prompts_dir=str(PROMPTS_DIR), **kwargs)
|
||
plugin._client = fake_client # type: ignore[assignment]
|
||
if count_tokens is not None:
|
||
plugin._count_tokens = count_tokens
|
||
return plugin
|
||
|
||
|
||
def _word_count(text: str) -> int:
|
||
return len(text.split())
|
||
|
||
|
||
def test_empty_transcript_returns_empty_string_without_llm_calls() -> None:
|
||
"""Пустой транскрипт → "" без единого обращения к LLM."""
|
||
fake_client = _FakeLlmClient()
|
||
plugin = _make_plugin(fake_client, count_tokens=_word_count)
|
||
|
||
result = plugin.summarize("")
|
||
|
||
assert result == ""
|
||
assert fake_client.prompts == []
|
||
|
||
|
||
def test_single_chunk_returns_map_result_without_reduce_call() -> None:
|
||
"""Транскрипт, помещающийся в один чанк, → результат map без reduce-вызова."""
|
||
fake_client = _FakeLlmClient(responses=["итоговое резюме одного чанка"])
|
||
plugin = _make_plugin(fake_client, count_tokens=_word_count, chunk_minutes=20)
|
||
transcript = "[Аня 00:00] привет всем\n[Боря 00:03] привет как дела"
|
||
|
||
result = plugin.summarize(transcript)
|
||
|
||
assert result == "итоговое резюме одного чанка"
|
||
assert len(fake_client.prompts) == 1
|
||
assert "привет всем" in fake_client.prompts[0]
|
||
assert "привет как дела" in fake_client.prompts[0]
|
||
# промпт собран из утверждённого map-шаблона (проверяем характерный заголовок)
|
||
assert "## Ключевые тезисы" in fake_client.prompts[0]
|
||
|
||
|
||
def test_multiple_chunks_are_mapped_then_reduced_once() -> None:
|
||
"""Несколько чанков → map каждого чанка, затем один reduce-вызов над частичными резюме."""
|
||
fake_client = _FakeLlmClient(
|
||
responses=["частичное резюме A", "частичное резюме B", "итоговое резюме"]
|
||
)
|
||
plugin = _make_plugin(fake_client, count_tokens=_word_count, chunk_minutes=1)
|
||
lines = [f"[Аня {minute:02d}:00] " + " ".join(["слово"] * 5) for minute in range(0, 4, 2)]
|
||
# Раздвигаем фразы на разные "чанки" по времени (chunk_minutes=1 закрывает
|
||
# чанк уже на второй фразе, т.к. между ними больше минуты).
|
||
transcript = "\n".join(lines)
|
||
|
||
result = plugin.summarize(transcript)
|
||
|
||
assert result == "итоговое резюме"
|
||
assert len(fake_client.prompts) == 3
|
||
map_prompts, reduce_prompt = fake_client.prompts[:2], fake_client.prompts[2]
|
||
assert all("## Ключевые тезисы" in p for p in map_prompts)
|
||
assert "частичное резюме A" in reduce_prompt
|
||
assert "частичное резюме B" in reduce_prompt
|
||
assert "Частичные резюме:" in reduce_prompt
|
||
|
||
|
||
def test_hierarchical_reduce_when_partial_summaries_exceed_token_budget() -> None:
|
||
"""Много крупных частичных резюме не влезают в бюджет одного reduce → иерархический reduce."""
|
||
# 10 чанков ⇒ 10 map-вызовов; каждое частичное резюме — 1000 "токенов"
|
||
# (по фейковому счётчику слов), бюджет reduce — 6000 → одной группой не
|
||
# свести, нужна как минимум одна промежуточная стадия.
|
||
big_word = "слово " * 1000
|
||
|
||
class _BigFakeLlmClient(_FakeLlmClient):
|
||
def complete(self, prompt: str, *, max_tokens: int | None = None) -> str:
|
||
self.prompts.append(prompt)
|
||
self.max_tokens_seen.append(max_tokens)
|
||
call_no = len(self.prompts)
|
||
if "Транскрипт:" in prompt:
|
||
return big_word.strip()
|
||
return f"промежуточное резюме {call_no}"
|
||
|
||
fake_client = _BigFakeLlmClient()
|
||
plugin = _make_plugin(fake_client, count_tokens=_word_count, chunk_minutes=1)
|
||
|
||
lines = [f"[Аня {minute:02d}:00] реплика номер {minute}" for minute in range(0, 40, 2)]
|
||
transcript = "\n".join(lines)
|
||
|
||
result = plugin.summarize(transcript)
|
||
|
||
map_calls = [p for p in fake_client.prompts if "Транскрипт:" in p]
|
||
reduce_calls = [p for p in fake_client.prompts if "Частичные резюме:" in p]
|
||
|
||
assert len(map_calls) == 20 # по фразе на чанк (chunk_minutes=1, реплики каждые 2 мин)
|
||
assert len(reduce_calls) >= 2 # хотя бы одна промежуточная группа + финальный reduce
|
||
assert result.startswith("промежуточное резюме")
|
||
|
||
|
||
def test_map_prompt_uses_approved_template_placeholder() -> None:
|
||
"""Плейсхолдер `{transcript_chunk}` подставляется через .replace, не .format."""
|
||
fake_client = _FakeLlmClient(responses=["резюме"])
|
||
plugin = _make_plugin(fake_client, count_tokens=_word_count)
|
||
transcript = "[Аня 00:00] текст с фигурными скобками {как тут} и всё равно ок"
|
||
|
||
result = plugin.summarize(transcript)
|
||
|
||
assert result == "резюме"
|
||
assert "текст с фигурными скобками {как тут} и всё равно ок" in fake_client.prompts[0]
|
||
|
||
|
||
def test_map_and_reduce_use_different_max_tokens() -> None:
|
||
"""map и reduce вызываются с разными `max_tokens` (раздельные per-tier лимиты, ADR-004)."""
|
||
fake_client = _FakeLlmClient(
|
||
responses=["частичное резюме A", "частичное резюме B", "итоговое резюме"]
|
||
)
|
||
plugin = _make_plugin(
|
||
fake_client,
|
||
count_tokens=_word_count,
|
||
chunk_minutes=1,
|
||
max_tokens_map=1024,
|
||
max_tokens_reduce=2048,
|
||
)
|
||
lines = [f"[Аня {minute:02d}:00] " + " ".join(["слово"] * 5) for minute in range(0, 4, 2)]
|
||
transcript = "\n".join(lines)
|
||
|
||
plugin.summarize(transcript)
|
||
|
||
assert fake_client.max_tokens_seen == [1024, 1024, 2048]
|
||
|
||
|
||
def test_max_tokens_map_and_reduce_default_to_max_tokens_when_not_set() -> None:
|
||
"""Без явных `max_tokens_map`/`max_tokens_reduce` оба используют общий `max_tokens`
|
||
(обратная совместимость с конфигурацией)."""
|
||
fake_client = _FakeLlmClient(responses=["резюме"])
|
||
plugin = _make_plugin(fake_client, count_tokens=_word_count, max_tokens=777)
|
||
|
||
plugin.summarize("[Аня 00:00] короткая фраза")
|
||
|
||
assert fake_client.max_tokens_seen == [777]
|
||
assert plugin.max_tokens_map == 777
|
||
assert plugin.max_tokens_reduce == 777
|
||
|
||
|
||
def test_summarize_closes_llm_client_after_use() -> None:
|
||
"""Плагин закрывает HTTP-клиент LLM по завершении `summarize` (освобождение пула соединений)."""
|
||
fake_client = _FakeLlmClient(responses=["резюме"])
|
||
plugin = _make_plugin(fake_client, count_tokens=_word_count)
|
||
|
||
plugin.summarize("[Аня 00:00] короткая фраза")
|
||
|
||
assert fake_client.closed is True
|
||
assert plugin._client is None
|
||
|
||
|
||
def test_reduce_forces_pairwise_progress_when_summaries_exceed_budget_individually() -> None:
|
||
"""Если каждое частичное резюме само превышает бюджет reduce (например, из-за
|
||
неудачно большого `max_tokens` в конфиге), группировка по бюджету не сокращает
|
||
список — plugin обязан гарантировать прогресс (принудительное попарное
|
||
объединение), а не зависать в бесконечном цикле."""
|
||
|
||
class _HugeFakeLlmClient(_FakeLlmClient):
|
||
def complete(self, prompt: str, *, max_tokens: int | None = None) -> str:
|
||
self.prompts.append(prompt)
|
||
self.max_tokens_seen.append(max_tokens)
|
||
call_no = len(self.prompts)
|
||
# Защита от регресса: если бы `_reduce` зациклился без прогресса,
|
||
# число вызовов росло бы неограниченно — превращаем зависание
|
||
# в быстрый явный отказ теста вместо вечного цикла.
|
||
if call_no > 50:
|
||
raise AssertionError("нет прогресса в _reduce — похоже на бесконечный цикл")
|
||
if "Транскрипт:" in prompt:
|
||
return f"огромное резюме {call_no}"
|
||
return f"свод {call_no}"
|
||
|
||
fake_client = _HugeFakeLlmClient()
|
||
# Фейковый счётчик токенов всегда возвращает значение >= бюджета reduce —
|
||
# ни одно частичное резюме "не помещается", группировка по бюджету
|
||
# вырождается в синглтоны на каждом уровне.
|
||
plugin = _make_plugin(fake_client, count_tokens=lambda _text: 7000, chunk_minutes=1)
|
||
|
||
lines = [f"[Аня {minute:02d}:00] реплика {minute}" for minute in range(0, 16, 2)]
|
||
transcript = "\n".join(lines)
|
||
|
||
result = plugin.summarize(transcript)
|
||
|
||
assert result.startswith("свод")
|
||
reduce_calls = [p for p in fake_client.prompts if "Частичные резюме:" in p]
|
||
assert len(reduce_calls) >= 1
|
||
|
||
|
||
def test_prompt_templates_are_loaded_lazily_and_cached(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
"""Файл промпта читается с диска один раз (лениво) и переиспользуется при повторных вызовах."""
|
||
read_calls = {"count": 0}
|
||
original_read_text = Path.read_text
|
||
|
||
def counting_read_text(self: Path, *args: Any, **kwargs: Any) -> str:
|
||
read_calls["count"] += 1
|
||
result: str = original_read_text(self, *args, **kwargs)
|
||
return result
|
||
|
||
monkeypatch.setattr(Path, "read_text", counting_read_text)
|
||
|
||
plugin = QwenLocal(prompts_dir=str(PROMPTS_DIR))
|
||
plugin._map_prompt_template()
|
||
plugin._map_prompt_template()
|
||
plugin._reduce_prompt_template()
|
||
plugin._reduce_prompt_template()
|
||
|
||
assert read_calls["count"] == 2 # по одному чтению на файл: map и reduce
|