Coverage for agentos/tools/policy_engine.py: 39%
210 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 12:29 +0800
1"""
2PolicyEngine — declarative rule engine with conditions, actions, and priorities.
4Supports:
5 - Rule DSL: WHEN {condition} THEN {action} WITH priority {int}
6 - Conditions: equals, contains, regex, gt/lt/gte/lte, in_set, exists, custom callable
7 - Actions: set, log, reject, allow, call, chain
8 - Priority ordering + first-match / all-match modes
9 - Dynamic rule add/remove at runtime
10 - JSON serialization for persistence
11"""
13from __future__ import annotations
15import json
16import logging
17import re
18from dataclasses import dataclass, field
19from enum import Enum
20from typing import Any
22logger = logging.getLogger(__name__)
25# ============================================================================
26# Condition Operators
27# ============================================================================
30class Op(Enum):
31 EQ = "eq"
32 NE = "ne"
33 GT = "gt"
34 GTE = "gte"
35 LT = "lt"
36 LTE = "lte"
37 CONTAINS = "contains"
38 REGEX = "regex"
39 IN_SET = "in_set"
40 EXISTS = "exists"
41 CUSTOM = "custom"
43 def evaluate(self, actual: Any, expected: Any) -> bool:
44 if self == Op.EQ:
45 return actual == expected
46 elif self == Op.NE:
47 return actual != expected
48 elif self == Op.GT:
49 return actual > expected
50 elif self == Op.GTE:
51 return actual >= expected
52 elif self == Op.LT:
53 return actual < expected
54 elif self == Op.LTE:
55 return actual <= expected
56 elif self == Op.CONTAINS:
57 return expected in actual if isinstance(actual, (str, list, set, tuple)) else False
58 elif self == Op.REGEX:
59 return bool(re.search(expected, actual)) if isinstance(actual, str) else False
60 elif self == Op.IN_SET:
61 return actual in expected if isinstance(expected, (list, set, tuple)) else False
62 elif self == Op.EXISTS:
63 return actual is not None
64 elif self == Op.CUSTOM:
65 # expected is a callable
66 return bool(expected(actual))
67 return False
70# ============================================================================
71# Condition
72# ============================================================================
75@dataclass
76class Condition:
77 """A single condition: field OP value."""
79 field: str
80 op: Op
81 value: Any
83 def evaluate(self, context: dict[str, Any]) -> bool:
84 actual = context.get(self.field)
85 return self.op.evaluate(actual, self.value)
87 def to_dict(self) -> dict[str, Any]:
88 result: dict[str, Any] = {"field": self.field, "op": self.op.value}
89 if self.op != Op.EXISTS:
90 result["value"] = self.value if self.op != Op.CUSTOM else "<callable>"
91 return result
93 @classmethod
94 def from_dict(cls, d: dict[str, Any]) -> Condition:
95 return cls(field=d["field"], op=Op(d["op"]), value=d.get("value"))
98# ============================================================================
99# Actions
100# ============================================================================
103class ActionType(Enum):
104 SET = "set"
105 LOG = "log"
106 REJECT = "reject"
107 ALLOW = "allow"
108 CALL = "call"
109 CHAIN = "chain"
111 def execute(
112 self,
113 context: dict[str, Any],
114 params: Any,
115 ) -> dict[str, Any] | None:
116 result: dict[str, Any] | None = None
117 if self == ActionType.SET:
118 if isinstance(params, dict):
119 for k, v in params.items():
120 context[k] = v
121 elif self == ActionType.LOG:
122 logger.info("policy_engine: %s", params)
123 elif self == ActionType.REJECT:
124 result = {"action": "reject", "reason": str(params)}
125 elif self == ActionType.ALLOW:
126 result = {"action": "allow"}
127 elif self == ActionType.CALL:
128 if callable(params):
129 params(context)
130 elif self == ActionType.CHAIN:
131 if isinstance(params, list):
132 for sub_action in params:
133 if isinstance(sub_action, Action):
134 sub_action.execute(context)
135 return result
137 @classmethod
138 def from_str(cls, s: str) -> ActionType:
139 return ActionType(s.lower())
142# ============================================================================
143# Action
144# ============================================================================
147@dataclass
148class Action:
149 """Action definition."""
151 type: ActionType
152 params: Any = None
154 def execute(self, context: dict[str, Any]) -> dict[str, Any] | None:
155 return self.type.execute(context, self.params)
157 def to_dict(self) -> dict[str, Any]:
158 result = {"type": self.type.value}
159 if self.params is not None:
160 result["params"] = self.params if not callable(self.params) else "<callable>"
161 return result
163 @classmethod
164 def from_dict(cls, d: dict[str, Any]) -> Action:
165 return cls(type=ActionType.from_str(d["type"]), params=d.get("params"))
168# ============================================================================
169# Rule
170# ============================================================================
173@dataclass
174class Rule:
175 """A single rule: when all conditions match, execute actions."""
177 name: str
178 conditions: list[Condition] = field(default_factory=list)
179 actions: list[Action] = field(default_factory=list)
180 priority: int = 0
181 enabled: bool = True
182 description: str = ""
184 def matches(self, context: dict[str, Any]) -> bool:
185 if not self.enabled:
186 return False
187 return all(c.evaluate(context) for c in self.conditions)
189 def evaluate(self, context: dict[str, Any]) -> dict[str, Any] | None:
190 """Evaluate rule against context. Returns action result if triggered."""
191 if not self.matches(context):
192 return None
193 results = []
194 for action in self.actions:
195 r = action.execute(context)
196 if r:
197 results.append(r)
198 return results[-1] if results else None
200 def to_dict(self) -> dict[str, Any]:
201 return {
202 "name": self.name,
203 "conditions": [c.to_dict() for c in self.conditions],
204 "actions": [a.to_dict() for a in self.actions],
205 "priority": self.priority,
206 "enabled": self.enabled,
207 "description": self.description,
208 }
210 @classmethod
211 def from_dict(cls, d: dict[str, Any]) -> Rule:
212 return cls(
213 name=d["name"],
214 conditions=[Condition.from_dict(c) for c in d.get("conditions", [])],
215 actions=[Action.from_dict(a) for a in d.get("actions", [])],
216 priority=d.get("priority", 0),
217 enabled=d.get("enabled", True),
218 description=d.get("description", ""),
219 )
222# ============================================================================
223# Match Mode
224# ============================================================================
227class MatchMode(Enum):
228 FIRST = "first" # Stop after first matching rule
229 ALL = "all" # Evaluate all rules, collect results
232# ============================================================================
233# PolicyEngine
234# ============================================================================
237class PolicyEngine:
238 """Declarative rule engine.
240 Usage:
241 pe = PolicyEngine()
243 pe.add_rule(
244 name="admin-access",
245 conditions=[
246 Condition(field="role", op=Op.EQ, value="admin"),
247 ],
248 actions=[Action(type=ActionType.ALLOW)],
249 priority=100,
250 )
252 pe.add_rule(
253 name="rate-limit",
254 conditions=[
255 Condition(field="requests_per_min", op=Op.GT, value=100),
256 ],
257 actions=[Action(type=ActionType.REJECT, params="rate limit exceeded")],
258 priority=50,
259 )
261 result = pe.evaluate({"role": "admin", "requests_per_min": 150})
262 # → {"action": "allow"} (higher priority matches first)
263 """
265 def __init__(self, mode: MatchMode = MatchMode.FIRST):
266 self._rules: list[Rule] = []
267 self._mode = mode
268 self._rule_names: set[str] = set()
270 # ---------- CRUD ----------
272 def add_rule(
273 self,
274 name: str,
275 conditions: list[Condition] | None = None,
276 actions: list[Action] | None = None,
277 priority: int = 0,
278 enabled: bool = True,
279 description: str = "",
280 ) -> Rule:
281 if name in self._rule_names:
282 raise ValueError(f"Rule '{name}' already exists")
283 rule = Rule(
284 name=name,
285 conditions=conditions or [],
286 actions=actions or [],
287 priority=priority,
288 enabled=enabled,
289 description=description,
290 )
291 self._rules.append(rule)
292 self._rule_names.add(name)
293 self._sort()
294 return rule
296 def remove_rule(self, name: str) -> bool:
297 before = len(self._rules)
298 self._rules = [r for r in self._rules if r.name != name]
299 self._rule_names.discard(name)
300 return len(self._rules) < before
302 def get_rule(self, name: str) -> Rule | None:
303 for r in self._rules:
304 if r.name == name:
305 return r
306 return None
308 def enable_rule(self, name: str) -> bool:
309 rule = self.get_rule(name)
310 if rule:
311 rule.enabled = True
312 return True
313 return False
315 def disable_rule(self, name: str) -> bool:
316 rule = self.get_rule(name)
317 if rule:
318 rule.enabled = False
319 return True
320 return False
322 def _sort(self) -> None:
323 self._rules.sort(key=lambda r: -r.priority)
325 # ---------- evaluation ----------
327 def evaluate(self, context: dict[str, Any]) -> dict[str, Any] | None:
328 """Evaluate rules against context."""
329 results = []
330 for rule in self._rules:
331 if not rule.enabled:
332 continue
333 if rule.matches(context):
334 result = rule.evaluate(context)
335 results.append({"rule": rule.name, "result": result})
336 if self._mode == MatchMode.FIRST:
337 break
338 if not results:
339 return None
340 if self._mode == MatchMode.FIRST:
341 return results[0]
342 return {"matches": results}
344 def evaluate_all(self, context: dict[str, Any]) -> list[dict[str, Any]]:
345 """Evaluate all rules, return list of matches (ignores mode)."""
346 matches = []
347 for rule in self._rules:
348 if rule.matches(context):
349 result = rule.evaluate(context)
350 if result:
351 matches.append({"rule": rule.name, "result": result})
352 return matches
354 # ---------- serialization ----------
356 def to_dict(self) -> dict[str, Any]:
357 return {"rules": [r.to_dict() for r in self._rules], "mode": self._mode.value}
359 def to_json(self) -> str:
360 return json.dumps(self.to_dict(), indent=2, ensure_ascii=False)
362 @classmethod
363 def from_dict(cls, d: dict[str, Any]) -> PolicyEngine:
364 pe = cls(mode=MatchMode(d.get("mode", "first")))
365 for rule_d in d.get("rules", []):
366 pe.add_rule(
367 name=rule_d["name"],
368 conditions=[Condition.from_dict(c) for c in rule_d.get("conditions", [])],
369 actions=[Action.from_dict(a) for a in rule_d.get("actions", [])],
370 priority=rule_d.get("priority", 0),
371 enabled=rule_d.get("enabled", True),
372 description=rule_d.get("description", ""),
373 )
374 return pe
376 @classmethod
377 def from_json(cls, json_str: str) -> PolicyEngine:
378 return cls.from_dict(json.loads(json_str))
380 @property
381 def rules(self) -> list[Rule]:
382 return list(self._rules)