arborist/aborist/qa/keys.py

"""The 8-dim Merkle-AGI v9.8 cache_key.

v9.8 invariant: no answer is reused unless all eight match and the
record is live (not failed/stale/quarantined):

    1. source_root              — content fingerprint of the document
    2. question_hash             — SHA-256 of normalized question text
    3. model_profile_hash        — model_id + revision + quantization
    4. conversation_hash         — full canonical OpenAI messages array
    5. governance_policy_hash    — sampling/policy parameters dict
    6. schema_version            — aborist DB schema version
    7. canonicalization_version  — text normalization rules
    8. chunking_version          — chunker name & parameters

Bumping ANY of these eight dimensions yields a distinct cache_key, so
prior records cannot be served. This is the runtime drift detection
the providence whitepaper compresses into "cache_key = source_root +
':' + question_hash" — that's a simplification; the rigorous form is
all eight dimensions hashed together.
"""

from __future__ import annotations

import hashlib
import json

from aborist.document import canonicalize


def _sha256(s: str) -> str:
    return hashlib.sha256(s.encode("utf-8")).hexdigest()


def _canonical_json(obj) -> str:
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


QUESTION_DEDUP_MODES = ("strict", "equivalence_class")
DEFAULT_QUESTION_DEDUP = "equivalence_class"

# Lookup-time fidelity. Decoupled from write-time `question_dedup`:
# write determines under which `cache_key` a record lands; fidelity
# determines which `cache_key`s a lookup will check.
#
#   strict             only the cache_key matching the agent's policy
#                      is checked. No fallback. Audit-grade behavior.
#   equivalence_class  primary cache_key checked first; if miss AND the
#                      OTHER dedup mode produces a different cache_key,
#                      the alternate is also checked. Lets a fast-cache
#                      agent reuse records written under either mode.
FIDELITY_MODES = ("strict", "equivalence_class")
DEFAULT_FIDELITY = "equivalence_class"


def canonical_question(
    question: str, *, mode: str = DEFAULT_QUESTION_DEDUP
) -> str:
    """Canonical form of `question` for the given dedup ``mode``.

    Two modes:

    - ``"equivalence_class"`` (default): four-step canonicalization —
      ``canonicalize()`` (NFC + ws-collapse + strip ends), then
      lowercase, then trailing-punctuation strip, then standalone-article
      filter (``the``, ``a``, ``an``). All variants of "Who is THE
      Batman?" / "who is batman" / "who is X." collapse to one form.
      The default for chat-style agents that prefer fast cache hits.

    - ``"strict"``: only ``canonicalize()`` — NFC + ws-collapse + strip
      ends. Case-sensitive, punctuation-sensitive, article-sensitive.
      Maximum granularity. The choice for audit-grade agents that want
      every distinct phrasing to get its own answer.

    Exposed as a function so callers can dedup BEFORE hashing — e.g.
    inject the canonical form into the user message used for
    ``conversation_hash``, while still sending the verbatim question to
    the LLM. Without this split ``"who is batman"``, ``"who is
    batman?"``, and ``"who is the batman?"`` collapse on
    ``question_hash`` (under equivalence_class) but each hits
    ``conversation_hash`` differently, missing cache.

    The choice of mode flows through ``policy["question_dedup"]`` into
    ``governance_policy_hash`` so two agents under different modes
    write records under different ``cache_key``s — they coexist in
    parallel namespaces, never collide.
    """
    if mode not in QUESTION_DEDUP_MODES:
        raise ValueError(
            f"question dedup mode must be one of {QUESTION_DEDUP_MODES}, got {mode!r}"
        )
    canon = canonicalize(question)
    if mode == "strict":
        return canon
    canon = canon.lower().rstrip(_QUESTION_TRAILING_STRIP)
    tokens = [t for t in canon.split() if t not in _QUESTION_ARTICLE_STRIP]
    return " ".join(tokens)


def question_hash(
    question: str, *, mode: str = DEFAULT_QUESTION_DEDUP
) -> str:
    """SHA-256 of the dedup-mode-canonicalized question.

    See ``canonical_question`` for what each mode does. The hash is the
    SHA-256 of the canonical form. Bumping
    ``_QUESTION_TRAILING_STRIP`` or ``_QUESTION_ARTICLE_STRIP`` (the
    equivalence-class strip sets) orphans prior cache records whose
    canonical question contained newly-stripped tokens; they live as
    history but won't be re-hit on lookup.

    Equivalence class examples (mode="equivalence_class")::

        "who is X"           |
        "who is X?"          |
        "Who Is X."          | -> same question_hash
        "who is the X"       |
        "who is a X"         |
        "who is an X？"      |  (CJK question mark)

    Strict mode (mode="strict") distinguishes all of those.

    What's IN the trailing-strip set: ``.?!,;:`` (ASCII), ``？！。、``
    (CJK full-width), ``…`` (ellipsis). Pairs like ``"`` ``'`` ``)``
    ``]`` ``}`` are NOT — naive one-sided stripping breaks balance.
    Apostrophes aren't either — ``X's`` is a different question from
    ``X``.
    """
    return _sha256(canonical_question(question, mode=mode))


# Trailing punctuation that carries no semantic difference at the end
# of a question. Order doesn't matter (rstrip walks char-by-char from
# the right). Repeats handled trivially: ``X???`` → ``X``.
#
# ASCII:  . ? ! , ; :
# CJK:    ？ U+FF1F  full-width question mark
#         ！ U+FF01  full-width exclamation
#         。 U+3002  ideographic full stop
#         、 U+3001  ideographic comma
# Other:  … U+2026  horizontal ellipsis
_QUESTION_TRAILING_STRIP = ".?!,;:？！。、…"

# English articles stripped as standalone tokens after lowercasing. The
# question equivalence class treats "the foo" and "foo" as the same
# question — fox's 2026-04-29 catch: `who is the batman` & `who is
# batman` produced different cache records under earlier rules. Tokens
# are matched as EXACT lowercase strings, so substrings like "thesis"
# (contains "the") stay untouched.
#
# Conservative on purpose: only ASCII English articles. "El", "la",
# "los", "le", "les", "der", "die", "das" etc. are not stripped today.
# Adding them when needed flows through the same equivalence-class
# expansion the trailing-punctuation set went through.
_QUESTION_ARTICLE_STRIP = frozenset({"the", "a", "an"})


def model_profile_hash(
    model_id: str, revision: str = "", quantization: str = ""
) -> str:
    """SHA-256 of model identity. Bumping any field bumps the cache key."""
    return _sha256(f"{model_id}|{revision}|{quantization}")


def conversation_hash(messages: list[dict]) -> str:
    """SHA-256 of canonical JSON of the full OpenAI messages array.

    Order matters: a 6-turn dialogue arriving at the same final question
    produces a different hash than a single-turn ask.
    """
    return _sha256(_canonical_json(messages))


def governance_policy_hash(policy: dict) -> str:
    """SHA-256 of canonical JSON of the sampling/policy dict.

    Includes temperature, top_p, max_tokens, and the system prompt — any
    of those changing means the answer is governed differently and the
    cache must miss.
    """
    return _sha256(_canonical_json(policy))


def cache_key(
    source_root: str,
    question_hash_value: str,
    model_profile_hash_value: str,
    conversation_hash_value: str,
    governance_policy_hash_value: str,
    schema_version: str,
    canonicalization_version: str,
    chunking_version: str,
) -> str:
    """SHA-256 of all 8 dimensions joined with '|'.

    Any drift in any dimension produces a distinct cache_key.
    """
    return _sha256(
        "|".join(
            [
                source_root,
                question_hash_value,
                model_profile_hash_value,
                conversation_hash_value,
                governance_policy_hash_value,
                schema_version,
                canonicalization_version,
                chunking_version,
            ]
        )
    )