api/logging_run.py

"""Generate view (miru's Logging Mode): bulk generation with probability logging.

Single-user server-side run state (tracer + originating inputs) lives in
``state.get_logging_session()`` — the streaming handler updates it so Stop and
Continue can reach the live tracer, and Continue is invalidated when the model
changes (object identity, like upstream).
"""

from __future__ import annotations

import time
from collections.abc import Iterator

from api.helpers import ChatValidationError, ui_sampling_params
from api.interactive import _reset_tracer_for_mode
from api.models import model_manager
from api.serialize import error_payload, fig_json
from api.state import get_logging_session
from miru_tracer.core.logging_config import get_logger
from miru_tracer.core.tracer import LLMTracer
from miru_tracer.visualization.plots import (
    get_generation_stats,
    plot_probability_visualizations,
)

logger = get_logger(__name__)


def _finalize(tracer, params, heatmap_ranks, prob_mode) -> dict:
    """Stats, plots, and export for a finished (or stopped) run."""
    stats = get_generation_stats(tracer.history)
    ranks = int(heatmap_ranks) if heatmap_ranks else 10
    if tracer.history:
        ranks = min(ranks, len(tracer.history[0].top_k_tokens))
    figures = plot_probability_visualizations(
        tracer.history,
        top_k=ranks,
        probability_mode=prob_mode,
        temperature=params.temperature,
    )
    return {
        "stats": stats,
        "fig_heatmap": fig_json(figures[0]) if figures else None,
        "fig_confidence": fig_json(figures[1]) if len(figures) > 1 else None,
        "export": tracer.export_to_dict(params),
    }


def _stream_run(
    tracer, params, originals, max_new_tokens, log_top_k, log_full_probs,
    stop_at_eos, heatmap_ranks, prob_mode, start_time,
) -> Iterator[dict]:
    """Shared streaming loop for generate and continue."""
    tokens_before = len(tracer.history)
    for step_data in tracer.generate_stream(
        max_new_tokens=int(max_new_tokens),
        params=params,
        log_top_k=max(int(log_top_k or 10), 1),
        log_full_probs=bool(log_full_probs),
        stop_at_eos=bool(stop_at_eos),
    ):
        yield {
            "ok": True,
            "type": "progress",
            "text": tracer.get_generated_text(),
            "progress": {
                "step": len(tracer.history),
                "new_tokens": len(tracer.history) - tokens_before,
                "last_token": step_data.token_text,
            },
        }
    final = _finalize(tracer, params, heatmap_ranks, prob_mode)
    logger.info(
        f"Generation finished: {len(tracer.history)} tokens in "
        f"{time.time() - start_time:.2f}s"
    )
    yield {
        "ok": True,
        "type": "final",
        "text": tracer.get_generated_text(),
        "originals": list(originals),
        **final,
    }


def logging_generate(
    mode: str, prompt: str, chat_json: str, raw_text: str,
    thinking: str, think_prefill: str,
    max_new_tokens: int, strategy: str, temperature: float,
    top_k: int, top_p: float,
    stop_at_eos: bool, log_top_k: int, heatmap_ranks: int,
    log_full_probs: bool, prob_mode: str,
) -> Iterator[dict]:
    model = model_manager.get_model()
    tokenizer = model_manager.get_tokenizer()
    device = model_manager.get_device()
    run = get_logging_session()

    if model is None or tokenizer is None:
        yield error_payload("No model loaded. Load one in the Model view first.")
        return
    if max_new_tokens is None or max_new_tokens < 1:
        yield error_payload("Maximum new tokens must be at least 1")
        return

    start_time = time.time()
    try:
        params = ui_sampling_params(strategy, temperature, top_k, top_p)
        tracer = LLMTracer(model, tokenizer, device)
        _reset_tracer_for_mode(
            tracer, mode, prompt, chat_json, raw_text, thinking, think_prefill
        )
        originals = (mode, prompt, chat_json, raw_text, thinking, think_prefill)
        with run.lock:
            run.tracer = tracer  # reachable by Stop mid-run
            run.originals = originals
        logger.info(
            f"Logging generation started: mode={mode}, "
            f"max_tokens={max_new_tokens}, strategy={strategy}"
        )
        yield from _stream_run(
            tracer, params, originals, max_new_tokens, log_top_k,
            log_full_probs, stop_at_eos, heatmap_ranks, prob_mode, start_time,
        )
    except ChatValidationError as e:
        with run.lock:
            run.tracer = None
            run.originals = None
        yield error_payload(str(e))
    except Exception as e:
        logger.error(f"Generation error: {e}", exc_info=True)
        with run.lock:
            run.tracer = None
            run.originals = None
        yield error_payload(f"Error during generation:\n\n{e}", trace=True)


def logging_continue(
    max_new_tokens: int, strategy: str, temperature: float,
    top_k: int, top_p: float,
    stop_at_eos: bool, log_top_k: int, heatmap_ranks: int,
    log_full_probs: bool, prob_mode: str,
) -> Iterator[dict]:
    """Resume the same tracer; history accumulates across runs."""
    run = get_logging_session()
    with run.lock:
        tracer = run.tracer
        originals = run.originals
    if tracer is None:
        yield error_payload("No previous generation to continue from.")
        return
    if max_new_tokens is None or max_new_tokens < 1:
        yield error_payload("Maximum new tokens must be at least 1")
        return
    if model_manager.get_model() is None or tracer.model is not model_manager.get_model():
        yield error_payload(
            "Model has been unloaded or changed.\n\nPlease start a new generation."
        )
        return

    start_time = time.time()
    try:
        params = ui_sampling_params(strategy, temperature, top_k, top_p)
        logger.info(f"Continuing generation: max_tokens={max_new_tokens}")
        yield from _stream_run(
            tracer, params, originals, max_new_tokens, log_top_k,
            log_full_probs, stop_at_eos, heatmap_ranks, prob_mode, start_time,
        )
    except Exception as e:
        logger.error(f"Continuation error: {e}", exc_info=True)
        yield error_payload(f"Error during continuation:\n\n{e}", trace=True)


def logging_stop() -> dict:
    """Cooperative stop; the running stream finalizes the partial run itself."""
    run = get_logging_session()
    with run.lock:
        tracer = run.tracer
    if tracer is None:
        return error_payload("No generation running.")
    tracer.request_stop()
    logger.info(f"Stop requested after {len(tracer.history)} tokens")
    return {"ok": True}


def logging_replot(heatmap_ranks: int, prob_mode: str, temperature: float) -> dict:
    """Re-render plots with new display settings — no regeneration."""
    run = get_logging_session()
    with run.lock:
        tracer = run.tracer
    if tracer is None or not tracer.history:
        return {"ok": True, "fig_heatmap": None, "fig_confidence": None}
    try:
        ranks = min(
            int(heatmap_ranks) if heatmap_ranks else 10,
            len(tracer.history[0].top_k_tokens),
        )
        figures = plot_probability_visualizations(
            tracer.history,
            top_k=ranks,
            probability_mode=prob_mode,
            temperature=float(temperature),
        )
        return {
            "ok": True,
            "fig_heatmap": fig_json(figures[0]) if figures else None,
            "fig_confidence": fig_json(figures[1]) if len(figures) > 1 else None,
        }
    except Exception as e:
        logger.error(f"Error refreshing visualizations: {e}", exc_info=True)
        return error_payload(str(e))