Production grade bot with re-index button

src/config.py

@@ -0,0 +1,9 @@
+# config.py
+INDEX_DIR = "src/index"
+FAISS_INDEX_PATH = f"{INDEX_DIR}/faiss.index"
+DOC_STORE_PATH = f"{INDEX_DIR}/documents.pkl"
+
+HF_DATASET = "OnlyTheTruth03/ott"
+HF_SPLIT = "train"
+
+TOP_K = 4

src/ingest.py

@@ -1,57 +1,48 @@
-# src/ingest.py
+# ingest.py
 import os
 import pickle
-import faiss
 from datasets import load_dataset
 from pypdf import PdfReader
 from sentence_transformers import SentenceTransformer
+import faiss
 
-DATASET_NAME = "OnlyTheTruth03/ott"
-INDEX_DIR = "src/index"
+from config import INDEX_DIR, FAISS_INDEX_PATH, DOC_STORE_PATH, HF_DATASET, HF_SPLIT
 
-os.makedirs(INDEX_DIR, exist_ok=True)
 
-embedder = SentenceTransformer("all-MiniLM-L6-v2")
+def build_index():
+    os.makedirs(INDEX_DIR, exist_ok=True)
 
+    dataset = load_dataset(HF_DATASET, split=HF_SPLIT)
 
-def build_index():
-    print("📥 Loading HF dataset...")
-    dataset = load_dataset(DATASET_NAME, split="train")
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
     documents = []
+    embeddings = []
 
     for row in dataset:
-        # HF auto-parquet PDF object
-        pdf_obj = row[dataset.column_names[0]]
+        pdf_obj = row["pdf"]  # HF auto column
+        reader = PdfReader(pdf_obj)
 
-        # ✅ THIS is the key line
-        pdf_path = pdf_obj.path
-
-        print(f"📄 Reading PDF from: {pdf_path}")
-        reader = PdfReader(pdf_path)
-
-        for page_no, page in enumerate(reader.pages, start=1):
+        for page_no, page in enumerate(reader.pages):
             text = page.extract_text()
             if not text:
                 continue
 
             documents.append({
-                "text": text.strip(),
-                "page": page_no
+                "text": text,
+                "page": page_no + 1,
+                "source": "dataset_pdf"
             })
+            embeddings.append(text)
 
-    if not documents:
-        raise RuntimeError("❌ No text extracted from PDF")
-
-    texts = [d["text"] for d in documents]
-    embeddings = embedder.encode(texts).astype("float32")
+    vectors = embedder.encode(embeddings, show_progress_bar=True)
 
-    index = faiss.IndexFlatL2(embeddings.shape[1])
-    index.add(embeddings)
+    index = faiss.IndexFlatL2(vectors.shape[1])
+    index.add(vectors)
 
-    faiss.write_index(index, f"{INDEX_DIR}/faiss.index")
+    faiss.write_index(index, FAISS_INDEX_PATH)
 
-    with open(f"{INDEX_DIR}/documents.pkl", "wb") as f:
+    with open(DOC_STORE_PATH, "wb") as f:
         pickle.dump(documents, f)
 
     print("✅ FAISS index built successfully")

src/rag.py

@@ -1,69 +1,61 @@
-# src/rag.py
+# rag.py
 import os
 import pickle
 import faiss
-import numpy as np
-from sentence_transformers import SentenceTransformer
 from groq import Groq
-from ingest import build_index  # 👈 important
+from sentence_transformers import SentenceTransformer
 
-INDEX_DIR = "src/index"
-TOP_K = 4
-
-embedder = SentenceTransformer("all-MiniLM-L6-v2")
+from config import FAISS_INDEX_PATH, DOC_STORE_PATH, TOP_K
 
 GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
 if not GROQ_API_KEY:
-    raise RuntimeError("❌ GROQ_API_KEY not set in Hugging Face Secrets")
+    raise RuntimeError("❌ GROQ_API_KEY not set")
 
 client = Groq(api_key=GROQ_API_KEY)
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
 
 def load_index():
-    index_path = f"{INDEX_DIR}/faiss.index"
-    docs_path = f"{INDEX_DIR}/documents.pkl"
+    if not os.path.exists(FAISS_INDEX_PATH):
+        raise RuntimeError("❌ FAISS index not found. Run ingestion first.")
 
-    # 🔥 AUTO INGEST IF MISSING
-    if not os.path.exists(index_path) or not os.path.exists(docs_path):
-        print("⚠️ FAISS index missing. Running ingestion...")
-        build_index()
+    index = faiss.read_index(FAISS_INDEX_PATH)
 
-    index = faiss.read_index(index_path)
-
-    with open(docs_path, "rb") as f:
+    with open(DOC_STORE_PATH, "rb") as f:
         documents = pickle.load(f)
 
     return index, documents
 
 
-def retrieve(query, top_k=TOP_K):
+def retrieve(query):
     index, documents = load_index()
 
-    query_embedding = embedder.encode([query]).astype("float32")
-    distances, indices = index.search(query_embedding, top_k)
-
-    results = []
-    for idx in indices[0]:
-        if idx == -1:
-            continue
-        results.append(documents[idx])
+    q_vec = embedder.encode([query])
+    _, indices = index.search(q_vec, TOP_K)
 
-    return results
+    return [documents[i] for i in indices[0]]
 
 
 def ask_llm(query, contexts):
     context_text = "\n\n".join(
-        f"[p.{c['page']}]\n{c['text']}"
-        for c in contexts
+        f"(Page {c['page']}): {c['text']}" for c in contexts
     )
 
+    prompt = f"""
+Answer the question using only the context below.
+If the answer is not found, say so.
+
+Context:
+{context_text}
+
+Question:
+{query}
+"""
+
     response = client.chat.completions.create(
-        model="llama-3.1-8b-instant",
-        messages=[
-            {"role": "system", "content": "You are an astrology tutor."},
-            {"role": "user", "content": f"{context_text}\n\nQuestion: {query}"}
-        ],
-        temperature=0.2
+        model="llama-3.1-70b-versatile",
+        messages=[{"role": "user", "content": prompt}],
+        temperature=0.2,
     )
 
     return response.choices[0].message.content

src/streamlit_app.py

@@ -1,16 +1,25 @@
-# src/streamlit_app.py
+# streamlit_app.py
 import streamlit as st
+from ingest import build_index
 from rag import retrieve, ask_llm
 
-st.set_page_config(page_title="OTT Astrology Bot", layout="wide")
-st.title("🪐 OTT Astrology Assistant")
+st.title("📘 OTT Knowledge Bot")
 
-query = st.text_input("Ask your astrology question:")
+if st.button("🔁 Build / Rebuild Index"):
+    with st.spinner("Building index..."):
+        build_index()
+    st.success("Index built successfully")
+
+query = st.text_input("Ask a question")
 
 if query:
-    with st.spinner("Thinking..."):
+    with st.spinner("Searching..."):
         contexts = retrieve(query)
         answer = ask_llm(query, contexts)
 
-    st.markdown("### 🪐 Answer")
+    st.markdown("### ✅ Answer")
     st.write(answer)
+
+    with st.expander("📚 Sources"):
+        for c in contexts:
+            st.markdown(f"- Page {c['page']}")