chatbot_rag.py

from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain.chains import RetrievalQA
import traceback  # ✅ added


def build_qa():
    """Builds and returns the RAG QA pipeline."""
    print("🚀 Starting QA pipeline...")

    # 1. Embeddings
    print("🔹 Loading embeddings...")
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

    # 2. Load vector DB
    print("🔹 Loading Chroma DB...")
    vectorstore = Chroma(
        persist_directory="db",
        collection_name="rag-docs",
        embedding_function=embeddings,
    )
    print("📂 Docs in DB:", vectorstore._collection.count())

    # 3. LLM
    print("🔹 Loading LLM...")
    model_id = "microsoft/phi-3-mini-4k-instruct"
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",
        torch_dtype="auto"
    )
    print("✅ LLM loaded.")

    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=256,
        temperature=0.2,
    )
    llm = HuggingFacePipeline(pipeline=pipe)

    # 4. QA Chain
    print("🔹 Building RetrievalQA...")
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    qa = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        return_source_documents=False
    )

    print("✅ QA pipeline ready.")
    return qa


# Build at import time (so it's ready when app runs)
try:
    qa_pipeline = build_qa()
except Exception as e:
    qa_pipeline = None
    print("❌ Failed to build QA pipeline:", e)
    traceback.print_exc()   # ✅ added: full error details


def get_answer(query: str) -> str:
    """Takes user query and returns chatbot response."""
    if qa_pipeline is None:
        return "⚠️ QA pipeline not initialized."
    return qa_pipeline.run(query)