app/training/sft_dataset.py

"""SFT dataset helpers."""

from __future__ import annotations

from app.common.types import CandidateAction, PolyGuardState


def build_sft_example(state: PolyGuardState, candidates: list[CandidateAction], target_candidate_id: str) -> dict:
    return {
        "prompt": {
            "patient_id": state.patient.patient_id,
            "medications": [m.model_dump(mode="json") for m in state.patient.medications],
            "candidates": [c.model_dump(mode="json") for c in candidates],
        },
        "target_candidate_id": target_candidate_id,
    }