chroma/answer.py

"""Answer a question about ONE agenda item in a single LLM call.

Where :mod:`chroma.report` runs a map-reduce over *every* chunk (the right tool for
summarizing a whole meeting), this module is built for the much narrower "ask the
packet" case: the packet has already been sliced down to one item's section, so the
relevant text is small. We assemble the most relevant context once and answer it in a
**single** completion -- far faster than N chunk calls, and more coherent because the
model sees the context whole instead of pre-digested bullet notes.

Scope is still chosen *structurally* upstream (the item->pages slice); the optional
relevance retrieval here only ranks chunks *within* that slice when it is too big to
feed wholesale -- it never decides which item the user is asking about.

The LLM call is injectable via ``complete`` (defaults to
:func:`chroma.llm.chat_complete`) so callers can swap in a fake (or a llama.cpp
streaming completer) for testing.
"""

from __future__ import annotations

import os
from typing import Callable

from .store import AgendaStore

Completer = Callable[..., str]

# How much packet text (characters) to feed the model in one shot. Sized to leave
# room under the local GGUF's context window (LLAMA_N_CTX=8192 tokens) for the
# system/user prompt wrappers and the answer itself.
DEFAULT_CHAR_BUDGET = int(os.getenv("ANSWER_CHAR_BUDGET", "16000"))
# How many chunks to pull from the relevance query before packing to the budget.
DEFAULT_N_RESULTS = int(os.getenv("ANSWER_N_RESULTS", "16"))
# Ceiling on a single-pass item answer's length (the thoroughness slider scales up to
# this). Raise ANSWER_MAX_TOKENS for longer answers — pair it with a bigger LLAMA_N_CTX
# and GPU_DURATION_ANSWER so the generation fits the context + ZeroGPU window.
ANSWER_MAX_TOKENS = int(os.getenv("ANSWER_MAX_TOKENS", "1536"))

ANSWER_SYSTEM = (
    "You are a civic-records analyst answering a resident's question about ONE item on "
    "a public meeting's agenda, using ONLY the packet text you are given. "
    "Write the answer in clean GitHub-flavored markdown:\n"
    "- Open with a one- or two-sentence direct answer in **bold**.\n"
    "- Then use short `##` sub-headings with `-` bullet lists for the specifics that matter.\n"
    "- Quote concrete dollar amounts, dates, deadlines, vote/motion language, and agenda "
    "item numbers verbatim when they appear.\n"
    "- Never invent facts or use outside knowledge. If the packet does not address the "
    "question, say so plainly in a single sentence.\n"
    "No preamble, no \"Based on the packet…\", no restating the question, no closing summary."
)
ANSWER_PROMPT = (
    "Agenda item: {title}\n\n"
    "Resident's question:\n{query}\n\n"
    "Packet text for this item (its supporting documents):\n"
    '"""\n{context}\n"""\n\n'
    "Write the markdown answer now."
)


def _default_completer() -> Completer:
    from .llm import chat_complete

    return chat_complete


def _invoke(complete: Completer, prompt: str, system: str, max_tokens: int) -> str:
    """Call ``complete``, degrading gracefully if it ignores system/max_tokens."""
    try:
        return complete(prompt, system=system, max_tokens=max_tokens)
    except TypeError:
        pass
    try:
        return complete(prompt, system=system)
    except TypeError:
        return complete(prompt)


def budget_for(max_sections: int) -> tuple[int, int]:
    """Map the UI "thoroughness" knob onto ``(char_budget, answer_max_tokens)``.

    Higher thoroughness feeds more packet context and allows a longer answer. The
    char budget is capped so the prompt still fits the local model's context window.
    """
    n = int(max_sections or 0)
    char_budget = min(DEFAULT_CHAR_BUDGET, max(6000, n * 320))
    # Scale so the UI's max thoroughness (~120) reaches ANSWER_MAX_TOKENS.
    max_tokens = min(ANSWER_MAX_TOKENS, max(640, round(n * ANSWER_MAX_TOKENS / 120)))
    return char_budget, max_tokens


def _split_docs(documents: list[dict]) -> tuple[str, list[dict]]:
    """Separate the focusing item-title doc from the packet body docs.

    The item-report loader (:func:`webapp.backend._item_documents`) prepends a tiny
    ``…:item`` doc holding the selected item's title; the real content is the
    ``…:packet-slice`` / ``…:packet`` / ``…:agenda`` doc(s).
    """
    title = ""
    body: list[dict] = []
    for d in documents:
        doc_id = str(d.get("doc_id", ""))
        text = (d.get("text") or "").strip()
        if doc_id.endswith(":item"):
            # text looks like "Selected agenda item: <title>"
            title = text.split(":", 1)[1].strip() if ":" in text else text
        elif text:
            body.append(d)
    return title, body


def pick_engine(documents: list[dict], *, char_budget: int = DEFAULT_CHAR_BUDGET) -> str:
    """Choose ``"single"`` or ``"mapreduce"`` for these item documents.

    Single-pass when the body already fits one context window, or when the item was
    structurally located (a ``…:packet-slice`` doc) and is only modestly over budget —
    there, relevance retrieval still surfaces the chunks that matter. Map-reduce only
    when the text is *many* times the window (a located slice that big, or the
    un-located full-packet fallback), where one-shot retrieval would silently drop
    most of the content.
    """
    located = any(str(d.get("doc_id", "")).endswith(":packet-slice") for d in documents)
    _, body = _split_docs(documents)
    body_chars = sum(len(d.get("text") or "") for d in body)
    if body_chars <= char_budget:
        return "single"
    if located and body_chars <= 2 * char_budget:
        return "single"
    return "mapreduce"


def resolve_engine(documents: list[dict], engine: str = "auto", *,
                   char_budget: int = DEFAULT_CHAR_BUDGET) -> str:
    """Honor an explicit ``engine`` choice, else fall back to :func:`pick_engine`.

    ``engine`` is the user's selection from the UI: ``"single"`` (semantic search /
    one-shot answer) or ``"mapreduce"`` (read every section). Anything else
    (``"auto"``, ``""``) defers to the size-based heuristic.
    """
    if engine in ("single", "mapreduce"):
        return engine
    return pick_engine(documents, char_budget=char_budget)


def assemble_context(
    documents: list[dict],
    query: str,
    *,
    store: AgendaStore | None = None,
    char_budget: int = DEFAULT_CHAR_BUDGET,
    n_results: int = DEFAULT_N_RESULTS,
) -> tuple[str, dict]:
    """Build the single-pass context string from the item's documents.

    If the packet body fits ``char_budget`` it is used whole (no embedding needed --
    the fastest path). Otherwise the body is chunked + embedded into an ephemeral
    Chroma collection, the chunks most relevant to ``query`` are retrieved, packed up
    to the budget, and re-sorted into reading order for coherence.

    Returns ``(context, info)`` where ``info`` carries ``title`` plus
    ``{retrieved, used_chunks, total_chunks, context_chars}`` for UI notes / eval.
    """
    title, body = _split_docs(documents)
    info: dict = {"title": title, "retrieved": False, "used_chunks": 0,
                  "total_chunks": 0, "context_chars": 0}
    if not body:
        return "", info

    body_text = "\n\n".join((d.get("text") or "").strip() for d in body).strip()
    if len(body_text) <= char_budget:
        info["context_chars"] = len(body_text)
        return body_text, info

    # Too big for one shot: rank chunks by relevance, pack to budget, restore order.
    store = store if store is not None else AgendaStore(path=None)
    store.add_documents(body)
    all_chunks = store.get_chunks()
    info["total_chunks"] = len(all_chunks)
    hits = store.query(query, n_results=min(n_results, max(1, len(all_chunks))))

    selected: list[dict] = []
    used = 0
    for h in hits:
        text = h.get("text") or ""
        if not text:
            continue
        if used and used + len(text) > char_budget:
            continue
        selected.append(h)
        used += len(text)
        if used >= char_budget:
            break

    def _order(c: dict) -> tuple:
        m = c.get("metadata") or {}
        return (str(m.get("doc_id", "")), m.get("chunk_index", 0))

    selected.sort(key=_order)
    context = "\n\n".join((c.get("text") or "").strip() for c in selected).strip()
    info.update(retrieved=True, used_chunks=len(selected), context_chars=len(context))
    return context, info


def answer_question(
    documents: list[dict],
    query: str,
    *,
    complete: Completer | None = None,
    char_budget: int = DEFAULT_CHAR_BUDGET,
    n_results: int = DEFAULT_N_RESULTS,
    max_tokens: int = 1024,
    progress: Callable[[float, str], None] | None = None,
) -> str:
    """Answer ``query`` about one agenda item in a single LLM call (non-streaming).

    Used by the remote backend and the eval harness; the ZeroGPU path streams the
    same prompt via :func:`webapp.local_llm.gpu_answer`.
    """
    query = (query or "").strip()
    if not query:
        return "_Enter a question to ask about this item._"
    if not documents:
        return "_No agenda text was available for this item (the packet may be empty or failed to download)._"

    complete = complete or _default_completer()
    if progress:
        progress(0.1, "Finding the most relevant pages…")
    context, info = assemble_context(
        documents, query, char_budget=char_budget, n_results=n_results)
    if not context.strip():
        return "_No extractable text found for this item._"

    if progress:
        progress(0.45, "Writing the answer…")
    out = _invoke(
        complete,
        ANSWER_PROMPT.format(title=info["title"], query=query, context=context),
        ANSWER_SYSTEM,
        max_tokens,
    )
    if progress:
        progress(1.0, "Done")
    return (out or "").strip()