Lotty/src/backend/api/v1/guardrails/endpoints.py

from http import HTTPStatus
from uuid import UUID

from django.http import Http404, HttpRequest
from ninja import Router

from api.v1.guardrails.schemas import (
    ExperimentGuardrailCheckResultOut,
    GuardrailCheckResultOut,
    GuardrailCreateIn,
    GuardrailOut,
    GuardrailTriggerOut,
    GuardrailUpdateIn,
)
from apps.experiments.models import Experiment
from apps.guardrails.models import Guardrail
from apps.guardrails.services import (
    check_all_running_experiments,
    check_experiment_guardrails,
    guardrail_create,
    guardrail_delete,
    guardrail_list,
    guardrail_trigger_list,
    guardrail_update,
)
from apps.metrics.services import metric_definition_get
from apps.users.auth.bearer import jwt_bearer

router = Router(tags=["guardrails"], auth=jwt_bearer)


@router.post(
    "/guardrails/check-all",
    response={HTTPStatus.OK: GuardrailCheckResultOut},
    summary="Check guardrails across all running experiments",
)
def check_all_guardrails(
    request: HttpRequest,
) -> tuple[int, GuardrailCheckResultOut]:
    results = check_all_running_experiments()
    return HTTPStatus.OK, GuardrailCheckResultOut(**results)


@router.post(
    "/experiments/{experiment_id}/guardrails/check",
    response={HTTPStatus.OK: ExperimentGuardrailCheckResultOut},
    summary="Manually trigger guardrail check for experiment",
)
def check_guardrails_for_experiment(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[int, ExperimentGuardrailCheckResultOut]:
    try:
        experiment = (
            Experiment.objects.select_related("flag")
            .prefetch_related(
                "variants",
                "guardrails__metric",
            )
            .get(pk=experiment_id)
        )
    except Experiment.DoesNotExist:
        raise Http404 from Experiment.DoesNotExist
    triggers = check_experiment_guardrails(experiment)
    return HTTPStatus.OK, ExperimentGuardrailCheckResultOut(
        experiment_id=str(experiment.pk),
        triggered=len(triggers),
        triggers=list(triggers),
    )


@router.get(
    "/experiments/{experiment_id}/guardrails",
    response={HTTPStatus.OK: list[GuardrailOut]},
    summary="List experiment guardrails",
)
def list_guardrails(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[int, list[GuardrailOut]]:
    try:
        experiment = Experiment.objects.get(pk=experiment_id)
    except Experiment.DoesNotExist:
        raise Http404 from Experiment.DoesNotExist
    grs = guardrail_list(experiment)
    return HTTPStatus.OK, [GuardrailOut.from_guardrail(g) for g in grs]


@router.post(
    "/experiments/{experiment_id}/guardrails",
    response={HTTPStatus.CREATED: GuardrailOut},
    summary="Create a guardrail for experiment",
)
def create_guardrail(
    request: HttpRequest,
    experiment_id: UUID,
    payload: GuardrailCreateIn,
) -> tuple[int, GuardrailOut]:
    try:
        experiment = Experiment.objects.get(pk=experiment_id)
    except Experiment.DoesNotExist:
        raise Http404 from Experiment.DoesNotExist
    metric = metric_definition_get(payload.metric_id)
    if not metric:
        raise Http404
    g = guardrail_create(
        experiment=experiment,
        metric=metric,
        threshold=payload.threshold,
        observation_window_minutes=payload.observation_window_minutes,
        action=payload.action,
    )
    g = Guardrail.objects.select_related("metric").get(pk=g.pk)
    return HTTPStatus.CREATED, GuardrailOut.from_guardrail(g)


@router.get(
    "/experiments/{experiment_id}/guardrail-triggers",
    response={HTTPStatus.OK: list[GuardrailTriggerOut]},
    summary="List guardrail triggers (audit)",
)
def list_guardrail_triggers(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[int, list[GuardrailTriggerOut]]:
    try:
        experiment = Experiment.objects.get(pk=experiment_id)
    except Experiment.DoesNotExist:
        raise Http404 from Experiment.DoesNotExist
    triggers = guardrail_trigger_list(experiment)
    return HTTPStatus.OK, [
        GuardrailTriggerOut.model_validate(trigger) for trigger in triggers
    ]


@router.patch(
    "/experiments/{experiment_id}/guardrails/{guardrail_id}",
    response={HTTPStatus.OK: GuardrailOut},
    summary="Update a guardrail",
)
def update_guardrail(
    request: HttpRequest,
    experiment_id: UUID,
    guardrail_id: UUID,
    payload: GuardrailUpdateIn,
) -> tuple[int, GuardrailOut]:
    try:
        g = Guardrail.objects.select_related("metric", "experiment").get(
            pk=guardrail_id,
            experiment_id=experiment_id,
        )
    except Guardrail.DoesNotExist:
        raise Http404 from Guardrail.DoesNotExist
    g = guardrail_update(
        guardrail=g,
        **payload.model_dump(exclude_unset=True),
    )
    g = Guardrail.objects.select_related("metric").get(pk=g.pk)
    return HTTPStatus.OK, GuardrailOut.from_guardrail(g)


@router.delete(
    "/experiments/{experiment_id}/guardrails/{guardrail_id}",
    response={HTTPStatus.NO_CONTENT: None},
    summary="Delete a guardrail",
)
def delete_guardrail_endpoint(
    request: HttpRequest,
    experiment_id: UUID,
    guardrail_id: UUID,
) -> tuple[int, None]:
    try:
        g = Guardrail.objects.select_related("experiment").get(
            pk=guardrail_id,
            experiment_id=experiment_id,
        )
    except Guardrail.DoesNotExist:
        raise Http404 from Guardrail.DoesNotExist
    guardrail_delete(guardrail=g)
    return HTTPStatus.NO_CONTENT, None