app/models/baselines/beam_search_planner.py

"""Constrained beam search baseline."""

from __future__ import annotations

from app.common.types import CandidateAction, PolyGuardAction


def choose_beam_search(candidates: list[CandidateAction], beam_width: int = 3) -> PolyGuardAction:
    legal = [c for c in candidates if c.legality_precheck]
    if not legal:
        legal = candidates
    topk = sorted(legal, key=lambda c: (c.estimated_safety_delta + c.burden_delta), reverse=True)[:beam_width]
    chosen = topk[0]
    return PolyGuardAction(
        mode=chosen.mode,
        action_type=chosen.action_type,
        target_drug=chosen.target_drug,
        replacement_drug=chosen.replacement_drug,
        dose_bucket=chosen.dose_bucket,
        taper_days=chosen.taper_days,
        monitoring_plan=chosen.monitoring_plan,
        candidate_id=chosen.candidate_id,
        confidence=0.74,
        rationale_brief=f"Beam-search({beam_width}) top candidate.",
    )