Update app.py

app.py

@@ -1,32 +1,41 @@
 import os
 import threading
 from flask import Flask, render_template, request, jsonify
-from rss_processor import fetch_rss_feeds, process_and_store_articles, download_from_hf_hub, upload_to_hf_hub, clean_text
+from rss_processor import fetch_rss_feeds, process_and_store_articles, download_from_hf_hub, upload_to_hf_hub, clean_text, LOCAL_DB_DIR
 import logging
 import time
 from datetime import datetime
 import hashlib
-import glob
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
 
 app = Flask(__name__)
 
-# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Global flag to track background loading
-loading_complete = True  # Start as True to allow initial rendering
+loading_complete = True
 last_update_time = time.time()
-last_data_hash = None  # Track the hash of the last data to detect changes
+last_data_hash = None
 
 def get_embedding_model():
-    """Returns a singleton instance of the embedding model to avoid reloading."""
     if not hasattr(get_embedding_model, "model"):
         get_embedding_model.model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     return get_embedding_model.model
 
+def get_vector_db():
+    if not os.path.exists(LOCAL_DB_DIR):
+        return None
+    try:
+        return Chroma(
+            persist_directory=LOCAL_DB_DIR,
+            embedding_function=get_embedding_model(),
+            collection_name="news_articles"
+        )
+    except Exception as e:
+        logger.error(f"Failed to load vector DB: {e}")
+        return None
+
 def load_feeds_in_background():
     global loading_complete, last_update_time
     try:
@@ -42,196 +51,117 @@ def load_feeds_in_background():
     finally:
         loading_complete = True
 
-def get_all_docs_from_dbs():
-    """Aggregate documents and metadata from all Chroma DB folders."""
-    all_docs = {'documents': [], 'metadatas': []}
-    seen_ids = set()
-    embedding_function = get_embedding_model()
-
-    for db_path in glob.glob("chroma_db_*"):
-        if not os.path.isdir(db_path):
-            continue
-        try:
-            temp_vector_db = Chroma(
-                persist_directory=db_path,
-                embedding_function=embedding_function,
-                collection_name="news_articles"
-            )
-            # Skip empty databases
-            if temp_vector_db._collection.count() == 0:
-                continue
-
-            db_data = temp_vector_db.get(include=['documents', 'metadatas'])
-            if db_data.get('documents') and db_data.get('metadatas'):
-                for doc, meta in zip(db_data['documents'], db_data['metadatas']):
-                    # Use a more robust unique identifier
-                    doc_id = f"{meta.get('title', 'No Title')}|{meta.get('link', '')}|{meta.get('published', 'Unknown Date')}"
-                    if doc_id not in seen_ids:
-                        seen_ids.add(doc_id)
-                        all_docs['documents'].append(doc)
-                        all_docs['metadatas'].append(meta)
-        except Exception as e:
-            logger.error(f"Error loading DB {db_path}: {e}")
-    
-    return all_docs
+def get_all_docs_from_db():
+    vector_db = get_vector_db()
+    if not vector_db or vector_db._collection.count() == 0:
+        return {'documents': [], 'metadatas': []}
+    return vector_db.get(include=['documents', 'metadatas'])
 
 def compute_data_hash(categorized_articles):
-    """Compute a hash of the current articles to detect changes."""
-    if not categorized_articles:
-        return ""
-    # Create a sorted string representation of the articles for consistent hashing
+    if not categorized_articles: return ""
     data_str = ""
     for cat, articles in sorted(categorized_articles.items()):
         for article in sorted(articles, key=lambda x: x["published"]):
             data_str += f"{cat}|{article['title']}|{article['link']}|{article['published']}|"
     return hashlib.sha256(data_str.encode('utf-8')).hexdigest()
 
+def process_docs_into_articles(docs_data):
+    enriched_articles = []
+    seen_keys = set()
+    for doc, meta in zip(docs_data['documents'], docs_data['metadatas']):
+        if not meta: continue
+        title = meta.get("title", "No Title")
+        link = meta.get("link", "")
+        description = meta.get("original_description", "No Description")
+        published = meta.get("published", "Unknown Date").strip()
+
+        key = f"{title}|{link}|{published}"
+        if key not in seen_keys:
+            seen_keys.add(key)
+            try:
+                published_iso = datetime.strptime(published, "%Y-%m-%d %H:%M:%S").isoformat()
+            except (ValueError, TypeError):
+                published_iso = "1970-01-01T00:00:00"
+            
+            enriched_articles.append({
+                "title": title,
+                "link": link,
+                "description": description,
+                "category": meta.get("category", "Uncategorized"),
+                "published": published_iso,
+                "image": meta.get("image", "svg"),
+            })
+    return enriched_articles
+
 @app.route('/')
 def index():
     global loading_complete, last_update_time, last_data_hash
 
-    # Check if any DB exists; if not, download from Hugging Face
-    if not glob.glob("chroma_db_*"):
-        logger.info("No Chroma DBs found, downloading from Hugging Face Hub...")
+    if not os.path.exists(LOCAL_DB_DIR):
+        logger.info(f"No Chroma DB found at '{LOCAL_DB_DIR}', downloading from Hugging Face Hub...")
         download_from_hf_hub()
 
-    # Start background RSS feed update
     loading_complete = False
     threading.Thread(target=load_feeds_in_background, daemon=True).start()
 
-    # Load existing data immediately
     try:
-        all_docs = get_all_docs_from_dbs()
-        total_docs = len(all_docs['documents'])
-        logger.info(f"Total articles across all DBs at startup: {total_docs}")
-        if not all_docs.get('metadatas'):
-            logger.info("No articles in any DB yet")
+        all_docs = get_all_docs_from_db()
+        if not all_docs['metadatas']:
+            logger.info("No articles in the DB yet")
             return render_template("index.html", categorized_articles={}, has_articles=False, loading=True)
 
-        # Process and categorize articles with deduplication
-        enriched_articles = []
-        seen_keys = set()
-        for doc, meta in zip(all_docs['documents'], all_docs['metadatas']):
-            if not meta:
-                continue
-            title = meta.get("title", "No Title")
-            link = meta.get("link", "")
-            description = meta.get("original_description", "No Description")
-            published = meta.get("published", "Unknown Date").strip()
-
-            title = clean_text(title)
-            link = clean_text(link)
-            description = clean_text(description)
-
-            description_hash = hashlib.sha256(description.encode('utf-8')).hexdigest()
-            key = f"{title}|{link}|{published}|{description_hash}"
-            if key not in seen_keys:
-                seen_keys.add(key)
-                try:
-                    published = datetime.strptime(published, "%Y-%m-%d %H:%M:%S").isoformat() if "Unknown" not in published else published
-                except (ValueError, TypeError):
-                    published = "1970-01-01T00:00:00"
-                enriched_articles.append({
-                    "title": title,
-                    "link": link,
-                    "description": description,
-                    "category": meta.get("category", "Uncategorized"),
-                    "published": published,
-                    "image": meta.get("image", "svg"),
-                })
-
+        enriched_articles = process_docs_into_articles(all_docs)
         enriched_articles.sort(key=lambda x: x["published"], reverse=True)
-
+        
         categorized_articles = {}
         for article in enriched_articles:
             cat = article["category"]
-            if cat not in categorized_articles:
-                categorized_articles[cat] = []
-            categorized_articles[cat].append(article)
-
-        categorized_articles = dict(sorted(categorized_articles.items(), key=lambda x: x[0].lower()))
+            categorized_articles.setdefault(cat, []).append(article)
+        
+        categorized_articles = dict(sorted(categorized_articles.items()))
 
         for cat in categorized_articles:
-            categorized_articles[cat] = sorted(categorized_articles[cat], key=lambda x: x["published"], reverse=True)[:10]
-            if len(categorized_articles[cat]) >= 2:
-                logger.debug(f"Category {cat} top 2: {categorized_articles[cat][0]['title']} | {categorized_articles[cat][1]['title']}")
+            categorized_articles[cat] = categorized_articles[cat][:10]
 
-        # Compute initial data hash
         last_data_hash = compute_data_hash(categorized_articles)
-
-        logger.info(f"Displaying articles at startup: {sum(len(articles) for articles in categorized_articles.values())} total")
-        return render_template("index.html", 
-                              categorized_articles=categorized_articles, 
-                              has_articles=True, 
-                              loading=True)
+        
+        return render_template("index.html", categorized_articles=categorized_articles, has_articles=True, loading=True)
     except Exception as e:
-        logger.error(f"Error retrieving articles at startup: {e}")
+        logger.error(f"Error retrieving articles at startup: {e}", exc_info=True)
         return render_template("index.html", categorized_articles={}, has_articles=False, loading=True)
 
 @app.route('/search', methods=['POST'])
 def search():
     query = request.form.get('search')
     if not query:
-        logger.info("Empty search query received")
         return jsonify({"categorized_articles": {}, "has_articles": False, "loading": False})
 
     try:
         logger.info(f"Performing semantic search for: '{query}'")
+        vector_db = get_vector_db()
+        if not vector_db:
+            return jsonify({"categorized_articles": {}, "has_articles": False, "loading": False})
+        
+        results = vector_db.similarity_search_with_relevance_scores(query, k=50)
         
-        embedding_function = get_embedding_model()
         enriched_articles = []
         seen_keys = set()
-        db_paths = glob.glob("chroma_db_*")
-
-        if not db_paths:
-            logger.warning("No Chroma DBs found for search.")
-            return jsonify({"categorized_articles": {}, "has_articles": False, "loading": False})
-
-        all_search_results = []
-        for db_path in db_paths:
-            if not os.path.isdir(db_path): continue
-            try:
-                vector_db = Chroma(
-                    persist_directory=db_path,
-                    embedding_function=embedding_function,
-                    collection_name="news_articles"
-                )
-                if vector_db._collection.count() > 0:
-                    results = vector_db.similarity_search_with_relevance_scores(query, k=20)
-                    all_search_results.extend(results)
-            except Exception as e:
-                logger.error(f"Error searching in DB {db_path}: {e}")
-
-        # Sort all results by relevance score (higher is better)
-        all_search_results.sort(key=lambda x: x[1], reverse=True)
-
-        # Process and deduplicate top results
-        for doc, score in all_search_results:
+        for doc, score in results:
             meta = doc.metadata
-            title = clean_text(meta.get("title", "No Title"))
-            link = clean_text(meta.get("link", ""))
-            description = clean_text(meta.get("original_description", "No Description"))
-            published = meta.get("published", "Unknown Date").strip()
-
-            description_hash = hashlib.sha256(description.encode('utf-8')).hexdigest()
-            key = f"{title}|{link}|{published}|{description_hash}"
-            
+            title = meta.get("title", "No Title")
+            link = meta.get("link", "")
+            key = f"{title}|{link}|{meta.get('published', '')}"
             if key not in seen_keys:
                 seen_keys.add(key)
                 enriched_articles.append({
-                    "title": meta.get("title", "No Title"),
-                    "link": meta.get("link", ""),
+                    "title": title,
+                    "link": link,
                     "description": meta.get("original_description", "No Description"),
                     "category": meta.get("category", "Uncategorized"),
-                    "published": published,
+                    "published": meta.get("published", "Unknown Date"),
                     "image": meta.get("image", "svg"),
                 })
 
-        logger.info(f"Found {len(enriched_articles)} unique articles from semantic search.")
-        if not enriched_articles:
-            return jsonify({"categorized_articles": {}, "has_articles": False, "loading": False})
-
-        # Categorize the articles
         categorized_articles = {}
         for article in enriched_articles:
             cat = article["category"]
@@ -246,141 +176,58 @@ def search():
         logger.error(f"Semantic search error: {e}", exc_info=True)
         return jsonify({"categorized_articles": {}, "has_articles": False, "loading": False}), 500
 
+@app.route('/get_all_articles/<category>')
+def get_all_articles(category):
+    try:
+        all_docs = get_all_docs_from_db()
+        enriched_articles = process_docs_into_articles(all_docs)
+        
+        category_articles = [
+            article for article in enriched_articles if article["category"] == category
+        ]
+        
+        category_articles.sort(key=lambda x: x["published"], reverse=True)
+        return jsonify({"articles": category_articles, "category": category})
+    except Exception as e:
+        logger.error(f"Error fetching all articles for category {category}: {e}")
+        return jsonify({"articles": [], "category": category}), 500
 
 @app.route('/check_loading')
 def check_loading():
     global loading_complete, last_update_time
-    if loading_complete:
-        return jsonify({"status": "complete", "last_update": last_update_time})
-    return jsonify({"status": "loading"}), 202
+    return jsonify({"status": "complete" if loading_complete else "loading", "last_update": last_update_time})
 
 @app.route('/get_updates')
 def get_updates():
     global last_update_time, last_data_hash
     try:
-        all_docs = get_all_docs_from_dbs()
-        if not all_docs.get('metadatas'):
-            return jsonify({"articles": [], "last_update": last_update_time, "has_updates": False})
-
-        enriched_articles = []
-        seen_keys = set()
-        for doc, meta in zip(all_docs['documents'], all_docs['metadatas']):
-            if not meta:
-                continue
-            title = meta.get("title", "No Title")
-            link = meta.get("link", "")
-            description = meta.get("original_description", "No Description")
-            published = meta.get("published", "Unknown Date").strip()
-
-            title = clean_text(title)
-            link = clean_text(link)
-            description = clean_text(description)
-
-            description_hash = hashlib.sha256(description.encode('utf-8')).hexdigest()
-            key = f"{title}|{link}|{published}|{description_hash}"
-            if key not in seen_keys:
-                seen_keys.add(key)
-                try:
-                    published = datetime.strptime(published, "%Y-%m-%d %H:%M:%S").isoformat() if "Unknown" not in published else published
-                except (ValueError, TypeError):
-                    published = "1970-01-01T00:00:00"
-                enriched_articles.append({
-                    "title": title,
-                    "link": link,
-                    "description": description,
-                    "category": meta.get("category", "Uncategorized"),
-                    "published": published,
-                    "image": meta.get("image", "svg"),
-                })
-
-        enriched_articles.sort(key=lambda x: x["published"], reverse=True)
+        all_docs = get_all_docs_from_db()
+        if not all_docs['metadatas']:
+            return jsonify({"articles": {}, "last_update": last_update_time, "has_updates": False})
+        
+        enriched_articles = process_docs_into_articles(all_docs)
         categorized_articles = {}
         for article in enriched_articles:
             cat = article["category"]
-            if cat not in categorized_articles:
-                categorized_articles[cat] = []
-            key = f"{article['title']}|{article['link']}|{article['published']}"
-            if key not in [f"{a['title']}|{a['link']}|{a['published']}" for a in categorized_articles[cat]]:
-                categorized_articles[cat].append(article)
-
+            categorized_articles.setdefault(cat, []).append(article)
+        
         for cat in categorized_articles:
-            unique_articles = []
-            seen_cat_keys = set()
-            for article in sorted(categorized_articles[cat], key=lambda x: x["published"], reverse=True):
-                key = f"{clean_text(article['title'])}|{clean_text(article['link'])}|{article['published']}"
-                if key not in seen_cat_keys:
-                    seen_cat_keys.add(key)
-                    unique_articles.append(article)
-            categorized_articles[cat] = unique_articles[:10]
+            categorized_articles[cat].sort(key=lambda x: x["published"], reverse=True)
+            categorized_articles[cat] = categorized_articles[cat][:10]
 
-        # Compute hash of new data
         current_data_hash = compute_data_hash(categorized_articles)
-
-        # Compare with last data hash to determine if there are updates
         has_updates = last_data_hash != current_data_hash
+        
         if has_updates:
             logger.info("New RSS data detected, sending updates to frontend")
             last_data_hash = current_data_hash
-            return jsonify({
-                "articles": categorized_articles,
-                "last_update": last_update_time,
-                "has_updates": True
-            })
+            return jsonify({"articles": categorized_articles, "last_update": last_update_time, "has_updates": True})
         else:
-            logger.info("No new RSS data, skipping update")
-            return jsonify({
-                "articles": {},
-                "last_update": last_update_time,
-                "has_updates": False
-            })
+            return jsonify({"articles": {}, "last_update": last_update_time, "has_updates": False})
     except Exception as e:
         logger.error(f"Error fetching updates: {e}")
         return jsonify({"articles": {}, "last_update": last_update_time, "has_updates": False}), 500
 
-@app.route('/get_all_articles/<category>')
-def get_all_articles(category):
-    try:
-        all_docs = get_all_docs_from_dbs()
-        if not all_docs.get('metadatas'):
-            return jsonify({"articles": [], "category": category})
-
-        enriched_articles = []
-        seen_keys = set()
-        for doc, meta in zip(all_docs['documents'], all_docs['metadatas']):
-            if not meta or meta.get("category") != category:
-                continue
-            title = meta.get("title", "No Title")
-            link = meta.get("link", "")
-            description = meta.get("original_description", "No Description")
-            published = meta.get("published", "Unknown Date").strip()
-
-            title = clean_text(title)
-            link = clean_text(link)
-            description = clean_text(description)
-
-            description_hash = hashlib.sha256(description.encode('utf-8')).hexdigest()
-            key = f"{title}|{link}|{published}|{description_hash}"
-            if key not in seen_keys:
-                seen_keys.add(key)
-                try:
-                    published = datetime.strptime(published, "%Y-%m-%d %H:%M:%S").isoformat() if "Unknown" not in published else published
-                except (ValueError, TypeError):
-                    published = "1970-01-01T00:00:00"
-                enriched_articles.append({
-                    "title": title,
-                    "link": link,
-                    "description": description,
-                    "category": meta.get("category", "Uncategorized"),
-                    "published": published,
-                    "image": meta.get("image", "svg"),
-                })
-
-        enriched_articles.sort(key=lambda x: x["published"], reverse=True)
-        return jsonify({"articles": enriched_articles, "category": category})
-    except Exception as e:
-        logger.error(f"Error fetching all articles for category {category}: {e}")
-        return jsonify({"articles": [], "category": category}), 500
-
 @app.route('/card')
 def card_load():
     return render_template("card.html")