Update app2.py

app2.py

@@ -4,6 +4,7 @@ import hashlib
 import os
 from typing import List, Tuple, Dict, Any, Optional
 import datetime
+import signal
 import feedparser
 import aiohttp
 import gradio as gr
@@ -13,6 +14,9 @@ from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.orm import declarative_base, sessionmaker
 from urllib.parse import urljoin
 import logging
+from sqlalchemy.orm import Session
+from sqlalchemy.future import select
+
 Base = declarative_base()
 
 class Article(Base):
@@ -253,11 +257,8 @@ async def main():
     try:
         engine, Session = await create_db_engine("sqlite:///monitoring.db")
         db_session = Session()
-    except SQLAlchemyError as e:
-        logger.error(f"Failed to connect to database: {e}")
-        return
 
-    demo = gr.Blocks()
+        demo = gr.Blocks()
 
     with demo:
         gr.Markdown("# Website Monitor and Chatbot")
@@ -318,23 +319,40 @@ async def main():
         send_button.click(
             chatbot_response,
             inputs=[message_input, chatbot_interface],
-            outputs=[chatbot_interface, message_input])
-
-        
+            outputs=[chatbot_interface, message_input])        
 
-    asyncio.create_task(periodic_update_with_error_handling(db_session))
-    asyncio.create_task(update_db_status(db_status))
-
-    feed_updater = gr.Timer(300)  # 5 minutes
+    # Set up the timer
+        feed_updater = gr.Timer(interval=300)
         feed_updater.tick(fn=update_feed_content,
                          outputs=feed_content)
-    try:
+
+        # Create background tasks
+        update_tasks = [
+            asyncio.create_task(periodic_update_with_error_handling(db_session)),
+            asyncio.create_task(update_db_status(db_status))
+        ]
+
+        # Launch the demo
         await demo.launch()
+
+        # Wait for background tasks to complete
+        for task in update_tasks:
+            task.cancel()
+        await asyncio.gather(*update_tasks, return_exceptions=True)
+
+    except Exception as e:
+        logger.error(f"Error in main: {e}")
     finally:
         if db_session:
             await db_session.close()
+        if engine:
             engine.dispose()
 
+def signal_handler():
+    for task in asyncio.all_tasks():
+        task.cancel()
 
 if __name__ == "__main__":
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
     asyncio.run(main())