Update rss_processor.py

rss_processor.py

@@ -4,7 +4,7 @@ from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.docstore.document import Document
 import logging
-from huggingface_hub import HfApi, login, snapshot_download
+from huggingface_hub import HfApi, login
 import shutil
 import json
 from datetime import datetime
@@ -12,24 +12,34 @@ import dateutil.parser
 import hashlib
 import re
 
+# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Constants
+MAX_ARTICLES_PER_FEED = 10
 LOCAL_DB_DIR = "chroma_db"
+FEEDS_FILE = "rss_feeds.json"
 COLLECTION_NAME = "news_articles"
 HF_API_TOKEN = os.getenv("HF_TOKEN")
 REPO_ID = "broadfield-dev/news-rag-db"
-FEEDS_FILE = "rss_feeds.json"
 
+# Initialize Hugging Face API
 login(token=HF_API_TOKEN)
 hf_api = HfApi()
 
-def get_embedding_model():
-    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
+# Initialize embedding model (global, reusable)
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# 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 clean_text(text):
+    """Clean text by removing HTML tags and extra whitespace."""
     if not text or not isinstance(text, str):
         return ""
     text = re.sub(r'<.*?>', '', text)
@@ -51,25 +61,27 @@ def fetch_rss_feeds():
         for feed_info in feeds:
             feed_url = feed_info.get("url")
             if not feed_url:
+                logger.warning(f"Skipping feed with no URL in category '{category}'")
                 continue
 
             try:
-                logger.info(f"Fetching '{feed_info.get('name', feed_url)}' from category '{category}'")
-                # Add a User-Agent to prevent getting blocked
-                feed = feedparser.parse(feed_url, agent="RSSNewsBot/1.0 (+http://huggingface.co/spaces/broadfield-dev/RSS_News)")
-
+                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}")
                     continue
-                
-                for entry in feed.entries[:10]: # Process max 10 entries per feed
+                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", ""))
                     
-                    cleaned_title = clean_text(title)
-                    cleaned_link = clean_text(link)
-                    
+                    title = clean_text(title)
+                    link = clean_text(link)
+                    description = clean_text(description)
+
                     published = "Unknown Date"
                     for date_field in ["published", "updated", "created", "pubDate"]:
                         if date_field in entry:
@@ -77,104 +89,153 @@ 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):
+                            except (ValueError, TypeError) as e:
+                                logger.debug(f"Failed to parse {date_field} '{entry[date_field]}': {e}")
                                 continue
 
-                    key = f"{cleaned_title}|{cleaned_link}|{published}"
+                    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"
-                        if 'media_content' in entry and entry.media_content:
-                            image = entry.media_content[0].get('url', 'svg')
-                        elif 'media_thumbnail' in entry and entry.media_thumbnail:
-                            image = entry.media_thumbnail[0].get('url', '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,
                             "link": link,
                             "description": description,
                             "published": published,
-                            "category": category, # Directly use category from JSON
+                            "category": category,  # Use JSON category directly
                             "image": image,
                         })
+                        article_count += 1
             except Exception as e:
                 logger.error(f"Error fetching {feed_url}: {e}")
-
     logger.info(f"Total articles fetched: {len(articles)}")
     return articles
 
-def process_and_store_articles(articles):
-    vector_db = Chroma(
-        persist_directory=LOCAL_DB_DIR,
-        embedding_function=get_embedding_model(),
-        collection_name=COLLECTION_NAME
-    )
-    
-    try:
-        existing_ids = set(vector_db.get(include=[])["ids"])
-    except Exception:
-        existing_ids = set()
+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"
 
-    docs_to_add = []
-    ids_to_add = []
-    
+def process_and_store_articles(articles):
+    documents = []
+    existing_ids = set(vector_db.get()["ids"])  # Load existing IDs once
     for article in articles:
-        cleaned_title = clean_text(article["title"])
-        cleaned_link = clean_text(article["link"])
-        doc_id = f"{cleaned_title}|{cleaned_link}|{article['published']}"
-        
-        if doc_id in existing_ids:
-            continue
-
-        metadata = {
-            "title": article["title"],
-            "link": article["link"],
-            "original_description": article["description"],
-            "published": article["published"],
-            "category": article["category"],
-            "image": article["image"],
-        }
-        doc = Document(page_content=clean_text(article["description"]), metadata=metadata)
-        docs_to_add.append(doc)
-        ids_to_add.append(doc_id)
-        existing_ids.add(doc_id)
+        try:
+            title = clean_text(article["title"])
+            link = clean_text(article["link"])
+            description = clean_text(article["description"])
+            published = article["published"]
+            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}")
+                continue
+            metadata = {
+                "title": article["title"],
+                "link": article["link"],
+                "original_description": article["description"],
+                "published": article["published"],
+                "category": article["category"],
+                "image": article["image"],
+            }
+            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
+        except Exception as e:
+            logger.error(f"Error processing article {article['title']}: {e}")
     
-    if docs_to_add:
-        vector_db.add_documents(documents=docs_to_add, ids=ids_to_add)
-        vector_db.persist()
-        logger.info(f"Added {len(docs_to_add)} new articles to DB. Total in DB: {vector_db._collection.count()}")
+    if documents:
+        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'])}")
+        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:
-            snapshot_download(
-                repo_id=REPO_ID,
-                repo_type="dataset",
-                local_dir=".",
-                local_dir_use_symlinks=False,
-                allow_patterns=f"{LOCAL_DB_DIR}/**",
-                token=HF_API_TOKEN
-            )
+            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.warning(f"Could not download DB from Hub (this is normal on first run): {e}")
+            logger.error(f"Error downloading from Hugging Face Hub: {e}")
+    else:
+        logger.info("Local Chroma DB exists, loading existing data.")
 
 def upload_to_hf_hub():
     if os.path.exists(LOCAL_DB_DIR):
         try:
-            hf_api.upload_folder(
-                folder_path=LOCAL_DB_DIR,
-                path_in_repo=LOCAL_DB_DIR,
-                repo_id=REPO_ID,
-                repo_type="dataset",
-                token=HF_API_TOKEN,
-                commit_message="Update RSS news database"
-            )
+            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}")
 
 if __name__ == "__main__":
-    download_from_hf_hub()
+    download_from_hf_hub()  # Ensure DB is initialized
     articles = fetch_rss_feeds()
-    if articles:
-        process_and_store_articles(articles)
-        upload_to_hf_hub()
+    process_and_store_articles(articles)
+    upload_to_hf_hub()