src/ui/cache.py

"""Persistent sample-level inference cache and orchestration."""
from __future__ import annotations

import store
from data import get_record
from inference import default_model
import gradio as gr
from ui.parsing import _tag
from viz.renderers import (
    _color_reason,
    _pred_table,
    _reason_panel,
    _reason_placeholder,
    _status_line,
    _tldr_card,
)


def _output_tuple(
    results: dict,
    reason_html: str,
    physics_html: str,
    chemistry_html: str,
    status_html: str,
    inchikey: str,
    guest_name: str,
    model: str,
    prompt_version: str,
) -> tuple:
    return (
        _pred_table(results),
        reason_html,
        physics_html,
        chemistry_html,
        status_html,
        inchikey,
        guest_name,
        model,
        prompt_version,
    )


def _reason_from_results(results: dict) -> str:
    combined = results.get("combined") or {}
    if combined.get("error"):
        return _reason_panel(
            _reason_placeholder(f"Combined trajectory failed: {combined['error']}", kind="error")
        )
    text = combined.get("text") or ""
    if text:
        return _reason_panel(_color_reason(_tag(text, "think") or text))
    return _reason_panel(
        _reason_placeholder("Combined reasoning trace will appear here shortly…", kind="running")
    )


def _card_from_results(results: dict, mode: str, kind: str) -> str:
    row = results.get(mode)
    if row is None:
        return _tldr_card("__running__", kind=kind)
    if row.get("error"):
        return _tldr_card("__error__", kind=kind)
    return _tldr_card(row.get("tldr") or "", kind=kind)


async def predict(
    guest_name: str,
    model: str | None = None,
    prompt_mode: str = "Default",
    custom_text: str = "",
    force_rerun: bool = False,
    feature_preset: str = "Recommended",
    custom_features: list[str] | None = None,
    progress=gr.Progress(),
):
    """Async generator yielding prediction UI updates.

    Output order must match `run_btn.click(outputs=[...])`:
        pred_html, reason_html, physics_card, chem_card, status_html,
        inchikey_st, guest_st, model_st, promptver_st
    """
    model = model or default_model()
    if not guest_name:
        yield _output_tuple(
            {},
            _reason_panel(_reason_placeholder("Pick a guest and click Run prediction.")),
            _tldr_card("", kind="physics"),
            _tldr_card("", kind="chemistry"),
            _status_line("No guest selected.", kind="error"),
            "",
            "",
            model,
            "",
        )
        return

    progress(0.0, desc="Loading record...")
    try:
        rec = get_record(guest_name)
    except Exception as exc:  # noqa: BLE001 - surfaced to UI
        yield _output_tuple(
            {},
            _reason_panel(_reason_placeholder(f"Could not load features: {exc}", kind="error")),
            _tldr_card("", kind="physics"),
            _tldr_card("", kind="chemistry"),
            _status_line(f"Dataset error: {exc}", kind="error"),
            "",
            guest_name,
            model,
            "",
        )
        return

    inchikey = str(rec.get("inchikey", ""))

    # READ-ONLY: show the stored (batch) prediction for this guest+model. The
    # dashboard never runs live inference — that wrote fresh stochastic values
    # into the canonical store and diverged from the batch data. Look up by
    # guest+model regardless of the UI prompt hash, so the batch "v5" row matches
    # (the UI's _prompt_version is a different hashed key that never would).
    sample = store.get_sample_any(inchikey, store.norm_model(model))
    if sample is not None:
        results = sample.get("results") or {}
        yield _output_tuple(
            results,
            _reason_from_results(results),
            _card_from_results(results, "physics", "physics"),
            _card_from_results(results, "chemistry", "chemistry"),
            _status_line("Stored prediction (read-only). Review below.", kind="done"),
            inchikey,
            guest_name,
            model,
            sample.get("prompt_version", ""),
        )
        return

    # No stored prediction for this guest — the board shows batch results only.
    yield _output_tuple(
        {},
        _reason_panel(_reason_placeholder(
            "This guest has no stored prediction yet. The dashboard displays stored "
            "batch predictions only; live re-runs are disabled.", kind="error")),
        _tldr_card("", kind="physics"),
        _tldr_card("", kind="chemistry"),
        _status_line("No stored prediction for this guest.", kind="error"),
        inchikey,
        guest_name,
        model,
        "",
    )