server/batch_predict.py

from __future__ import annotations

from typing import Any

from clip_quality_env.rubric import RubricState


def predict_clip_label(
    clip: dict[str, Any],
    rubric: RubricState,
    reasoning_context: str = "",
) -> tuple[str, str, float]:
    """
    Deterministic label prediction for a single clip.

    Returns:
        (label, reasoning, confidence)
    """
    label = rubric.derive_label(clip)
    dominant_features = rubric.get_dominant_features(clip)

    reasoning_parts: list[str] = []
    for feature in dominant_features:
        value = clip.get(feature)
        if not isinstance(value, (int, float)):
            continue
        status = rubric.get_feature_status(feature, float(value))
        reasoning_parts.append(f"{feature}={value:.3g} ({status})")

    signal_summary = (
        ", ".join(reasoning_parts) if reasoning_parts else "no dominant numeric signals"
    )
    context = reasoning_context.strip()
    if context:
        reasoning = f"{context} | Signals: {signal_summary}"
    else:
        reasoning = f"Label {label} based on: {signal_summary}"

    confidence = 0.85 if label != "BORDERLINE" else 0.55
    return label, reasoning, confidence


def batch_predict(
    clips: list[dict[str, Any]],
    rubric: RubricState,
    reasoning_context: str = "",
) -> list[dict[str, Any]]:
    """
    Predict labels for all clips in the episode.

    Returns:
        list of dictionaries with clip_id, label, reasoning, confidence.
    """
    results: list[dict[str, Any]] = []
    for clip in clips:
        label, reasoning, confidence = predict_clip_label(
            clip=clip,
            rubric=rubric,
            reasoning_context=reasoning_context,
        )
        results.append(
            {
                "clip_id": str(clip.get("clip_id", "")),
                "label": label,
                "reasoning": reasoning,
                "confidence": confidence,
            }
        )
    return results