app/agents/dosing_agent.py

"""Dosing analysis agent."""

from __future__ import annotations

from app.common.types import PolyGuardState
from app.knowledge.drug_catalog import DRUG_CLASSES
from app.models.dosing.dose_policy_features import build_dose_features
from app.models.dosing.infer import infer_dosing_quality
from app.models.dosing.pkpd_state import PKPDState
from app.models.dosing.surrogate_pkpd import step_pkpd


class DosingAgent:
    name = "DosingAgent"

    def run(self, state: PolyGuardState) -> dict:
        sensitive_classes = {"anticoagulant", "sedative", "glucose_lowering"}
        dose_sensitive = [
            m.drug
            for m in state.patient.medications
            if DRUG_CLASSES.get(m.drug) in sensitive_classes
        ][:3]
        analyses: list[dict] = []
        for drug in dose_sensitive:
            feats = build_dose_features(state.patient, drug)
            base_state = PKPDState(
                effect_level=min(1.0, 0.35 + feats["adherence"] * 0.45),
                toxicity_level=min(1.0, 0.08 + feats["organ_stress"] * 0.4),
                underdose_risk=max(0.0, 1.0 - (0.35 + feats["adherence"] * 0.45)),
                organ_stress=feats["organ_stress"],
                interaction_load=feats["interaction_load"],
            )
            lower = infer_dosing_quality(step_pkpd(base_state, dose_delta=-0.2, organ_factor=feats["organ_stress"], interaction_factor=feats["interaction_load"]))
            hold = infer_dosing_quality(step_pkpd(base_state, dose_delta=0.0, organ_factor=feats["organ_stress"], interaction_factor=feats["interaction_load"]))
            higher = infer_dosing_quality(step_pkpd(base_state, dose_delta=0.2, organ_factor=feats["organ_stress"], interaction_factor=feats["interaction_load"]))
            analyses.append(
                {
                    "drug": drug,
                    "features": feats,
                    "options": {
                        "reduce": lower,
                        "hold": hold,
                        "increase": higher,
                    },
                }
            )
        return {
            "dose_sensitive_drugs": dose_sensitive,
            "dosing_active": bool(dose_sensitive),
            "recommend_mode": "DOSE_OPT" if dose_sensitive else "REGIMEN_OPT",
            "analyses": analyses,
        }