tinymind-native-colab-handoff/bundle/evaluation/evidence_ledger.py

"""Evidence Ledger v2: exact token archive with integrity checks."""

from __future__ import annotations

import hashlib
import json
import struct
from pathlib import Path
from typing import Iterable


class LedgerCorruptionError(RuntimeError):
    pass


def encode_tokens(tokens: Iterable[int]) -> bytes:
    return b"".join(struct.pack("<q", int(token)) for token in tokens)


def decode_tokens(payload: bytes) -> list[int]:
    if len(payload) % 8 != 0:
        raise LedgerCorruptionError("token payload length is not divisible by 8")
    return [value[0] for value in struct.iter_unpack("<q", payload)]


def _sha256(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def _merkle_root(hashes: list[str]) -> str:
    if not hashes:
        return _sha256(b"")
    level = [bytes.fromhex(h) for h in hashes]
    while len(level) > 1:
        if len(level) % 2 == 1:
            level.append(level[-1])
        level = [hashlib.sha256(level[i] + level[i + 1]).digest() for i in range(0, len(level), 2)]
    return level[0].hex()


class EvidenceLedgerV2:
    """Content-addressed exact token ledger.

    Chunks are stored as binary int64 token payloads. The model may regenerate
    KV from retrieved chunks, but exact recall is guaranteed by this ledger.
    """

    def __init__(self, root: str | Path, chunk_tokens: int = 8192):
        self.root = Path(root)
        self.chunk_tokens = int(chunk_tokens)
        self.chunks_dir = self.root / "chunks"
        self.manifest_path = self.root / "manifest.json"
        self.semantic_path = self.root / "semantic_index.json"

    def ingest(
        self,
        tokens: Iterable[int],
        passkeys: dict[int, int] | None = None,
        semantic_spans: dict[str, tuple[int, int]] | None = None,
    ) -> dict:
        self.chunks_dir.mkdir(parents=True, exist_ok=True)
        passkeys = passkeys or {}
        semantic_spans = semantic_spans or {}
        chunks: list[dict] = []
        token_buffer: list[int] = []
        start = 0
        total_tokens = 0
        stream_hasher = hashlib.sha256()

        def flush() -> None:
            nonlocal token_buffer, start
            if not token_buffer:
                return
            payload = encode_tokens(token_buffer)
            digest = _sha256(payload)
            path = self.chunks_dir / f"{digest}.bin"
            path.write_bytes(payload)
            chunks.append(
                {
                    "chunk_id": len(chunks),
                    "start": start,
                    "length": len(token_buffer),
                    "sha256": digest,
                    "path": str(path),
                }
            )
            start += len(token_buffer)
            token_buffer = []

        for token in tokens:
            value = int(token)
            encoded = struct.pack("<q", value)
            stream_hasher.update(encoded)
            token_buffer.append(value)
            total_tokens += 1
            if len(token_buffer) >= self.chunk_tokens:
                flush()
        flush()

        semantic_index = {
            key: {"start": int(span[0]), "length": int(span[1])}
            for key, span in semantic_spans.items()
        }
        self.semantic_path.write_text(json.dumps(semantic_index, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
        manifest = {
            "schema_version": "tinymind-evidence-ledger-v2",
            "total_tokens": total_tokens,
            "chunk_tokens": self.chunk_tokens,
            "chunk_count": len(chunks),
            "chunks": chunks,
            "stream_sha256": stream_hasher.hexdigest(),
            "merkle_root": _merkle_root([chunk["sha256"] for chunk in chunks]),
            "passkeys": {str(k): int(v) for k, v in passkeys.items()},
            "semantic_index_path": str(self.semantic_path),
            "kv_tokens_stored": 0,
        }
        self.manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
        return manifest

    def _manifest(self) -> dict:
        return json.loads(self.manifest_path.read_text(encoding="utf-8"))

    def _read_chunk(self, chunk: dict) -> list[int]:
        payload = Path(chunk["path"]).read_bytes()
        digest = _sha256(payload)
        if digest != chunk["sha256"]:
            raise LedgerCorruptionError(f"chunk {chunk['chunk_id']} hash mismatch")
        return decode_tokens(payload)

    def verify_integrity(self) -> dict:
        manifest = self._manifest()
        hashes: list[str] = []
        for chunk in manifest["chunks"]:
            payload = Path(chunk["path"]).read_bytes()
            digest = _sha256(payload)
            if digest != chunk["sha256"]:
                raise LedgerCorruptionError(f"chunk {chunk['chunk_id']} hash mismatch")
            hashes.append(digest)
        root = _merkle_root(hashes)
        if root != manifest["merkle_root"]:
            raise LedgerCorruptionError("merkle root mismatch")
        return {"passed": True, "merkle_root": root, "chunk_count": len(hashes)}

    def _chunk_for_position(self, position: int) -> dict:
        position = int(position)
        for chunk in self._manifest()["chunks"]:
            if int(chunk["start"]) <= position < int(chunk["start"]) + int(chunk["length"]):
                return chunk
        raise IndexError(f"position {position} outside ledger")

    def recall_position(self, position: int) -> dict:
        chunk = self._chunk_for_position(position)
        tokens = self._read_chunk(chunk)
        offset = int(position) - int(chunk["start"])
        return {"position": int(position), "token": tokens[offset], "chunk_id": int(chunk["chunk_id"])}

    def recall_span(self, start: int, length: int) -> dict:
        start = int(start)
        length = int(length)
        out: list[int] = []
        for position in range(start, start + length):
            out.append(self.recall_position(position)["token"])
        return {"start": start, "length": length, "tokens": out}

    def recall_hash(self, sha256: str) -> dict:
        for chunk in self._manifest()["chunks"]:
            if chunk["sha256"] == sha256:
                return {"sha256": sha256, "chunk_id": int(chunk["chunk_id"]), "tokens": self._read_chunk(chunk)}
        raise KeyError(sha256)

    def recall_semantic(self, key: str) -> dict:
        index = json.loads(self.semantic_path.read_text(encoding="utf-8"))
        row = index[key]
        recalled = self.recall_span(row["start"], row["length"])
        recalled["semantic_key"] = key
        return recalled