Update rss_processor.py

rss_processor.py

@@ -4,49 +4,89 @@ from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.docstore.document import Document
 import logging
-from huggingface_hub import HfApi, login
-import shutil
+from huggingface_hub import HfApi, login, snapshot_download
+from huggingface_hub.utils import HfHubHTTPError
 import json
-from datetime import datetime
 import dateutil.parser
 import hashlib
 import re
+from datetime import datetime
 
 # Setup logging
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Constants
-MAX_ARTICLES_PER_FEED = 1000
-LOCAL_DB_DIR = "chroma_db"
+# --- Constants ---
+# Local and repository configuration
+LOCAL_DB_DIR = "chroma_db_news"  # Using a more descriptive local directory name
 FEEDS_FILE = "rss_feeds.json"
 COLLECTION_NAME = "news_articles"
+REPO_ID = "broadfield-dev/news-rag-db" # Your Hugging Face Hub repo ID
+
+# RSS feed fetching configuration
+MAX_ARTICLES_PER_FEED = 1000
+
+# Hugging Face credentials
 HF_API_TOKEN = os.getenv("HF_TOKEN")
-REPO_ID = "broadfield-dev/news-rag-db"
+if not HF_API_TOKEN:
+    raise ValueError("Hugging Face API token not found. Please set the HF_TOKEN environment variable.")
 
+# --- Global Initializations ---
 # Initialize Hugging Face API
 login(token=HF_API_TOKEN)
 hf_api = HfApi()
 
-# Initialize embedding model (global, reusable)
+# Initialize embedding model once to be reused
+logger.info("Loading embedding model...")
 embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+logger.info("Embedding model loaded.")
 
-# Initialize vector DB with a specific collection name
-vector_db = Chroma(
-    persist_directory=LOCAL_DB_DIR,
-    embedding_function=embedding_model,
-    collection_name=COLLECTION_NAME
-)
+
+def setup_local_db():
+    """
+    Ensures the local database directory exists.
+    If it doesn't, it attempts to download from Hugging Face Hub.
+    If the Hub repo is empty or doesn't exist, it creates a new local directory.
+    This function handles the "build a new dataset if one does not exist" requirement.
+    """
+    if os.path.exists(LOCAL_DB_DIR):
+        logger.info(f"Local database found at '{LOCAL_DB_DIR}'.")
+        return
+
+    logger.info(f"Local database not found. Attempting to download from Hugging Face Hub repo: {REPO_ID}")
+    try:
+        # snapshot_download is the correct function for downloading a whole repository/folder
+        snapshot_download(
+            repo_id=REPO_ID,
+            repo_type="dataset",
+            local_dir=LOCAL_DB_DIR,
+            token=HF_API_TOKEN,
+        )
+        logger.info(f"Database downloaded successfully from {REPO_ID} to {LOCAL_DB_DIR}.")
+    except HfHubHTTPError as e:
+        # This error (e.g., 404 Not Found) is expected if the repo is new or empty.
+        logger.warning(
+            f"Failed to download from Hub (Repo might be new or empty): {e}. "
+            f"A new local database will be created at '{LOCAL_DB_DIR}'."
+        )
+        os.makedirs(LOCAL_DB_DIR, exist_ok=True)
+    except Exception as e:
+        logger.error(f"An unexpected error occurred during DB download: {e}")
+        logger.info(f"Creating a new local database at '{LOCAL_DB_DIR}'.")
+        os.makedirs(LOCAL_DB_DIR, exist_ok=True)
 
 def clean_text(text):
     """Clean text by removing HTML tags and extra whitespace."""
     if not text or not isinstance(text, str):
         return ""
+    # Remove HTML tags
     text = re.sub(r'<.*?>', '', text)
+    # Normalize whitespace
     text = ' '.join(text.split())
     return text.strip().lower()
 
 def fetch_rss_feeds():
+    """Fetches and parses articles from a list of RSS feeds in a JSON file."""
     articles = []
     seen_keys = set()
     
@@ -54,7 +94,7 @@ def fetch_rss_feeds():
         with open(FEEDS_FILE, 'r') as f:
             feed_categories = json.load(f)
     except FileNotFoundError:
-        logger.error(f"{FEEDS_FILE} not found. No feeds to process.")
+        logger.error(f"'{FEEDS_FILE}' not found. Please create it. No feeds to process.")
         return []
 
     for category, feeds in feed_categories.items():
@@ -68,19 +108,16 @@ def fetch_rss_feeds():
                 logger.info(f"Fetching {feed_url}")
                 feed = feedparser.parse(feed_url)
                 if feed.bozo:
-                    logger.warning(f"Parse error for {feed_url}: {feed.bozo_exception}")
+                    logger.warning(f"Feed parsing error for {feed_url}: {feed.bozo_exception}")
                     continue
-                article_count = 0
-                for entry in feed.entries:
-                    if article_count >= MAX_ARTICLES_PER_FEED:
-                        break
-                    title = entry.get("title", "No Title")
-                    link = entry.get("link", "")
-                    description = entry.get("summary", entry.get("description", ""))
+                
+                for entry in feed.entries[:MAX_ARTICLES_PER_FEED]:
+                    title = clean_text(entry.get("title", "No Title"))
+                    link = entry.get("link", "") # Don't clean link URL
+                    description = clean_text(entry.get("summary", entry.get("description", "")))
                     
-                    title = clean_text(title)
-                    link = clean_text(link)
-                    description = clean_text(description)
+                    if not description: # Skip articles without content
+                        continue
 
                     published = "Unknown Date"
                     for date_field in ["published", "updated", "created", "pubDate"]:
@@ -89,153 +126,135 @@ def fetch_rss_feeds():
                                 parsed_date = dateutil.parser.parse(entry[date_field])
                                 published = parsed_date.strftime("%Y-%m-%d %H:%M:%S")
                                 break
-                            except (ValueError, TypeError) as e:
-                                logger.debug(f"Failed to parse {date_field} '{entry[date_field]}': {e}")
+                            except (ValueError, TypeError):
                                 continue
 
+                    # Create a unique key to deduplicate articles before processing
                     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)
-                        image = "svg"
-                        for img_source in [
-                            lambda e: clean_text(e.get("media_content", [{}])[0].get("url")) if e.get("media_content") else "",
-                            lambda e: clean_text(e.get("media_thumbnail", [{}])[0].get("url")) if e.get("media_thumbnail") else "",
-                            lambda e: clean_text(e.get("enclosure", {}).get("url")) if e.get("enclosure") else "",
-                            lambda e: clean_text(next((lnk.get("href") for lnk in e.get("links", []) if lnk.get("type", "").startswith("image")), "")),
-                        ]:
-                            try:
-                                img = img_source(entry)
-                                if img and img.strip():
-                                    image = img
-                                    break
-                            except (IndexError, AttributeError, TypeError):
-                                continue
-
                         articles.append({
-                            "title": title,
+                            "title": entry.get("title", "No Title"),
                             "link": link,
                             "description": description,
                             "published": published,
-                            "category": category,  # Use JSON category directly
-                            "image": image,
+                            "category": category,
                         })
-                        article_count += 1
             except Exception as e:
-                logger.error(f"Error fetching {feed_url}: {e}")
-    logger.info(f"Total articles fetched: {len(articles)}")
+                logger.error(f"Error fetching or parsing feed {feed_url}: {e}")
+                
+    logger.info(f"Total unique articles fetched: {len(articles)}")
     return articles
 
-def categorize_feed(url):
-    """Categorize an RSS feed based on its URL."""
-    if not url or not isinstance(url, str):
-        logger.warning(f"Invalid URL provided for categorization: {url}")
-        return "Uncategorized"
-
-    url = url.lower().strip()  # Normalize the URL
-
-    logger.debug(f"Categorizing URL: {url}")  # Add debugging for visibility
-
-    if any(keyword in url for keyword in ["nature", "science.org", "arxiv.org", "plos.org", "annualreviews.org", "journals.uchicago.edu", "jneurosci.org", "cell.com", "nejm.org", "lancet.com"]):
-        return "Academic Papers"
-    elif any(keyword in url for keyword in ["reuters.com/business", "bloomberg.com", "ft.com", "marketwatch.com", "cnbc.com", "foxbusiness.com", "wsj.com", "bworldonline.com", "economist.com", "forbes.com"]):
-        return "Business"
-    elif any(keyword in url for keyword in ["investing.com", "cnbc.com/market", "marketwatch.com/market", "fool.co.uk", "zacks.com", "seekingalpha.com", "barrons.com", "yahoofinance.com"]):
-        return "Stocks & Markets"
-    elif any(keyword in url for keyword in ["whitehouse.gov", "state.gov", "commerce.gov", "transportation.gov", "ed.gov", "dol.gov", "justice.gov", "federalreserve.gov", "occ.gov", "sec.gov", "bls.gov", "usda.gov", "gao.gov", "cbo.gov", "fema.gov", "defense.gov", "hhs.gov", "energy.gov", "interior.gov"]):
-        return "Federal Government"
-    elif any(keyword in url for keyword in ["weather.gov", "metoffice.gov.uk", "accuweather.com", "weatherunderground.com", "noaa.gov", "wunderground.com", "climate.gov", "ecmwf.int", "bom.gov.au"]):
-        return "Weather"
-    elif any(keyword in url for keyword in ["data.worldbank.org", "imf.org", "un.org", "oecd.org", "statista.com", "kff.org", "who.int", "cdc.gov", "bea.gov", "census.gov", "fdic.gov"]):
-        return "Data & Statistics"
-    elif any(keyword in url for keyword in ["nasa", "spaceweatherlive", "space", "universetoday", "skyandtelescope", "esa"]):
-        return "Space"
-    elif any(keyword in url for keyword in ["sciencedaily", "quantamagazine", "smithsonianmag", "popsci", "discovermagazine", "scientificamerican", "newscientist", "livescience", "atlasobscura"]):
-        return "Science"
-    elif any(keyword in url for keyword in ["wired", "techcrunch", "arstechnica", "gizmodo", "theverge"]):
-        return "Tech"
-    elif any(keyword in url for keyword in ["horoscope", "astrostyle"]):
-        return "Astrology"
-    elif any(keyword in url for keyword in ["cnn_allpolitics", "bbci.co.uk/news/politics", "reuters.com/arc/outboundfeeds/newsletter-politics", "politico.com/rss/politics", "thehill"]):
-        return "Politics"
-    elif any(keyword in url for keyword in ["weather", "swpc.noaa.gov", "foxweather"]):
-        return "Earth Weather"
-    elif "vogue" in url:
-        return "Lifestyle"
-    elif any(keyword in url for keyword in ["phys.org", "aps.org", "physicsworld"]):
-        return "Physics"
-    else:
-        logger.warning(f"No matching category found for URL: {url}")
-        return "Uncategorized"
+def process_and_store_articles(articles, vector_db):
+    """Processes articles and stores them in the Chroma DB, avoiding duplicates."""
+    new_docs = []
+    new_doc_ids = []
+    
+    # Get all existing document IDs from the database once to check for duplicates
+    try:
+        existing_ids = set(vector_db.get(include=[])["ids"])
+        logger.info(f"Found {len(existing_ids)} existing documents in the database.")
+    except Exception as e:
+        logger.error(f"Could not retrieve existing IDs from DB. Assuming empty. Error: {e}")
+        existing_ids = set()
 
-def process_and_store_articles(articles):
-    documents = []
-    existing_ids = set(vector_db.get()["ids"])  # Load existing IDs once
     for article in articles:
         try:
+            # Recreate the same unique ID format for checking against the DB
             title = clean_text(article["title"])
-            link = clean_text(article["link"])
-            description = clean_text(article["description"])
+            link = article["link"]
             published = article["published"]
+            description = article["description"]
             description_hash = hashlib.sha256(description.encode('utf-8')).hexdigest()
             doc_id = f"{title}|{link}|{published}|{description_hash}"
+
             if doc_id in existing_ids:
-                logger.debug(f"Skipping duplicate in DB: {doc_id}")
+                logger.debug(f"Skipping duplicate article (ID already in DB): {title[:50]}...")
                 continue
+            
+            # Add to our in-memory set to avoid duplicates from the same batch
+            existing_ids.add(doc_id)
+
             metadata = {
                 "title": article["title"],
                 "link": article["link"],
-                "original_description": article["description"],
                 "published": article["published"],
                 "category": article["category"],
-                "image": article["image"],
+                # Store original description if needed, or keep it clean
+                # "original_description": article["description"], 
             }
-            doc = Document(page_content=description, metadata=metadata, id=doc_id)
-            documents.append(doc)
-            existing_ids.add(doc_id)  # Update in-memory set to avoid duplicates within this batch
+            # The Document object itself doesn't take an ID
+            doc = Document(page_content=description, metadata=metadata)
+            new_docs.append(doc)
+            new_doc_ids.append(doc_id)
+
         except Exception as e:
-            logger.error(f"Error processing article {article['title']}: {e}")
+            logger.error(f"Error processing article '{article.get('title', 'N/A')}': {e}")
     
-    if documents:
+    if new_docs:
+        logger.info(f"Adding {len(new_docs)} new documents to the database...")
         try:
-            vector_db.add_documents(documents)
-            vector_db.persist()
-            logger.info(f"Added {len(documents)} new articles to DB. Total documents: {len(vector_db.get()['ids'])}")
+            # Provide the list of documents and a parallel list of their unique IDs
+            vector_db.add_documents(documents=new_docs, ids=new_doc_ids)
+            vector_db.persist() # Save changes to disk
+            logger.info("Successfully added new documents and persisted the database.")
         except Exception as e:
-            logger.error(f"Error storing articles: {e}")
-
-def download_from_hf_hub():
-    if not os.path.exists(LOCAL_DB_DIR):
-        try:
-            hf_api.create_repo(repo_id=REPO_ID, repo_type="dataset", exist_ok=True, token=HF_API_TOKEN)
-            logger.info(f"Downloading Chroma DB from {REPO_ID}...")
-            hf_api.hf_hub_download(repo_id=REPO_ID, filename="chroma_db", local_dir=LOCAL_DB_DIR, repo_type="dataset", token=HF_API_TOKEN)
-        except Exception as e:
-            logger.error(f"Error downloading from Hugging Face Hub: {e}")
+            logger.error(f"Failed to add documents to Chroma DB: {e}")
     else:
-        logger.info("Local Chroma DB exists, loading existing data.")
+        logger.info("No new articles to add to the database.")
 
 def upload_to_hf_hub():
-    if os.path.exists(LOCAL_DB_DIR):
-        try:
-            logger.info(f"Uploading updated Chroma DB to {REPO_ID}...")
-            for root, _, files in os.walk(LOCAL_DB_DIR):
-                for file in files:
-                    local_path = os.path.join(root, file)
-                    remote_path = os.path.relpath(local_path, LOCAL_DB_DIR)
-                    hf_api.upload_file(
-                        path_or_fileobj=local_path,
-                        path_in_repo=remote_path,
-                        repo_id=REPO_ID,
-                        repo_type="dataset",
-                        token=HF_API_TOKEN
-                    )
-            logger.info(f"Database uploaded to: {REPO_ID}")
-        except Exception as e:
-            logger.error(f"Error uploading to Hugging Face Hub: {e}")
+    """Uploads the local Chroma DB directory to the Hugging Face Hub."""
+    if not os.path.exists(LOCAL_DB_DIR):
+        logger.warning(f"Local database directory '{LOCAL_DB_DIR}' not found. Nothing to upload.")
+        return
+    
+    try:
+        # Ensure the repo exists before uploading.
+        hf_api.create_repo(repo_id=REPO_ID, repo_type="dataset", exist_ok=True)
 
-if __name__ == "__main__":
-    download_from_hf_hub()  # Ensure DB is initialized
+        logger.info(f"Uploading database from '{LOCAL_DB_DIR}' to Hugging Face repo: {REPO_ID}...")
+        # upload_folder is the recommended way to upload a directory's contents.
+        hf_api.upload_folder(
+            folder_path=LOCAL_DB_DIR,
+            repo_id=REPO_ID,
+            repo_type="dataset",
+            commit_message=f"Update database - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
+        )
+        logger.info(f"Database successfully uploaded to {REPO_ID}.")
+    except Exception as e:
+        logger.error(f"Error uploading to Hugging Face Hub: {e}")
+
+
+def main():
+    """Main execution function to run the data pipeline."""
+    # 1. Ensure local DB exists by downloading from Hub or creating a new one
+    setup_local_db()
+
+    # 2. Initialize the vector DB object *after* the directory is guaranteed to exist
+    logger.info("Initializing Chroma vector database...")
+    vector_db = Chroma(
+        persist_directory=LOCAL_DB_DIR,
+        embedding_function=embedding_model,
+        collection_name=COLLECTION_NAME
+    )
+    logger.info("Chroma DB initialized.")
+
+    # 3. Fetch new articles from RSS feeds
     articles = fetch_rss_feeds()
-    process_and_store_articles(articles)
-    upload_to_hf_hub()
+
+    # 4. Process new articles and add them to the DB
+    if articles:
+        process_and_store_articles(articles, vector_db)
+    
+    # 5. Upload the potentially updated database back to the Hub
+    upload_to_hf_hub()
+    
+    logger.info("Script finished.")
+
+
+if __name__ == "__main__":
+    main()