app/models/baselines/imitation.py

"""Imitation baseline from logged actions."""

from __future__ import annotations

from app.common.types import CandidateAction, PolyGuardAction
from app.models.baselines.rules_only import choose_rules_only


def choose_imitation(candidates: list[CandidateAction], preferred_candidate_id: str | None = None) -> PolyGuardAction:
    if preferred_candidate_id:
        for c in candidates:
            if c.candidate_id == preferred_candidate_id:
                return PolyGuardAction(
                    mode=c.mode,
                    action_type=c.action_type,
                    target_drug=c.target_drug,
                    replacement_drug=c.replacement_drug,
                    dose_bucket=c.dose_bucket,
                    taper_days=c.taper_days,
                    monitoring_plan=c.monitoring_plan,
                    candidate_id=c.candidate_id,
                    confidence=0.7,
                    rationale_brief="Imitation-selected candidate from demonstration.",
                )
    return choose_rules_only(candidates)