Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import subprocess
-from flask import Flask, render_template, request, Response, jsonify
+from flask import Flask, render_template, request
 from rss_processor import fetch_rss_feeds, process_and_store_articles, vector_db
 import logging
 import time
@@ -24,28 +24,37 @@ def load_feeds_in_background():
 
 @app.route('/')
 def index():
-    # Show existing articles immediately, even if empty
-    stored_docs = vector_db.similarity_search("news", k=1000)  # Try to retrieve all available articles
-    logger.info(f"Found {len(stored_docs)} documents in vector DB")
-    # Use a set to ensure unique articles by title, link, and description hash
-    unique_articles = {}
-    for doc in stored_docs:
-        title = doc.metadata["title"]
-        link = doc.metadata["link"]
-        description = doc.metadata["original_description"]
-        desc_hash = hashlib.md5(description.encode()).hexdigest()
-        key = f"{title}|{link}|{desc_hash}"
-        if key not in unique_articles:
-            unique_articles[key] = {
-                "title": title,
-                "link": link,
-                "description": description,
-                "category": doc.metadata["category"],
-                "published": doc.metadata["published"],
-                "image": doc.metadata.get("image", "svg"),
-            }
-    enriched_articles = list(unique_articles.values())
-    logger.info(f"Enriched {len(enriched_articles)} unique articles for display")
+    # Show all existing articles immediately, even if empty
+    try:
+        # Get all documents from Chroma DB
+        all_docs = vector_db.get(include=['documents', 'metadatas'])
+        stored_docs = [
+            Document(page_content=doc['documents'], metadata=doc['metadatas'])
+            for doc in all_docs['documents']
+        ]
+        logger.info(f"Found {len(stored_docs)} documents in vector DB")
+        # Use a set to ensure unique articles by title, link, and full description hash
+        unique_articles = {}
+        for doc in stored_docs:
+            title = doc.metadata["title"]
+            link = doc.metadata["link"]
+            description = doc.metadata["original_description"]
+            desc_hash = hashlib.md5(description.encode()).hexdigest()
+            key = f"{title}|{link}|{desc_hash}"
+            if key not in unique_articles:
+                unique_articles[key] = {
+                    "title": title,
+                    "link": link,
+                    "description": description,
+                    "category": doc.metadata["category"],
+                    "published": doc.metadata["published"],
+                    "image": doc.metadata.get("image", "svg"),
+                }
+        enriched_articles = list(unique_articles.values())
+        logger.info(f"Enriched {len(enriched_articles)} unique articles for display")
+    except Exception as e:
+        logger.error(f"Error retrieving documents from vector DB: {e}")
+        enriched_articles = []  # Fallback if DB is empty or inaccessible
 
     # Start loading new feeds in the background
     subprocess.Popen(["python", "rss_processor.py", "load_feeds"])
@@ -57,14 +66,49 @@ def index():
             categorized_articles[cat] = []
         categorized_articles[cat].append(article)
 
-    return render_template("index.html", categorized_articles=categorized_articles, loading_new_feeds=True)
+    return render_template("index.html", categorized_articles=categorized_articles, loading_new_feeds=True, has_articles=bool(enriched_articles))
+
+@app.route('/search', methods=['POST'])
+def search():
+    query = request.form.get('search')
+    if query:
+        logger.info(f"Processing search query: {query}")
+        results = vector_db.similarity_search(query, k=10)
+        unique_search_articles = {}
+        for doc in results:
+            title = doc.metadata["title"]
+            link = doc.metadata["link"]
+            description = doc.metadata["original_description"]
+            desc_hash = hashlib.md5(description.encode()).hexdigest()
+            key = f"{title}|{link}|{desc_hash}"
+            if key not in unique_search_articles:
+                unique_search_articles[key] = {
+                    "title": title,
+                    "link": link,
+                    "description": description,
+                    "category": doc.metadata["category"],
+                    "published": doc.metadata["published"],
+                    "image": doc.metadata.get("image", "svg"),
+                }
+        enriched_articles = list(unique_search_articles.values())
+        logger.info(f"Search returned {len(enriched_articles)} unique results")
+
+        categorized_articles = {}
+        for article in enriched_articles:
+            cat = article["category"]
+            if cat not in categorized_articles:
+                categorized_articles[cat] = []
+            categorized_articles[cat].append(article)
+
+        return render_template("index.html", categorized_articles=categorized_articles, loading_new_feeds=True, has_articles=bool(enriched_articles))
+    return render_template("index.html", categorized_articles={}, loading_new_feeds=True, has_articles=False)
 
 @app.route('/check_feeds', methods=['GET'])
 def check_feeds():
     try:
         # Check if vector DB has new or updated documents
-        docs = vector_db.similarity_search("news", k=1)
-        if docs:
+        all_docs = vector_db.get(include=['documents', 'metadatas'])
+        if all_docs['documents']:
             logger.info("Feeds loaded successfully in vector DB")
             return jsonify({"status": "loaded"})
         return jsonify({"status": "loading"}), 202