haiku.rag/haiku_rag_slim/haiku/rag/capabilities/compaction.py

from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, cast

from pydantic_ai import RunContext
from pydantic_ai.capabilities import AbstractCapability

from haiku.rag.capabilities._base import RAGCapabilityBase
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.store.models.citation import Citation

CAPABILITY_ID = "haiku-rag-evidence-compaction"

CAPSULE_HEADER = (
    "[Evidence cited earlier in this conversation, kept so later questions can "
    "rely on it. Cite these chunk_ids directly when you use them.]"
)

RECEIPT = (
    "[Evidence retrieved for an earlier question, no longer shown. It does not "
    "count as cited for the current question.]"
)

ENTRY_SEPARATOR = "\n\n"


def group_label(position: int) -> str:
    """Name a group by its position among the groups, not by question number.

    A question identity is a message count, so a header built from it would present
    an index as a turn number, and an ordinal over the groups is not the
    conversation's ordinal either whenever a question in between cited nothing. The
    label claims only what it is: a grouping, newest first.
    """
    return f"[Cited evidence group {position}]"


def picture_label(chunk_id: str, self_ref: str) -> str:
    return (
        f"Page image retrieved from the knowledge base for cited evidence "
        f"[{chunk_id}] ({self_ref}). Not provided by the user."
    )


@dataclass(frozen=True)
class DiscoveredEvidence:
    """One evidence capability's records, as the compactor found them.

    Read-only and rebuilt per request: the compactor merges these into a view and
    persists nothing about evidence itself.
    """

    capability: str
    record: CapabilityEvidenceRecord
    citations: Mapping[str, Citation]
    tool_names: frozenset[str]


@dataclass(frozen=True)
class RetainedPicture:
    """A picture to re-attach, with the label that must accompany it.

    Addressed by owner, document and reference, because a reference such as
    ``#/pictures/0`` repeats across documents and capabilities. The label travels
    with it so it can never be emitted without its image.
    """

    capability: str
    chunk_id: str
    document_id: str
    self_ref: str
    label: str


@dataclass(frozen=True)
class Capsule:
    """Everything the compactor would insert, and nothing about where it goes."""

    text: str = ""
    pictures: tuple[RetainedPicture, ...] = ()


@dataclass(frozen=True)
class _Entry:
    capability: str
    chunk_id: str
    question: int
    citation: Citation

    def render(self) -> str:
        title = self.citation.document_title
        uri = self.citation.document_uri
        source = f'"{title}"' if title else uri
        if title and uri and uri != title:
            source = f"{source} ({uri})"
        return f"[{self.chunk_id}] Source: {source}\n{self.citation.content}"


def _eligible_entries(evidence: Sequence[DiscoveredEvidence]) -> list[_Entry]:
    """Cited evidence with content, newest citing question first.

    Evidence cited in several questions belongs to the most recent one, so it is
    rendered once and grouped where the model last used it.

    An occurrence and its canonical ``Citation`` are written by the same call, so a
    cited chunk without one is not a state this design produces. Rendering the rest
    regardless would quietly drop evidence an answer rested on, so it is reported.
    """
    entries = []
    for discovered in evidence:
        for chunk_id, occurrence in discovered.record.occurrences.items():
            if not occurrence.cited_in_questions:
                continue
            citation = discovered.citations.get(chunk_id)
            if citation is None:
                raise ValueError(
                    f"{discovered.capability} cited {chunk_id} in question(s) "
                    f"{occurrence.cited_in_questions} but has no citation record "
                    "for it, so its content cannot be retained."
                )
            entries.append(
                _Entry(
                    capability=discovered.capability,
                    chunk_id=chunk_id,
                    question=max(occurrence.cited_in_questions),
                    citation=citation,
                )
            )
    entries.sort(key=lambda entry: (-entry.question, entry.capability, entry.chunk_id))
    return entries


def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule:
    """Render every cited piece of evidence, grouped by the question that cited it.

    Everything cited is kept whole and everything else is dropped. There is no
    character budget: what a model can hold is the model's business, and a knob for
    it would only half-rescue models that fail on long conversations regardless.

    A host that needs earlier evidence pruned can compact further on top, on the wire
    only. Removing or reordering the stored history breaks the message counts that
    question identities and epochs are derived from, and the next record written is
    refused.

    Pure: no I/O and no message history, so what goes on the wire stays separable
    from what it should contain. Picture bytes are fetched by the caller, which is
    why a picture travels with its label rather than beside it.
    """
    entries = _eligible_entries(evidence)
    if not entries:
        return Capsule()

    lines = [CAPSULE_HEADER]
    pictures: list[RetainedPicture] = []
    seen: set[tuple[str, str, str]] = set()
    position = 0
    current_question: int | None = None
    for entry in entries:
        if entry.question != current_question:
            position += 1
            current_question = entry.question
            lines.append(group_label(position))
        lines.append(entry.render())
        for self_ref in entry.citation.picture_refs:
            # Overlapping chunks cite one figure, and a provider counts it twice.
            # Identity is owner plus document plus reference, so the same reference
            # in another document stays a different picture.
            identity = (entry.capability, entry.citation.document_id, self_ref)
            if identity in seen:
                continue
            seen.add(identity)
            pictures.append(
                RetainedPicture(
                    capability=entry.capability,
                    chunk_id=entry.chunk_id,
                    document_id=entry.citation.document_id,
                    self_ref=self_ref,
                    label=picture_label(entry.chunk_id, self_ref),
                )
            )
    return Capsule(text=ENTRY_SEPARATOR.join(lines), pictures=tuple(pictures))


@dataclass
class EvidenceCompactionCapability(AbstractCapability[Any]):
    """Rewrites the history from what the evidence capabilities recorded.

    Registering it is what turns compaction on: a host that leaves it out gets an
    untouched transcript, which is why it has no enable flag. It reads the evidence
    capabilities through the run's registry and holds no reference to any of them,
    so a host running one capability, both, or neither needs no wiring change.

    Registering two is rejected by pydantic-ai before the run starts, since they
    would share this capability's id.
    """

    def discover(self, ctx: RunContext[Any]) -> list[DiscoveredEvidence]:
        """Read what each evidence capability recorded, without writing anything.

        The registry holds the per-run instances, which are the ones carrying
        state; the registered objects never do. That includes a deferred capability
        the model has not loaded, whose record is simply empty.
        """
        discovered = []
        for capability in ctx.capabilities.values():
            if not isinstance(capability, RAGCapabilityBase):
                continue
            state = capability.state
            discovered.append(
                DiscoveredEvidence(
                    capability=capability.state_namespace,
                    record=cast(CapabilityEvidenceRecord, cast(Any, state).evidence),
                    citations=cast(Any, state).citation_index,
                    tool_names=frozenset(capability.evidence_tool_names()),
                )
            )
        return sorted(discovered, key=lambda evidence: evidence.capability)


def create_capability() -> EvidenceCompactionCapability:
    """Create the capability that compacts history from recorded evidence."""
    return EvidenceCompactionCapability(
        id=CAPABILITY_ID,
        description=(
            "Replaces earlier questions' evidence on the wire with a capsule of "
            "what was cited."
        ),
    )


__all__ = [
    "CAPABILITY_ID",
    "CAPSULE_HEADER",
    "RECEIPT",
    "Capsule",
    "DiscoveredEvidence",
    "EvidenceCompactionCapability",
    "RetainedPicture",
    "build_capsule",
    "create_capability",
    "group_label",
    "picture_label",
]