polyguard-rl/app/api/routes.py

"""API routes."""

from __future__ import annotations

from fastapi import APIRouter, Depends, HTTPException

from app.api.dependencies import get_service
from app.tools.medication_alternatives import build_alternatives_response
from app.api.schemas import (
    BatchInferRequest,
    EvidenceQueryRequest,
    MedicationAlternativesRequest,
    OrchestrateRequest,
    ResetRequest,
    StepCandidateRequest,
    StepRequest,
)
from app.api.service import APIService

router = APIRouter()


@router.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@router.post("/env/reset")
def env_reset(payload: ResetRequest, service: APIService = Depends(get_service)) -> dict:
    try:
        return service.reset(**payload.model_dump(mode="json"))
    except ValueError as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc


@router.post("/env/step")
def env_step(payload: StepRequest, service: APIService = Depends(get_service)) -> dict:
    return service.step(payload.model_dump(mode="json"))


@router.post("/env/step_candidate")
def env_step_candidate(payload: StepCandidateRequest, service: APIService = Depends(get_service)) -> dict:
    result = service.step_candidate(
        candidate_id=payload.candidate_id,
        confidence=payload.confidence,
        rationale_brief=payload.rationale_brief,
    )
    if result is None:
        raise HTTPException(status_code=404, detail=f"Candidate {payload.candidate_id!r} is not legal in this state.")
    return result


@router.get("/env/catalog")
def env_catalog(service: APIService = Depends(get_service)) -> dict:
    return service.catalog()


@router.get("/env/state")
def env_state(service: APIService = Depends(get_service)) -> dict:
    return service.env.get_state()


@router.get("/env/trace")
def env_trace(service: APIService = Depends(get_service)) -> list[dict]:
    return service.env.get_trace()


@router.get("/env/legal_actions")
def env_legal_actions(service: APIService = Depends(get_service)) -> list[dict]:
    return service.env.get_legal_actions()


@router.get("/env/reward_breakdown")
def env_reward_breakdown(service: APIService = Depends(get_service)) -> dict:
    return service.env.get_reward_breakdown()


@router.get("/env/uncertainty")
def env_uncertainty(service: APIService = Depends(get_service)) -> dict:
    return service.env.get_uncertainty_report().model_dump(mode="json")


@router.post("/agents/orchestrate")
def agents_orchestrate(
    payload: OrchestrateRequest = OrchestrateRequest(),
    service: APIService = Depends(get_service),
) -> dict:
    return service.orchestrate(coordination_mode=payload.coordination_mode)


@router.post("/policy/infer")
def policy_infer(service: APIService = Depends(get_service)) -> dict:
    return service.infer_policy()


@router.get("/policy/model_status")
def policy_model_status(service: APIService = Depends(get_service)) -> dict:
    return service.model_status()


@router.post("/policy/batch_infer")
def policy_batch_infer(
    payload: BatchInferRequest = BatchInferRequest(),
    service: APIService = Depends(get_service),
) -> list[dict]:
    return service.batch_infer(batch_size=payload.batch_size)


@router.post("/eval/run_baselines")
def eval_baselines(service: APIService = Depends(get_service)) -> dict:
    return service.run_baselines()


@router.post("/eval/run_policy")
def eval_run_policy(service: APIService = Depends(get_service)) -> dict:
    return service.run_policy_eval()


@router.post("/eval/run_dosing")
def eval_run_dosing(service: APIService = Depends(get_service)) -> dict:
    return service.run_dosing_eval()


@router.get("/metrics/training")
def metrics_training(service: APIService = Depends(get_service)) -> dict:
    return service.get_metrics()


@router.get("/cases/sample")
def cases_sample(service: APIService = Depends(get_service)) -> dict:
    return service.sample_case()


@router.get("/cases/search")
def cases_search(q: str, service: APIService = Depends(get_service)) -> list[dict]:
    return service.search_cases(q)


@router.post("/evidence/query")
def evidence_query(payload: EvidenceQueryRequest, service: APIService = Depends(get_service)) -> list[dict]:
    return service.evidence_query(query=payload.query, top_k=payload.top_k)


@router.post("/tools/medication_alternatives")
def medication_alternatives(payload: MedicationAlternativesRequest) -> dict:
    """OpenFDA class neighbors + optional external POST (env: POLYGUARD_MED_TOOL_URL / TOKEN)."""
    return build_alternatives_response(
        query_drug=payload.query_drug,
        regimen_drugs=payload.regimen_drugs,
        max_suggestions=payload.max_suggestions,
    )