Update src/rag_pipeline.py

src/rag_pipeline.py

@@ -1,50 +1,17 @@
 """
-rag_pipeline.py
----------------
-Ties together the retriever and the LLM to answer user questions.
-
-RAG in plain English
-  1. The user asks a question.
-  2. We convert the question to an embedding (same model used for documents).
-  3. FAISS finds the top-k document chunks whose embeddings are most similar.
-  4. We inject those chunks into a prompt as "context".
-  5. The LLM reads the context and generates an answer.
-
-This approach grounds the LLM's answer in YOUR documents rather than its
-general training data — which reduces hallucinations and keeps answers
-relevant.
-
-LLM options (set LLM_PROVIDER below):
-  "openai"  – uses OpenAI's API (requires OPENAI_API_KEY env variable)
-  "ollama"  – uses a locally running Ollama model (free, no API key needed)
+rag_pipeline.py - Connects retriever and LLM to answer questions.
 """
 
 import os
-from langchain_classic.chains import RetrievalQA
 from langchain_core.prompts import PromptTemplate
 from langchain_community.vectorstores import FAISS
 
-# ── configuration ─────────────────────────────────────────────────────────────
-
-# Choose your LLM backend: "groq", "openai", or "ollama"
-LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").strip().lower()
-
-# Groq settings (only used when LLM_PROVIDER = "groq")
-GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.1-8b-instant")
-
-# OpenAI settings (only used when LLM_PROVIDER = "openai")
-OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-3.5-turbo")
-
-# Ollama settings (only used when LLM_PROVIDER = "ollama")
-# Make sure Ollama is running: https://ollama.com/download
-OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
+LLM_PROVIDER  = os.getenv("LLM_PROVIDER", "groq")
+GROQ_MODEL    = os.getenv("GROQ_MODEL", "llama3-8b-8192")
+OPENAI_MODEL  = os.getenv("OPENAI_MODEL", "gpt-3.5-turbo")
+OLLAMA_MODEL  = os.getenv("OLLAMA_MODEL", "llama3")
 OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
 
-
-# ── prompt template ───────────────────────────────────────────────────────────
-
-# This prompt tells the LLM exactly what role to play and how to use
-# the retrieved context. Editing this string is a great way to experiment!
 RAG_PROMPT_TEMPLATE = """You are a helpful assistant that answers questions about
 Data Science interview preparation based only on the provided context.
 
@@ -64,20 +31,14 @@ PROMPT = PromptTemplate(
 )
 
 
-# ── LLM factory ───────────────────────────────────────────────────────────────
-
 def _get_llm():
-    """
-    Load and return the chosen LLM.
-
-    Reads LLM_PROVIDER and instantiates the right LangChain wrapper.
-    """
     if LLM_PROVIDER == "groq":
         api_key = os.getenv("GROQ_API_KEY")
         if not api_key:
             raise EnvironmentError(
-                "GROQ_API_KEY environment variable is not set.\n"
-                "Create a key at https://console.groq.com/keys and add it to .env."
+                "GROQ_API_KEY is not set.\n"
+                "Get your free key at: https://console.groq.com/keys\n"
+                "Then add to .env: GROQ_API_KEY=gsk_..."
             )
         from langchain_groq import ChatGroq
         print(f"  Using Groq model: {GROQ_MODEL}")
@@ -86,82 +47,43 @@ def _get_llm():
     elif LLM_PROVIDER == "openai":
         api_key = os.getenv("OPENAI_API_KEY")
         if not api_key:
-            raise EnvironmentError(
-                "OPENAI_API_KEY environment variable is not set.\n"
-                "Set it with:  export OPENAI_API_KEY='sk-...'\n"
-                "Or switch to Ollama by setting LLM_PROVIDER=ollama"
-            )
+            raise EnvironmentError("OPENAI_API_KEY is not set.")
         from langchain_openai import ChatOpenAI
         print(f"  Using OpenAI model: {OPENAI_MODEL}")
         return ChatOpenAI(model_name=OPENAI_MODEL, temperature=0)
 
     elif LLM_PROVIDER == "ollama":
         from langchain_ollama import ChatOllama
-        print(f"  Using Ollama model: {OLLAMA_MODEL} at {OLLAMA_BASE_URL}")
+        print(f"  Using Ollama model: {OLLAMA_MODEL}")
         return ChatOllama(model=OLLAMA_MODEL, base_url=OLLAMA_BASE_URL, temperature=0)
 
     else:
-        raise ValueError(
-            f"Unknown LLM_PROVIDER: '{LLM_PROVIDER}'. Choose 'groq', 'openai', or 'ollama'."
-        )
+        raise ValueError(f"Unknown LLM_PROVIDER: '{LLM_PROVIDER}'. Use 'groq', 'openai', or 'ollama'.")
 
 
-# ── pipeline builder ──────────────────────────────────────────────────────────
-
-def build_rag_chain(vector_store: FAISS, k: int = 4) -> RetrievalQA:
-    """
-    Build a RetrievalQA chain that answers questions from the vector store.
-
-    Parameters
-    ----------
-    vector_store : FAISS  – the indexed documents
-    k            : int    – number of chunks to retrieve per query
-
-    Returns
-    -------
-    RetrievalQA  – a ready-to-use QA chain
-    """
+def build_rag_chain(vector_store, k: int = 4):
+    from langchain.chains import RetrievalQA
     llm = _get_llm()
-
     retriever = vector_store.as_retriever(
         search_type="similarity",
         search_kwargs={"k": k},
     )
-
-    # RetrievalQA wires the retriever → prompt → LLM for us
     chain = RetrievalQA.from_chain_type(
         llm=llm,
-        chain_type="stuff",          # "stuff" = paste all chunks into one prompt
+        chain_type="stuff",
         retriever=retriever,
-        return_source_documents=True,  # so we can show users where the answer came from
+        return_source_documents=True,
         chain_type_kwargs={"prompt": PROMPT},
     )
-    print("  OK: RAG chain ready.")
+    print("  RAG chain ready.")
     return chain
 
 
-# ── query helper ──────────────────────────────────────────────────────────────
-
-def ask_question(chain: RetrievalQA, question: str) -> dict:
-    """
-    Run a question through the RAG chain.
-
-    Parameters
-    ----------
-    chain    : RetrievalQA  – built with build_rag_chain()
-    question : str          – the user's question
-
-    Returns
-    -------
-    dict with keys:
-      "answer"   : str              – the LLM's answer
-      "sources"  : list[Document]   – the retrieved chunks used as context
-    """
+def ask_question(chain, question: str) -> dict:
     if not question.strip():
         return {"answer": "Please enter a question.", "sources": []}
 
     result = chain.invoke({"query": question})
-
     return {
         "answer": result.get("result", "No answer returned."),
         "sources": result.get("source_documents", []),