Create rag_agent.py

rag_agent.py

@@ -0,0 +1,75 @@
+import os
+import google.generativeai as genai
+from dotenv import load_dotenv
+
+# --- CLOUD-READY KEY LOADING ---
+# 1. Try to load from environment (Cloud/Hugging Face)
+api_key = os.getenv("GEMINI_API_KEY")
+
+# 2. If not found, try loading from local .env file (Local Dev)
+if not api_key:
+    load_dotenv()
+    api_key = os.getenv("GEMINI_API_KEY")
+
+if not api_key:
+    # If still missing, just warn (prevents crash on startup, fails later if called)
+    print("⚠️ Warning: GEMINI_API_KEY not found. Agent functionality will fail.")
+else:
+    genai.configure(api_key=api_key)
+
+# --- MODEL CONFIGURATION ---
+print("📡 Connecting to Gemini 2.5 Flash...")
+try:
+    # Try the latest model first
+    model = genai.GenerativeModel('gemini-2.5-flash')
+except Exception as e:
+    print(f"⚠️ Model 2.5 not found, falling back to Pro. Error: {e}")
+    model = genai.GenerativeModel('gemini-pro')
+
+class SentinelAgent:
+    def __init__(self):
+        # Simulated Vector DB (Log Store)
+        self.system_logs = {
+            "CPU_SPIKE": "Log 10:42am - Process 'minerd' started using 99% CPU. Unknown user 'xmr_bot'.",
+            "MEMORY_LEAK": "Log 10:45am - OutOfMemoryError: Java Heap Space. Service 'PaymentGateway' crashed.",
+            "NETWORK_LAG": "Log 10:50am - DDOS detected from IP Block 192.168.x.x. Latency increased to 5000ms."
+        }
+
+    def investigate(self, anomaly_value, z_score):
+        print("🤖 Sentinel Agent analyzing logs...")
+        
+        # Simple heuristic retrieval
+        if anomaly_value > 100:
+            context = self.system_logs["CPU_SPIKE"]
+        elif anomaly_value > 80:
+            context = self.system_logs["MEMORY_LEAK"]
+        else:
+            context = self.system_logs["NETWORK_LAG"]
+
+        prompt = f"""
+        You are Sentinel, an Autonomous MLOps Agent.
+        
+        **ALERT:** Anomaly Detected!
+        - Metric Value: {anomaly_value}
+        - Deviation: {z_score:.2f} sigma
+        
+        **RETRIEVED LOGS:**
+        "{context}"
+        
+        **TASK:**
+        Identify the root cause and recommend a fix. Short and technical.
+        """
+        
+        try:
+            response = model.generate_content(prompt)
+            return response.text
+        except Exception as e:
+            return f"Agent Error: {str(e)}"
+
+# --- Quick Test Block ---
+if __name__ == "__main__":
+    agent = SentinelAgent()
+    print("🔥 Simulation: Testing Agent...")
+    report = agent.investigate(120, 4.5)
+    print("\n--- 📄 REPORT ---")
+    print(report)