llm_classifier.py

from config import client
from utils import clean_ai_response, parse_json_safe

def classify_alert(log_entry):
    prompt = f"""

    You are a SOC (Security Operations Center) AI assistant.

    Analyze the following security alert and classify it.



    Alert details:

    Timestamp: {log_entry.get('timestamp')}

    Source IP: {log_entry.get('source_ip')}

    Event: {log_entry.get('event')}

    Threat Intelligence: Reputation = {log_entry.get('ip_reputation')}, Note = {log_entry.get('intel_note')}

    GeoIP: {log_entry.get('geo_location')}



    Tasks:

    1. Categorize the attack type (e.g., brute force, malware, data exfiltration, benign).

    2. Assign a priority (High, Medium, Low).

    3. Suggest next action.



    Return response in JSON with keys: category, priority, action.

    """

    response = client.chat.completions.create(
        model="gemini-2.5-flash",
        messages=[{"role": "user", "content": prompt}],
    )

    raw_output = response.choices[0].message.content
    cleaned = clean_ai_response(raw_output)
    return parse_json_safe(cleaned)