api/analysis.py

"""Log analysis: visualize previously exported generation logs (schema v1+v2).

The frontend reads the uploaded file client-side and POSTs the parsed JSON.
"""

from __future__ import annotations

import numpy as np

from api.serialize import error_payload, fig_json
from miru_tracer.core.logging_config import get_logger
from miru_tracer.core.schema import parse_log
from miru_tracer.visualization.plots import plot_probability_visualizations

logger = get_logger(__name__)


def analyze_log(data: dict, heatmap_ranks: int, prob_mode: str) -> dict:
    if not isinstance(data, dict):
        return error_payload("No log data provided")

    def truncate(text, limit=200):
        return text[:limit] + "..." if len(text) > limit else text

    try:
        log = parse_log(data)
        metadata = {
            "schema_version": log.schema_version,
            "mode": log.mode,
            "prompt": truncate(log.prompt),
            "generated_text": truncate(log.generated_text),
            "timestamp": log.timestamp,
            "num_steps": log.num_steps,
            "sampling_params": log.sampling_params,
        }
        if not log.history:
            return {
                "ok": True,
                "metadata": metadata,
                "stats": "No history data found in log",
                "fig_heatmap": None,
                "fig_confidence": None,
            }

        probs = [step.probability for step in log.history]
        stats_text = (
            f"Mean: {np.mean(probs):.4f}\n"
            f"Std Dev: {np.std(probs):.4f}\n"
            f"Min: {np.min(probs):.4f}\n"
            f"Max: {np.max(probs):.4f}\n"
            f"Median: {np.median(probs):.4f}\n"
            f"Total steps: {len(log.history)}\n"
        )

        # Cap heatmap ranks at what was actually logged
        ranks = min(
            int(heatmap_ranks) if heatmap_ranks else 10,
            len(log.history[0].top_k_tokens),
        )
        figures = plot_probability_visualizations(
            log.history,
            top_k=ranks,
            probability_mode=prob_mode,
            temperature=log.temperature,
        )
        return {
            "ok": True,
            "metadata": metadata,
            "stats": stats_text,
            "fig_heatmap": fig_json(figures[0]) if figures else None,
            "fig_confidence": fig_json(figures[1]) if len(figures) > 1 else None,
        }
    except ValueError as e:
        # parse_log rejects files that aren't miru-tracer/Jacobina logs
        return error_payload(str(e))
    except Exception as e:
        logger.error(f"Log analysis error: {e}", exc_info=True)
        return error_payload(f"Error analyzing log:\n\n{e}", trace=True)