Kavin2615/meta-openenv-hackathon-demo

1"""2learning.py - RL-Compatible Learning Agent for AI Misuse Triage.3"""4import json5import os6from inference import RuleBasedAgent7 8class LearningAgent(RuleBasedAgent):9    """10    An extension of the RuleBasedAgent that supports simple reward-based 11    weight updating (lightweight RL approach).12    """13    def __init__(self, memory_file="agent_memory.json", log_file="training_log.jsonl"):14        super().__init__()15        self.memory_file = memory_file16        self.log_file = log_file17        18        # Initialize default weights for each rule index19        # We identify rules by their index in self._RULES20        self.rule_weights = {str(i): 1.0 for i in range(len(self._RULES))}21        self.load_memory()22 23    def load_memory(self):24        """Load saved rule weights from disk."""25        if os.path.exists(self.memory_file):26            try:27                with open(self.memory_file, "r") as f:28                    data = json.load(f)29                    # Update weights; keep defaults for newly added rules if any30                    for k, v in data.items():31                        self.rule_weights[k] = v32            except Exception as e:33                print(f"Warning: Failed to load agent memory: {e}")34 35    def save_memory(self):36        """Save rule weights to disk."""37        try:38            with open(self.memory_file, "w") as f:39                json.dump(self.rule_weights, f, indent=2)40        except Exception as e:41            print(f"Warning: Failed to save agent memory: {e}")42 43    def decide(self, observation: str) -> dict:44        """45        Evaluate observation using rule weights and return the highest scored decision.46        """47        obs_lower = observation.lower()48        49        best_score = -float('inf')50        best_decision = None51 52        # Check rule list and pick the one with the highest weight53        for idx, (keywords, risk, category, action) in enumerate(self._RULES):54            if any(kw.lower() in obs_lower for kw in keywords):55                weight = self.rule_weights.get(str(idx), 1.0)56                if weight > best_score:57                    best_score = weight58                    best_decision = {59                        "risk_label": risk,60                        "category": category,61                        "action": action,62                        "rationale": f"(Confidence Score: {weight:.2f}) Detected indicator keyword(s) matching '{category}' pattern.",63                        "rule_idx": str(idx) # Internal tracking for reward targeting64                    }65 66        if best_decision is None:67            # Fallback to benign patterns68            for pat in self._BENIGN_PATTERNS:69                if pat in obs_lower:70                    best_decision = {71                        "risk_label": "benign",72                        "category": "writing_assistance",73                        "action": "allow",74                        "rationale": "No harmful indicators detected. Content appears to be a routine informational or programming request.",75                        "rule_idx": "benign"76                    }77                    break78 79        if best_decision is None:80            # Absolute default behavior:81            # Instead of unconditionally defaulting to 'suspicious', we perform a 82            # FINAL SAFETY CHECK.83            ambiguous_signals = ["security", "credential", "login", "access", "bypass", "password"]84            if any(sig in obs_lower for sig in ambiguous_signals):85                best_decision = {86                    "risk_label": "suspicious",87                    "category": "other",88                    "action": "warn",89                    "rationale": "Could not classify with high confidence. Contains ambiguous terminology.",90                    "rule_idx": "default"91                }92            else:93                best_decision = {94                    "risk_label": "benign",95                    "category": "writing_assistance",96                    "action": "allow",97                    "rationale": "No harmful indicators detected. Content appears to be a routine request for assistance.",98                    "rule_idx": "safe_default"99                }100            101        return best_decision102 103    def update_policy(self, decision: dict, reward: float):104        """105        Update the weight of the rule that produced the decision based on the reward signal.106        Reward is expected to be positive (+1.0) for good behavior, or negative (-1.0) for bad.107        """108        rule_idx = decision.get("rule_idx")109        # We only update dynamic rules, not the fallback "benign" or "default" states110        if rule_idx and rule_idx not in ["benign", "default"]:111            current_weight = self.rule_weights.get(rule_idx, 1.0)112            # Lightweight QA-update: increase/decrease weight113            new_weight = current_weight + (0.5 * reward)114            # Cap the weight to prevent explosion115            new_weight = max(0.1, min(new_weight, 10.0))116            self.rule_weights[rule_idx] = new_weight117            self.save_memory()118 119    def log_episode(self, episode_data: dict):120        """121        Log the full episode context and reward into a JSONL file.122        """123        try:124            with open(self.log_file, "a") as f:125                f.write(json.dumps(episode_data) + "\n")126        except Exception as e:127            print(f"Warning: Failed to log episode: {e}")128