95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
"""
|
|
PolicyEngine — applies static rules + learned heuristics + live overrides
|
|
to scheduler decisions.
|
|
|
|
Constructor accepts an optional `db` (anything with an `events` dict-like
|
|
attribute). If db is None, the heuristics layer is skipped.
|
|
"""
|
|
|
|
|
|
class PolicyEngine:
|
|
|
|
def __init__(self, db=None):
|
|
self.db = db
|
|
self.rules = []
|
|
self.overrides = {}
|
|
|
|
# -----------------------------
|
|
# REGISTER RULE
|
|
# -----------------------------
|
|
def add_rule(self, rule):
|
|
self.rules.append(rule)
|
|
|
|
# -----------------------------
|
|
# SET LIVE OVERRIDE
|
|
# -----------------------------
|
|
def set_override(self, key, value):
|
|
self.overrides[key] = value
|
|
|
|
# -----------------------------
|
|
# CLEAR OVERRIDE
|
|
# -----------------------------
|
|
def clear_override(self, key):
|
|
if key in self.overrides:
|
|
del self.overrides[key]
|
|
|
|
# -----------------------------
|
|
# APPLY POLICY TO NODE SELECTION
|
|
# -----------------------------
|
|
def evaluate(self, action, target, node):
|
|
score_modifier = 0
|
|
for rule in self.rules:
|
|
try:
|
|
score_modifier += rule.apply(action, target, node)
|
|
except Exception:
|
|
pass
|
|
score_modifier += self._heuristics(action, target, node)
|
|
score_modifier += self._overrides(action, target, node)
|
|
return score_modifier
|
|
|
|
# -----------------------------
|
|
# HEURISTICS (LEARNING LAYER)
|
|
# -----------------------------
|
|
def _heuristics(self, action, target, node):
|
|
if self.db is None:
|
|
return 0
|
|
if not hasattr(self.db, "events"):
|
|
return 0
|
|
try:
|
|
history = self.db.events
|
|
except Exception:
|
|
return 0
|
|
|
|
score = 0
|
|
success_count = 0
|
|
# history may be dict-of-lists or a flat list
|
|
if isinstance(history, dict):
|
|
iterable = (e for sess in history.values() for e in (sess or []))
|
|
else:
|
|
iterable = history or []
|
|
|
|
for event in iterable:
|
|
if not isinstance(event, dict):
|
|
continue
|
|
if event.get("type") in ("action_end", "task_update"):
|
|
d = event.get("data", event)
|
|
if (d.get("action") == action.get("name")
|
|
and d.get("state") == "done"
|
|
and d.get("node") == node.get("name")):
|
|
success_count += 1
|
|
score += success_count * 2
|
|
return score
|
|
|
|
# -----------------------------
|
|
# LIVE OVERRIDES
|
|
# -----------------------------
|
|
def _overrides(self, action, target, node):
|
|
score = 0
|
|
forced = self.overrides.get("force_node")
|
|
if forced and node.get("name") == forced:
|
|
score += 1000
|
|
avoid = self.overrides.get("avoid_node")
|
|
if avoid and node.get("name") == avoid:
|
|
score -= 1000
|
|
return score
|