server/reward.py

from __future__ import annotations

from dataclasses import dataclass

from graders.health import HealthReport
from graders.security import SecurityReport


@dataclass
class RewardBreakdown:
    security_score: float
    health_score: float
    total_reward: float

    @classmethod
    def from_reports(
        cls,
        security: SecurityReport,
        health: HealthReport,
        initial_leaks: int | None = None,
    ) -> "RewardBreakdown":
        if initial_leaks is None:
            initial_leaks = security.finding_count
        if initial_leaks <= 0:
            security_score = 1.0
        else:
            fixed = max(initial_leaks - security.finding_count, 0)
            security_score = round(fixed / initial_leaks, 4)
        health_score = round(health.score, 4)
        total_reward = round(security_score * health_score, 4)
        return cls(
            security_score=security_score,
            health_score=health_score,
            total_reward=total_reward,
        )