dify-1.8.1/api/schedule/alert_check_task.py

"""
Alert Rule Check Task
Celery Beat task to check alert rules and send notifications
"""

import logging
import time
from datetime import datetime, timedelta, timezone

import click

import app
from extensions.ext_redis import redis_client
from services.alert_service import AlertService

logger = logging.getLogger(__name__)


@app.celery.task(queue="dataset")
def check_alert_rules_task():
    """
    Periodically check all enabled alert rules and trigger notifications if thresholds are exceeded.

    This task should be run every minute via Celery Beat.
    Each rule has its own check_interval_minutes that determines how often it should be checked.
    """
    logger.info(click.style("Start checking alert rules", fg="green"))
    start_at = time.perf_counter()

    try:
        # Get all enabled rules
        rules = AlertService.get_all_enabled_rules()

        logger.info(click.style(f"Found {len(rules)} enabled alert rules", fg="green"))

        for rule in rules:
            try:
                # Check if it's time to check this rule
                now = datetime.now(timezone.utc)
                last_triggered_str = rule.get("last_triggered_at", "")

                if last_triggered_str:
                    try:
                        last_triggered_at = datetime.fromisoformat(last_triggered_str)
                        # Calculate next check time
                        next_check_time = last_triggered_at + timedelta(minutes=int(rule["check_interval_minutes"]))
                        if now < next_check_time:
                            logger.debug(f"Skipping rule {rule['id']} - next check at {next_check_time}, now is {now}")
                            continue
                    except (ValueError, TypeError):
                        # Invalid timestamp, proceed with check
                        pass

                # Get current metric value
                metric_value = AlertService.get_metric_value(
                    rule["app_id"], rule["metric_type"], time_range_minutes=int(rule["check_interval_minutes"])
                )

                logger.debug(
                    f"Rule {rule['id']}: metric_type={rule['metric_type']}, value={metric_value}, "
                    f"threshold={rule['operator']} {rule['threshold_value']}"
                )

                # Check if threshold is exceeded
                is_triggered = AlertService.check_threshold(
                    metric_value, rule["operator"], float(rule["threshold_value"])
                )

                rule_key = AlertService._get_rule_key(rule["id"])

                if is_triggered:
                    # Increment consecutive trigger count
                    consecutive_count = int(rule.get("consecutive_trigger_count", 0)) + 1

                    logger.info(
                        f"Rule {rule['id']} threshold exceeded: {metric_value} {rule['operator']} {rule['threshold_value']} "
                        f"(consecutive: {consecutive_count}/{rule['consecutive_triggers']})"
                    )

                    # Update consecutive count
                    redis_client.hset(rule_key, "consecutive_trigger_count", str(consecutive_count))

                    # Check if we've reached the consecutive trigger threshold
                    if consecutive_count >= int(rule["consecutive_triggers"]):
                        logger.info(f"Rule {rule['id']} reached consecutive trigger threshold, sending notifications")
                        # Trigger alert and send notifications
                        AlertService.trigger_alert(rule, metric_value)
                        # Reset consecutive count after triggering (done in trigger_alert via update)
                        redis_client.hset(rule_key, "consecutive_trigger_count", "0")
                else:
                    # Reset consecutive count if not triggered
                    consecutive_count = int(rule.get("consecutive_trigger_count", 0))
                    if consecutive_count > 0:
                        logger.debug(f"Rule {rule['id']} no longer triggered, resetting consecutive count")
                        redis_client.hset(rule_key, "consecutive_trigger_count", "0")

                # Update last check time
                redis_client.hset(rule_key, "last_triggered_at", now.isoformat())
                redis_client.expire(rule_key, AlertService.RETENTION_SECONDS)

            except Exception:
                logger.exception(f"Error checking alert rule {rule.get('id', 'unknown')}")
                continue

        end_at = time.perf_counter()
        logger.info(click.style(f"Alert rule check completed: latency: {end_at - start_at}", fg="green"))

    except Exception:
        logger.exception("Error in alert check task")