knowline_layout.py

# knowline_layout.py: release copy of openjev.header_layout (OpenJev src/openjev/header_layout.py), generated by
# make_knowline_layout.py. Only the helper import differs (knowline_server instead of openjev.systemone.prompt).
# source sha256: 8dfa4571ac18c1f890bb1662ac8241b156417a49dc05d4ebec8257406a3a5011
"""Request layouts beyond the default "chat" layout of openjev.systemone.prompt: the header (the request's question list
before the state), read twice, and option rotations. Spec: docs/plans/2026-10-22-header-layout.md.

Pure Python. The only outside object is the model's tokenizer or HF processor (`template`): its chat template
(apply_chat_template) renders the turns and its encode() counts tokens for the read-twice cut. With every switch off,
render_request() / render() give byte-identical prompts to openjev.systemone.prompt.render(..., style="chat"), the layout
our training data and release serving use (tests/test_header_layout.py checks this).

    from openjev.header_layout import Layout, render, render_request
    text = render({"state": ..., "questions": {...}}, 0, tokenizer, Layout(header=True))          # one question
    r = render_request(request, tokenizer, Layout(header=True, read_twice=True), rotations=(0, 1))  # a whole request
    r["prompts"][qid][offset] -> (prompt text ending in "Answer:", answer keys in shown order)

Option rotations are a scoring technique: score the question at offsets first_pass_offsets(n), map each offset's label
distribution back with to_original(), average them with equal weight (combine_rotations), and when the two first-pass
offsets pick different options of a 3+ option question also score gated_offsets(n) and average all of them.
"""
import uuid
from dataclasses import dataclass

from knowline_server import ANSWER, DEFAULT_QUESTION, INSTRUCTION, find_labels, options_of, render_value, state_messages

HEADER_HEAD = "Questions that will be asked about the conversation or state below:\n"
HEADER_ITEM = "- "
READ_AGAIN = "Read the same input again before answering:\n"
TWICE_MAX_TOKENS = 3000
_KW = dict(tokenize=False, add_generation_prompt=True, enable_thinking=False)


@dataclass(frozen=True)
class Layout:
    header: bool = False             # list every question of the request before the state
    header_min_questions: int = 1    # requests with fewer questions keep the default layout (1 = every request)
    read_twice: bool = False         # write the whole input twice, joined by READ_AGAIN, ...
    twice_max_tokens: int = TWICE_MAX_TOKENS  # ... in text-only requests of at most this many tokens read once


# ------------------------------------------------------------------ pieces
def question_text(q):
    """The question stem, as the question block shows it."""
    return render_value(q["instructions"]) if q.get("instructions") is not None else DEFAULT_QUESTION


def header_text(questions):
    """'Questions that will be asked ...:\\n- <stem>\\n- <stem>\\n\\n' in request order (stems only, no options)."""
    return HEADER_HEAD + "\n".join(HEADER_ITEM + question_text(q) for q in questions.values()) + "\n\n"


def rotation(n, offset):
    """Original option indices in shown order for a cyclic rotation by `offset` (offset 0 = request order)."""
    offset %= n
    return list(range(offset, n)) + list(range(offset))


def question_block(q, labels, offset=0):
    """-> (block text, answer keys in shown order). Offset 0 is the default layout's block byte for byte."""
    keys, texts = options_of(q)
    order = rotation(len(keys), offset)
    lines = "".join(f"{labels[j]}. {texts[i]}\n" for j, i in enumerate(order))
    return f"Question: {question_text(q)}\nOptions:\n{lines.rstrip()}", [keys[i] for i in order]


def _prepend(msgs, text):
    """`text` before the first non-system turn: prefixed to that turn's content when it is a user turn, else its own user
    turn (without the trailing newlines) inserted there."""
    i = next((k for k, m in enumerate(msgs) if m.get("role") != "system"), len(msgs))
    if i < len(msgs) and msgs[i].get("role") == "user":
        c = msgs[i].get("content")
        c = [{"type": "text", "text": text}] + list(c) if isinstance(c, list) else text + (c or "")
        return msgs[:i] + [{**msgs[i], "content": c}] + msgs[i + 1:]
    return msgs[:i] + [{"role": "user", "content": text.rstrip("\n")}] + msgs[i:]


def _join(a, b):
    """Content of two adjacent turns of one role, folded into one turn with a blank line between them (for a text `b`
    after a text or part-list `a`, exactly the default layout's fold of the ask into the last user turn)."""
    if isinstance(a, list):
        return a + ([{"type": "text", "text": "\n\n" + b}] if isinstance(b, str) else [{"type": "text", "text": "\n\n"}] + list(b))
    if isinstance(b, list):
        return [{"type": "text", "text": f"{a}\n\n"}] + list(b)
    return f"{a}\n\n{b}"


def _merge_adjacent(msgs):
    out = []
    for m in msgs:
        if out and out[-1].get("role") == m.get("role") and m.get("role") in ("user", "system"):
            out[-1] = {**out[-1], "content": _join(out[-1].get("content"), m.get("content"))}
        else:
            out.append(dict(m))
    return out


def _conv(msgs, hdr, m1, m2):
    """[system turns] + [header] + state turns + ask(m1); with m2 also READ_AGAIN + the same non-system turns + ask(m2)."""
    if hdr:
        msgs = _prepend(msgs, hdr)
    conv = msgs + [{"role": "user", "content": INSTRUCTION + "\n\n" + m1}]
    if m2 is None:
        return conv
    again = [m for m in conv if m.get("role") != "system"]  # system turns appear once, at the start
    again[-1] = {"role": "user", "content": INSTRUCTION + "\n\n" + m2}
    return conv + _prepend(again, READ_AGAIN)


def _scaffold(template, state, hdr, twice):
    """-> (prefix, middle or None, ending, media). The question block goes at the markers: once at m1, twice at m1 and m2."""
    m1 = f"OPENJEV_{uuid.uuid4().hex}"
    m2 = f"OPENJEV_{uuid.uuid4().hex}" if twice else None
    msgs, media = state_messages(state)
    try:
        text = template.apply_chat_template(_conv(msgs, hdr, m1, m2), **_KW)
    except Exception:
        try:
            if not twice:  # the default layout's fold: the ask into the last user turn (strict-alternation templates)
                conv = _conv(msgs, hdr, m1, None)
                if len(conv) < 2 or conv[-2].get("role") != "user":
                    raise ValueError("last turn is not a user turn")
                conv = conv[:-2] + [{**conv[-2], "content": _join(conv[-2]["content"], conv[-1]["content"])}]
            else:
                conv = _merge_adjacent(_conv(msgs, hdr, m1, m2))
            text = template.apply_chat_template(conv, **_KW)
        except Exception:  # roles the template rejects: the whole chat as one user message, as the default layout does
            msgs, media = [{"role": "user", "content": render_value(state)}], []
            text = template.apply_chat_template(_conv(msgs, hdr, m1, m2), **_KW)
    if text.count(m1) != 1 or (m2 is not None and text.count(m2) != 1):
        raise ValueError("chat template dropped or duplicated the question slot")
    prefix, rest = text.split(m1)
    if m2 is None:
        return prefix, None, rest, media
    middle, ending = rest.split(m2)
    return prefix, middle, ending, media


_LABELS = {}


def labels_and_ids(template):
    """The default engine's label set and token ids: single-token labels A..Z, AA.. checked right after the real prompt
    ending (as openjev.systemone.engine.Engine finds them)."""
    key = id(template)
    if key not in _LABELS:
        prefix, _, ending, _ = _scaffold(template, "x", None, False)
        block, _ = question_block({"type": "noul"}, ["A", "B"])
        _LABELS[key] = (template, find_labels(getattr(template, "tokenizer", template), prefix + block + ending + ANSWER))
    return _LABELS[key][1]


def default_labels(template):
    return labels_and_ids(template)[0]


# ------------------------------------------------------------------ rendering
def render_request(request, template, layout=Layout(), rotations=(0,), labels=None):
    """-> {"prefix": shared prefix text, "twice": bool, "once_tokens": int | None, "media": [...],
           "prompts": {qid: {offset: (prompt text, answer keys in shown order)}}}.
    `request` is a /v1/systemone body ({"state", "questions"}). `rotations`: cyclic option offsets (see rotation()) for every
    question, or a function (question id, number of options) -> offsets, e.g. lambda qid, n: first_pass_offsets(n)."""
    state, questions = request.get("state", ""), request.get("questions") or {}
    if not questions:
        raise ValueError("no questions")
    labels = labels or default_labels(template)
    hdr = header_text(questions) if layout.header and len(questions) >= layout.header_min_questions else None
    prefix, _, ending, media = _scaffold(template, state, hdr, False)
    offs = rotations if callable(rotations) else (lambda qid, n: rotations)
    blocks = {qid: {o: question_block(q, labels, o) for o in dict.fromkeys(offs(qid, len(options_of(q)[0])))}
              for qid, q in questions.items()}
    blocks = {qid: bo for qid, bo in blocks.items() if bo}
    once = {qid: {o: (prefix + b + ending + ANSWER, keys) for o, (b, keys) in bo.items()} for qid, bo in blocks.items()}
    out = {"prefix": prefix, "twice": False, "once_tokens": None, "media": media, "prompts": once}
    if not layout.read_twice or media:  # read twice is text-only; a request with media is read once
        return out
    tok = getattr(template, "tokenizer", template)
    n = lambda s: len(tok.encode(s, add_special_tokens=False))
    p = n(prefix)
    first = {qid: (prefix + question_block(q, labels, 0)[0] + ending + ANSWER) for qid, q in questions.items()}
    out["once_tokens"] = p + sum(n(t) - p for t in first.values())  # shared prefix + every question row, read once
    if out["once_tokens"] > layout.twice_max_tokens:
        return out
    prefix2, middle, ending2, _ = _scaffold(template, state, hdr, True)
    out.update(prefix=prefix2, twice=True,
               prompts={qid: {o: (prefix2 + b + middle + b + ending2 + ANSWER, keys) for o, (b, keys) in bo.items()}
                        for qid, bo in blocks.items()})
    return out


def render(request, question_idx, template, layout=Layout(), rotation_offset=0, labels=None):
    """The prompt text of one question (by position in the request or by id), ending in "Answer:"."""
    qids = list((request.get("questions") or {}))
    qid = qids[question_idx] if isinstance(question_idx, int) else question_idx
    return render_request(request, template, layout, (rotation_offset,), labels)["prompts"][qid][rotation_offset][0]


# ------------------------------------------------------------------ rotations (scoring side)
def first_pass_offsets(n):
    """Request order and its rotation by n // 2 (a swap for two options)."""
    return list(dict.fromkeys([0, n // 2])) if n >= 2 else [0]


def gated_offsets(n):
    """Second pass for 3+ option questions whose first-pass offsets disagree: the other four-order rotations."""
    first = first_pass_offsets(n)
    return [o for o in dict.fromkeys([n // 4, 3 * n // 4]) if o not in first] if n >= 3 else []


def to_original(shown_probs, offset):
    """A label distribution in shown order -> the same distribution in request (original option) order."""
    n = len(shown_probs)
    out = [0.0] * n
    for j, i in enumerate(rotation(n, offset)):
        out[i] = shown_probs[j]
    return out


def argmax(p):
    return max(range(len(p)), key=p.__getitem__)


def combine_rotations(by_offset):
    """{offset: distribution in shown order} -> equal-weight mean in request order (each already tempered)."""
    rows = [to_original(p, o) for o, p in by_offset.items()]
    return [sum(r[i] for r in rows) / len(rows) for i in range(len(rows[0]))]


def needs_gate(by_offset):
    """True when the first-pass offsets of a 3+ option question pick different original options."""
    n = len(next(iter(by_offset.values())))
    first = first_pass_offsets(n)
    if n < 3 or any(o not in by_offset for o in first):
        return False
    return len({argmax(to_original(by_offset[o], o)) for o in first}) > 1