interactive-story/backend/app/knowledge/inject.py

"""M7: fitting retrieved knowledge into the prompt, and saying what it cost.

`retrieval.py` decides which passages are worth offering. This module decides
how many of them the prompt can actually afford, renders them with the framing
their class carries, and produces the provenance record the Insights panel and
the acceptance tests read.

It is pure. It takes a `retrieval.Result`, a budget and a token counter, and
returns text — no database, no session, no clock. That is what lets
`context/builder.py` import it without the import cycle a fuller dependency
would create, and it is why the whole budget arithmetic is testable without a
campaign.

## The pressure rules

`CONTEXT-AND-MEMORY.md` §29-31 and §37-40 of the design ask for four different
behaviours under pressure, and they are four different mechanisms here:

    always-included Canon   protected. Counted with the system block, before
                            any history is chosen. If it cannot fit alongside
                            the other protected sections and the reply reserve,
                            the turn fails with `ContextOverflow` rather than
                            sending a prompt known to overflow.
    retrieved Canon         bounded, and first in line for the retrieved budget.
    Reference               bounded, and capped at a share of it, so Reference
                            can never crowd out Canon.
    Inspiration             capped smallest, filled last, dropped first.

Every one of those is spent out of `KNOWLEDGE_SHARE` of what is left after the
protected context and the reply reserve are subtracted, so none of it can reach
the current state, the reader's input, the narrator rules or the output reserve.
Whatever is not spent returns to the story history rather than being lost.

## Rendering

Each passage arrives labelled with the file it came from, its heading trail and
its index, because that label is the provenance the reader inspects and it is
also what lets a narrator say where something came from. Hidden passages carry
`[narrator only]` on that same line — in the passage, not only in a preamble at
the top of the section, because a passage is read where it sits.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable

from . import classes
from .records import Candidate, Result

#: Share of the non-protected budget that retrieved knowledge may spend.
#:
#: Story cards already take up to 40% (`CARD_BUDGET_SHARE`), and the history is
#: what is left. A third is enough for several passages at the chunker's
#: typical size and leaves the majority of the window to the story itself,
#: which is the thing the reader came for.
KNOWLEDGE_SHARE = 0.33

#: What each class may take of the knowledge budget. Canon may take all of it;
#: the other two are capped so that they cannot, whatever they score.
CLASS_SHARE = {
    classes.CANON: 1.00,
    classes.REFERENCE: 0.50,
    classes.INSPIRATION: 0.25,
}

#: A ceiling on always-included Canon, as a share of the whole context budget.
#:
#: `always_include` is the one place a reader can put unbounded text into every
#: prompt, and it must not be allowed to consume the whole context window
#: (`IMPORTED-KNOWLEDGE-DESIGN.md` §32, `CONTEXT-AND-MEMORY.md` §29). It does
#: not fail silently either: what does not fit is
#: reported as dropped, with its token cost, in the same record everything else
#: appears in.
ALWAYS_SHARE = 0.20

#: The order classes are filled in, highest authority first.
FILL_ORDER = (classes.CANON, classes.REFERENCE, classes.INSPIRATION)


@dataclass
class Section:
    label: str
    text: str


@dataclass
class Plan:
    """A retrieval result, priced and ready to be cut to a budget."""

    result: Result
    count_tokens: Callable[[str], int]
    #: Sections for the system block: the untrusted-data rule and the Canon
    #: this campaign has marked as always in force.
    protected: list[Section] = field(default_factory=list)
    protected_tokens: int = 0
    _always_used: list[Candidate] = field(default_factory=list)
    _always_dropped: list[Candidate] = field(default_factory=list)
    _live_used: list[Candidate] = field(default_factory=list)
    _live_dropped: list[Candidate] = field(default_factory=list)
    _budget: int = 0
    _spent: int = 0


def plan(
    result: Result, count_tokens: Callable[[str], int], context_budget: int
) -> Plan:
    """Prices the protected half: the framing rule and always-included Canon.

    Called before the builder knows how much history it can afford, because the
    answer depends on this.
    """
    ready = Plan(result=result, count_tokens=count_tokens)
    if not result.candidates and not result.suppressed:
        return ready

    always = [c for c in result.candidates if c.always_include]
    others = [c for c in result.candidates if not c.always_include]

    # The rule is emitted whenever anything at all will be shown, including when
    # only always-included Canon survives. A framed section with no frame is the
    # failure mode this section exists to prevent.
    if not always and not others:
        return ready

    rule = classes.KNOWLEDGE_RULE
    if any(c.visibility == classes.HIDDEN for c in result.candidates):
        rule = f"{rule}\n{classes.HIDDEN_RULE}"
    ready.protected.append(Section(classes.SECTION_RULE, rule))

    if always:
        cap = max(0, int(context_budget * ALWAYS_SHARE))
        lines: list[str] = []
        spent = 0
        for candidate in always:
            rendered = render(candidate)
            cost = count_tokens(rendered) + count_tokens("\n\n")
            if spent + cost > cap:
                ready._always_dropped.append(candidate)
                continue
            lines.append(rendered)
            spent += cost
            ready._always_used.append(candidate)
        if lines:
            body = "\n\n".join([classes.ALWAYS_FRAMING] + lines)
            ready.protected.append(Section(classes.SECTION_ALWAYS_CANON, body))
    ready.protected_tokens = sum(count_tokens(s.text) for s in ready.protected)
    return ready


def select(ready: Plan, available: int) -> list[Section]:
    """Fills the retrieved-knowledge budget out of `available`. Returns sections.

    `available` is what the context builder has left for everything elastic, so
    only `KNOWLEDGE_SHARE` of it is spendable here — the remainder belongs to
    the story history and is left untouched.

    Classes are filled in authority order, each against its own cap and against
    what is left. A passage that does not fit is recorded as dropped rather than
    dropped silently: a reader asking "why is that not in the prompt?" gets
    "there was no budget for it", with the number.
    """
    ready._budget = budget = max(0, int(available * KNOWLEDGE_SHARE))
    candidates = [c for c in ready.result.candidates if not c.always_include]
    if not candidates or budget <= 0:
        ready._live_dropped.extend(candidates)
        return []

    separator_cost = ready.count_tokens("\n\n")
    sections: list[Section] = []
    spent = 0
    for classification in FILL_ORDER:
        members = [c for c in candidates if c.classification == classification]
        if not members:
            continue
        cap = min(budget - spent, int(budget * CLASS_SHARE[classification]))
        lines: list[str] = []
        used = 0
        for candidate in members:
            rendered = render(candidate)
            cost = ready.count_tokens(rendered) + separator_cost
            if used + cost > cap:
                ready._live_dropped.append(candidate)
                continue
            lines.append(rendered)
            used += cost
            ready._live_used.append(candidate)
        if lines:
            body = "\n\n".join([classes.CLASS_FRAMING[classification]] + lines)
            sections.append(Section(classes.CLASS_SECTIONS[classification], body))
            spent += used
    ready._spent = spent
    return sections


def render(candidate: Candidate) -> str:
    """One passage as the narrator sees it: a provenance line, then the text.

    The label is not decoration. It is what makes a claim in the prompt
    attributable — the difference between the narrator reading a fact and the
    narrator reading a fact *from a file the reader imported and classified* —
    and it is the same identification the inspector shows, so the two agree.
    """
    parts = [candidate.filename or candidate.title or "imported source"]
    if candidate.heading_path:
        parts.append(candidate.heading_path)
    parts.append(f"passage {candidate.chunk_index + 1}")
    label = " · ".join(parts)
    if candidate.visibility == classes.HIDDEN:
        label = f"{label} {classes.HIDDEN_MARKER}"
    return f"[{label}]\n{candidate.text}"


def report(ready: Plan) -> dict:
    """What the Insights panel and the tests read about this turn's knowledge.

    Everything needed to answer F05 and F06 for imported material: which source,
    which file, which class, which visibility, which passage, what it scored on
    each path and combined, how it was found, what it cost, and what was
    considered and set aside.

    This dict is written into the turn's context snapshot, and the rendered text
    goes with it. That is deliberate, and it is what
    `IMPORTED-KNOWLEDGE-DESIGN.md` §49-50 requires: a turn's evidence must
    survive the source being deleted, so the record holds the text rather than a
    pointer to a row that can go away.
    """
    result = ready.result
    return {
        "used": [_used(c, ready) for c in ready._always_used + ready._live_used],
        "dropped": [
            dict(_record(c), reason="over the knowledge budget")
            for c in ready._always_dropped + ready._live_dropped
        ],
        "suppressed": [
            dict(_record(c), duplicate_of=c.duplicate_of) for c in result.suppressed
        ],
        "terms": result.terms,
        "considered": result.considered,
        "generated": result.generated,
        "rejected": result.rejected,
        "semantic_floor": result.semantic_floor,
        "semantic_calibrated": result.semantic_calibrated,
        "embedding_model": result.embedding_model,
        "semantic_used": result.semantic_used,
        "semantic_note": result.semantic_note,
        "scan_truncated": result.scan_truncated,
        "budget": ready._budget,
        "spent": ready._spent,
        "protected_tokens": ready.protected_tokens,
    }


def _record(candidate: Candidate) -> dict:
    return candidate.as_record()


def _used(candidate: Candidate, ready: Plan) -> dict:
    """A used passage, with the text that was actually supplied."""
    rendered = render(candidate)
    return dict(
        _record(candidate),
        text=candidate.text,
        rendered=rendered,
        prompt_tokens=ready.count_tokens(rendered),
    )