Update rss_processor.py

rss_processor.py

@@ -4,6 +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
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -52,6 +53,12 @@ RSS_FEEDS = [
     "https://www.discovermagazine.com/rss",
     "https://www.atlasobscura.com/feeds/latest"
 ]
+HF_API_TOKEN = os.getenv("DEMO_HF_API_TOKEN", "YOUR_HF_API_TOKEN")
+REPO_ID = "broadfield-dev/news-rag-db"
+
+# Initialize Hugging Face API
+login(token=HF_API_TOKEN)
+hf_api = HfApi()
 
 # Initialize embedding model and vector DB
 embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
@@ -122,6 +129,39 @@ def process_and_store_articles(articles):
         except Exception as e:
             logger.error(f"Error storing articles: {e}")
 
+def download_from_hf_hub():
+    if os.path.exists(LOCAL_DB_DIR):
+        shutil.rmtree(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.download_repo(repo_id=REPO_ID, repo_type="dataset", local_dir=LOCAL_DB_DIR, token=HF_API_TOKEN)
+    except Exception as e:
+        logger.error(f"Error downloading from Hugging Face Hub: {e}")
+        raise
+
+def upload_to_hf_hub():
+    if 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"Uploading 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}")
+            raise
+
 if __name__ == "__main__":
     articles = fetch_rss_feeds()
-    process_and_store_articles(articles)
+    process_and_store_articles(articles)
+    upload_to_hf_hub()