Blame for server/utils/reward_kernel.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any

from pptx.presentation import Presentation as PptxPresentation

from .pptx_functions import PptxEditor
from .presentbench.scoring import (
    score_presentbench,
    score_presentbench_slide,
)
from .presentbench.spec_builder import build_presentbench_eval_spec
from .pptx_extraction import PptxExtractionService
from .reward_metrics import clamp_reward
from .reward_models import (
    EvalSpec,
    IntermediateSlideRewardResult,
    RewardResult,
    ExtractedSlide,
    SourcePack,
    TaskConstraints,
    to_serializable,
)
from .reward_prompts import build_task_spec
from .slidesgenbench.quizbank_service import QuizBankGenerationService
from .slidesgenbench.quantitative_judge import (
    QuantitativeQuizJudgeService,
)
from .slidesgenbench.scoring import score_slidesgenbench
from .slidesgenbench.spec_builder import build_slidesgenbench_eval_spec
from .slidesgenbench.rendered_aesthetics import (
    compute_intermediate_rendered_aesthetics_score,
)
from .slidesgenbench.text_layout import compute_slide_text_layout_scores

SPEC_VERSION = "1.0"
INTERMEDIATE_SLIDESGENBENCH_AESTHETIC_WEIGHT = 0.20
INTERMEDIATE_SLIDESGENBENCH_TEXT_LAYOUT_WEIGHT = 0.35

DEFAULT_REWARD_KERNEL_CONFIG: dict[str, Any] = {
    "branch_weights": {"presentbench": 0.45, "slidesgenbench": 0.55},
}


def _reward_details_payload(
    result: RewardResult | IntermediateSlideRewardResult,
) -> dict[str, Any]:
    return to_serializable(result)


def _write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")


def _persist_eval_spec(cache_dir: str, eval_spec: EvalSpec) -> dict[str, str]:
    root = Path(cache_dir) / eval_spec.spec_hash
    root.mkdir(parents=True, exist_ok=True)
    task_spec_path = root / "task_spec.json"
    presentbench_path = root / "presentbench_eval_spec.json"
    slidesgenbench_path = root / "slidesgenbench_eval_spec.json"
    scoring_config_path = root / "scoring_config.json"
    eval_spec_path = root / "eval_spec.json"

    _write_json(task_spec_path, to_serializable(eval_spec.task_spec))
    _write_json(presentbench_path, to_serializable(eval_spec.presentbench))
    _write_json(slidesgenbench_path, to_serializable(eval_spec.slidesgenbench))
    _write_json(scoring_config_path, eval_spec.scoring_config)
    _write_json(eval_spec_path, to_serializable(eval_spec))
    return {
        "cache_root": str(root),
        "task_spec": str(task_spec_path),
        "presentbench_eval_spec": str(presentbench_path),
        "slidesgenbench_eval_spec": str(slidesgenbench_path),
        "scoring_config": str(scoring_config_path),
        "eval_spec": str(eval_spec_path),
    }


def build_eval_spec(
    prompt: str,
    source_pack: SourcePack,
    task_constraints: TaskConstraints | None = None,
    *,
    quiz_bank_service: QuizBankGenerationService,
    cache_dir: str | None = None,
    mode: str = "eval",
) -> EvalSpec:
    if not source_pack.task_id:
        raise ValueError("source_pack.task_id is required")
    if not source_pack.documents:
        raise ValueError("source_pack.documents must not be empty")

    task_spec = build_task_spec(prompt, source_pack, task_constraints)
    slidesgenbench = build_slidesgenbench_eval_spec(
        task_spec,
        source_pack,
        quiz_bank_service=quiz_bank_service,
        mode=mode,
    )
    presentbench = build_presentbench_eval_spec(task_spec, mode=mode)
    payload = {
        "task_spec": to_serializable(task_spec),
        "presentbench": to_serializable(presentbench),
        "slidesgenbench": to_serializable(slidesgenbench),
        "scoring_config": {**DEFAULT_REWARD_KERNEL_CONFIG, "mode": mode},
        "spec_version": SPEC_VERSION,
    }
    spec_hash = hashlib.sha256(
        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
    ).hexdigest()
    eval_spec = EvalSpec(
        task_spec=task_spec,
        presentbench=presentbench,
        slidesgenbench=slidesgenbench,
        scoring_config={**DEFAULT_REWARD_KERNEL_CONFIG, "mode": mode},
        spec_version=SPEC_VERSION,
        spec_hash=spec_hash,
    )
    if cache_dir:
        artifact_paths = _persist_eval_spec(cache_dir, eval_spec)
        eval_spec.task_spec.metadata["cache_artifacts"] = artifact_paths
    return eval_spec


def _reward_result_for_failure(
    eval_spec: EvalSpec,
    *,
    error: Exception,
    mode: str,
) -> RewardResult:
    return RewardResult(
        reward_total=0.0,
        reward_breakdown={"R_total": 0.0, "R_pb": 0.0, "R_sg": 0.0},
        hard_caps={
            "C_fidelity_critical": 1.0,
            "C_blankness": 1.0,
            "C_hard": 0.0,
        },
        soft_penalties={},
        checklist_results=[],
        quiz_results=[],
        aesthetics_results={},
        artifacts={},
        metadata={
            "task_id": eval_spec.task_spec.task_id,
            "spec_hash": eval_spec.spec_hash,
            "spec_version": eval_spec.spec_version,
            "mode": mode,
            "used_mllm": False,
            "failure_counts": {"inspection": 1},
            "error": str(error),
        },
    )


def evaluate_presentation(
    eval_spec: EvalSpec,
    presentation: PptxEditor | PptxPresentation | str,
    *,
    use_mllm: bool = False,
    judge: Any | None = None,
    render_service: Any | None = None,
    inspection_service: PptxExtractionService | None = None,
    aesthetics_service: Any | None = None,
    quantitative_quiz_judge_service: QuantitativeQuizJudgeService,
    mode: str = "eval",
) -> RewardResult:
    del judge
    inspection_service = inspection_service or PptxExtractionService()

    try:
        presentation_extraction = inspection_service.inspect_presentation(presentation)
    except Exception as error:
        return _reward_result_for_failure(eval_spec, error=error, mode=mode)

    rendered_presentation = None
    render_error: str | None = None
    if render_service and hasattr(render_service, "render_presentation"):
        try:
            rendered_presentation = render_service.render_presentation(presentation)
        except Exception as error:
            render_error = str(error)

    presentbench_result = score_presentbench(
        eval_spec.task_spec,
        presentation_extraction,
        eval_spec.presentbench,
        aesthetics_service=aesthetics_service,
    )
    slidesgenbench_result = score_slidesgenbench(
        eval_spec.task_spec,
        presentation_extraction,
        eval_spec.slidesgenbench,
        quantitative_quiz_judge_service=quantitative_quiz_judge_service,
        rendered_presentation=rendered_presentation,
    )

    branch_weights = eval_spec.scoring_config["branch_weights"]
    reward_total = clamp_reward(
        branch_weights["presentbench"] * presentbench_result.reward_total
        + branch_weights["slidesgenbench"] * slidesgenbench_result.reward_total
    )

    return RewardResult(
        reward_total=reward_total,
        reward_breakdown={
            "R_total": reward_total,
            "R_pb": presentbench_result.reward_total,
            "R_sg": slidesgenbench_result.reward_total,
            **presentbench_result.reward_breakdown,
            **slidesgenbench_result.reward_breakdown,
        },
        hard_caps=presentbench_result.hard_caps,
        soft_penalties=presentbench_result.soft_penalties,
        checklist_results=presentbench_result.checklist_results,
        quiz_results=slidesgenbench_result.quiz_results,
        aesthetics_results=(
            slidesgenbench_result.aesthetics_results
            if slidesgenbench_result.aesthetics_results
            else presentbench_result.aesthetics_results
        ),
        artifacts={
            "presentation_digest": presentation_extraction.metadata.get(
                "presentation_digest"
            ),
            "rendered_presentation": to_serializable(rendered_presentation)
            if rendered_presentation is not None
            else None,
        },
        metadata={
            "task_id": eval_spec.task_spec.task_id,
            "spec_hash": eval_spec.spec_hash,
            "presentbench_spec_hash": eval_spec.presentbench.spec_hash,
            "slidesgenbench_spec_hash": eval_spec.slidesgenbench.spec_hash,
            "spec_version": eval_spec.spec_version,
            "mode": mode,
            "slide_count": presentation_extraction.slide_count,
            "presentbench_item_count": len(eval_spec.presentbench.checklist),
            "slidesgenbench_question_count": len(eval_spec.slidesgenbench.quiz_bank),
            "judge_call_count": 0,
            "failure_counts": {
                "inspection": 0,
                "judge": 0,
                "render": 0 if render_error is None else 1,
            },
            "used_mllm": use_mllm,
            "inspection_mode": presentation_extraction.metadata.get("inspection_mode"),
            "presentation_digest": presentation_extraction.metadata.get(
                "presentation_digest"
            ),
            "render_backend": getattr(rendered_presentation, "backend", None),
            "render_error": render_error,
        },
    )


def evaluate_slide(
    eval_spec: EvalSpec,
    slide_index: int,
    *,
    presentation: PptxEditor | PptxPresentation | None = None,
    slide_extraction: ExtractedSlide | None = None,
    use_mllm: bool = False,
    judge: Any | None = None,
    render_service: Any | None = None,
    inspection_service: PptxExtractionService | None = None,
    aesthetics_service: Any | None = None,
    previous_slide_extractions: list[ExtractedSlide] | None = None,
    mode: str = "eval",
) -> IntermediateSlideRewardResult:
    del judge, aesthetics_service
    inspection_service = inspection_service or PptxExtractionService()
    try:
        if slide_extraction is None:
            if presentation is None:
                raise ValueError("presentation or slide_extraction is required")
            slide_extraction = inspection_service.inspect_slide(
                slide_index,
                presentation=presentation,
            )
    except Exception as error:
        return IntermediateSlideRewardResult(
            slide_index=slide_index,
            reward_total=clamp_reward(0.0),
            reward_breakdown={
                "R_slide": 0.0,
                "S_prompt_alignment": 0.0,
                "S_local_completeness": 0.0,
                "S_local_correctness": 0.0,
                "S_local_fidelity": 0.0,
                "S_local_usability": 0.0,
            },
            hard_caps={
                "C_slide_fidelity_critical": 1.0,
                "C_slide_blankness": 1.0,
                "C_slide_hard": 0.0,
            },
            metadata={
                "slide_index": slide_index,
                "spec_hash": eval_spec.spec_hash,
                "mode": mode,
                "error": str(error),
            },
        )

    presentbench_result = score_presentbench_slide(
        eval_spec.task_spec,
        eval_spec.presentbench,
        slide_index,
        slide_extraction,
        previous_slide_extractions=previous_slide_extractions,
        use_mllm=use_mllm,
        mode=mode,
    )

    rendered_aesthetics: dict[str, Any] = {}
    text_layout_scores = compute_slide_text_layout_scores(
        slide_extraction,
        config=eval_spec.slidesgenbench.scoring_config.get("text_layout_config"),
        weights=eval_spec.slidesgenbench.scoring_config.get("text_layout_weights"),
    )
    render_error: str | None = None
    if (
        presentation is not None
        and render_service
        and hasattr(render_service, "render_presentation")
    ):
        try:
            rendered_presentation = render_service.render_presentation(presentation)
            rendered_slide_map = {
                slide.slide_index: slide for slide in rendered_presentation.slide_images
            }
            current_rendered_slide = rendered_slide_map.get(slide_index)
            if current_rendered_slide is None:
                raise ValueError(
                    f"missing rendered slide for slide_index {slide_index}"
                )

            previous_slide = None
            previous_rendered_slide = None
            if slide_index > 1:
                previous_rendered_slide = rendered_slide_map.get(slide_index - 1)
                if previous_rendered_slide is None:
                    raise ValueError(
                        f"missing rendered previous slide for slide_index {slide_index - 1}"
                    )
                if previous_slide_extractions:
                    previous_slide = previous_slide_extractions[-1]
                else:
                    previous_slide = inspection_service.inspect_slide(
                        slide_index - 1,
                        presentation=presentation,
                    )

            rendered_aesthetics = compute_intermediate_rendered_aesthetics_score(
                current_slide=slide_extraction,
                current_rendered_slide=current_rendered_slide,
                previous_slide=previous_slide,
                previous_rendered_slide=previous_rendered_slide,
                slide_width_in=float(
                    slide_extraction.metadata.get("slide_width_in") or 0.0
                ),
                slide_height_in=float(
                    slide_extraction.metadata.get("slide_height_in") or 0.0
                ),
                metric_weights=eval_spec.slidesgenbench.scoring_config.get(
                    "aesthetic_weights"
                ),
                harmony_config=eval_spec.slidesgenbench.scoring_config.get(
                    "harmony_config"
                ),
                rhythm_config=eval_spec.slidesgenbench.scoring_config.get(
                    "rhythm_config"
                ),
            )
        except Exception as error:
            render_error = str(error)

    s_slide_aesthetic = float(rendered_aesthetics.get("aesthetic", 0.0))
    s_slide_text_layout = float(text_layout_scores.get("text_layout", 0.0))
    aesthetic_weight = INTERMEDIATE_SLIDESGENBENCH_AESTHETIC_WEIGHT
    text_layout_weight = INTERMEDIATE_SLIDESGENBENCH_TEXT_LAYOUT_WEIGHT
    presentbench_weight = max(0.0, 1.0 - aesthetic_weight - text_layout_weight)
    reward_pre_cap = (
        presentbench_weight * presentbench_result.reward_total
        + aesthetic_weight * s_slide_aesthetic
        + text_layout_weight * s_slide_text_layout
    )
    text_layout_hard_cap = float(text_layout_scores.get("hard_cap", 1.0))
    reward_total = clamp_reward(reward_pre_cap * text_layout_hard_cap)
    reward_breakdown = {
        **presentbench_result.reward_breakdown,
        "R_slide_presentbench": presentbench_result.reward_total,
        "R_slide_pre_text_layout_cap": reward_pre_cap,
        "S_slide_aesthetic": s_slide_aesthetic,
        "S_slide_harmony": float(rendered_aesthetics.get("harmony", 0.0)),
        "S_slide_engagement": float(rendered_aesthetics.get("engagement", 0.0)),
        "S_slide_usability_rendered": float(rendered_aesthetics.get("usability", 0.0)),
        "S_slide_rhythm": float(rendered_aesthetics.get("rhythm", 0.0)),
        "S_slide_text_layout": s_slide_text_layout,
        "S_slide_text_bounds": float(text_layout_scores.get("text_bounds", 0.0)),
        "S_slide_text_density": float(text_layout_scores.get("text_density", 0.0)),
        "S_slide_text_overlap": float(text_layout_scores.get("text_overlap", 0.0)),
        "C_slide_text_layout_hard": text_layout_hard_cap,
        "R_slide": reward_total,
    }
    metadata = {
        **presentbench_result.metadata,
        "render_backend": rendered_aesthetics.get("backend"),
        "render_error": render_error,
        "used_rendered_slide_aesthetics": bool(rendered_aesthetics),
        "intermediate_slidesgenbench_aesthetic_weight": aesthetic_weight,
        "intermediate_slidesgenbench_text_layout_weight": text_layout_weight,
        "text_layout": text_layout_scores,
    }
    return IntermediateSlideRewardResult(
        slide_index=presentbench_result.slide_index,
        reward_total=reward_total,
        reward_breakdown=reward_breakdown,
        hard_caps=presentbench_result.hard_caps,
        soft_penalties=presentbench_result.soft_penalties,
        checklist_results=presentbench_result.checklist_results,
        aesthetics_results=rendered_aesthetics,
        artifacts=presentbench_result.artifacts,
        metadata=metadata,
    )


def compute_presentation_reward(
    prompt: str,
    source_pack: SourcePack,
    presentation: PptxEditor | PptxPresentation | str,
    *,
    task_constraints: TaskConstraints | None = None,
    use_mllm: bool = False,
    judge: Any | None = None,
    render_service: Any | None = None,
    inspection_service: PptxExtractionService | None = None,
    aesthetics_service: Any | None = None,
    quiz_bank_service: QuizBankGenerationService,
    quantitative_quiz_judge_service: QuantitativeQuizJudgeService,
    cache_dir: str | None = None,
    mode: str = "eval",
) -> RewardResult:
    eval_spec = build_eval_spec(
        prompt,
        source_pack,
        task_constraints,
        cache_dir=cache_dir,
        mode=mode,
        quiz_bank_service=quiz_bank_service,
    )
    return evaluate_presentation(
        eval_spec,
        presentation,
        use_mllm=use_mllm,
        judge=judge,
        render_service=render_service,
        inspection_service=inspection_service,
        aesthetics_service=aesthetics_service,
        quantitative_quiz_judge_service=quantitative_quiz_judge_service,
        mode=mode,
    )


def compute_intermediate_slide_reward(
    prompt: str,
    source_pack: SourcePack,
    *,
    slide_index: int,
    presentation: PptxEditor | PptxPresentation | None = None,
    slide_extraction: ExtractedSlide | None = None,
    task_constraints: TaskConstraints | None = None,
    use_mllm: bool = False,
    judge: Any | None = None,
    render_service: Any | None = None,
    inspection_service: PptxExtractionService | None = None,
    aesthetics_service: Any | None = None,
    previous_slide_extractions: list[ExtractedSlide] | None = None,
    quiz_bank_service: QuizBankGenerationService,
    cache_dir: str | None = None,
    mode: str = "eval",
) -> IntermediateSlideRewardResult:
    eval_spec = build_eval_spec(
        prompt,
        source_pack,
        task_constraints,
        cache_dir=cache_dir,
        mode=mode,
        quiz_bank_service=quiz_bank_service,
    )
    return evaluate_slide(
        eval_spec,
        slide_index,
        presentation=presentation,
        slide_extraction=slide_extraction,
        use_mllm=use_mllm,
        judge=judge,
        render_service=render_service,
        inspection_service=inspection_service,
        aesthetics_service=aesthetics_service,
        previous_slide_extractions=previous_slide_extractions,
        mode=mode,
    )


__all__ = [
    "build_eval_spec",
    "compute_intermediate_slide_reward",
    "compute_presentation_reward",
    "evaluate_presentation",
    "evaluate_slide",
]