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.dict(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