Spaces:
Sleeping
Sleeping
| import os | |
| import feedparser | |
| 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.utils import HfHubHTTPError | |
| import json | |
| from datetime import datetime | |
| import dateutil.parser | |
| import hashlib | |
| import re | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) | |
| MAX_ARTICLES_PER_FEED = 1000 | |
| 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" | |
| 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") | |
| vector_db = None | |
| def clean_text(text): | |
| if not text or not isinstance(text, str): | |
| return "" | |
| text = re.sub(r'<.*?>', '', text) | |
| text = ' '.join(text.split()) | |
| return text.strip().lower() | |
| 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}'") | |
| 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") | |
| 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)}") | |
| 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 vector_db: | |
| logger.error("Vector database is not initialized. Cannot process articles.") | |
| return | |
| documents = [] | |
| doc_ids = [] | |
| try: | |
| existing_ids = set(vector_db.get(include=[])["ids"]) | |
| logger.info(f"Found {len(existing_ids)} existing document IDs in the database.") | |
| except Exception: | |
| existing_ids = set() | |
| logger.info("No existing documents found or error retrieving them. Starting fresh.") | |
| for article in articles: | |
| try: | |
| title = clean_text(article["title"]) | |
| link = article["link"] | |
| description = 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: | |
| continue | |
| metadata = { | |
| "title": article["title"], | |
| "link": article["link"], | |
| "published": article["published"], | |
| "category": article["category"], | |
| "image": article["image"], | |
| } | |
| doc = Document(page_content=description, metadata=metadata) | |
| documents.append(doc) | |
| doc_ids.append(doc_id) | |
| existing_ids.add(doc_id) | |
| except Exception as e: | |
| logger.error(f"Error processing article {article['title']}: {e}") | |
| if documents: | |
| try: | |
| vector_db.add_documents(documents=documents, ids=doc_ids) | |
| vector_db.persist() | |
| logger.info(f"Added {len(documents)} new articles to DB. Total documents now: {len(vector_db.get()['ids'])}") | |
| except Exception as e: | |
| logger.error(f"Error storing articles: {e}") | |
| else: | |
| logger.info("No new articles to add.") | |
| def download_from_hf_hub(): | |
| if os.path.exists(LOCAL_DB_DIR): | |
| logger.info(f"Local database directory '{LOCAL_DB_DIR}' already exists. Skipping download.") | |
| return | |
| logger.info(f"Attempting to download database from Hugging Face Hub repo: {REPO_ID}") | |
| try: | |
| snapshot_download( | |
| repo_id=REPO_ID, | |
| repo_type="dataset", | |
| local_dir=LOCAL_DB_DIR, | |
| token=HF_API_TOKEN, | |
| ) | |
| logger.info(f"Database successfully downloaded to '{LOCAL_DB_DIR}'.") | |
| except HfHubHTTPError as e: | |
| logger.warning(f"Failed to download from Hub (repo may be new or empty): {e}. Building new dataset locally.") | |
| os.makedirs(LOCAL_DB_DIR, exist_ok=True) | |
| except Exception as e: | |
| logger.error(f"An unexpected error occurred during download: {e}. Creating new local directory.") | |
| os.makedirs(LOCAL_DB_DIR, exist_ok=True) | |
| def upload_to_hf_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: | |
| hf_api.create_repo(repo_id=REPO_ID, repo_type="dataset", exist_ok=True) | |
| logger.info(f"Uploading updated Chroma DB to {REPO_ID}...") | |
| hf_api.upload_folder( | |
| folder_path=LOCAL_DB_DIR, | |
| repo_id=REPO_ID, | |
| repo_type="dataset", | |
| commit_message=f"Update database - {datetime.now().isoformat()}" | |
| ) | |
| logger.info(f"Database uploaded successfully to Hugging Face Hub.") | |
| except Exception as e: | |
| logger.error(f"Error uploading to Hugging Face Hub: {e}") | |
| if __name__ == "__main__": | |
| download_from_hf_hub() | |
| global vector_db | |
| vector_db = Chroma( | |
| persist_directory=LOCAL_DB_DIR, | |
| embedding_function=embedding_model, | |
| collection_name=COLLECTION_NAME | |
| ) | |
| articles = fetch_rss_feeds() | |
| if articles: | |
| process_and_store_articles(articles) | |
| upload_to_hf_hub() | |
| logger.info("Script finished.") |