207 lines
6.5 KiB
Python
207 lines
6.5 KiB
Python
"""预警规则引擎 (MVP-1 / P2)
|
||
|
||
YAML 驱动的预警规则引擎，对 ``CandidateEvent`` 列表应用规则，
|
||
满足条件的事件会被升级为 ``AlertEvent``。
|
||
|
||
规则 YAML 示例 (config/rules/fire.yaml)::
|
||
|
||
name: fire_critical
|
||
event_type: fire
|
||
enabled: true
|
||
min_confidence: 0.6
|
||
severity: critical # 覆盖默认严重性 (可选)
|
||
description: 检测到火焰，立即触发预警
|
||
|
||
规则条件支持:
|
||
|
||
- ``min_confidence``: 置信度阈值
|
||
- ``allowed_sources``: 允许的来源 (source_id 白名单，None 表示不限制)
|
||
- ``required_labels``: 检测项 label 必须包含其中之一
|
||
- ``min_bbox_area``: 边界框最小面积
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
import yaml
|
||
|
||
from models.event_schemas import (
|
||
AlertEvent,
|
||
CandidateEvent,
|
||
EventType,
|
||
SeverityLevel,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class AlertRule:
|
||
"""单条预警规则。"""
|
||
|
||
name: str
|
||
event_type: EventType
|
||
enabled: bool = True
|
||
min_confidence: float = 0.0
|
||
severity: Optional[SeverityLevel] = None
|
||
allowed_sources: Optional[List[str]] = None
|
||
required_labels: Optional[List[str]] = None
|
||
min_bbox_area: int = 0
|
||
description: str = ""
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: Dict[str, Any]) -> "AlertRule":
|
||
return cls(
|
||
name=str(data["name"]),
|
||
event_type=EventType(data["event_type"]),
|
||
enabled=bool(data.get("enabled", True)),
|
||
min_confidence=float(data.get("min_confidence", 0.0)),
|
||
severity=SeverityLevel(data["severity"]) if data.get("severity") else None,
|
||
allowed_sources=data.get("allowed_sources"),
|
||
required_labels=data.get("required_labels"),
|
||
min_bbox_area=int(data.get("min_bbox_area", 0)),
|
||
description=str(data.get("description", "")),
|
||
)
|
||
|
||
def matches(self, event: CandidateEvent) -> bool:
|
||
"""判断候选事件是否命中规则。"""
|
||
|
||
if not self.enabled:
|
||
return False
|
||
if event.event_type != self.event_type:
|
||
return False
|
||
if event.confidence < self.min_confidence:
|
||
return False
|
||
if self.allowed_sources and event.source_id not in self.allowed_sources:
|
||
return False
|
||
if self.required_labels:
|
||
labels = {event.detection.class_name, event.detection.label}
|
||
if not any(lbl in labels for lbl in self.required_labels):
|
||
return False
|
||
if self.min_bbox_area > 0 and event.detection.bbox.area < self.min_bbox_area:
|
||
return False
|
||
return True
|
||
|
||
|
||
@dataclass
|
||
class _RuleStats:
|
||
loaded: int = 0
|
||
enabled: int = 0
|
||
files: List[str] = field(default_factory=list)
|
||
|
||
|
||
class AlertRuleEngine:
|
||
"""预警规则引擎。
|
||
|
||
支持从单个 YAML 文件或目录批量加载规则。
|
||
"""
|
||
|
||
def __init__(self, rules: Optional[List[AlertRule]] = None) -> None:
|
||
self.rules: List[AlertRule] = list(rules or [])
|
||
self._stats = _RuleStats(
|
||
loaded=len(self.rules),
|
||
enabled=sum(1 for r in self.rules if r.enabled),
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 加载
|
||
# ------------------------------------------------------------------
|
||
|
||
@classmethod
|
||
def from_directory(cls, rules_dir: str | Path) -> "AlertRuleEngine":
|
||
"""从目录加载所有 ``*.yaml`` / ``*.yml`` 规则文件。"""
|
||
|
||
engine = cls()
|
||
engine.load_directory(rules_dir)
|
||
return engine
|
||
|
||
def load_directory(self, rules_dir: str | Path) -> int:
|
||
path = Path(rules_dir)
|
||
if not path.exists():
|
||
logger.warning("规则目录不存在: %s", path)
|
||
return 0
|
||
|
||
count = 0
|
||
for file in sorted(path.glob("*.y*ml")):
|
||
try:
|
||
count += self.load_file(file)
|
||
except Exception as exc: # noqa: BLE001
|
||
logger.error("加载规则文件失败 %s: %s", file, exc)
|
||
logger.info("规则引擎加载完成: 共 %d 条规则 (来自 %s)", count, path)
|
||
return count
|
||
|
||
def load_file(self, rule_file: str | Path) -> int:
|
||
path = Path(rule_file)
|
||
with path.open("r", encoding="utf-8") as fp:
|
||
data = yaml.safe_load(fp)
|
||
|
||
if data is None:
|
||
return 0
|
||
|
||
# 单条规则 dict 或 列表 或 {"rules": [...]} 三种格式都支持
|
||
if isinstance(data, dict) and "rules" in data:
|
||
rule_items = data["rules"]
|
||
elif isinstance(data, list):
|
||
rule_items = data
|
||
else:
|
||
rule_items = [data]
|
||
|
||
added = 0
|
||
for item in rule_items:
|
||
rule = AlertRule.from_dict(item)
|
||
self.rules.append(rule)
|
||
self._stats.loaded += 1
|
||
if rule.enabled:
|
||
self._stats.enabled += 1
|
||
added += 1
|
||
self._stats.files.append(str(path))
|
||
return added
|
||
|
||
# ------------------------------------------------------------------
|
||
# 评估
|
||
# ------------------------------------------------------------------
|
||
|
||
def evaluate(self, events: List[CandidateEvent]) -> List[AlertEvent]:
|
||
"""对一批候选事件执行规则评估，返回触发的预警事件。"""
|
||
|
||
alerts: List[AlertEvent] = []
|
||
for event in events:
|
||
for rule in self.rules:
|
||
if rule.matches(event):
|
||
alerts.append(self._build_alert(event, rule))
|
||
event.triggered_rules.append(rule.name)
|
||
break # 命中一条即可，避免重复
|
||
return alerts
|
||
|
||
@staticmethod
|
||
def _build_alert(event: CandidateEvent, rule: AlertRule) -> AlertEvent:
|
||
severity = rule.severity or event.severity
|
||
return AlertEvent(
|
||
event_type=event.event_type,
|
||
severity=severity,
|
||
confidence=event.confidence,
|
||
source_id=event.source_id,
|
||
detections=[event.detection],
|
||
rule_name=rule.name,
|
||
metadata={"description": rule.description},
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 自省
|
||
# ------------------------------------------------------------------
|
||
|
||
@property
|
||
def stats(self) -> Dict[str, Any]:
|
||
return {
|
||
"loaded": self._stats.loaded,
|
||
"enabled": self._stats.enabled,
|
||
"files": list(self._stats.files),
|
||
}
|
||
|
||
|
||
__all__ = ["AlertRule", "AlertRuleEngine"]
|