app/api/schemas.py

"""API schemas."""

from __future__ import annotations

from typing import Any, Optional

from pydantic import BaseModel, ConfigDict, Field

from app.common.enums import ActionType, DecisionMode, Difficulty, DoseBucket, SubEnvironment


class StrictSchema(BaseModel):
    model_config = ConfigDict(extra="forbid")


class ResetRequest(StrictSchema):
    task_id: Optional[str] = None
    seed: Optional[int] = None
    difficulty: Optional[Difficulty] = None
    sub_environment: Optional[SubEnvironment] = None
    scenario_id: Optional[str] = None
    patient_id: Optional[str] = None


class StepRequest(StrictSchema):
    mode: DecisionMode
    action_type: ActionType
    target_drug: Optional[str] = None
    replacement_drug: Optional[str] = None
    dose_bucket: DoseBucket
    taper_days: Optional[int] = None
    monitoring_plan: Optional[str] = None
    evidence_query: Optional[str] = None
    new_drug_name: Optional[str] = None
    candidate_components: list[str] = Field(default_factory=list)
    candidate_id: str
    confidence: float
    rationale_brief: str


class StepCandidateRequest(StrictSchema):
    candidate_id: str
    confidence: float
    rationale_brief: str


class OrchestrateRequest(StrictSchema):
    coordination_mode: Optional[str] = None


class BatchInferRequest(StrictSchema):
    batch_size: int = 4


class EvidenceQueryRequest(StrictSchema):
    query: str
    top_k: int = 5