src/qa.py

"""
qa.py — Phi-2 FAST + ReRank (stable) — Prefer semantic ranking, neighbor-fill last-resort
---------------------------------------------------------------------------------------
- Uses intfloat/e5-small-v2 for embeddings
- Uses microsoft/phi-2 for generation
- Re-ranks candidate pool from FAISS then picks top_k by true cosine similarity
- Neighbor expansion only if not enough high-sim items
- Logs chunk indices + similarity scores for debugging
"""

import os
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch

print("✅ qa.py (Phi-2 FAST + ReRank stable) loaded from:", __file__)

# ---------------------------
# Cache
# ---------------------------
CACHE_DIR = "/tmp/hf_cache"
os.makedirs(CACHE_DIR, exist_ok=True)
os.environ.update({
    "HF_HOME": CACHE_DIR,
    "TRANSFORMERS_CACHE": CACHE_DIR,
    "HF_DATASETS_CACHE": CACHE_DIR,
    "HF_MODULES_CACHE": CACHE_DIR
})
print(f"✅ Using Hugging Face cache at {CACHE_DIR}")

# ---------------------------
# Embeddings
# ---------------------------
try:
    _query_model = SentenceTransformer("intfloat/e5-small-v2", cache_folder=CACHE_DIR)
    print("✅ Loaded embedding model: intfloat/e5-small-v2")
except Exception as e:
    print(f"⚠️ Embedding load failed ({e}), falling back to MiniLM")
    _query_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", cache_folder=CACHE_DIR)

# ---------------------------
# Phi-2 model
# ---------------------------
MODEL_NAME = "microsoft/phi-2"
print(f"✅ Loading LLM: {MODEL_NAME}")

_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
_model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    cache_dir=CACHE_DIR,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.bfloat16,
    low_cpu_mem_usage=True,
).to("cpu")

_answer_model = pipeline(
    "text-generation",
    model=_model,
    tokenizer=_tokenizer,
    device=-1,
    model_kwargs={"torch_dtype": torch.bfloat16, "low_cpu_mem_usage": True},
)
print("✅ Phi-2 text-generation pipeline ready (optimized).")

# ---------------------------
# Prompts
# ---------------------------
STRICT_PROMPT = (
    "You are an enterprise documentation assistant.\n"
    "Use ONLY the CONTEXT chunks below to answer the QUESTION.\n"
    "Cite the chunk number(s) you used, e.g. [Chunk 3].\n"
    "If the document does not contain the answer, reply exactly:\n"
    "\"I don't know based on the provided document.\"\n\n"
    "Context:\n{context}\n\nQuestion: {query}\nAnswer:"
)

REASONING_PROMPT = (
    "You are an expert enterprise assistant with reasoning capacity.\n"
    "Prefer the provided CONTEXT but you may cautiously infer when reasonable.\n"
    "If you infer, say so and prefer facts from the document.\n"
    "If the document lacks the answer, say:\n"
    "\"I don't know based on the provided document.\"\n\n"
    "Context:\n{context}\n\nQuestion: {query}\nAnswer:"
)

# ---------------------------
# Retrieval: FAISS -> rerank -> neighbor fill (last resort)
# ---------------------------
def retrieve_chunks(query: str, index, chunks: list, top_k: int = 3, min_similarity: float = 0.55, candidate_multiplier: int = 4):
    """
    Steps:
    1. Encode query (E5 style).
    2. Run FAISS search for k*candidate_multiplier candidates.
    3. Re-embed those candidate texts and compute cosine similarity with query embedding.
    4. Sort by similarity and pick top_k where similarity >= min_similarity.
    5. If fewer than top_k passed threshold, fill remaining slots by:
       - selecting neighboring chunks around the *highest-scoring* chunk(s),
         but only if absolutely necessary (keeps noise low).
    Returns: ordered list of chunks (strings)
    Also prints indices + similarity scores for debugging.
    """

    if not index or not chunks:
        return []

    try:
        # 1. encode query
        q_emb = _query_model.encode(
            [f"query: {query.strip()}"],
            convert_to_numpy=True,
            normalize_embeddings=True
        )[0]

        # 2. FAISS initial retrieval (get a larger candidate pool)
        num_candidates = max(top_k * candidate_multiplier, top_k + 2)
        distances, indices = index.search(np.array([q_emb]).astype("float32"), num_candidates)
        candidate_indices = [int(i) for i in indices[0] if i >= 0]

        # protective dedupe and clamp
        candidate_indices = list(dict.fromkeys(candidate_indices))  # preserve order, unique

        # 3. Re-embed candidate texts and compute true cosine similarity
        candidate_texts = [chunks[i] for i in candidate_indices]
        # Encode passages (passage prefix helps alignment)
        doc_embs = _query_model.encode(
            [f"passage: {c}" for c in candidate_texts],
            convert_to_numpy=True,
            normalize_embeddings=True
        )
        sims = cosine_similarity([q_emb], doc_embs)[0]

        # Pair up indices and sims and sort descending
        paired = [(candidate_indices[i], float(sims[i])) for i in range(len(candidate_indices))]
        paired_sorted = sorted(paired, key=lambda x: x[1], reverse=True)

        # Debug print: top candidates and their similarity
        print("🔎 Candidate ranking (index : sim):")
        for idx, sim in paired_sorted[: min(len(paired_sorted), top_k * 3)]:
            print(f"  - Chunk {idx} : {sim:.4f}")

        # 4. Pick those meeting threshold
        selected = [idx for idx, sim in paired_sorted if sim >= min_similarity]

        # Preserve order by similarity
        selected = selected[:top_k]

        # 5. If not enough, fill by neighbors around highest-scoring items
        if len(selected) < top_k:
            needed = top_k - len(selected)
            # pick highest scoring indices as anchor(s)
            anchors = [idx for idx, _ in paired_sorted[:3]]  # top 3 anchors
            expanded = []
            for a in anchors:
                # neighbors ordered by proximity: a, a-1, a+1, a-2, a+2 ...
                if a not in expanded:
                    expanded.append(a)
                offset = 1
                while len(expanded) < top_k and offset < 5:
                    for cand in (a - offset, a + offset):
                        if 0 <= cand < len(chunks) and cand not in expanded:
                            expanded.append(cand)
                            if len(expanded) >= top_k:
                                break
                    offset += 1
                if len(expanded) >= top_k:
                    break
            # final selected: first maintain previously selected, then add neighbors from expanded preserving order
            final_order = []
            for idx, _sim in paired_sorted:
                if idx in selected and idx not in final_order:
                    final_order.append(idx)
            for idx in expanded:
                if idx not in final_order:
                    final_order.append(idx)
            selected = final_order[:top_k]

        # final chunk strings (ordered by selected list)
        final_chunks = [chunks[i] for i in selected]

        print(f"✅ retrieve_chunks: returning {len(final_chunks)} chunks (top_k={top_k}, min_sim={min_similarity})")
        print(f"  chunk indices: {selected}")

        # Also return the indices? (if you want to display chunk numbers in UI, you can)
        return final_chunks

    except Exception as e:
        print(f"⚠️ Retrieval error: {e}")
        return []


# ---------------------------
# Answer generation
# ---------------------------
def generate_answer(query: str, retrieved_chunks: list, reasoning_mode: bool = False):
    """
    - reasoning_mode=False => strict factual, deterministic
    - reasoning_mode=True  => allow cautious inference (slower / longer)
    """
    if not retrieved_chunks:
        return "Sorry, I couldn’t find relevant information in the document."

    # Add chunk headings so model can cite them if needed
    context_lines = []
    for i, chunk in enumerate(retrieved_chunks, start=1):
        # Use [Chunk i] markers — LLM will echo them when asked to cite sources
        context_lines.append(f"[Chunk {i}]: {chunk.strip()}")
    context = "\n".join(context_lines)

    prompt = (REASONING_PROMPT if reasoning_mode else STRICT_PROMPT).format(
        context=context, query=query
    )

    try:
        # deterministic in strict mode
        if reasoning_mode:
            max_new_tokens = 220
            temp = 0.6
            do_sample = True
        else:
            max_new_tokens = 140
            temp = 0.0
            do_sample = False

        result = _answer_model(
            prompt,
            max_new_tokens=max_new_tokens,
            temperature=temp,
            do_sample=do_sample,
            early_stopping=True,
            pad_token_id=_tokenizer.eos_token_id,
        )

        text = result[0].get("generated_text", "").strip()
        # remove the prompt echo if present
        if "Answer:" in text:
            out = text.split("Answer:")[-1].strip()
        else:
            out = text

        # Enforce exact fallback phrase if model tries to paraphrase missing-answer
        if not reasoning_mode and ("i don't know" in out.lower() or "not present" in out.lower()):
            return "I don't know based on the provided document."

        return out

    except Exception as e:
        print(f"⚠️ Generation failed: {e}")
        return "⚠️ Error: Could not generate an answer."

# ---------------------------
# Local debug main
# ---------------------------
if __name__ == "__main__":
    from vectorstore import build_faiss_index

    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.",
        "Appendix: Communication user creation steps are explained later in this guide."
    ]
    embeddings = [
        _query_model.encode([f"passage: {c}"], convert_to_numpy=True, normalize_embeddings=True)[0]
        for c in dummy_chunks
    ]
    index = build_faiss_index(embeddings)

    query = "How do I create a communication user?"
    retrieved = retrieve_chunks(query, index, dummy_chunks, top_k=3, min_similarity=0.55)
    print("🔍 Retrieved:", retrieved)
    print("💬 Answer:", generate_answer(query, retrieved, reasoning_mode=False))