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

"""Folder summary generator (phase 94, task 01).

``ls`` becomes a drill-down tree (phase 94, ``00_phase.md``): the LLM
lists the synced projects, then drills into folders. Each level's
listing shows the folder's **subtree summary** — a 1–3 sentence
plain-text description of what the folder's documents cover, generated
at SYNC time by the aipi ``lite`` model (the ``KB_OVERVIEW_MODE``
contract of ``app.rag.overview``: change-gated by the caller,
fail-soft — an old summary is better than none — and E2E-mockable via
the ``FOLDER_SUMMARY_MODE`` system-prompt marker).

ONE concept end to end (:func:`group_by_folder`): a folder row's
documents are its **recursive subtree** — every document whose path
equals the folder or starts with ``folder + "/"`` — exactly the set the
``ls`` count rule (``00_phase.md``) counts. A document under ``a/b/``
therefore contributes to the ``a``, ``a/b``, and ``""`` (source root)
groups alike: each level's listing shows its own accurate subtree
summary, and the number next to a folder is the number of documents its
summary describes.

Storage: ``folder_summaries`` (migration 0017) — PK
``(source, folder_path)``; ``folder_path = ""`` is the SOURCE ROOT
(the top-level source summary). Rows exist only for folders with
≥ 2 documents (the :data:`MIN_DOCS_PER_FOLDER` rule — a
single-document folder is fully described by its one file line, so no
``lite`` burn); after a changed sync, rows whose folder dropped below
2 documents are pruned (a pruned/renamed folder's summary would
otherwise go stale), while rows for folders that still have
≥ 2 documents persist (an unchanged folder's summary is still true).

Chat turns never generate folder summaries — the agent's ``ls`` output
(phase 94, task 03) only reads the stored rows. Generation is the
caller's job at sync time (phase 94, task 02), and :func:`generate_
folder_summaries` only flushes — the sync path owns the transaction
(the phase-53 ``bump_sources_version`` convention).

Self-heal (phase 96): an exhausted one-shot retry can still leave a
candidate folder without a row. :func:`missing_folder_summaries`
names those gaps (the generator's exact candidate computation minus
the stored keys — one concept, as with :func:`group_by_folder`), and
:func:`generate_folder_summaries`'s ``only_missing=True`` fills
EXACTLY those on the next sync — existing rows stay
byte-identical (text AND ``updated_at``), the prune pass still runs.
"""
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.models import Document, FolderSummary
from app.rag.llm import LLMError
from app.rag.retriever import TRUNCATION_MARKER

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

#: System-prompt marker for folder-summary generation — the E2E mock
#: LLM keys on it (same convention as ``SUMMARY_MODE`` /
#: ``KB_OVERVIEW_MODE``, PLAN §6).
FOLDER_SUMMARY_MODE = "FOLDER_SUMMARY_MODE"

#: Locked instruction for the ``lite`` model (phase 94): the folder
#: summary is the drill-down ``ls``'s per-level picture, so it must be
#: a short plain-text description of what the folder's documents
#: cover — 1–3 sentences (a folder is a skim, not a read: the document
#: summary's 3–6 sentences would blow up a 20-folder listing),
#: strictly grounded in the listed titles/paths/summary lines.
FOLDER_SUMMARY_INSTRUCTION = (
    "From the document list below, write a 1-3 sentence plain-text "
    "summary of what this folder's documents cover, in natural "
    "language. Do not use markdown. Do not invent anything that is not "
    "in the list."
)

#: Full system prompt: marker first (the mock's key), then the
#: instruction (the ``app.rag.overview.SYSTEM_PROMPT`` shape).
SYSTEM_PROMPT = f"{FOLDER_SUMMARY_MODE}: {FOLDER_SUMMARY_INSTRUCTION}"

#: First line of the user message — the folder the summary describes:
#: ``<source>`` for the source root, ``<source>/<folder_path>`` for a
#: folder. The deterministic E2E mock keys on it to name the folder in
#: its canned reply (``tests/e2e/mock_llm.py``), so the stored row is
#: a pure function of the request.
FOLDER_HEADER_PREFIX = "Folder: "

#: A folder is summarized only while its recursive subtree holds at
#: least this many documents — a single-document folder is fully
#: described by its one file line, so no ``lite`` burn (phase 94
#: ``00_phase.md`` scope rule; the prune rule applies the same count).
MIN_DOCS_PER_FOLDER = 2

#: One document row for the grouping/prompting:
#: ``(source, path, title, summary)`` — the ``app.rag.overview``
#: ``build_overview_prompt`` row shape (``summary`` is the stored
#: lite-written text or ``None`` for markdown docs / the fail-soft
#: path).
DocRow = tuple[str, str, str, str | None]


class FolderSummaryLLM(Protocol):
    """The one-shot chat surface the folder-summary generator needs.

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

    settings: Settings

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


def folder_of(path: str) -> str:
    """The directory prefix before the last ``/`` (``""`` for root-level files).

    The single notion of "which folder owns this document" used by the
    whole module: ``folder_of("a.md") == ""``,
    ``folder_of("a/b.md") == "a"``, ``folder_of("a/b/c.md") == "a/b"``.
    Iterating :func:`folder_of` over its own result walks a path's
    folder prefixes from nearest to farthest, ending at ``""``.
    """
    idx = path.rfind("/")
    return path[:idx] if idx >= 0 else ""


def group_by_folder(rows: Sequence[DocRow]) -> dict[tuple[str, str], list[DocRow]]:
    """Group document rows by the folders their recursive subtree fills.

    The ONE concept of this module, documented in the module docstring:
    a folder row's documents are its **recursive subtree** — every
    document whose path equals the folder or starts with
    ``folder + "/"`` — exactly the set the ``ls`` count rule counts.
    Each row therefore lands in the ``""`` (source root) group AND in
    the group of every folder prefix of its path: a doc under
    ``a/b/`` contributes to the ``a``, ``a/b``, and ``""`` groups.
    Per source the candidate rows are thus ``""`` (all of the source's
    documents — the top-level source summary) plus every distinct
    folder prefix of an indexed path.

    *rows* are ``(source, path, title, summary)`` tuples (the
    :data:`DocRow` shape, e.g. straight from the ``documents``
    catalogue query). Returns ``{(source, folder_path): [rows]}`` —
    group lists keep the input (catalogue) order, so downstream prompt
    building is deterministic. Groups of ANY size (≥ 1) are returned;
    the ≥ 2 :data:`MIN_DOCS_PER_FOLDER` rule is applied by
    :func:`generate_folder_summaries`, not here.

    """
    # Pass 1: the distinct TRUE folder prefixes of the catalogue — a
    # folder is a slash-boundary prefix of at least one indexed path
    # (the ``00_phase.md`` candidate definition: ``""`` + every distinct
    # folder prefix, per source).
    folders: set[tuple[str, str]] = set()
    for row in rows:
        source, path = row[0], row[1]
        folder = folder_of(path)
        while folder:
            folders.add((source, folder))
            folder = folder_of(folder)

    # Pass 2: every row lands in the source-root group, in the group of
    # every ancestor folder (the ``startswith folder + "/"`` arm of the
    # count rule), and — when its path IS one of the source's folder
    # prefixes (a file sharing its name with a directory) — in that
    # folder's group too (the ``path == folder`` arm). A row therefore
    # belongs to its folder's group iff its path equals the folder or
    # starts with ``folder + "/"`` — exactly the set the ``ls`` count
    # rule counts, for EVERY group.
    groups: dict[tuple[str, str], list[DocRow]] = {}
    for row in rows:
        source, path = row[0], row[1]
        groups.setdefault((source, ""), []).append(row)
        if (source, path) in folders:
            groups.setdefault((source, path), []).append(row)
        folder = folder_of(path)
        while folder:
            groups.setdefault((source, folder), []).append(row)
            folder = folder_of(folder)
    return groups


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
    folder summary. Blank/whitespace-only summaries yield ``''`` as
    well (the ``app.rag.overview`` helper, mirrored locally so each
    lite-mode module stays self-contained).
    """
    if not summary:
        return ""
    for line in summary.splitlines():
        stripped = line.strip()
        if stripped:
            return stripped
    return ""


def _folder_label(source: str, folder_path: str) -> str:
    """The user-message folder identifier (the ``FOLDER_HEADER_PREFIX`` tail).

    ``<source>`` for the source root (``folder_path = ""``),
    ``<source>/<folder_path>`` for a folder — the canonical folder
    identity the E2E mock echoes into its canned reply.
    """
    return source if not folder_path else f"{source}/{folder_path}"


def build_folder_summary_prompt(
    source: str,
    folder_path: str,
    docs: Sequence[DocRow],
    max_chars: int | None = None,
) -> tuple[str, str]:
    """The ``(system, user)`` message pair for one folder-summary call.

    * ``system`` — :data:`SYSTEM_PROMPT`: the ``FOLDER_SUMMARY_MODE``
      marker + the locked instruction.
    * ``user`` — the folder header line
      (``FOLDER_HEADER_PREFIX + <source>[/<folder_path>]`` — the line
      the E2E mock parses to name the folder), then one line per
      document, ``path — title — {first line of summary}``, joined
      with newlines, in the given (catalogue) order. The summary field
      is omitted when the document has none (markdown docs and the
      fail-soft path — no dangling dash, the overview convention). The
      whole message is capped at *max_chars* (default
      ``BOR_FOLDER_SUMMARY_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).
    """
    lines = [FOLDER_HEADER_PREFIX + _folder_label(source, folder_path)]
    for _source, path, title, summary in docs:
        line = f"{path} — {title}"
        first = _first_summary_line(summary)
        if first:
            line += f" — {first}"
        lines.append(line)
    user = "\n".join(lines)
    limit = (
        max_chars
        if max_chars is not None
        else get_settings().folder_summary_input_max_chars
    )
    if len(user) > limit:
        user = user[:limit] + "\n" + TRUNCATION_MARKER
    return SYSTEM_PROMPT, user


async def summarize_folder(
    source: str, folder_path: str, docs: Sequence[DocRow], llm: FolderSummaryLLM
) -> str:
    """One-shot ``lite`` summary of one folder's recursive subtree.

    Builds the prompt from the folder's documents (each ``path``,
    ``title``, first summary line), makes ONE :meth:`LLMClient.chat`
    call against ``llm.settings.llm_summary_model`` (the ``lite``
    model — no new model management), and returns the model's text
    trimmed. Raises :class:`LLMError` when the reply is empty after
    trimming (the client already rejects empty content; re-asserted
    defensively, the summarizer's rule — a silent summary must never be
    stored), and propagates any :class:`LLMError` the client raises
    (the generator's fail-soft path catches it per folder).
    """
    system, user = build_folder_summary_prompt(source, folder_path, docs)
    raw = await llm.chat(
        [{"role": "system", "content": system}, {"role": "user", "content": user}],
        model=llm.settings.llm_summary_model,
    )
    summary = raw.strip()
    if not summary:
        label = _folder_label(source, folder_path)
        raise LLMError(
            f"folder summary model returned empty content for {label} — "
            "refusing to store a silent summary"
        )
    return summary


def _upsert(db: Session, source: str, folder_path: str, summary: str) -> None:
    """Insert or update the row for one folder (fresh UTC stamp)."""
    now = datetime.now(UTC)
    row = db.get(FolderSummary, (source, folder_path))
    if row is None:
        db.add(
            FolderSummary(
                source=source, folder_path=folder_path, summary=summary, updated_at=now
            )
        )
    else:
        row.summary = summary
        row.updated_at = now


def _catalog_rows(db: Session) -> list[DocRow]:
    """The document catalogue in catalogue order (one query).

    ``(source, path, title, summary)`` ordered by ``(source, path)`` —
    the EXACT query :func:`generate_folder_summaries` runs, so the gap
detector and the fill can never disagree about what the catalogue
    contains (phase 96, task 02: one concept, as with
    :func:`group_by_folder`).
    """
    result = db.execute(
        select(Document.source, Document.path, Document.title, Document.summary)
        .order_by(Document.source, Document.path)
    ).all()
    return [
        (source, path, title, summary)
        for source, path, title, summary in result
    ]


def _candidates(rows: Sequence[DocRow]) -> dict[tuple[str, str], list[DocRow]]:
    """The generator's candidate map (recursive subtree ≥ 2 docs).

    The :func:`group_by_folder` groups filtered by
    :data:`MIN_DOCS_PER_FOLDER` — the EXACT set a full regeneration
    would cover, so a "missing" folder can never disagree with what a
    regeneration would (re)generate.
    """
    groups = group_by_folder(rows)
    return {
        key: docs for key, docs in groups.items() if len(docs) >= MIN_DOCS_PER_FOLDER
    }


def _stored_keys(db: Session) -> set[tuple[str, str]]:
    """The ``(source, folder_path)`` keys that already hold a row.

    One bounded key select — no count scan (the phase-94
    table-emptiness probe's house pattern).
    """
    return {
        (source, folder_path)
        for source, folder_path in db.execute(
            select(FolderSummary.source, FolderSummary.folder_path)
        ).all()
    }


def missing_folder_summaries(db: Session) -> list[tuple[str, str]]:
    """The candidate folders whose stored row is ABSENT (phase 96, 02).

    The unchanged-sync self-heal gate: a one-shot reply can still
    exhaust its retries and leave a candidate folder without a row, and
    the phase-94 change gate only regenerated on a KB change — so the
gap persisted until the next KB change. This function names the gap:
    the generator's candidate folders (the EXACT candidate computation
    — one catalogue query in catalogue order, :func:`group_by_folder`,
    recursive subtree ≥ :data:`MIN_DOCS_PER_FOLDER`) minus the stored
    keys (one bounded select). The gate probe and the
    :func:`generate_folder_summaries` fill both key off this one
    concept.

    Returns the missing keys sorted by ``(source, folder_path)``;
    ``[]`` when there is no gap — including an empty catalogue over an
    empty table (no candidates, no gap). A single-document folder is
    never listed (it is not a candidate), and a stored row for a folder
    that dropped below the minimum is NOT missing (it is stale — the
    prune pass owns it).
    """
    return sorted(
        key for key in _candidates(_catalog_rows(db)) if key not in _stored_keys(db)
    )


async def generate_folder_summaries(
    db: Session,
    llm: FolderSummaryLLM,
    *,
    skip: bool = False,
    only_missing: bool = False,
) -> dict[str, int]:
    """Regenerate the stored folder summaries for the current catalogue.

    The sync-path orchestrator (phase 94, task 02 calls it
    change-gated, like the KB overview; ``--limit`` debug runs pass
    ``skip=True``). Steps, in order:

    1. ``skip=True`` → a no-op: the zero stats dict is returned, the
       LLM is never called, and no rows are touched.
    2. Group the document catalogue by :func:`group_by_folder` (the
       recursive-subtree concept) and keep the candidate folders —
       the ones whose recursive subtree holds
       :data:`MIN_DOCS_PER_FOLDER` (≥ 2) documents. Single-document
       folders get no row (their one file line IS their summary).
    3. For each candidate (deterministic ``(source, folder_path)``
       order) call :func:`summarize_folder` and UPSERT — fail-soft PER
       FOLDER: one folder's :class:`LLMError` is logged and counted,
       its previous row (if any) is kept, and the remaining folders
       still land (a ``lite`` outage must never fail the sync — the KB
       is the product, the summaries are auxiliary). With
       ``only_missing=True`` the iteration is restricted to the
       candidates that have NO stored row (the same gap
       :func:`missing_folder_summaries` reports, derived from the SAME
       catalogue pass — one bounded stored-key select, no second
       catalogue query): existing rows stay byte-identical (summary
       text AND ``updated_at`` — never re-stamped, even a stale-looking
       one; staleness is the changed-KB regeneration's job) and no
       ``lite`` call is burned for a folder that already has a summary
       — the unchanged-sync self-heal fill (phase 96, task 02).
    4. DELETE rows whose folder no longer has ≥ 2 documents — a
       pruned/renamed folder's summary goes stale and is dropped.
       Rows for folders that still qualify persist (regenerated in
       step 3 — an unchanged folder's summary is still true). The
       prune pass runs in BOTH modes: under ``only_missing`` on an
       unchanged catalogue it is a no-op (the invariant kept), and it
       still drops rows whose folder fell below the minimum.

    Only flushes — the CALLER commits (the phase-53
    ``bump_sources_version`` convention: the sync path owns the
    transaction, so a failed sync rolls the summaries back with it).

    Returns the small stats dict ``{"generated", "failed", "pruned"}``
    for the caller's summary-line logging (PLAN §9 ample logging) —
    the caller logs the mode, not the generator.
    """
    stats = {"generated": 0, "failed": 0, "pruned": 0}
    if skip:
        return stats

    candidates = _candidates(_catalog_rows(db))
    keys = sorted(candidates)
    if only_missing:
        stored = _stored_keys(db)
        keys = [key for key in keys if key not in stored]

    for key in keys:
        source, folder_path = key
        docs = candidates[key]
        try:
            summary = await summarize_folder(source, folder_path, docs, llm)
        except LLMError as e:
            stats["failed"] += 1
            logger.error(
                "folder summary failed for %s/%s — %s", source, folder_path, e
            )
            continue
        _upsert(db, source, folder_path, summary)
        stats["generated"] += 1

    existing = db.execute(
        select(FolderSummary.source, FolderSummary.folder_path)
    ).all()
    for source, folder_path in existing:
        if (source, folder_path) not in candidates:
            row = db.get(FolderSummary, (source, folder_path))
            if row is not None:
                db.delete(row)
                stats["pruned"] += 1

    db.flush()
    logger.info(
        "folder_summaries: generated=%d failed=%d pruned=%d",
        stats["generated"],
        stats["failed"],
        stats["pruned"],
    )
    return stats