388 lines
13 KiB
Python
388 lines
13 KiB
Python
"""Unit: KB overview generator (phase 31, task 02).
|
|
|
|
The prompt tests are pure (no DB): ``KB_OVERVIEW_MODE`` system prompt,
|
|
per-document user lines (source/path/title/first summary line), and the
|
|
input cap with the shared ``[…truncated…]`` marker. The loader and
|
|
regenerator tests run against the local compose Postgres (preferred —
|
|
real single-row upsert), skipping with clear instructions when the stack
|
|
is not up — same pattern as ``test_importer.py``.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.config import Settings, get_settings
|
|
from app.models import Document, KbOverview
|
|
from app.rag.llm import LLMError
|
|
from app.rag.overview import (
|
|
KB_OVERVIEW_MODE,
|
|
OVERVIEW_INSTRUCTION,
|
|
SYSTEM_PROMPT,
|
|
build_overview_prompt,
|
|
load_kb_overview,
|
|
regenerate_overview,
|
|
)
|
|
from app.rag.retriever import TRUNCATION_MARKER
|
|
|
|
REPLY = "- Homelab\n - Kubernetes (k3s)\n- Deployments\n - Borg backups"
|
|
|
|
|
|
class _FakeLLM:
|
|
"""Duck-typed stand-in for ``LLMClient`` (``chat`` + ``settings``).
|
|
|
|
Records the messages and the ``model`` kwarg it was called with;
|
|
returns a canned reply or raises (e.g. :class:`LLMError`).
|
|
"""
|
|
|
|
def __init__(self, reply: str = REPLY, fail: Exception | None = None) -> None:
|
|
self._reply = reply
|
|
self._fail = fail
|
|
self.settings = Settings(_env_file=None) # pyright: ignore[reportCallIssue]
|
|
self.messages: list[dict[str, str]] = []
|
|
self.model: str | None = None
|
|
self.calls = 0
|
|
|
|
async def chat(
|
|
self, messages: list[dict[str, str]], model: str | None = None
|
|
) -> str:
|
|
self.calls += 1
|
|
self.messages = list(messages)
|
|
self.model = model
|
|
if self._fail is not None:
|
|
raise self._fail
|
|
return self._reply
|
|
|
|
|
|
# ---------- build_overview_prompt: system ----------
|
|
|
|
|
|
def test_system_prompt_has_marker_and_locked_instruction() -> None:
|
|
assert SYSTEM_PROMPT.startswith(KB_OVERVIEW_MODE)
|
|
assert OVERVIEW_INSTRUCTION in SYSTEM_PROMPT
|
|
for fragment in (
|
|
"compact plain-text outline",
|
|
"basic categories and topics",
|
|
"Group by source",
|
|
"use `-` bullet lines",
|
|
"~1500 characters",
|
|
"no markdown headings",
|
|
"no topics not present in the list",
|
|
):
|
|
assert fragment in SYSTEM_PROMPT
|
|
system, _ = build_overview_prompt([])
|
|
assert system == SYSTEM_PROMPT
|
|
assert KB_OVERVIEW_MODE in system # the marker the E2E mock keys on
|
|
|
|
|
|
# ---------- build_overview_prompt: user lines ----------
|
|
|
|
|
|
def test_user_lines_carry_source_path_title_and_first_summary_line() -> None:
|
|
rows = [
|
|
("Homelab", "kubernetes/k3s.md", "K3s Cluster", None),
|
|
(
|
|
"Deployments",
|
|
"backups/borg.yaml",
|
|
"Borg Backups",
|
|
"Backups run nightly via the borg schedule.\nSource: Deployments/backups/borg.yaml",
|
|
),
|
|
]
|
|
system, user = build_overview_prompt(rows)
|
|
assert system == SYSTEM_PROMPT
|
|
assert user == (
|
|
"Homelab — kubernetes/k3s.md — K3s Cluster\n"
|
|
"Deployments — backups/borg.yaml — Borg Backups — "
|
|
"Backups run nightly via the borg schedule."
|
|
)
|
|
|
|
|
|
def test_user_line_omits_summary_field_when_absent() -> None:
|
|
_, user = build_overview_prompt([("Homelab", "a.md", "Title A", None)])
|
|
assert user == "Homelab — a.md — Title A"
|
|
assert not user.endswith(" — ") # no dangling dash
|
|
|
|
|
|
def test_user_line_omits_summary_field_when_blank() -> None:
|
|
_, user = build_overview_prompt([("Homelab", "a.md", "Title A", " \n\t ")])
|
|
assert user == "Homelab — a.md — Title A"
|
|
|
|
|
|
def test_user_line_uses_only_first_summary_line() -> None:
|
|
"""Multi-line summaries contribute their first line only (the
|
|
model-written lead sentence; the ``Source: …`` pointer is last)."""
|
|
_, user = build_overview_prompt(
|
|
[
|
|
(
|
|
"Homelab",
|
|
"a.yaml",
|
|
"Title A",
|
|
"First line.\nSecond line.\nSource: Homelab/a.yaml",
|
|
)
|
|
]
|
|
)
|
|
assert user == "Homelab — a.yaml — Title A — First line."
|
|
assert "Second line" not in user
|
|
assert "Source:" not in user
|
|
|
|
|
|
def test_zero_rows_yield_empty_user() -> None:
|
|
_, user = build_overview_prompt([])
|
|
assert user == ""
|
|
|
|
|
|
# ---------- build_overview_prompt: input cap ----------
|
|
|
|
|
|
def test_user_prompt_truncated_with_marker_when_over_custom_cap() -> None:
|
|
rows = [("S", f"p{i}.md", f"T{i}", None) for i in range(10)]
|
|
_, full = build_overview_prompt(rows, max_chars=10_000)
|
|
cap = 25
|
|
_, user = build_overview_prompt(rows, max_chars=cap)
|
|
assert user == full[:cap] + "\n" + TRUNCATION_MARKER
|
|
assert user.endswith(TRUNCATION_MARKER)
|
|
assert len(user) > cap # the marker makes the cut visible past the cap
|
|
|
|
|
|
def test_user_prompt_at_exact_cap_not_truncated() -> None:
|
|
rows = [("S", "p.md", "T", None)] # "S — p.md — T" = 12 chars
|
|
_, user = build_overview_prompt(rows, max_chars=12)
|
|
assert user == "S — p.md — T"
|
|
assert TRUNCATION_MARKER not in user
|
|
|
|
|
|
def test_user_prompt_truncated_at_default_cap() -> None:
|
|
"""No explicit cap → ``BOR_OVERVIEW_INPUT_MAX_CHARS`` (read from the
|
|
live settings, so the test holds for any configured value)."""
|
|
cap = get_settings().overview_input_max_chars
|
|
rows = [("S", f"p{i}.md", "T", None) for i in range(3_000)]
|
|
_, user = build_overview_prompt(rows)
|
|
assert user.endswith(TRUNCATION_MARKER)
|
|
body = user.removesuffix("\n" + TRUNCATION_MARKER)
|
|
assert len(body) == cap # cut exactly at the cap, marker on its own line
|
|
assert "p2999.md" not in body # the overflow never reaches the model
|
|
|
|
|
|
# ---------- load_kb_overview ----------
|
|
|
|
|
|
def test_load_kb_overview_no_row_empty(db: Session) -> None:
|
|
db.execute(text("DELETE FROM kb_overview"))
|
|
db.commit()
|
|
assert load_kb_overview(db) == ""
|
|
|
|
|
|
def test_load_kb_overview_returns_trimmed_content(db: Session) -> None:
|
|
db.execute(text("DELETE FROM kb_overview"))
|
|
db.add(KbOverview(id=1, content=f" {REPLY} \n"))
|
|
db.commit()
|
|
assert load_kb_overview(db) == REPLY
|
|
|
|
|
|
def test_load_kb_overview_whitespace_only_row_empty(db: Session) -> None:
|
|
db.execute(text("DELETE FROM kb_overview"))
|
|
db.add(KbOverview(id=1, content=" \n\t "))
|
|
db.commit()
|
|
assert load_kb_overview(db) == ""
|
|
|
|
|
|
# ---------- regenerate_overview ----------
|
|
|
|
|
|
def _add_doc(
|
|
db: Session, source: str, path: str, title: str, summary: str | None = None
|
|
) -> Document:
|
|
doc = Document(
|
|
id=uuid.uuid4(),
|
|
source=source,
|
|
path=path,
|
|
full_path=f"/tmp/{path}",
|
|
title=title,
|
|
content="body",
|
|
content_hash="0" * 64,
|
|
summary=summary,
|
|
)
|
|
db.add(doc)
|
|
db.commit()
|
|
return doc
|
|
|
|
|
|
def _truncate_documents(db: Session) -> None:
|
|
db.execute(text("TRUNCATE chunks, documents"))
|
|
db.commit()
|
|
|
|
|
|
def _overview_row(db: Session) -> KbOverview | None:
|
|
return db.get(KbOverview, 1)
|
|
|
|
|
|
@pytest.fixture()
|
|
def clean_overview(db: Session):
|
|
db.execute(text("DELETE FROM kb_overview"))
|
|
db.commit()
|
|
yield
|
|
db.execute(text("DELETE FROM kb_overview"))
|
|
db.commit()
|
|
|
|
|
|
def test_regenerate_happy_path_creates_single_row_id_1(
|
|
db: Session, clean_overview, caplog: pytest.LogCaptureFixture
|
|
) -> None:
|
|
_truncate_documents(db)
|
|
try:
|
|
_add_doc(db, "Homelab", "b.md", "B")
|
|
_add_doc(db, "Homelab", "a.md", "A", "A summary.\nSource: Homelab/a.md")
|
|
llm = _FakeLLM()
|
|
with caplog.at_level(logging.INFO, logger="app.rag.overview"):
|
|
ok = asyncio.run(regenerate_overview(llm, db))
|
|
assert ok is True
|
|
assert llm.calls == 1
|
|
row = db.get(KbOverview, 1)
|
|
assert row is not None, "the single row must be upserted"
|
|
assert row.id == 1
|
|
assert row.content == REPLY
|
|
assert row.updated_at is not None
|
|
age = datetime.now(UTC) - row.updated_at
|
|
assert age.total_seconds() < 300, "updated_at must be a fresh UTC timestamp"
|
|
assert f"overview: regenerated docs=2 chars={len(REPLY)}" in caplog.text
|
|
finally:
|
|
_truncate_documents(db)
|
|
|
|
|
|
def test_regenerate_updates_existing_row_in_place(db: Session, clean_overview) -> None:
|
|
_truncate_documents(db)
|
|
past = datetime(2020, 1, 1, 12, 0, 0, tzinfo=UTC)
|
|
db.add(KbOverview(id=1, content="old outline", updated_at=past))
|
|
db.commit()
|
|
try:
|
|
_add_doc(db, "Homelab", "a.md", "A")
|
|
ok = asyncio.run(regenerate_overview(_FakeLLM(), db))
|
|
assert ok is True
|
|
count = db.scalar(text("SELECT count(*) FROM kb_overview"))
|
|
assert count == 1, "still exactly one row after the re-generation"
|
|
row = db.get(KbOverview, 1)
|
|
assert row is not None
|
|
assert row.content == REPLY # replaced, not appended
|
|
assert row.updated_at is not None and row.updated_at > past
|
|
finally:
|
|
_truncate_documents(db)
|
|
|
|
|
|
def test_regenerate_orders_documents_by_source_path(db: Session, clean_overview) -> None:
|
|
_truncate_documents(db)
|
|
try:
|
|
# Inserted out of order — the model must see a deterministic order.
|
|
_add_doc(db, "Deployments", "z.md", "Z")
|
|
_add_doc(db, "Homelab", "b.md", "B")
|
|
_add_doc(db, "Homelab", "a.md", "A")
|
|
llm = _FakeLLM()
|
|
assert asyncio.run(regenerate_overview(llm, db)) is True
|
|
finally:
|
|
_truncate_documents(db)
|
|
|
|
user = next(m["content"] for m in llm.messages if m["role"] == "user")
|
|
assert user == (
|
|
"Deployments — z.md — Z\nHomelab — a.md — A\nHomelab — b.md — B"
|
|
)
|
|
|
|
|
|
def test_regenerate_calls_the_configured_summary_model(
|
|
db: Session, clean_overview
|
|
) -> None:
|
|
_truncate_documents(db)
|
|
try:
|
|
_add_doc(db, "Homelab", "a.md", "A")
|
|
llm = _FakeLLM()
|
|
assert asyncio.run(regenerate_overview(llm, db)) is True
|
|
finally:
|
|
_truncate_documents(db)
|
|
|
|
assert llm.model == llm.settings.llm_summary_model # the ``lite`` default
|
|
assert llm.model == "lite"
|
|
assert [m["role"] for m in llm.messages] == ["system", "user"]
|
|
assert KB_OVERVIEW_MODE in llm.messages[0]["content"]
|
|
assert "A" in llm.messages[1]["content"]
|
|
|
|
|
|
def test_regenerate_zero_docs_leaves_existing_row_untouched(
|
|
db: Session, clean_overview
|
|
) -> None:
|
|
db.execute(text("DELETE FROM kb_overview"))
|
|
db.add(KbOverview(id=1, content="old outline"))
|
|
db.commit()
|
|
_truncate_documents(db)
|
|
llm = _FakeLLM()
|
|
ok = asyncio.run(regenerate_overview(llm, db))
|
|
assert ok is False
|
|
assert llm.calls == 0, "no KB → no wasted lite-model call"
|
|
row = db.get(KbOverview, 1)
|
|
assert row is not None and row.content == "old outline" # untouched
|
|
|
|
|
|
def test_regenerate_zero_docs_without_row_creates_nothing(db: Session, clean_overview) -> None:
|
|
_truncate_documents(db)
|
|
llm = _FakeLLM()
|
|
ok = asyncio.run(regenerate_overview(llm, db))
|
|
assert ok is False
|
|
assert llm.calls == 0
|
|
assert _overview_row(db) is None, "no row must be invented for an empty KB"
|
|
|
|
|
|
def test_regenerate_llm_error_leaves_previous_row_intact(
|
|
db: Session, clean_overview, caplog: pytest.LogCaptureFixture
|
|
) -> None:
|
|
db.execute(text("DELETE FROM kb_overview"))
|
|
db.add(KbOverview(id=1, content="old outline"))
|
|
db.commit()
|
|
_truncate_documents(db)
|
|
try:
|
|
_add_doc(db, "Homelab", "a.md", "A")
|
|
llm = _FakeLLM(fail=LLMError("simulated lite-model failure"))
|
|
with caplog.at_level(logging.ERROR, logger="app.rag.overview"):
|
|
ok = asyncio.run(regenerate_overview(llm, db))
|
|
assert ok is False, "fail-soft: the import must not be dragged down"
|
|
row = db.get(KbOverview, 1)
|
|
assert row is not None and row.content == "old outline", (
|
|
"the previous outline stays — an old outline is better than none"
|
|
)
|
|
assert "overview: regeneration failed —" in caplog.text
|
|
assert "simulated lite-model failure" in caplog.text
|
|
finally:
|
|
_truncate_documents(db)
|
|
|
|
|
|
def test_regenerate_llm_error_without_row_creates_nothing(
|
|
db: Session, clean_overview
|
|
) -> None:
|
|
_truncate_documents(db)
|
|
try:
|
|
_add_doc(db, "Homelab", "a.md", "A")
|
|
llm = _FakeLLM(fail=LLMError("simulated lite-model failure"))
|
|
ok = asyncio.run(regenerate_overview(llm, db))
|
|
finally:
|
|
_truncate_documents(db)
|
|
|
|
assert ok is False
|
|
assert _overview_row(db) is None
|
|
|
|
|
|
def test_regenerate_opens_own_session_when_none(db: Session, clean_overview) -> None:
|
|
"""``session=None`` → a private ``SessionLocal`` session is opened,
|
|
committed, and closed (the import-script call path)."""
|
|
_truncate_documents(db)
|
|
try:
|
|
_add_doc(db, "Homelab", "a.md", "A")
|
|
ok = asyncio.run(regenerate_overview(_FakeLLM()))
|
|
finally:
|
|
_truncate_documents(db)
|
|
|
|
assert ok is True
|
|
row = db.get(KbOverview, 1) # fresh lookup via the test session
|
|
assert row is not None and row.content == REPLY
|