203 lines
8.0 KiB
Python
203 lines
8.0 KiB
Python
"""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 (markdown docs and the fail-soft path — 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()
|