app/agents/supervisor_agent.py

"""Supervisor agent."""

from __future__ import annotations

from app.common.enums import DecisionMode
from app.common.types import PolyGuardState
from app.models.policy.uncertainty import estimate_uncertainty


class SupervisorAgent:
    name = "SupervisorAgent"

    def run(self, state: PolyGuardState, dosing_active: bool) -> dict:
        uncertainty = estimate_uncertainty(state)
        if uncertainty > 0.72:
            mode = DecisionMode.REVIEW
        elif state.sub_environment.value == "PRECISION_DOSING":
            mode = DecisionMode.DOSE_OPT
        elif dosing_active:
            mode = DecisionMode.DOSE_OPT
        else:
            mode = DecisionMode.REGIMEN_OPT
        return {"mode": mode.value, "uncertainty": uncertainty, "sub_environment": state.sub_environment.value}