cmc

rag.py

@@ -1,30 +1,29 @@
-# rag.py – bullet-proof and container-safe
+# rag.py — production-grade, Parquet-ready, FAISS-cached RAG pipeline
 from __future__ import annotations
-import os, re, json, requests, traceback
-from functools import lru_cache
+import os, re, json, pickle
 from typing import List, Tuple
+from functools import lru_cache
+import requests
+from supabase import create_client
 
+# LangChain + HF
+from datasets import load_dataset
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_huggingface import HuggingFaceEndpoint
 from langchain_core.prompts import PromptTemplate
 from langchain.chains import RetrievalQA
-from supabase import create_client
 
 # ------------------------------------------------------------------ CONFIG
-DATASET_API = "https://datasets-server.huggingface.co/rows"
-DATASET     = "NimrodDev/LD_Events2"
-CONFIG      = "default"
-SPLIT       = "train"
-LIMIT       = 500
-LLM_MODEL   = "microsoft/DialoGPT-medium"
-
+DATASET_ID  = "NimrodDev/LD_Events2"
+HF_TOKEN    = os.getenv("HF_TOKEN")
 SUPABASE_URL = os.getenv("SUPABASE_URL")
 SUPABASE_KEY = os.getenv("SUPABASE_KEY")
-HF_TOKEN     = os.getenv("HF_TOKEN")
+LLM_MODEL   = "mistralai/Mistral-7B-Instruct-v0.3"  # stronger model
+CACHE_DIR   = os.getenv("HF_HOME", "/tmp/hf_cache")
+FAISS_PATH  = "/tmp/faiss_index.pkl"
 
-# ------------------------------------------------------------------ SAFE CACHE
-CACHE_DIR = os.getenv("HF_HOME", "/tmp/hf_cache")
 os.makedirs(CACHE_DIR, exist_ok=True)
 os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
 os.environ["HF_HOME"] = CACHE_DIR
@@ -32,14 +31,13 @@ os.environ["HF_HUB_CACHE"] = CACHE_DIR
 
 supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
 
-# ------------------------------------------------------------------ INTENTS
+# ------------------------------------------------------------------ INTENT DETECTION
 GREETING_RE  = re.compile(r"\b(hi|hello|hey|good morning|good afternoon|good evening)\b", re.I)
 THANKS_RE    = re.compile(r"\b(thank|thanks|appreciate)\b", re.I)
 BYE_RE       = re.compile(r"\b(bye|goodbye|see you|later)\b", re.I)
 MONEY_RE     = re.compile(r"\b(price|cost|budget|cheap|expensive|money|usd|ksh|payment|deposit)\b", re.I)
 COMPLAIN_RE  = re.compile(r"\b(complain|bad|terrible|awful|disappointed|angry|slow|rude)\b", re.I)
 
-# ------------------------------------------------------------------ FALLBACKS
 FALLBACKS = {
     "LD Events": {
         "greeting": "Hello! 👋 I’m Amina, your assistant for *LD Events* (weddings, graduations, corporate events) "
@@ -83,72 +81,67 @@ def _detect_intent(text: str) -> str:
 def _fallback_answer(company: str, intent: str) -> str:
     return FALLBACKS[company].get(intent, FALLBACKS[company]["default"])
 
-# ------------------------------------------------------------------ DATA FETCH
+# ------------------------------------------------------------------ DATA FETCH (Parquet)
 @lru_cache(maxsize=1)
 def get_texts() -> List[str]:
     try:
-        url = f"{DATASET_API}?dataset={DATASET}&config={CONFIG}&split={SPLIT}&offset=0&length={LIMIT}"
-        r = requests.get(url, timeout=60)
-        r.raise_for_status()
-        rows = r.json()["rows"]
-        texts = [row["row"]["text"] for row in rows if row["row"].get("text")]
-        print(f"✓ Fetched {len(texts)} texts from {DATASET}")
+        print("🔍 Loading dataset from Parquet branch...")
+        ds = load_dataset(DATASET_ID, split="train", revision="refs/convert/parquet")
+        texts = [str(row["text"]) for row in ds if row.get("text")]
+        print(f"✅ Loaded {len(texts)} rows from {DATASET_ID}")
         return texts
     except Exception as e:
-        print(f"⚠ Dataset fetch failed: {e} – using empty corpus")
+        print(f"⚠️ Failed to load Parquet dataset: {e}")
         return []
 
-# ------------------------------------------------------------------ EMBEDDINGS
+# ------------------------------------------------------------------ VECTORSTORE
 @lru_cache(maxsize=1)
 def get_vectorstore() -> FAISS:
-    texts = get_texts()
-    embeddings = None
-
-    try:
-        # Primary attempt: use Inference API embeddings
-        from langchain_huggingface import HuggingFaceInferenceAPIEmbeddings
-        embeddings = HuggingFaceInferenceAPIEmbeddings(
-            api_key=HF_TOKEN,
-            model_name="sentence-transformers/all-MiniLM-L6-v2"
-        )
-        print("✅ Using HuggingFaceInferenceAPIEmbeddings")
-    except ImportError:
-        # Fallback: local embeddings (container-safe)
+    if os.path.exists(FAISS_PATH):
         try:
-            from langchain_huggingface import HuggingFaceEmbeddings
-            embeddings = HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-MiniLM-L6-v2",
-                cache_folder=CACHE_DIR
-            )
-            print("⚙️ Using local HuggingFaceEmbeddings fallback")
-        except Exception as e:
-            traceback.print_exc()
-            raise RuntimeError(f"❌ Failed to load embeddings: {e}")
+            with open(FAISS_PATH, "rb") as f:
+                vs = pickle.load(f)
+                print("📦 Loaded cached FAISS index.")
+                return vs
+        except Exception:
+            print("⚠️ Cached FAISS corrupted, rebuilding...")
 
+    texts = get_texts()
     if not texts:
-        print("⚠ No dataset texts available; creating dummy FAISS index.")
-        return FAISS.from_texts([""], embeddings)
+        print("⚠️ No dataset found; creating dummy FAISS index.")
+        return FAISS.from_texts(["No context available."], HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2"))
 
-    splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=50)
-    docs = splitter.create_documents(texts, metadatas=[{"source": DATASET}] * len(texts))
-    print(f"🔹 Building FAISS index with {len(docs)} chunks")
-    return FAISS.from_documents(docs, embeddings)
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=100)
+    docs = splitter.create_documents(texts)
+    vs = FAISS.from_documents(docs, embeddings)
+
+    with open(FAISS_PATH, "wb") as f:
+        pickle.dump(vs, f)
+    print("✅ FAISS index created and cached.")
+    return vs
 
 # ------------------------------------------------------------------ LLM
 @lru_cache(maxsize=1)
 def get_llm():
     return HuggingFaceEndpoint(
         repo_id=LLM_MODEL,
+        huggingfacehub_api_token=HF_TOKEN,
         temperature=0.1,
-        max_new_tokens=150,
-        huggingfacehub_api_token=HF_TOKEN
+        max_new_tokens=200
     )
 
 PROMPT = PromptTemplate.from_template("""
 You are Amina, assistant for {company}.
-Use only the context below. If unsure, say: "A human agent will follow up."
-Context: {context}
-Question: {question}
+Answer clearly and politely using the context below.
+If context lacks information, say: "Let me get back to you on that."
+
+Context:
+{context}
+
+Question:
+{question}
+
 Answer:
 """)
 
@@ -157,40 +150,35 @@ def ask_question(phone: str, question: str) -> Tuple[str, List]:
     intent  = _detect_intent(question)
     company = _company_from_text(question)
 
+    # Quick intents skip RAG
     if intent in ("greeting", "thanks", "bye"):
         answer = _fallback_answer(company, intent)
         _save_chat(phone, question, answer)
         return answer, []
 
-    try:
-        vs = get_vectorstore()
-        docs = vs.similarity_search(question, k=3)
-    except Exception as e:
-        print(f"❌ Vector store retrieval failed: {e}")
-        docs = []
-
-    if not docs or docs[0].page_content.strip() == "":
-        answer = _fallback_answer(company, intent if intent in ("money", "complain") else "default")
-        _save_chat(phone, question, answer)
-        return answer, []
+    vs = get_vectorstore()
+    retriever = vs.as_retriever(search_kwargs={"k": 4})
+    qa = RetrievalQA.from_chain_type(
+        llm=get_llm(),
+        retriever=retriever,
+        chain_type_kwargs={"prompt": PROMPT},
+        return_source_documents=True
+    )
 
     try:
-        qa = RetrievalQA.from_chain_type(
-            llm=get_llm(),
-            retriever=vs.as_retriever(search_kwargs={"k": 3}),
-            return_source_documents=True,
-            chain_type_kwargs={"prompt": PROMPT}
-        )
         result = qa({"query": question, "company": company})
-        answer = result.get("result", "").strip() or _fallback_answer(company, "default")
-        docs_used = result.get("source_documents", [])
+        answer = result.get("result", "").strip()
+        docs = result.get("source_documents", [])
     except Exception as e:
-        traceback.print_exc()
-        answer = _fallback_answer(company, "default")
-        docs_used = []
+        print(f"❌ QA error: {e}")
+        answer, docs = "", []
+
+    # Fallback logic
+    if not answer or len(answer.split()) < 3:
+        answer = _fallback_answer(company, intent if intent in ("money", "complain") else "default")
 
     _save_chat(phone, question, answer)
-    return answer, docs_used
+    return answer, docs
 
 # ------------------------------------------------------------------ SUPABASE LOGGING
 def _save_chat(phone: str, q: str, a: str) -> None:
@@ -198,4 +186,4 @@ def _save_chat(phone: str, q: str, a: str) -> None:
         supabase.table("chat_memory").insert({"user_phone": phone, "role": "user", "message": q}).execute()
         supabase.table("chat_memory").insert({"user_phone": phone, "role": "assistant", "message": a}).execute()
     except Exception as e:
-        print(f"⚠ Chat log save failed: {e}")
+        print(f"⚠️ Supabase logging failed: {e}")