src/solomon/binding.py

"""Serving identity: which readout is served, and proof that it is the one that was measured.

Split out of `solomon/service.py` unchanged. Three things live here:

  * the frozen selection file (readout mode, served design, where the binding is),
  * the two-letter prompt identity, and
  * `RuntimeBinding` -- the one safety property this release keeps. A binding carries the runtime identity
    of the model that was scored, and every key of it must equal the live engine's or the service refuses
    to load. The binding also carries the calibration, so a temperature cannot drift from the stack it was
    fitted against, and it refuses a calibration fitted on a different model.
"""
import copy
import hashlib
import importlib
import json
from pathlib import Path

from solomon.calibration import NoTemperature, ROOT, digest, load_temperature
from solomon.engine_contract import BOOLEAN_TASK, LABEL

SELECTION = ROOT / 'serving/selection.json'
DEFAULT_BINDING = ROOT / 'serving/serving-binding.json'
DESIGN_FREEZE = ROOT / 'serving/design-freeze.json'
CONTRACT = 'solomon-v1'
MODES = ('four_collapsed', 'two_letter')
BINDING_SCHEMA = 'solomon-serving-binding-v1'
# ------------------------------------------------------------------ selection and prompts
def load_selection(path=SELECTION):
    """The frozen selection file -> {'mode', 'design', 'prompts', 'binding', 'sha256'}.

    Accepted keys: 'readout', 'readout_mode' or 'mode' (four_collapsed | two_letter).
    Optional: 'design' {single_choice, ordered} override, 'two_letter_prompts' (inline headers dict or
    'module:ATTRIBUTE'), 'binding' (path to the Solomon serving-binding.json: absolute, or relative to the
    selection file's own directory, else to the project root).
    """
    path = Path(path)
    if not path.exists():
        raise ValueError(f'Solomon selection not found: {path} (pass mode= explicitly for development use)')
    raw = json.loads(path.read_text())
    mode = raw.get('readout') or raw.get('readout_mode') or raw.get('mode')
    if mode not in MODES:
        raise ValueError('selection.json names no readout mode (readout / readout_mode / mode)')
    return {'mode': mode, 'design': raw.get('design'), 'prompts': raw.get('two_letter_prompts'), 'binding': raw.get('binding'),
            'evidence_head': raw.get('evidence_head'),
            'sha256': hashlib.sha256(path.read_bytes()).hexdigest(), 'path': str(path), 'dir': str(path.resolve().parent)}


def default_design(mode):
    """four_collapsed serves the frozen design (single R, ordered S); two_letter was trained on S rows only."""
    if mode == 'two_letter':
        return {'single_choice': 'S', 'ordered': 'S'}
    frozen = json.loads(DESIGN_FREEZE.read_text()) if DESIGN_FREEZE.exists() else {'single_choice': 'R', 'ordered': 'S'}
    return {'single_choice': frozen['single_choice'], 'ordered': frozen['ordered']}


class TwoLetterPrompts:
    """Header swap four-state -> two-letter. headers: {'boolean': str, 'label': str, optional 'entity': str}.

    The four-state block is HEADER + rest (question / label / answer cue); only HEADER changes, so a two-letter block
    is byte-identical to what training used iff the training builder does the same swap (checked by `check_parity`).
    """
    FOUR = {'boolean': BOOLEAN_TASK, 'label': LABEL}

    def __init__(self, headers):
        if not isinstance(headers, dict) or not {'boolean', 'label'} <= set(headers) or set(headers) - {'boolean', 'label', 'entity'}:
            raise ValueError('two-letter prompts need headers for boolean and label (optional entity)')
        if any(not isinstance(v, str) or not v.strip() for v in headers.values()):
            raise ValueError('two-letter headers must be nonempty strings')
        if any(v.startswith(four) or four.startswith(v) for v in headers.values() for four in self.FOUR.values()):
            raise ValueError('two-letter header must differ from the four-state headers')
        self.headers = dict(headers)
        self.sha256 = digest(self.headers)

    def header(self, task, kind):
        return self.headers.get('entity', self.headers['boolean']) if (kind == 'boolean' and task == 'entity') else self.headers[kind]

    def rewrite(self, task, block):
        """Four-state block -> (two-letter block, 2), or None when the block is not a four-state Noul block."""
        for kind, four in self.FOUR.items():
            if block.startswith(four) and (kind == 'label') == (task == 'multilabel') and task in ('boolean', 'entity', 'multilabel'):
                return self.header(task, kind) + block[len(four):], 2
        return None

    def check_parity(self, rows):
        """rows: two-letter training/eval rows {'task', 'block'} (Noul tasks). Returns mismatching row indices."""
        bad = []
        for i, row in enumerate(rows):
            task, block = row['task'], row['block']
            kind = 'label' if task == 'multilabel' else 'boolean'
            head = self.header(task, kind)
            if not block.startswith(head) or self.rewrite(task, self.FOUR[kind] + block[len(head):]) != (block, 2):
                bad.append(i)
        return bad


def load_prompts(spec):
    """spec: None (solomon/prompts_two_letter.py, the headers the two-letter readout was trained on), inline headers dict, or 'module:ATTRIBUTE'."""
    if spec is None:
        from solomon import prompts_two_letter as p9
        loaded = TwoLetterPrompts({'boolean': p9.BOOLEAN_TASK2, 'label': p9.LABEL2})
        loaded.sha256 = hashlib.sha256(Path(p9.__file__).read_bytes()).hexdigest()  # = the prompts_sha256 a two_letter binding pins
        return loaded
    if isinstance(spec, str):
        module, _, attr = spec.partition(':')
        spec = getattr(importlib.import_module(module), attr or 'TWO_LETTER_HEADERS')
    return spec if isinstance(spec, TwoLetterPrompts) else TwoLetterPrompts(spec)


class UnboundServing:
    """Development / no-binding-file path. Still answers everything; it simply claims no bound identity."""
    status = 'unbound'
    schema = None
    sha256 = None

    def __init__(self, reason='unbound: no Solomon serving binding is loaded'):
        self.reason = reason
        self.bound_keys = []
        self.calibration = NoTemperature()
        self.head_routing = None

    def describe(self):
        return {'status': self.status, 'schema': None, 'sha256': None, 'bound_runtime_keys': [],
                'reason': self.reason, 'calibration': self.calibration.describe()}


class RuntimeBinding:
    """The one safety property Solomon v1.1 keeps: the stack being served is the stack that was measured.

    A `solomon-serving-binding-v1` artifact carries the runtime identity of the scored model (15 keys for this
    release: every ServiceEngine identity key except the recomputed `fingerprint`). Every one of them must equal the
    live engine identity or construction raises, so the service refuses to load on drift. It carries NO thresholds,
    NO per-task status, NO temperatures and NO fitted correctness head: there is nothing here that could gate an
    answer, because nothing gates an answer.
    """
    status = 'bound'
    schema = BINDING_SCHEMA
    # Without these two the binding would not pin the model at all; a binding that omits them is refused.
    REQUIRED = ('adapter_sha256', 'trained_heads_sha256')
    # Keys that must never appear in a serving binding: they are the machinery of the abstention path (v1.1).
    # `temperatures` is deliberately NOT here -- a per-task scalar is a calibration parameter, not a decision, and
    # carrying it inside the binding is what stops it drifting from the stack it was fitted against.
    REFUSED = ('tasks', 'thresholds', 'confidence', 'correctness', 'abstention', 'policy', 'status')
    # A temperature is fitted on ONE model's logits and is meaningless on another, even though the numbers would
    # look identical. So a non-trivial calibration must record the model it was fitted on, and that model must be
    # the one this binding serves.
    CALIBRATION_BOUND_KEYS = ('adapter_sha256', 'trained_heads_sha256')

    def __init__(self, binding, *, mode, runtime, design, prompts_sha256=None):
        if not isinstance(binding, dict) or binding.get('schema') != BINDING_SCHEMA:
            raise ValueError('Solomon serving binding schema required')
        if binding.get('sha256') != digest({k: v for k, v in binding.items() if k != 'sha256'}):
            raise ValueError('Solomon serving binding checksum mismatch')
        present = [k for k in self.REFUSED if k in binding]
        if present:
            raise ValueError('serving binding carries removed abstention machinery: ' + ', '.join(present))
        if binding.get('contract') not in (None, CONTRACT):
            raise ValueError('Solomon serving binding was issued for a different serving contract: ' + str(binding.get('contract')))
        if binding.get('readout') != mode:
            raise ValueError('Solomon serving binding was built for a different readout mode')
        bound = binding.get('runtime') or {}
        if not isinstance(bound, dict) or any(k not in bound for k in self.REQUIRED):
            raise ValueError('serving binding must bind at least ' + ', '.join(self.REQUIRED))
        drift = sorted(k for k, v in bound.items() if runtime.get(k) != v)
        if drift:
            raise ValueError('Solomon serving binding runtime identity mismatch: ' + ', '.join(drift))
        if binding.get('design') is not None and binding['design'] != design:
            raise ValueError('Solomon serving binding design mismatch')
        if mode == 'two_letter' and binding.get('prompts_sha256') != prompts_sha256:
            raise ValueError('Solomon serving binding prompt identity mismatch')
        routing = binding.get('head_routing')
        if routing is not None and (not isinstance(routing, dict) or any(not isinstance(k, str) or not isinstance(v, str)
                                                                          for k, v in routing.items())):
            raise ValueError('serving binding head_routing must map head_key -> head_key')
        self.head_routing = dict(routing) if routing is not None else None
        self.binding = copy.deepcopy(binding)
        self.sha256 = binding['sha256']
        self.bound_keys = sorted(bound)
        # Covered by this binding's checksum and by the identity check above, so the temperatures cannot drift
        # from the runtime they were measured on.
        self.calibration = load_temperature(binding.get('temperatures'))
        provenance = binding.get('provenance') or {}
        if self.calibration.file_sha256 is None:      # embedded copy: take the file hash from the provenance
            self.calibration.file_sha256 = provenance.get('calibration_file_sha256')
        self._check_calibration_model(bound, provenance)

    def _check_calibration_model(self, bound, provenance):
        """Refuse a calibration inherited from a different model.

        The 15-key identity check already proves the SERVING stack is the scored one. It cannot see this: a
        temperature fitted on model A and copied onto a binding for model B has the same numbers and passes every
        value check there is. The only thing that distinguishes them is which model's logits the scalars were
        fitted against, so that has to be recorded and matched.
        """
        if self.calibration.is_identity():
            return                                    # T = 1.0 everywhere is a no-op; safe on any model
        fitted = provenance.get('calibration_fitted_on') or {}
        missing = [k for k in self.CALIBRATION_BOUND_KEYS if not fitted.get(k)]
        if missing:
            raise ValueError(
                'serving binding carries temperatures but no record of the model they were fitted on (missing '
                + ', '.join('provenance.calibration_fitted_on.' + k for k in missing) + '). A temperature fitted '
                'on one model is meaningless on another, so an unattributed calibration is refused. Refit on this '
                'model, or serve at T = 1.0.')
        for key in self.CALIBRATION_BOUND_KEYS:
            if bound.get(key) != fitted[key]:
                raise ValueError(
                    f'calibration was fitted on {key} {fitted[key]}, binding carries {key} {bound.get(key)}. '
                    'Temperatures fitted on one model do not transfer to another even when the numbers match; '
                    'refit on this model against the error-rich fit panel, or serve it at T = 1.0.')

    def describe(self):
        return {'status': self.status, 'schema': self.schema, 'sha256': self.sha256,
                'bound_runtime_keys': list(self.bound_keys), 'calibration': self.calibration.describe(),
                **({'head_routing': dict(self.head_routing)} if self.head_routing is not None else {})}


def load_binding(path, *, mode, runtime, design, prompts_sha256=None):
    """None / missing file -> UnboundServing; a present but invalid or drifted file fails closed (raises)."""
    if isinstance(path, (UnboundServing, RuntimeBinding)) or hasattr(path, 'describe'):
        return path
    if isinstance(path, dict):
        return RuntimeBinding(path, mode=mode, runtime=runtime, design=design, prompts_sha256=prompts_sha256)
    path = Path(path) if path is not None else DEFAULT_BINDING
    if not path.is_absolute():
        path = ROOT / path
    if not path.exists():
        return UnboundServing(f'unbound: no Solomon serving binding at {path.relative_to(ROOT) if path.is_relative_to(ROOT) else path}')
    return RuntimeBinding(json.loads(path.read_text()), mode=mode, runtime=runtime, design=design, prompts_sha256=prompts_sha256)