brain-of-reese/app/rag/overview.py

"""KB overview generator (phase 31, task 02).

Builds the ``KB_OVERVIEW_MODE`` prompt from the document catalogue
(source, path, title, first summary line per document), calls the aipi
``lite`` model through the one-shot ``LLMClient.chat`` (phase 30, same
model as document summaries — no new model management), and stores the
plain-text outline in the single ``kb_overview`` row (migration 0005).

Chat turns never generate an overview — they read the stored row
(:func:`load_kb_overview`, one indexed PK lookup) and inject it into the
system prompt of every turn as the ``<knowledge_base>`` section
(phase 31, task 03). Generation happens at import time (phase 31, task
04) and is **best-effort and change-gated**: callers invoke it only when
an import added/updated at least one document (or no row exists yet),
and a ``lite`` failure logs and leaves the previous outline intact — an
old outline is better than none (phase locked decisions).

The ``KB_OVERVIEW_MODE`` marker follows the ``SUMMARY_MODE`` /
``DEFLECT_MODE`` convention: the deterministic E2E mock LLM keys on it
in the system prompt (wired in phase 31, task 05).
"""
from __future__ import annotations

import logging
from collections.abc import Sequence
from datetime import UTC, datetime
from typing import Protocol

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.config import Settings, get_settings
from app.db import SessionLocal
from app.models import Document, KbOverview
from app.rag.llm import LLMError
from app.rag.retriever import TRUNCATION_MARKER

logger = logging.getLogger("app.rag.overview")

#: System-prompt marker for KB overview generation — the E2E mock LLM keys
#: on it (same convention as ``SUMMARY_MODE`` / ``DEFLECT_MODE``).
KB_OVERVIEW_MODE = "KB_OVERVIEW_MODE"

#: Locked instruction for the ``lite`` model (phase 31): the overview is
#: the agent's *a priori* picture of the knowledge base, so it must be a
#: compact plain-text outline of basic categories and topics — strictly
#: grounded in the document list, nothing invented.
OVERVIEW_INSTRUCTION = (
    "From the document list below, write a compact plain-text outline of "
    "the basic categories and topics this knowledge base covers. Group by "
    "source where useful, use `-` bullet lines, at most ~1500 characters, "
    "no markdown headings, and no topics not present in the list."
)

#: Full system prompt: marker first (the mock's key), then the instruction.
SYSTEM_PROMPT = f"{KB_OVERVIEW_MODE}: {OVERVIEW_INSTRUCTION}"


class OverviewLLM(Protocol):
    """The one-shot chat surface the overview generator needs.

    :class:`app.rag.llm.LLMClient` satisfies it; unit tests pass a
    duck-typed fake (``chat`` + ``settings``) instead — same pattern as
    the summarizer's ``SummaryLLM`` protocol.
    """

    settings: Settings

    async def chat(
        self, messages: list[dict[str, str]], model: str | None = None
    ) -> str: ...


def _first_summary_line(summary: str | None) -> str:
    """First line of a stored summary, stripped; ``''`` when absent.

    The stored summary (phase 30) ends in the code-appended
    ``Source: …`` pointer line, so its first line is the model-written
    lead sentence — the best one-line picture of the document for the
    outline. Blank/whitespace-only summaries yield ``''`` as well.
    """
    if not summary:
        return ""
    for line in summary.splitlines():
        stripped = line.strip()
        if stripped:
            return stripped
    return ""


def build_overview_prompt(
    rows: Sequence[tuple[str, str, str, str | None]], max_chars: int | None = None
) -> tuple[str, str]:
    """The ``(system, user)`` message pair for one overview call.

    Each row is ``(source, path, title, summary)``.

    * ``system`` — :data:`SYSTEM_PROMPT`: the ``KB_OVERVIEW_MODE`` marker
      + the locked instruction.
    * ``user`` — one line per document,
      ``source — path — title — {first line of summary}``, joined with
      newlines. The summary field is omitted when the document has no
      summary (pre-phase-30 rows, the fail-soft path, and pre-backfill
      NULL rows — no dangling dash). The list is capped at *max_chars*
      (default
      ``BOR_OVERVIEW_INPUT_MAX_CHARS``): overflow is cut exactly at the
      cap and the shared ``[…truncated…]`` marker is appended on its own
      line, so the model never sees more than the cap and the cut is
      visible (summarizer convention, phase 30).
    """
    limit = max_chars if max_chars is not None else get_settings().overview_input_max_chars
    lines: list[str] = []
    for source, path, title, summary in rows:
        line = f"{source} — {path} — {title}"
        first = _first_summary_line(summary)
        if first:
            line += f" — {first}"
        lines.append(line)
    user = "\n".join(lines)
    if len(user) > limit:
        user = user[:limit] + "\n" + TRUNCATION_MARKER
    return SYSTEM_PROMPT, user


def load_kb_overview(db: Session) -> str:
    """The stored KB outline for prompt injection (phase 31, task 03).

    One indexed PK lookup (the single row, ``id = 1``) returning
    ``content`` trimmed. Returns ``""`` when the row is missing or
    empty — the chat path then omits the ``<knowledge_base>`` section
    entirely, keeping the prompt byte-identical to the pre-phase text
    (phase 15 convention).
    """
    row = db.get(KbOverview, 1)
    if row is None:
        return ""
    return (row.content or "").strip()


async def regenerate_overview(
    llm: OverviewLLM, session: Session | None = None
) -> bool:
    """Generate a fresh KB outline with the ``lite`` model and upsert it.

    Best-effort and change-gated by the caller (phase locked decisions):
    the import script (phase 31, task 04) and the admin sync (phase 32)
    invoke this only when an import added/updated at least one document
    (or no row exists yet). *session* may be supplied (the import
    script's own session, tests); a private one is opened and closed
    otherwise — same convention as ``import_sources``.

    * Zero documents → the existing row (if any) is left untouched and
      ``False`` is returned — no KB, no outline, no wasted model call.
    * Happy path → the model's text is upserted into the single row
      (``id = 1``, fresh UTC ``updated_at``), committed, and
      ``overview: regenerated docs=… chars=…`` is logged; returns
      ``True``. Idempotent by construction: one row, always ``id = 1``.
    * :class:`LLMError` → ``overview: regeneration failed — …`` is
      logged and ``False`` is returned with the previous row intact —
      an old outline is better than none.
    """
    owns_session = session is None
    if session is None:
        session = SessionLocal()
    try:
        result = session.execute(
            select(Document.source, Document.path, Document.title, Document.summary)
            .order_by(Document.source, Document.path)
        ).all()
        rows: list[tuple[str, str, str, str | None]] = [
            (source, path, title, summary)
            for source, path, title, summary in result
        ]
        if not rows:
            logger.info(
                "overview: no documents in the knowledge base — "
                "leaving any existing outline untouched"
            )
            return False

        system, user = build_overview_prompt(rows)
        try:
            text = await llm.chat(
                [{"role": "system", "content": system}, {"role": "user", "content": user}],
                model=llm.settings.llm_summary_model,
            )
        except LLMError as e:
            logger.error("overview: regeneration failed — %s", e)
            return False

        now = datetime.now(UTC)
        row = session.get(KbOverview, 1)
        if row is None:
            session.add(KbOverview(id=1, content=text, updated_at=now))
        else:
            row.content = text
            row.updated_at = now
        session.commit()
        logger.info("overview: regenerated docs=%d chars=%d", len(rows), len(text))
        return True
    finally:
        if owns_session:
            session.close()