Update rss_processor.py

rss_processor.py

@@ -5,35 +5,29 @@ from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.docstore.document import Document
 import logging
 from huggingface_hub import HfApi, login, snapshot_download
-from huggingface_hub.utils import HfHubHTTPError
-import json
+import shutil
+import rss_feeds
 from datetime import datetime
 import dateutil.parser
 import hashlib
 import re
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-MAX_ARTICLES_PER_FEED = 1000
 LOCAL_DB_DIR = "chroma_db"
-FEEDS_FILE = "rss_feeds.json"
+RSS_FEEDS = rss_feeds.RSS_FEEDS
 COLLECTION_NAME = "news_articles"
 HF_API_TOKEN = os.getenv("HF_TOKEN")
 REPO_ID = "broadfield-dev/news-rag-db"
 
-if not HF_API_TOKEN:
-    raise ValueError("HF_TOKEN environment variable not set.")
-
 login(token=HF_API_TOKEN)
 hf_api = HfApi()
 
-embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-
-'''def get_embedding_model():
+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'''
+    return get_embedding_model.model
 
 def clean_text(text):
     if not text or not isinstance(text, str):
@@ -45,85 +39,62 @@ def clean_text(text):
 def fetch_rss_feeds():
     articles = []
     seen_keys = set()
-    
-    try:
-        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.")
-        return []
-
-    for category, feeds in feed_categories.items():
-        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}'")
+    for feed_url in RSS_FEEDS:
+        try:
+            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
-
-            try:
-                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
-                article_count = 0
-                for entry in feed.entries:
-                    if article_count >= MAX_ARTICLES_PER_FEED:
-                        break
-                    
-                    title_raw = 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)
-
-                    if not clean_desc_val:
-                        continue
-
-                    published = "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")
+            article_count = 0
+            for entry in feed.entries:
+                if article_count >= 10:
+                    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)
+                
+                published = "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")
+                            break
+                        except (ValueError, TypeError):
+                            continue
+
+                key = f"{cleaned_title}|{cleaned_link}|{published}"
+                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 "",
+                    ]:
+                        try:
+                            img = img_source(entry)
+                            if img and img.strip():
+                                image = img
                                 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
-            except Exception as e:
-                logger.error(f"Error fetching {feed_url}: {e}")
-    logger.info(f"Total unique articles fetched: {len(articles)}")
+                        except (IndexError, AttributeError, TypeError):
+                            continue
+
+                    articles.append({
+                        "title": title,
+                        "link": link,
+                        "description": description,
+                        "published": published,
+                        "category": categorize_feed(feed_url),
+                        "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 categorize_feed(url):
@@ -167,8 +138,8 @@ def categorize_feed(url):
 def process_and_store_articles(articles):
     vector_db = Chroma(
         persist_directory=LOCAL_DB_DIR,
-        #embedding_function=get_embedding_model(),
-        embedding_function=embedding_model,
+        embedding_function=get_embedding_model(),
+        #embedding_function=embedding_model,
         collection_name=COLLECTION_NAME
     )