from huggingface_hub import hf_hub_download
import os

faiss_dir = os.path.join(os.path.dirname(__file__), "dify_faiss_index")
os.makedirs(faiss_dir, exist_ok=True)

faiss_path = os.path.join(faiss_dir, "index.faiss")
pkl_path = os.path.join(faiss_dir, "index.pkl")

if not os.path.exists(faiss_path):
    hf_hub_download(repo_id="k01010/k01010_dify-faiss-index", filename="index.faiss", local_dir=faiss_dir)
if not os.path.exists(pkl_path):
    hf_hub_download(repo_id="k01010/k01010_dify-faiss-index", filename="index.pkl", local_dir=faiss_dir)

from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from transformers import pipeline

# Load FAISS vector store and QA pipeline ONCE at module level
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-MiniLM-L6-cos-v1")
index_path = os.path.join(os.path.dirname(__file__), "dify_faiss_index")
vector_store = FAISS.load_local(index_path, embeddings, allow_dangerous_deserialization=True)
qa = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")

# Main RAG answer function
def answer_question(question, top_k=4):
    retriever = vector_store.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.get_relevant_documents(question)
    context = " ".join([doc.page_content for doc in docs])
    result = qa(question=question, context=context)
    return {
        "answer": result["answer"],
        "score": result["score"],
        "context": context,
        "sources": [getattr(doc, "metadata", {}) for doc in docs]
    }

if __name__ == "__main__":
    # Simple CLI for testing
    while True:
        q = input("Ask a question (or 'exit'): ")
        if q.lower() == "exit":
            break
        out = answer_question(q)
        print(f"Answer: {out['answer']}\nScore: {out['score']:.2f}\nSources: {out['sources']}")