Update src/qa.py

src/qa.py

@@ -3,21 +3,34 @@ qa.py — Retrieval + Generation Layer
 -------------------------------------
 Handles:
 • Query embedding (SentenceTransformer / E5-compatible)
-• Chunk retrieval (FAISS with neighborhood merging + re-ranking)
-• Answer generation (OpenAI GPT-4o-mini or fallback to Flan-T5)
+• Chunk retrieval (FAISS + cosine re-ranking)
+• Answer generation (OpenAI GPT-4o-mini or FLAN-T5 fallback)
 Optimized for Hugging Face Spaces & Streamlit.
 """
 
 import os
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from vectorstore import search_faiss
 from sklearn.metrics.pairwise import cosine_similarity
+from vectorstore import search_faiss
 
-print("✅ qa.py loaded from:", __file__)
+# ==========================================================
+# 1️⃣ Load OpenAI if key available
+# ==========================================================
+USE_OPENAI = bool(os.getenv("OPENAI_API_KEY"))
+
+if USE_OPENAI:
+    from openai import OpenAI
+    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+    print("✅ Using OpenAI GPT-4o-mini for answer generation")
+else:
+    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+    print("⚙️ No OpenAI key found — using fallback FLAN-T5 model")
+
+print("✅ qa.py loaded successfully")
 
 # ==========================================================
-# 1️⃣ Hugging Face Cache Setup
+# 2️⃣ Hugging Face Cache Setup (Safe for Spaces)
 # ==========================================================
 CACHE_DIR = "/tmp/hf_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
@@ -29,39 +42,21 @@ os.environ.update({
 })
 
 # ==========================================================
-# 2️⃣ OpenAI Integration (with safe fallback)
-# ==========================================================
-# ⚠️ TEMPORARY: You can hardcode your key here for testing
-os.environ["OPENAI_API_KEY"] = "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"
-
-USE_OPENAI = bool(os.getenv("OPENAI_API_KEY"))
-
-if USE_OPENAI:
-    try:
-        from openai import OpenAI
-        client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-        print("✅ Using OpenAI GPT-4o-mini for answer generation")
-    except Exception as e:
-        print(f"⚠️ OpenAI client initialization failed: {e}")
-        USE_OPENAI = False
-
-# ==========================================================
-# 3️⃣ Query Embedding Model
+# 3️⃣ Embedding Model (E5 for better retrieval)
 # ==========================================================
 try:
     _query_model = SentenceTransformer("intfloat/e5-small-v2", cache_folder=CACHE_DIR)
-    print("✅ Loaded query model: intfloat/e5-small-v2")
+    print("✅ Loaded embedding model: intfloat/e5-small-v2")
 except Exception as e:
-    print(f"⚠️ Query model load failed ({e}), using fallback MiniLM.")
+    print(f"⚠️ Failed to load e5-small-v2 ({e}), switching to MiniLM.")
     _query_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", cache_folder=CACHE_DIR)
+    print("✅ Loaded fallback: all-MiniLM-L6-v2")
 
 # ==========================================================
-# 4️⃣ Fallback LLM (if no OpenAI key or quota exhausted)
+# 4️⃣ Fallback Model (FLAN-T5)
 # ==========================================================
 if not USE_OPENAI:
-    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
     MODEL_NAME = "google/flan-t5-base"
-    print(f"⚙️ Using fallback model: {MODEL_NAME}")
     _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
     _model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
     _answer_model = pipeline("text2text-generation", model=_model, tokenizer=_tokenizer, device=-1)
@@ -71,8 +66,8 @@ if not USE_OPENAI:
 # ==========================================================
 PROMPT_TEMPLATE = """
 You are an enterprise knowledge assistant.
-Use ONLY the CONTEXT below to answer the QUESTION clearly and factually.
-If the context doesn’t contain the answer, reply exactly:
+Use ONLY the context below to answer the question clearly, precisely, and factually.
+If the context doesn’t contain the answer, say exactly:
 "I don't know based on the provided document."
 
 ---
@@ -86,10 +81,10 @@ Answer:
 """
 
 # ==========================================================
-# 6️⃣ Chunk Retrieval Function
+# 6️⃣ Chunk Retrieval
 # ==========================================================
 def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5):
-    """Retrieve top-K relevant chunks, merge nearby ones, and re-rank by cosine similarity."""
+    """Retrieve top-K relevant chunks and re-rank by semantic similarity."""
     if not index or not chunks:
         return []
 
@@ -100,12 +95,14 @@ def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5):
             normalize_embeddings=True
         )[0]
 
+        # Retrieve more and then re-rank
         distances, indices = index.search(np.array([query_emb]).astype("float32"), top_k * 2)
         merged_chunks = []
         for idx in indices[0]:
             neighbors = [chunks[i] for i in range(max(0, idx - 1), min(len(chunks), idx + 2))]
             merged_chunks.append(" ".join(neighbors))
 
+        # Re-rank by cosine similarity
         chunk_vecs = np.array([
             _query_model.encode([c], convert_to_numpy=True, normalize_embeddings=True)[0]
             for c in merged_chunks
@@ -120,46 +117,38 @@ def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5):
         return []
 
 # ==========================================================
-# 7️⃣ Answer Generation Function
+# 7️⃣ Answer Generation
 # ==========================================================
 def generate_answer(query: str, retrieved_chunks: list):
-    """Generate factual, complete answers using OpenAI (or Flan-T5 fallback)."""
+    """Generate factual answer using OpenAI GPT-4o-mini (preferred) or FLAN fallback."""
     if not retrieved_chunks:
         return "Sorry, I couldn’t find relevant information in the document."
 
-    context = "\n\n".join([
-        f"[Chunk {i+1}]: {chunk.strip()}"
-        for i, chunk in enumerate(retrieved_chunks)
-    ])
+    # Merge retrieved chunks
+    context = "\n\n".join(
+        [f"[Chunk {i+1}]: {chunk.strip()}" for i, chunk in enumerate(retrieved_chunks)]
+    )
     prompt = PROMPT_TEMPLATE.format(context=context, query=query)
 
     try:
         if USE_OPENAI:
-            completion = client.chat.completions.create(
-                model="gpt-3.5-turbo",
+            response = client.chat.completions.create(
+                model="gpt-4o-mini",
                 messages=[
-                    {"role": "system", "content": "You are a precise enterprise document assistant."},
+                    {"role": "system", "content": "You are a precise enterprise assistant that answers only from the provided context."},
                     {"role": "user", "content": prompt},
                 ],
                 temperature=0.4,
                 max_tokens=800,
             )
-            return completion.choices[0].message.content.strip()
+            return response.choices[0].message.content.strip()
 
         else:
             result = _answer_model(prompt, max_new_tokens=600, do_sample=False, temperature=0.3)
-            answer = result[0]["generated_text"].strip()
-            return answer
+            return result[0]["generated_text"].strip()
 
     except Exception as e:
         print(f"⚠️ Generation failed: {e}")
-        # Auto fallback to Flan-T5 if OpenAI fails mid-session
-        if USE_OPENAI:
-            try:
-                result = _answer_model(prompt, max_new_tokens=600, do_sample=False, temperature=0.3)
-                return result[0]["generated_text"].strip()
-            except Exception as e2:
-                print(f"⚠️ Fallback model also failed: {e2}")
         return "⚠️ Error: Could not generate an answer at the moment."
 
 # ==========================================================