formscout/agents/judge.py

"""
JudgeAgent — VLM-based final scorer with rationale, compensation tags, pain detection.

Input:  BiomechFeatures, ScoreResult (rubric candidate), MovementResult, keyframes
Output: JudgeResult(score, rationale, compensation_tags, corrective_hint, needs_human)
Failure: returns JudgeResult(needs_human=True, score=None) when uncertain.
Model:  Qwen3-VL-8B-Instruct via llama.cpp (8B params, Apache-2.0).
Gated:  No.

Safety: NEVER auto-scores pain. If any indication of pain/clearing test,
        sets needs_human=True and score=None.
"""
from __future__ import annotations

import json
import logging
from pathlib import Path

from formscout import config
from formscout.types import (
    BiomechFeatures, ScoreResult, MovementResult,
    IngestResult, JudgeResult,
)
from formscout.serving import get_vlm_client

logger = logging.getLogger(__name__)

_PROMPT_PATH = Path(__file__).parent / "prompts" / "c2_judge.md"


class JudgeAgent:
    """VLM judge that produces the final FMS score with rationale."""

    def __init__(self):
        self._client = get_vlm_client()
        self._system_prompt = _PROMPT_PATH.read_text(encoding="utf-8")

    def run(
        self,
        features: BiomechFeatures,
        rubric_score: ScoreResult,
        movement: MovementResult,
        ingest: IngestResult | None = None,
    ) -> JudgeResult:
        """
        Produce final score. Falls back to rubric score if VLM unavailable.
        """
        if not config.ENABLE_JUDGE:
            return self._fallback_from_rubric(rubric_score, features)

        if not self._client.available:
            logger.warning("JudgeAgent: VLM unavailable, using rubric score as final")
            return self._fallback_from_rubric(rubric_score, features)

        # Build context for the judge
        context = {
            "test": features.test_name,
            "side": features.side,
            "view": features.view,
            "features": {"angles": features.angles, "alignments": features.alignments},
            "candidate_score": rubric_score.score,
            "candidate_confidence": rubric_score.confidence,
            "exemplars": [],  # Phase 3: populated by RetrievalAgent
        }

        prompt = f"{self._system_prompt}\n\n{json.dumps(context, indent=2)}"

        # Optionally include keyframes
        images = None
        if ingest and ingest.frames:
            images = self._encode_keyframes(ingest.frames)

        result = self._client.complete(prompt, images=images, max_tokens=512, temperature=0.1)
        if result.get("fallback"):
            # transformers backend couldn't load/run — use the deterministic rubric
            return self._fallback_from_rubric(rubric_score, features)
        return self._parse_response(result)

    def _encode_keyframes(self, frames: list) -> list[str]:
        """Encode 3 keyframes for VLM context."""
        import cv2
        import base64

        n = len(frames)
        indices = [0, n // 2, n - 1] if n >= 3 else list(range(n))
        encoded = []
        for idx in indices:
            _, buf = cv2.imencode(".jpg", frames[idx], [cv2.IMWRITE_JPEG_QUALITY, 70])
            encoded.append(base64.b64encode(buf.tobytes()).decode())
        return encoded

    def _parse_response(self, result: dict) -> JudgeResult:
        """Parse VLM JSON response into JudgeResult."""
        if "error" in result:
            return JudgeResult(
                score=None, rationale=f"VLM error: {result['error']}",
                compensation_tags=[], corrective_hint="",
                confidence=0.0, needs_human=True,
            )

        needs_human = result.get("needs_human", False)
        score = result.get("score") if not needs_human else None
        if score is not None:
            score = max(0, min(3, int(score)))

        return JudgeResult(
            score=score,
            rationale=result.get("rationale", ""),
            compensation_tags=result.get("compensation_tags", []),
            corrective_hint=result.get("corrective_hint", ""),
            confidence=float(result.get("confidence", 0.5)),
            needs_human=needs_human,
        )

    def _fallback_from_rubric(self, rubric: ScoreResult, features: BiomechFeatures) -> JudgeResult:
        """When VLM is unavailable, promote the rubric score as the final score."""
        return JudgeResult(
            score=rubric.score,
            rationale=f"[rubric-only] {rubric.rationale}",
            compensation_tags=[],
            corrective_hint="",
            confidence=rubric.confidence * 0.8,
            needs_human=rubric.needs_human,
            notes="VLM unavailable — rubric score used as final",
        )