Update rss_processor.py

rss_processor.py

@@ -1,19 +1,17 @@
 import os
 import feedparser
 from chromadb import PersistentClient
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.docstore.document import Document
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_core.documents import Document
 import logging
 from huggingface_hub import HfApi, login, snapshot_download
-import shutil
-import rss_feeds
 from datetime import datetime
 import dateutil.parser
 import hashlib
 import json
 import re
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 LOCAL_DB_DIR = "chroma_db"
@@ -22,8 +20,17 @@ COLLECTION_NAME = "news_articles"
 HF_API_TOKEN = os.getenv("HF_TOKEN")
 REPO_ID = "broadfield-dev/news-rag-db"
 MAX_ARTICLES_PER_FEED = 1000
-login(token=HF_API_TOKEN)
-hf_api = HfApi()
+
+def initialize_hf_api():
+    if not HF_API_TOKEN:
+        logger.error("Hugging Face API token (HF_TOKEN) not set.")
+        raise ValueError("HF_TOKEN environment variable is not set.")
+    try:
+        login(token=HF_API_TOKEN)
+        return HfApi()
+    except Exception as e:
+        logger.error(f"Failed to login to Hugging Face Hub: {e}")
+        raise
 
 def get_embedding_model():
     if not hasattr(get_embedding_model, "model"):
@@ -35,11 +42,10 @@ def clean_text(text):
         return ""
     text = re.sub(r'<.*?>', '', text)
     text = ' '.join(text.split())
-    return text.strip().lower()
+    return text.strip()
 
 def fetch_rss_feeds():
     articles = []
-    seen_keys = set()
     
     try:
         with open(FEEDS_FILE, 'r') as f:
@@ -61,104 +67,56 @@ def fetch_rss_feeds():
                 if feed.bozo:
                     logger.warning(f"Parse 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_raw = entry.get("title", "No Title")
+                
+                for entry in feed.entries[:MAX_ARTICLES_PER_FEED]:
+                    title = entry.get("title", "No Title")
                     link = entry.get("link", "")
-                    description = entry.get("summary", entry.get("description", ""))
-                    
-                    clean_title_val = clean_text(title_raw)
-                    clean_desc_val = clean_text(description)
+                    description_raw = entry.get("summary", entry.get("description", ""))
+                    description = clean_text(description_raw)
 
-                    if not clean_desc_val:
+                    if not description:
                         continue
 
-                    published = "Unknown Date"
+                    published_str = "Unknown Date"
                     for date_field in ["published", "updated", "created", "pubDate"]:
                         if date_field in entry:
                             try:
                                 parsed_date = dateutil.parser.parse(entry[date_field])
-                                published = parsed_date.strftime("%Y-%m-%d %H:%M:%S")
+                                published_str = parsed_date.isoformat()
                                 break
                             except (ValueError, TypeError):
                                 continue
-
-                    description_hash = hashlib.sha256(clean_desc_val.encode('utf-8')).hexdigest()
-                    key = f"{clean_title_val}|{link}|{published}|{description_hash}"
                     
-                    if key not in seen_keys:
-                        seen_keys.add(key)
-                        
-                        image = "svg"
-                        for img_source in [
-                            lambda e: e.get("media_content", [{}])[0].get("url") if e.get("media_content") else "",
-                            lambda e: e.get("media_thumbnail", [{}])[0].get("url") if e.get("media_thumbnail") else "",
-                            lambda e: e.get("enclosure", {}).get("url") if e.get("enclosure") else "",
-                            lambda e: 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_raw,
-                            "link": link,
-                            "description": clean_desc_val,
-                            "published": published,
-                            "category": category,
-                            "image": image,
-                        })
-                        article_count += 1
+                    image = "svg"
+                    image_sources = [
+                        lambda e: e.get("media_content", [{}])[0].get("url") if e.get("media_content") else None,
+                        lambda e: e.get("media_thumbnail", [{}])[0].get("url") if e.get("media_thumbnail") else None,
+                        lambda e: e.get("enclosure", {}).get("url") if e.get("enclosure") and e.get("enclosure", {}).get('type', '').startswith('image') else None,
+                        lambda e: next((lnk.get("href") for lnk in e.get("links", []) if lnk.get("type", "").startswith("image")), None),
+                    ]
+                    for source_func in image_sources:
+                        try:
+                            img_url = source_func(entry)
+                            if img_url and isinstance(img_url, str) and img_url.strip():
+                                image = img_url
+                                break
+                        except (IndexError, AttributeError, TypeError):
+                            continue
+                    
+                    articles.append({
+                        "title": title,
+                        "link": link,
+                        "description": description,
+                        "published": published_str,
+                        "category": category,
+                        "image": image,
+                    })
             except Exception as e:
-                logger.error(f"Error fetching {feed_url}: {e}")
-    logger.info(f"Total unique articles fetched: {len(articles)}")
+                logger.error(f"Error fetching or parsing {feed_url}: {e}")
+    
+    logger.info(f"Total articles fetched: {len(articles)}")
     return articles
 
-def categorize_feed(url):
-    if not url or not isinstance(url, str):
-        logger.warning(f"Invalid URL provided for categorization: {url}")
-        return "Uncategorized"
-    url = url.lower().strip()
-    logger.debug(f"Categorizing URL: {url}")
-    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):
     if not os.path.exists(LOCAL_DB_DIR):
         os.makedirs(LOCAL_DB_DIR)
@@ -169,17 +127,19 @@ def process_and_store_articles(articles):
     try:
         existing_ids = set(collection.get(include=[])["ids"])
         logger.info(f"Loaded {len(existing_ids)} existing document IDs from {LOCAL_DB_DIR}.")
-    except Exception as e:
-        logger.info(f"No existing DB found or error loading IDs: {e}. Starting fresh.")
+    except Exception:
+        logger.info("No existing DB found or it is empty. Starting fresh.")
         existing_ids = set()
 
-    docs_to_add = []
+    contents_to_add = []
+    metadatas_to_add = []
     ids_to_add = []
-    
+
     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 not article.get('link'):
+            continue
+        
+        doc_id = hashlib.sha256(article['link'].encode('utf-8')).hexdigest()
         
         if doc_id in existing_ids:
             continue
@@ -187,32 +147,35 @@ def process_and_store_articles(articles):
         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)
+        
+        contents_to_add.append(article["description"])
+        metadatas_to_add.append(metadata)
         ids_to_add.append(doc_id)
     
-    if docs_to_add:
+    if ids_to_add:
+        logger.info(f"Found {len(ids_to_add)} new articles to add to the database.")
         try:
-            embeddings = get_embedding_model()
-            for doc, doc_id in zip(docs_to_add, ids_to_add):
-                collection.add(
-                    documents=[doc.page_content],
-                    metadatas=[doc.metadata],
-                    ids=[doc_id],
-                    embeddings=[embeddings.embed_query(doc.page_content)]
-                )
-            client.persist()
-            logger.info(f"Added {len(docs_to_add)} new articles to DB. Total in DB: {collection.count()}")
+            embedding_model = get_embedding_model()
+            embeddings_to_add = embedding_model.embed_documents(contents_to_add)
+            
+            collection.add(
+                embeddings=embeddings_to_add,
+                documents=contents_to_add,
+                metadatas=metadatas_to_add,
+                ids=ids_to_add
+            )
+            logger.info(f"Successfully added {len(ids_to_add)} new articles to DB. Total in DB: {collection.count()}")
         except Exception as e:
-            logger.error(f"Error storing articles: {e}")
+            logger.error(f"Error storing articles in ChromaDB: {e}", exc_info=True)
+    else:
+        logger.info("No new articles to add to the database.")
 
 def download_from_hf_hub():
-    if not os.path.exists(LOCAL_DB_DIR):
+    if not os.path.exists(os.path.join(LOCAL_DB_DIR, "chroma.sqlite3")):
         try:
             logger.info(f"Downloading Chroma DB from {REPO_ID} to {LOCAL_DB_DIR}...")
             snapshot_download(
@@ -220,16 +183,16 @@ def download_from_hf_hub():
                 repo_type="dataset",
                 local_dir=".",
                 local_dir_use_symlinks=False,
-                allow_patterns=f"{LOCAL_DB_DIR}/**",
+                allow_patterns=[f"{LOCAL_DB_DIR}/**"],
                 token=HF_API_TOKEN
             )
             logger.info("Finished downloading DB.")
         except Exception as e:
             logger.warning(f"Could not download from Hugging Face Hub (this is normal on first run): {e}")
     else:
-        logger.info("Local Chroma DB exists, loading existing data.")
+        logger.info(f"Local Chroma DB found at '{LOCAL_DB_DIR}', skipping download.")
 
-def upload_to_hf_hub():
+def upload_to_hf_hub(hf_api):
     if os.path.exists(LOCAL_DB_DIR):
         try:
             logger.info(f"Uploading updated Chroma DB '{LOCAL_DB_DIR}' to {REPO_ID}...")
@@ -238,9 +201,25 @@ def upload_to_hf_hub():
                 path_in_repo=LOCAL_DB_DIR,
                 repo_id=REPO_ID,
                 repo_type="dataset",
-                token=HF_API_TOKEN,
-                commit_message=f"Update RSS news database {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
+                commit_message=f"Update RSS news database {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
+                ignore_patterns=["*.bak", "*.tmp"]
             )
             logger.info(f"Database folder '{LOCAL_DB_DIR}' uploaded to: {REPO_ID}")
         except Exception as e:
-            logger.error(f"Error uploading to Hugging Face Hub: {e}")
+            logger.error(f"Error uploading to Hugging Face Hub: {e}", exc_info=True)
+
+def main():
+    try:
+        hf_api = initialize_hf_api()
+        download_from_hf_hub()
+        articles_to_process = fetch_rss_feeds()
+        if articles_to_process:
+            process_and_store_articles(articles_to_process)
+            upload_to_hf_hub(hf_api)
+        else:
+            logger.info("No articles fetched, skipping database processing and upload.")
+    except Exception as e:
+        logger.critical(f"An unhandled error occurred in main execution: {e}", exc_info=True)
+
+if __name__ == "__main__":
+    main()