api/serialize.py

"""Serializers: core objects -> JSON-safe payloads for the frontend."""

from __future__ import annotations

import json
import traceback

from miru_tracer.core.interventions import Intervention
from miru_tracer.core.sampling import SamplingParams, select_token
from miru_tracer.core.tracer import NextTokenDistribution


def fig_json(fig) -> dict | None:
    """A Plotly figure as a JSON-safe dict (None passes through)."""
    if fig is None:
        return None
    return json.loads(fig.to_json())


def error_payload(message: str, *, trace: bool = False) -> dict:
    """Uniform error shape; optionally append the traceback like the app did."""
    if trace:
        message = f"{message}\n\nTraceback:\n{traceback.format_exc()}"
    return {"ok": False, "error": message}


def candidates_payload(dist: NextTokenDistribution) -> list[dict]:
    """Next-token candidates table: rank/id/token/probability rows."""
    return [
        {
            "rank": rank,
            "token_id": token_id,
            "text": text_raw,
            "prob": round(prob, 6),
            "raw_prob": round(raw_prob, 6),
        }
        for rank, (token_id, prob, raw_prob, text_raw) in enumerate(
            zip(
                dist.top_k_tokens,
                dist.top_k_probs,
                dist.top_k_raw_probs,
                dist.top_k_texts_raw,
                strict=True,
            )
        )
    ]


def render_state(session_id: str, tracer, status: str, params: SamplingParams, log_topk) -> dict:
    """The canonical Interactive-mode payload every successful handler returns.

    The preview_id is exactly the token ``step()`` would commit; for sampling
    this draws the sample now — the next Step locks it in.
    """
    dist = tracer.peek(
        top_k=max(int(log_topk or 10), 1), temperature=params.temperature
    )
    preview_id = select_token(dist.raw_logits, params)
    return {
        "ok": True,
        "status": status,
        "session_id": session_id,
        "text": tracer.get_full_text() if tracer.history else "",
        "step": len(tracer.history),
        "candidates": candidates_payload(dist),
        "preview_id": int(preview_id),
        "eos": False,
    }


def intervention_to_dict(iv: Intervention, tokenizer=None) -> dict:
    return {
        "kind": iv.kind,
        "layer": iv.layer,
        "token_id": iv.token_id,
        "strength": iv.strength,
        "token_id_to": iv.token_id_to,
        "basis": iv.basis,
        "description": iv.describe(tokenizer),
    }


def intervention_rows_payload(rows: list[dict], tokenizer=None) -> list[dict]:
    return [
        {
            "enabled": row.get("enabled", True),
            **intervention_to_dict(row["intervention"], tokenizer),
        }
        for row in rows or []
    ]