Update src/qa.py

src/qa.py

@@ -3,8 +3,8 @@ qa.py — Retrieval + Generation Layer
 -------------------------------------
 Handles:
 • Query embedding (SentenceTransformer / E5-compatible)
-• Chunk retrieval (FAISS)
-• Answer generation (OpenAI or Flan-T5 fallback)
+• Chunk retrieval (FAISS with neighborhood merging + re-ranking)
+• Answer generation (OpenAI GPT-4o-mini or fallback to Flan-T5)
 Optimized for Hugging Face Spaces & Streamlit.
 """
 
@@ -14,16 +14,6 @@ from sentence_transformers import SentenceTransformer
 from vectorstore import search_faiss
 from sklearn.metrics.pairwise import cosine_similarity
 
-# Optional: use OpenAI if API 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 for answer generation")
-else:
-    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-    print("⚙️ Using fallback FLAN-T5 model (local)")
-
 print("✅ qa.py loaded from:", __file__)
 
 # ==========================================================
@@ -39,7 +29,24 @@ os.environ.update({
 })
 
 # ==========================================================
-# 2️⃣ Query Embedding Model
+# 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
 # ==========================================================
 try:
     _query_model = SentenceTransformer("intfloat/e5-small-v2", cache_folder=CACHE_DIR)
@@ -49,16 +56,18 @@ except Exception as e:
     _query_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", cache_folder=CACHE_DIR)
 
 # ==========================================================
-# 3️⃣ Fallback LLM (if no OpenAI key)
+# 4️⃣ Fallback LLM (if no OpenAI key or quota exhausted)
 # ==========================================================
 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)
 
 # ==========================================================
-# 4️⃣ Prompt Template
+# 5️⃣ Prompt Template
 # ==========================================================
 PROMPT_TEMPLATE = """
 You are an enterprise knowledge assistant.
@@ -77,7 +86,7 @@ Answer:
 """
 
 # ==========================================================
-# 5️⃣ Chunk Retrieval Function
+# 6️⃣ Chunk Retrieval Function
 # ==========================================================
 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."""
@@ -111,10 +120,10 @@ def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5):
         return []
 
 # ==========================================================
-# 6️⃣ Answer Generation Function
+# 7️⃣ Answer Generation Function
 # ==========================================================
 def generate_answer(query: str, retrieved_chunks: list):
-    """Generate factual, complete answers using OpenAI or FLAN."""
+    """Generate factual, complete answers using OpenAI (or Flan-T5 fallback)."""
     if not retrieved_chunks:
         return "Sorry, I couldn’t find relevant information in the document."
 
@@ -133,7 +142,7 @@ def generate_answer(query: str, retrieved_chunks: list):
                     {"role": "user", "content": prompt},
                 ],
                 temperature=0.4,
-                max_tokens=600,
+                max_tokens=800,
             )
             return completion.choices[0].message.content.strip()
 
@@ -144,11 +153,17 @@ def generate_answer(query: str, retrieved_chunks: list):
 
     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."
 
-
 # ==========================================================
-# 7️⃣ Local Test
+# 8️⃣ Local Test
 # ==========================================================
 if __name__ == "__main__":
     dummy_chunks = [