Update stock_vector_db.py

stock_vector_db.py

@@ -1,106 +1,76 @@
 import os
-import json
-from typing import List, Optional
-from datetime import datetime
-from langchain.vectorstores import FAISS
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.schema import Document
+import shutil
+from huggingface_hub import hf_hub_download, Repository
+from langchain_community.vectorstores.faiss import FAISS
 import streamlit as st
 
-class StockVectorDB:
-    def __init__(self, index_path: str = "vector_index", log_path: str = "vector_log.json", embedding_model: Optional[str] = None):
-        self.index_path = index_path
-        self.log_path = log_path
-        model = embedding_model or "sentence-transformers/all-MiniLM-L6-v2"
-        self.embedding_model = HuggingFaceEmbeddings(model_name=model)
+class HFVectorDB:
+    def __init__(self, hf_repo_id, hf_token, local_index_dir="/tmp/vector_index", embedding_model=None):
+        self.hf_repo_id = hf_repo_id
+        self.hf_token = hf_token
+        self.local_index_dir = local_index_dir
+        self.embedding_model = embedding_model
 
-        # Load FAISS index if it exists
-        if os.path.exists(index_path):
-            st.write(f"🔁 Loading existing FAISS index from '{index_path}'")
-            self.index = FAISS.load_local(index_path, self.embedding_model)
-        else:
-            st.write(f"🆕 No FAISS index found. Will create when documents are added.")
-            self.index = None  # Delay creation
-
-        # Load or initialize log
-        if os.path.exists(self.log_path):
-            with open(self.log_path, "r") as f:
-                try:
-                    self.log_data = json.load(f)
-                except json.JSONDecodeError:
-                    self.log_data = []
-        else:
-            self.log_data = []
+        os.makedirs(self.local_index_dir, exist_ok=True)
+        self.index = None
 
-    def store_top_picks(self, top_picks: List[dict], date: Optional[datetime] = None):
-        date = date or datetime.now()
-        formatted_date = date.strftime("%Y-%m-%d")
-        docs = []
+        # Download index files from HF repo (if exist)
+        self._download_index_files()
+        self._load_index()
 
-        for stock in top_picks:
-            try:
-                content = f"{stock['ticker']} {stock['company']} is {stock['sentiment']} due to: {stock['critical_news']}. Impact: {stock['impact_summary']}. Action: {stock['action']}. Reason: {stock['reason']}"
-                metadata = {
-                    "date": formatted_date,
-                    "ticker": stock["ticker"],
-                    "company": stock["company"],
-                    "sentiment": stock["sentiment"],
-                    "action": stock["action"]
-                }
-                docs.append(Document(page_content=content, metadata=metadata))
+    def _download_index_files(self):
+        try:
+            faiss_path = hf_hub_download(repo_id=self.hf_repo_id, filename="index.faiss", use_auth_token=self.hf_token)
+            pkl_path = hf_hub_download(repo_id=self.hf_repo_id, filename="index.pkl", use_auth_token=self.hf_token)
+            shutil.copy(faiss_path, os.path.join(self.local_index_dir, "index.faiss"))
+            shutil.copy(pkl_path, os.path.join(self.local_index_dir, "index.pkl"))
+            st.write("✅ Downloaded FAISS index files from HF repo")
+        except Exception as e:
+            st.write(f"⚠️ Could not download FAISS index files: {e}")
 
-                self.log_data.append({
-                    "ticker": stock["ticker"],
-                    "company": stock["company"],
-                    "date": formatted_date,
-                    "sentiment": stock["sentiment"],
-                    "action": stock["action"],
-                    "reason": stock["reason"]
-                })
-            except KeyError as e:
-                st.write(f"❌ Skipping malformed stock data: {e}")
-                continue
+    def _load_index(self):
+        try:
+            self.index = FAISS.load_local(self.local_index_dir, self.embedding_model)
+            st.write("✅ Loaded FAISS index from local")
+        except Exception:
+            st.write("ℹ️ No local FAISS index found, starting with empty index")
+            self.index = None
 
-        if docs:
-            if self.index is None:
-                self.index = FAISS.from_documents(docs, self.embedding_model)
-            else:
-                self.index.add_documents(docs)
-
-            self.save_index()
-            self.save_log()
-            st.write(f"✅ Stored {len(docs)} documents for {formatted_date}")
+    def save_index(self):
+        if self.index is not None:
+            self.index.save_local(self.local_index_dir)
+            st.write("✅ Saved FAISS index locally")
+            self._upload_index_files()
         else:
-            st.write("⚠️ No valid documents to store.")
+            st.write("⚠️ No FAISS index to save")
 
-    def save_index(self):
-        if self.index:
-            self.index.save_local(self.index_path)
+    def _upload_index_files(self):
+        repo_local_path = "/tmp/hf_dataset_repo_clone"
+        if os.path.exists(repo_local_path):
+            shutil.rmtree(repo_local_path)
 
-    def save_log(self):
-        with open(self.log_path, "w") as f:
-            json.dump(self.log_data, f, indent=2)
+        repo = Repository(local_dir=repo_local_path, clone_from=self.hf_repo_id, use_auth_token=self.hf_token)
 
-    def search(self, query: str, k: int = 5):
-        if not self.index:
-            st.write("⚠️ FAISS index is empty.")
-            return []
+        shutil.copy(os.path.join(self.local_index_dir, "index.faiss"), os.path.join(repo_local_path, "index.faiss"))
+        shutil.copy(os.path.join(self.local_index_dir, "index.pkl"), os.path.join(repo_local_path, "index.pkl"))
 
-        print(f"🔍 Searching for: '{query}' (top {k})")
-        results = self.index.similarity_search(query, k=k)
-        for res in results:
-            st.write(f"\n📌 Ticker: {res.metadata.get('ticker')} | Sentiment: {res.metadata.get('sentiment')} | Date: {res.metadata.get('date')}")
-            st.write(res.page_content)
-            st.write("-" * 80)
-        return results
+        repo.git_add(auto_lfs_track=True)
+        repo.git_commit("Update FAISS index files")
+        repo.git_push()
+        st.write("✅ Uploaded FAISS index files to HF repo")
 
-    def backup(self, backup_dir: str = "vector_backups"):
-        if not self.index:
-            st.write("⚠️ No index to backup.")
-            return
+    def add_documents(self, docs):
+        if self.index is None:
+            self.index = FAISS.from_documents(docs, self.embedding_model)
+            st.write("✅ Created new FAISS index")
+        else:
+            self.index.add_documents(docs)
+            st.write("✅ Added documents to FAISS index")
 
-        os.makedirs(backup_dir, exist_ok=True)
-        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-        backup_path = os.path.join(backup_dir, f"vector_index_{ts}")
-        self.index.save_local(backup_path)
-        st.write(f"📦 Backup saved to {backup_path}")
+        self.save_index()
+
+    def similarity_search(self, query, k=5):
+        if self.index is None:
+            st.write("⚠️ No index found, returning empty results")
+            return []
+        return self.index.similarity_search(query, k=k)