Upload 4 files

notify_slack.py

@@ -0,0 +1,27 @@
+import requests
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+SLACK_WEBHOOK_URL = os.getenv("SLACK_WEBHOOK_URL")
+
+def send_slack_alert(log_message, s3_link=None):
+    if not SLACK_WEBHOOK_URL:
+        print(" Slack webhook URL not found in .env")
+        return
+
+    text = f"* Incident Detected!*\n```{log_message}```"
+    if s3_link:
+        text += f"\n [Runbook]({s3_link})"
+
+    payload = {
+        "text": text,
+        "mrkdwn": True
+    }
+
+    response = requests.post(SLACK_WEBHOOK_URL, json=payload)
+
+    if response.status_code == 200:
+        print(" Slack alert sent.")
+    else:
+        print(f" Slack alert failed: {response.status_code} - {response.text}")

requirements.txt

@@ -1,3 +1,5 @@
-altair
-pandas
-streamlit
+streamlit
+torch
+transformers
+boto3
+python-dotenv

streamlit_app.py

@@ -0,0 +1,116 @@
+import streamlit as st
+import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+from upload_to_s3 import upload_file_to_s3
+from notify_slack import send_slack_alert
+
+# Load fine-tuned model
+model_path = "model/codementor-flan"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
+
+# Fallback rule-based classifier
+def fallback_label(log):
+    log_lower = log.lower()
+    if "login" in log_lower and "failed" in log_lower:
+        return "SECURITY"
+    elif "error" in log_lower or "failed" in log_lower:
+        return "ERROR"
+    elif "timeout" in log_lower or "not responding" in log_lower:
+        return "CRITICAL"
+    elif "cpu" in log_lower or "memory" in log_lower:
+        return "WARNING"
+    else:
+        return "INFO"
+
+# Hybrid classifier using Flan-T5 with fallback
+def classify_log(log):
+    prompt = f"""Classify this log message as one of: INFO, WARNING, ERROR, CRITICAL, SECURITY.
+
+Log: {log}
+
+Label:"""
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=3)
+    prediction = tokenizer.decode(outputs[0], skip_special_tokens=True).strip().upper()
+
+    valid_labels = {"INFO", "WARNING", "ERROR", "CRITICAL", "SECURITY"}
+    if prediction in valid_labels:
+        return prediction, "Model"
+    else:
+        return fallback_label(log), "Fallback"
+
+# Runbook generator in markdown format
+def generate_runbook(log_text, label):
+    prompt = f"""
+You are an expert SRE. Create a step-by-step runbook in markdown format for the following {label} log.
+
+Log message: "{log_text}"
+
+Include:
+1. Summary
+2. Possible causes
+3. Troubleshooting steps
+4. Mitigation actions
+5. Responsible team or escalation
+
+Only output valid markdown text.
+"""
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=300)
+    runbook_text = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+    safe_filename = log_text.lower().replace(" ", "_").replace(":", "").replace('"', '').replace(".", "")
+    runbook_path = f"runbooks/runbook_{safe_filename}.md"
+    os.makedirs("runbooks", exist_ok=True)
+    with open(runbook_path, "w", encoding="utf-8") as f:
+        f.write(runbook_text)
+
+    return runbook_path, runbook_text
+
+# Streamlit UI
+st.set_page_config(page_title="WatchTowerAI", layout="centered")
+st.title("WatchTowerAI - Log Classification and Runbook Generator")
+
+# Input: text or file
+log_input = st.text_input("Enter a log message")
+uploaded_file = st.file_uploader("Or upload a .log file", type=["txt", "log"])
+
+# Trigger button
+if st.button("Classify and Generate Runbook"):
+    logs = []
+
+    if log_input:
+        logs.append(log_input.strip())
+
+    if uploaded_file:
+        content = uploaded_file.read().decode("utf-8")
+        logs.extend([line.strip() for line in content.splitlines() if line.strip()])
+
+    if not logs:
+        st.warning("Please enter a log message or upload a file.")
+    else:
+        for log in logs:
+            with st.spinner(f"Processing: {log}"):
+                label, source = classify_log(log)
+                st.markdown(f"**Classification:** `{label}` ({source})")
+                st.markdown(f"**Log:** {log}")
+
+                if label in {"CRITICAL", "SECURITY"}:
+                    runbook_path, runbook_md = generate_runbook(log, label)
+
+                    s3_path = runbook_path.replace("\\", "/")
+                    success = upload_file_to_s3(runbook_path, "watchtowerai-artifacts", s3_path)
+
+                    if success:
+                        s3_url = f"https://watchtowerai-artifacts.s3.ap-south-1.amazonaws.com/{s3_path}"
+                        send_slack_alert(log, s3_url)
+                        st.success("Slack alert sent with S3 link.")
+                        st.markdown(f"[View Runbook on S3]({s3_url})")
+
+                    st.download_button("Download Runbook", runbook_md, file_name=os.path.basename(runbook_path))
+                else:
+                    st.info("No runbook generated for this log.")

upload_to_s3.py

@@ -0,0 +1,27 @@
+import boto3
+from botocore.exceptions import BotoCoreError, ClientError
+import os
+from dotenv import load_dotenv
+
+# Load AWS credentials from .env
+load_dotenv()
+
+aws_access_key = os.getenv("AWS_ACCESS_KEY_ID")
+aws_secret_key = os.getenv("AWS_SECRET_ACCESS_KEY")
+region = os.getenv("AWS_REGION", "ap-south-1")
+
+s3_client = boto3.client(
+    "s3",
+    aws_access_key_id=aws_access_key,
+    aws_secret_access_key=aws_secret_key,
+    region_name=region
+)
+
+def upload_file_to_s3(local_path, bucket, s3_key):
+    try:
+        s3_client.upload_file(local_path, bucket, s3_key)
+        print(f" Uploaded: {local_path} → s3://{bucket}/{s3_key}")
+        return True
+    except (BotoCoreError, ClientError) as e:
+        print(f" Upload failed: {e}")
+        return False