ai_triage.py

import json
import os
from dotenv import load_dotenv
from groq import Groq
from triage import check_ip_reputation, log_decision, rule_based_triage

load_dotenv()
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

SYSTEM_PROMPT = """You are a NOC/SOC L1 triage analyst. Given an alert and a rule-based
suggestion, decide the FINAL severity (P1, P2, P3, or P4), who to escalate to, and a short
reason. Agree with the rule-based suggestion unless the alert details clearly justify a
different call. Respond ONLY with valid JSON, no markdown, no preamble, in this exact shape:
{"severity": "P1", "reason": "...", "escalate_to": "...", "agrees_with_rules": true}"""


def ai_triage_alert(alert):
    """Runs rule-based triage first, then asks the LLM to confirm or override with reasoning."""
    reputation = check_ip_reputation(alert["ip"]) if "ip" in alert else None
    rule_severity, rule_reason, rule_escalate = rule_based_triage(
        alert, reputation
    )

    user_prompt = f"""
Alert: {json.dumps(alert)}
IP reputation score (0-100, higher=worse, null if no IP): {reputation}
Rule-based suggestion: severity={rule_severity}, reason="{rule_reason}", escalate_to="{rule_escalate}"
Confirm or override this suggestion with your own judgment.
"""
    try:
        response = client.chat.completions.create(
            model="openai/gpt-oss-120b",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.2,
            max_tokens=600,
            reasoning_effort="medium",      # keeps reasoning short so more budget goes to the JSON answer
            response_format={"type": "json_object"}
        )
        raw = response.choices[0].message.content.strip()
        # Strip accidental markdown fences if the model adds them
        raw = raw.replace("```json", "").replace("```", "").strip()
        if not raw:
             raise ValueError("Empty response from model")
        ai_result = json.loads(raw)

        result = {
            "severity": ai_result["severity"],
            "reason": ai_result["reason"],
            "escalate_to": ai_result["escalate_to"],
            "ip_reputation_score": reputation,
            "rule_suggested": rule_severity,
            "ai_agreed_with_rules": ai_result.get("agrees_with_rules", None),
        }
        log_decision(alert, result, source="ai")
        return result

    except Exception as e:
        # AI failed or returned bad JSON -> fall back to rules, never crash the pipeline
        print(f"AI layer failed ({e}), falling back to rule-based result.")
        result = {
            "severity": rule_severity,
            "reason": rule_reason,
            "escalate_to": rule_escalate,
            "ip_reputation_score": reputation,
            "rule_suggested": rule_severity,
            "ai_agreed_with_rules": None,
        }
        log_decision(alert, result, source="rule_fallback")
        return result


if __name__ == "__main__":
    test_alerts = [
        {
            "type": "failed_logins",
            "host": "10.0.0.5",
            "count": 47,
            "ip": "185.220.101.4",
        },
        {"type": "host_down", "host": "web-server-02"},
        {"type": "disk_space", "host": "db-server-01", "percent_used": 94},
    ]

    for alert in test_alerts:
        result = ai_triage_alert(alert)
        print(f"\nALERT: {alert}")
        print(
            f" -> AI Severity: {result['severity']} (rules suggested {result['rule_suggested']})"
        )
        print(f" -> Reason: {result['reason']}")
        print(f" -> Escalate to: {result['escalate_to']}")
        print(f" -> Agreed with rules: {result['ai_agreed_with_rules']}")