ria-toolkit-oss/src/ria_toolkit_oss/server/models.py

"""Pydantic request and response models for the RT-OSS HTTP server."""

from __future__ import annotations

from pydantic import BaseModel

# ---------------------------------------------------------------------------
# Orchestrator
# ---------------------------------------------------------------------------


class DeployRequest(BaseModel):
    config: dict


class DeployResponse(BaseModel):
    campaign_id: str


class CampaignStatusResponse(BaseModel):
    campaign_id: str
    status: str
    config_name: str
    progress: int
    total_steps: int
    started_at: float
    ended_at: float | None = None
    result: dict | None = None
    error: str | None = None


class CancelResponse(BaseModel):
    campaign_id: str
    cancelled: bool


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------


class SdrConfig(BaseModel):
    device: str
    center_freq: float
    sample_rate: float
    gain: float | str = "auto"


class LoadModelRequest(BaseModel):
    model_path: str
    label_map: dict[str, int]  # class_name -> class_index


class LoadModelResponse(BaseModel):
    loaded: bool
    model_path: str
    num_classes: int


class StartInferenceRequest(BaseModel):
    sdr_config: SdrConfig


class StartInferenceResponse(BaseModel):
    running: bool


class StopInferenceResponse(BaseModel):
    stopped: bool


class ConfigureRequest(BaseModel):
    """Partial SDR reconfiguration — only supplied fields are updated."""

    center_freq: float | None = None
    sample_rate: float | None = None
    gain: float | str | None = None


class ConfigureResponse(BaseModel):
    configured: bool


class InferenceStatusResponse(BaseModel):
    """Latest inference result as returned by GET /inference/status.

    When ``idle`` is True the radio is scanning but no signal was detected.
    ``device_id`` is the raw prediction label from the model's label map.
    The frontend is responsible for mapping device_id to a human name and
    determining whether the device is authorized.
    """

    timestamp: float
    idle: bool = False
    device_id: str | None = None  # prediction label; None when idle
    confidence: float = 0.0
    snr_db: float = 0.0