dify-1.8.1/api/controllers/console/app/alert.py

"""Alert rule management API controllers."""

from flask import request
from flask_login import current_user
from flask_restx import Resource, reqparse
from werkzeug.exceptions import Forbidden, NotFound

from controllers.console import api
from controllers.console.app.wraps import get_app_model
from controllers.console.wraps import account_initialization_required, setup_required
from libs.login import login_required
from services.alert_service import AlertService


class AlertRuleListApi(Resource):
    """List and create alert rules for an app."""

    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    def get(self, app_model):
        """Get all alert rules for an app."""
        # Check permission
        if app_model.tenant_id != current_user.current_tenant_id:
            raise Forbidden()

        rules = AlertService.get_alert_rules(str(app_model.id))
        return {"data": rules}

    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    def post(self, app_model):
        """Create a new alert rule."""
        # Check permission
        if app_model.tenant_id != current_user.current_tenant_id:
            raise Forbidden()

        parser = reqparse.RequestParser()
        parser.add_argument("name", type=str, required=True, location="json")
        parser.add_argument("description", type=str, required=False, location="json", default="")
        parser.add_argument("is_enabled", type=bool, required=False, location="json", default=True)
        parser.add_argument("metric_type", type=str, required=True, location="json")
        parser.add_argument("operator", type=str, required=True, location="json")
        parser.add_argument("threshold_value", type=float, required=True, location="json")
        parser.add_argument("check_interval_minutes", type=int, required=False, location="json", default=5)
        parser.add_argument("consecutive_triggers", type=int, required=False, location="json", default=2)
        parser.add_argument("notification_channels", type=dict, required=True, location="json")
        args = parser.parse_args()

        # Validate metric_type
        valid_metric_types = ["error_rate", "avg_response_time", "response_time_p95", "response_time_p99"]
        if args["metric_type"] not in valid_metric_types:
            return {
                "code": "invalid_param",
                "message": f"metric_type must be one of: {', '.join(valid_metric_types)}",
            }, 400

        # Validate operator
        valid_operators = [">", "<", ">=", "<=", "="]
        if args["operator"] not in valid_operators:
            return {"code": "invalid_param", "message": f"operator must be one of: {', '.join(valid_operators)}"}, 400

        # Validate notification_channels
        if not args["notification_channels"]:
            return {"code": "invalid_param", "message": "notification_channels cannot be empty"}, 400

        # Validate check_interval_minutes
        if args["check_interval_minutes"] < 1 or args["check_interval_minutes"] > 1440:
            return {"code": "invalid_param", "message": "check_interval_minutes must be between 1 and 1440"}, 400

        # Validate consecutive_triggers
        if args["consecutive_triggers"] < 1 or args["consecutive_triggers"] > 10:
            return {"code": "invalid_param", "message": "consecutive_triggers must be between 1 and 10"}, 400

        try:
            rule = AlertService.create_alert_rule(
                app_id=str(app_model.id), data=args, created_by=str(current_user.id)
            )
            return rule, 201
        except ValueError as e:
            return {"code": "invalid_param", "message": str(e)}, 400


class AlertRuleDetailApi(Resource):
    """Get, update, or delete a specific alert rule."""

    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    def get(self, app_model, rule_id):
        """Get a specific alert rule."""
        # Check permission
        if app_model.tenant_id != current_user.current_tenant_id:
            raise Forbidden()

        rule = AlertService.get_alert_rule(rule_id)
        if not rule:
            raise NotFound("Alert rule not found")

        # Verify rule belongs to this app
        if rule.get("app_id") != str(app_model.id):
            raise Forbidden()

        return rule

    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    def patch(self, app_model, rule_id):
        """Update an alert rule."""
        # Check permission
        if app_model.tenant_id != current_user.current_tenant_id:
            raise Forbidden()

        rule = AlertService.get_alert_rule(rule_id)
        if not rule:
            raise NotFound("Alert rule not found")

        # Verify rule belongs to this app
        if rule.get("app_id") != str(app_model.id):
            raise Forbidden()

        parser = reqparse.RequestParser()
        parser.add_argument("name", type=str, required=False, location="json")
        parser.add_argument("description", type=str, required=False, location="json")
        parser.add_argument("is_enabled", type=bool, required=False, location="json")
        parser.add_argument("metric_type", type=str, required=False, location="json")
        parser.add_argument("operator", type=str, required=False, location="json")
        parser.add_argument("threshold_value", type=float, required=False, location="json")
        parser.add_argument("check_interval_minutes", type=int, required=False, location="json")
        parser.add_argument("consecutive_triggers", type=int, required=False, location="json")
        parser.add_argument("notification_channels", type=dict, required=False, location="json")
        args = parser.parse_args()

        # Remove None values
        update_data = {k: v for k, v in args.items() if v is not None}

        # Validate metric_type if provided
        if "metric_type" in update_data:
            valid_metric_types = ["error_rate", "avg_response_time", "response_time_p95", "response_time_p99"]
            if update_data["metric_type"] not in valid_metric_types:
                return {
                    "code": "invalid_param",
                    "message": f"metric_type must be one of: {', '.join(valid_metric_types)}",
                }, 400

        # Validate operator if provided
        if "operator" in update_data:
            valid_operators = [">", "<", ">=", "<=", "="]
            if update_data["operator"] not in valid_operators:
                return {
                    "code": "invalid_param",
                    "message": f"operator must be one of: {', '.join(valid_operators)}",
                }, 400

        # Validate check_interval_minutes if provided
        if "check_interval_minutes" in update_data:
            if update_data["check_interval_minutes"] < 1 or update_data["check_interval_minutes"] > 1440:
                return {"code": "invalid_param", "message": "check_interval_minutes must be between 1 and 1440"}, 400

        # Validate consecutive_triggers if provided
        if "consecutive_triggers" in update_data:
            if update_data["consecutive_triggers"] < 1 or update_data["consecutive_triggers"] > 10:
                return {"code": "invalid_param", "message": "consecutive_triggers must be between 1 and 10"}, 400

        try:
            updated_rule = AlertService.update_alert_rule(rule_id, update_data)
            return updated_rule
        except ValueError as e:
            return {"code": "invalid_param", "message": str(e)}, 400

    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    def delete(self, app_model, rule_id):
        """Delete an alert rule."""
        # Check permission
        if app_model.tenant_id != current_user.current_tenant_id:
            raise Forbidden()

        rule = AlertService.get_alert_rule(rule_id)
        if not rule:
            raise NotFound("Alert rule not found")

        # Verify rule belongs to this app
        if rule.get("app_id") != str(app_model.id):
            raise Forbidden()

        try:
            AlertService.delete_alert_rule(rule_id)
            return {"result": "success"}
        except ValueError as e:
            return {"code": "invalid_param", "message": str(e)}, 400


class AlertRuleLogsApi(Resource):
    """Get alert logs for a specific rule."""

    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    def get(self, app_model, rule_id):
        """Get alert logs for a rule."""
        # Check permission
        if app_model.tenant_id != current_user.current_tenant_id:
            raise Forbidden()

        rule = AlertService.get_alert_rule(rule_id)
        if not rule:
            raise NotFound("Alert rule not found")

        # Verify rule belongs to this app
        if rule.get("app_id") != str(app_model.id):
            raise Forbidden()

        parser = reqparse.RequestParser()
        parser.add_argument("limit", type=int, required=False, location="args", default=100)
        args = parser.parse_args()

        # Validate limit
        limit = min(args["limit"], 1000)  # Cap at 1000

        logs = AlertService.get_alert_logs(rule_id, limit=limit)
        return {"data": logs}


class AlertNotificationConfigApi(Resource):
    """Get and save notification channel configuration for an app."""

    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    def get(self, app_model):
        """Get notification channel configuration."""
        # Check permission
        if app_model.tenant_id != current_user.current_tenant_id:
            raise Forbidden()

        config = AlertService.get_notification_config(str(app_model.id))
        return config

    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    def post(self, app_model):
        """Save notification channel configuration."""
        # Check permission
        if app_model.tenant_id != current_user.current_tenant_id:
            raise Forbidden()

        parser = reqparse.RequestParser()
        parser.add_argument("email", type=dict, required=False, location="json", default={})
        parser.add_argument("dingtalk", type=dict, required=False, location="json", default={})
        parser.add_argument("wecom", type=dict, required=False, location="json", default={})
        args = parser.parse_args()

        try:
            config = AlertService.save_notification_config(str(app_model.id), args)
            return config
        except ValueError as e:
            return {"code": "invalid_param", "message": str(e)}, 400


# Register API endpoints
api.add_resource(AlertRuleListApi, "/apps/<uuid:app_id>/alert-rules")
api.add_resource(AlertRuleDetailApi, "/apps/<uuid:app_id>/alert-rules/<string:rule_id>")
api.add_resource(AlertRuleLogsApi, "/apps/<uuid:app_id>/alert-rules/<string:rule_id>/logs")
api.add_resource(AlertNotificationConfigApi, "/apps/<uuid:app_id>/alert-notification-config")