Update src/qa.py

src/qa.py

@@ -3,22 +3,31 @@ qa.py — Retrieval + Generation Layer
 -------------------------------------
 Handles:
 • Query embedding (SentenceTransformer / E5-compatible)
-• Chunk retrieval (FAISS with neighborhood merging + re-ranking)
-• Answer generation (Flan-T5, tuned for factual completeness)
+• Chunk retrieval (FAISS)
+• Answer generation (OpenAI or Flan-T5 fallback)
 Optimized for Hugging Face Spaces & Streamlit.
 """
 
 import os
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 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__)
 
 # ==========================================================
-# 1️⃣ Hugging Face Cache Setup (Safe for Spaces)
+# 1️⃣ Hugging Face Cache Setup
 # ==========================================================
 CACHE_DIR = "/tmp/hf_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
@@ -33,111 +42,128 @@ os.environ.update({
 # 2️⃣ Query Embedding Model
 # ==========================================================
 try:
-    _query_model = SentenceTransformer(
-        "intfloat/e5-small-v2",
-        cache_folder=CACHE_DIR
-    )
+    _query_model = SentenceTransformer("intfloat/e5-small-v2", cache_folder=CACHE_DIR)
     print("✅ Loaded query model: intfloat/e5-small-v2")
 except Exception as e:
-    print(f"⚠️ Query model load failed ({e}), falling back to MiniLM.")
-    _query_model = SentenceTransformer(
-        "sentence-transformers/all-MiniLM-L6-v2",
-        cache_folder=CACHE_DIR
-    )
-    print("✅ Loaded fallback model: all-MiniLM-L6-v2")
+    print(f"⚠️ Query model load failed ({e}), using fallback MiniLM.")
+    _query_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", cache_folder=CACHE_DIR)
 
 # ==========================================================
-# 3️⃣ LLM for Answer Generation (OpenAI GPT with Flan fallback)
+# 3️⃣ Fallback LLM (if no OpenAI key)
 # ==========================================================
-from openai import OpenAI
-client = None
-
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-if OPENAI_API_KEY:
-    client = OpenAI(api_key=OPENAI_API_KEY)
-    LLM_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
-    print(f"✅ Using OpenAI model: {LLM_MODEL}")
-else:
-    # Fallback to Flan if no API key is provided
+if not USE_OPENAI:
     MODEL_NAME = "google/flan-t5-base"
-    print(f"⚠️ No OpenAI key found. 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
-    )
-
+    _answer_model = pipeline("text2text-generation", model=_model, tokenizer=_tokenizer, device=-1)
 
 # ==========================================================
-# 6️⃣ Answer Generation Function (GPT or Flan fallback)
+# 4️⃣ Prompt Template
 # ==========================================================
-def generate_answer(query: str, retrieved_chunks: list):
-    """
-    Generates grounded, context-only answers.
-    Uses GPT (preferred) or Flan-T5 (fallback) for response synthesis.
-    """
-    if not retrieved_chunks:
-        return "Sorry, I couldn’t find relevant information in the document."
-
-    # Combine retrieved chunks
-    context = "\n\n".join([
-        f"[Chunk {i+1}]: {chunk.strip()}" for i, chunk in enumerate(retrieved_chunks)
-    ])
-
-    # --- PROMPT TEMPLATE ---
-    system_prompt = """You are an enterprise knowledge assistant.
-Use ONLY the provided context to answer the user's question accurately.
-If the answer is not explicitly in the context, reply exactly:
+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:
 "I don't know based on the provided document."
-Be factual, concise, and structured when relevant.
-"""
 
-    user_prompt = f"""
+---
 Context:
 {context}
-
+---
 Question:
 {query}
-
+---
 Answer:
 """
 
-    # --- Use OpenAI GPT if key available ---
-    if client:
-        try:
-            response = client.chat.completions.create(
-                model=LLM_MODEL,
+# ==========================================================
+# 5️⃣ 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."""
+    if not index or not chunks:
+        return []
+
+    try:
+        query_emb = _query_model.encode(
+            [f"query: {query.strip()}"],
+            convert_to_numpy=True,
+            normalize_embeddings=True
+        )[0]
+
+        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))
+
+        chunk_vecs = np.array([
+            _query_model.encode([c], convert_to_numpy=True, normalize_embeddings=True)[0]
+            for c in merged_chunks
+        ])
+        scores = cosine_similarity(np.array([query_emb]), chunk_vecs)[0]
+        sorted_indices = np.argsort(scores)[::-1]
+
+        return [merged_chunks[i] for i in sorted_indices[:top_k]]
+
+    except Exception as e:
+        print(f"⚠️ Retrieval error: {e}")
+        return []
+
+# ==========================================================
+# 6️⃣ Answer Generation Function
+# ==========================================================
+def generate_answer(query: str, retrieved_chunks: list):
+    """Generate factual, complete answers using OpenAI or FLAN."""
+    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)
+    ])
+    prompt = PROMPT_TEMPLATE.format(context=context, query=query)
+
+    try:
+        if USE_OPENAI:
+            completion = client.chat.completions.create(
+                model="gpt-4o-mini",
                 messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt},
+                    {"role": "system", "content": "You are a precise enterprise document assistant."},
+                    {"role": "user", "content": prompt},
                 ],
-                temperature=0.2,        # factual, low creativity
-                max_tokens=500,
-                presence_penalty=0,
-                frequency_penalty=0
+                temperature=0.4,
+                max_tokens=600,
             )
-            answer = response.choices[0].message.content.strip()
+            return completion.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
-        except Exception as e:
-            print(f"⚠️ OpenAI generation failed: {e}")
-            return "⚠️ Error: Could not generate an answer at the moment."
 
-    # --- Otherwise, use Flan-T5 fallback ---
-    try:
-        result = _answer_model(
-            PROMPT_TEMPLATE.format(context=context, query=query),
-            max_new_tokens=600,
-            do_sample=False,
-            temperature=0.3,
-            repetition_penalty=1.1
-        )
-        answer = result[0]["generated_text"].strip()
-        if "I don't know" in answer:
-            return "I don't know based on the provided document."
-        return answer
     except Exception as e:
-        print(f"⚠️ Flan generation failed: {e}")
+        print(f"⚠️ Generation failed: {e}")
         return "⚠️ Error: Could not generate an answer at the moment."
+
+
+# ==========================================================
+# 7️⃣ Local Test
+# ==========================================================
+if __name__ == "__main__":
+    dummy_chunks = [
+        "Step 1: Open the dashboard and navigate to reports.",
+        "Step 2: Click 'Export' to download a CSV summary.",
+        "Step 3: Review the generated report in your downloads folder."
+    ]
+    from vectorstore import build_faiss_index
+
+    index = build_faiss_index([
+        _query_model.encode([f"passage: {chunk}"], convert_to_numpy=True, normalize_embeddings=True)[0]
+        for chunk in dummy_chunks
+    ])
+
+    query = "What are the steps to export a report?"
+    retrieved = retrieve_chunks(query, index, dummy_chunks)
+    print("🔍 Retrieved:", retrieved)
+    print("💬 Answer:", generate_answer(query, retrieved))