Update app.py

app.py

@@ -1,49 +1,44 @@
+from langchain_mongodb import MongoDBAtlasVectorSearch
+from langchain_huggingface import HuggingFaceEmbeddings
+# from dotenv import load_dotenv
 import os
-import gradio as gr
-from langchain_community.vectorstores import MongoDBAtlasVectorSearch
-from langchain_community.embeddings import HuggingFaceEmbeddings
 import pymongo
 import logging
 import nest_asyncio
 from langchain.docstore.document import Document
 import redis
-import asyncio
 import threading
-import time
+import asyncio
+import gradio as gr
 
-# Config
-nest_asyncio.apply()
-logging.basicConfig(level=logging.INFO)
+# config
+# nest_asyncio.apply()
+logging.basicConfig(level = logging.INFO)
 database = "AlertSimAndRemediation"
 collection = "alert_embed"
 stream_name = "alerts"
 
-# Environment variables
-MONGO_URI = os.getenv('MONGO_URI')
-REDIS_HOST = os.getenv('REDIS_HOST')
-REDIS_PWD = os.getenv('REDIS_PWD')
+# Global variables to store alert information
+latest_alert = "No alerts yet."
+alert_count = 0
 
-# Embedding model
+# embedding model
 embedding_args = {
-    "model_name": "BAAI/bge-large-en-v1.5",
-    "model_kwargs": {"device": "cpu"},
-    "encode_kwargs": {"normalize_embeddings": True}
+    "model_name" : "BAAI/bge-large-en-v1.5",
+    "model_kwargs" : {"device": "cpu"},
+    "encode_kwargs" : {"normalize_embeddings": True}
 }
 embedding_model = HuggingFaceEmbeddings(**embedding_args)
 
-# MongoDB connection
-connection = pymongo.MongoClient(MONGO_URI)
+# Mongo Connection
+connection = pymongo.MongoClient(os.environ["MONGO_URI"])
 alert_collection = connection[database][collection]
 
 # Redis connection
-r = redis.Redis(host=REDIS_HOST, password=REDIS_PWD, port=16652)
-
-# Global variables to store alert information
-latest_alert = "No alerts yet."
-alert_count = 0
+r = redis.Redis(host=os.environ['REDIS_HOST'], password=os.environ['REDIS_PWD'], port=16652)
 
 # Preprocessing
-def create_textual_description(entry_data):
+async def create_textual_description(entry_data):
     entry_dict = {k.decode(): v.decode() for k, v in entry_data.items()}
 
     category = entry_dict["Category"]
@@ -59,7 +54,7 @@ def create_textual_description(entry_data):
     return description, entry_dict
 
 # Saving alert doc
-def save(entry):
+async def save(entry):
     vector_search = MongoDBAtlasVectorSearch.from_documents(
         documents=[Document(
             page_content=entry["content"],
@@ -69,55 +64,65 @@ def save(entry):
         collection=alert_collection,
         index_name="alert_index",
     )
-    logging.info("Alert stored successfully!")
+    logging.info("Alerts stored successfully!")
 
 # Listening to alert stream
-def listen_to_alerts():
+async def listen_to_alerts(r):
+    global latest_alert, alert_count
+    try:
+        last_id = '$'
+
+        while True:
+            entries = r.xread({stream_name: last_id}, block=0, count=None)
+
+            if entries:
+                stream, new_entries = entries[0]
+
+                for entry_id, entry_data in new_entries:
+                    description, entry_dict = await create_textual_description(entry_data)
+                    await save({
+                        "content" : description,
+                        "metadata" : entry_dict
+                    })
+                    print(description)
+                    latest_alert = description
+                    alert_count += 1
+                    # Update the last ID read
+                    last_id = entry_id
+                    await asyncio.sleep(1) 
+
+    except KeyboardInterrupt:
+        print("Exiting...")
+
+def run_alert_listener():
+    asyncio.run(listen_to_alerts(r))
+
+# Start the alert listener thread
+alert_thread = threading.Thread(target=run_alert_listener)
+alert_thread.start()
+
+# gradio interface
+# Gradio interface
+def get_latest_alert():
     global latest_alert, alert_count
-    last_id = '$'
-
-    while True:
-        entries = r.xread({stream_name: last_id}, block=1000, count=None)
-
-        if entries:
-            stream, new_entries = entries[0]
-
-            for entry_id, entry_data in new_entries:
-                description, entry_dict = create_textual_description(entry_data)
-                save({
-                    "content": description,
-                    "metadata": entry_dict
-                })
-                latest_alert = description
-                alert_count += 1
-                last_id = entry_id
-
-# Start listening to alerts in a separate thread
-threading.Thread(target=listen_to_alerts, daemon=True).start()
-
-# Function to get current stats
-def get_current_stats():
     return latest_alert, f"Total Alerts: {alert_count}"
 
-# Gradio interface
-def create_interface():
-    with gr.Blocks() as iface:
-        gr.Markdown("# Alert Monitoring Service")
-        with gr.Row():
-            latest_alert_md = gr.Markdown("Waiting for alerts...")
-        with gr.Row():
-            alert_count_md = gr.Markdown("Total Alerts: 0")
-        
-        def update_stats():
-            while True:
-                time.sleep(1)  # Update every second
-                yield get_current_stats()
-        
-        iface.load(update_stats, None, [latest_alert_md, alert_count_md], every=1)
+with gr.Blocks() as app:
+    gr.Markdown("# Alert Dashboard 🔔")
+    
+    with gr.Row():
+        latest_alert_box = gr.Textbox(label="Latest Alert", lines=3, interactive=False)
+        alert_count_box = gr.Textbox(label="Alert Count", interactive=False)
+    
+    refresh_button = gr.Button("Refresh")
     
-    return iface
+    refresh_button.click(get_latest_alert, inputs=None, outputs=[latest_alert_box, alert_count_box])
+    
+    app.load(get_latest_alert, inputs=None, outputs=[latest_alert_box, alert_count_box])
+
+    # Auto-refresh every 5 seconds
+    app.load(get_latest_alert, inputs=None, outputs=[latest_alert_box, alert_count_box], every=5)
 
 # Launch the app
-if __name__ == "__main__":
-    iface = create_interface()
-    iface.queue().launch()
+# if __name__ == "__main__":
+app.launch()