updated chain

chain/qa_chain.py

@@ -12,7 +12,7 @@ load_dotenv()
 OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
 MODEL              = "openai/gpt-oss-120b"
 MAX_TOKENS         = 2048
-MAX_HISTORY        = 10      # keep last 10 exchanges (20 messages) to avoid token overflow
+MAX_HISTORY        = 10      # keep last 10 exchanges to avoid token overflow
 
 # ── OpenRouter Client ─────────────────────────────────────────────────────────
 client = OpenAI(
@@ -25,9 +25,8 @@ SYSTEM_PROMPT = """You are Study Saathi — a friendly and smart study assistant
 You help students understand their Operating Systems notes.
 
 Rules you must follow:
-- Answer ONLY from the provided context. Never use outside knowledge.
-- If the answer is not in the context, say: "Yeh topic notes mein nahi mila."
-- If user says that you have to give explaination, try not to use bullet points, use simple language, examples in plain text.
+- Answer ONLY from the provided context or conversation history. Never use outside knowledge.
+- If the answer is not in the context or history, say: "Yeh topic notes mein nahi mila."
 - You have access to the full conversation history. Use it to answer follow-up
   questions like "translate the above", "explain more", "give examples of that", etc.
 - Explain in simple Roman Urdu, Urdu, or English — based on what the user uses.
@@ -37,10 +36,25 @@ Rules you must follow:
 - For MCQ answer keys: always return them in a markdown table.
 """
 
+# ── Follow-up detection ───────────────────────────────────────────────────────
+FOLLOWUP_PATTERNS = [
+    r'\babove\b', r'\bupar\b', r'\bwoh\b', r'\bise\b', r'\busse\b',
+    r'\btranslate\b', r'\btarjuma\b', r'\bsummari[sz]e\b',
+    r'\bexplain more\b', r'\baur explain\b', r'\baur batao\b',
+    r'\bexpand\b', r'\brepeat\b', r'\bdobara\b', r'\bphir se\b',
+    r'\bsimplify\b', r'\brewrite\b', r'\bconvert\b',
+    r'\bthis\b', r'\bthat\b', r'\byeh\b', r'\bwahi\b',
+    r'\bprevious\b', r'\blast\b',
+]
+
+def is_followup(query: str) -> bool:
+    """Returns True if the query is a follow-up that doesn't need RAG."""
+    q = query.lower()
+    return any(re.search(p, q) for p in FOLLOWUP_PATTERNS)
+
 # ── Build Context Prompt ──────────────────────────────────────────────────────
 def build_context_prompt(query: str, context: str) -> str:
-    return f"""Use the following context from the student's notes to help answer the question.
-If the question is a follow-up (e.g. translate, summarize, explain more), use the conversation history.
+    return f"""Use the following context from the student's notes to help answer.
 
 --- CONTEXT START ---
 {context}
@@ -49,6 +63,13 @@ If the question is a follow-up (e.g. translate, summarize, explain more), use th
 Student's Question: {query}
 """
 
+def build_followup_prompt(query: str) -> str:
+    return f"""This is a follow-up question. Use the conversation history above to answer.
+Do NOT search for new context — work only from what was already discussed.
+
+Student's Follow-up: {query}
+"""
+
 # ── Detect MCQ Request ────────────────────────────────────────────────────────
 def extract_mcq_count(query: str):
     match = re.search(r'(\d+)\s*(mcq|question|mcqs|questions)', query.lower())
@@ -56,7 +77,6 @@ def extract_mcq_count(query: str):
 
 # ── Trim history to avoid token overflow ─────────────────────────────────────
 def trim_history(history: list) -> list:
-    # keep only last MAX_HISTORY pairs (user+assistant = 2 messages per pair)
     max_msgs = MAX_HISTORY * 2
     if len(history) > max_msgs:
         return history[-max_msgs:]
@@ -65,36 +85,41 @@ def trim_history(history: list) -> list:
 # ── Main Chain ────────────────────────────────────────────────────────────────
 def run_chain(query: str, topic: str = None, history: list = []) -> str:
     """
-    Full RAG chain with conversation memory:
-    1. Retrieve chunks from Pinecone
-    2. Format context
-    3. Build messages: system + trimmed history + new user message
-    4. Send to gpt-oss-120b via OpenRouter
-    5. Return response
+    Full RAG chain with conversation memory.
+    - Follow-up questions skip RAG and use history only.
+    - New questions retrieve chunks from Pinecone.
     """
-    mcq_count = extract_mcq_count(query)
+    mcq_count  = extract_mcq_count(query)
+    followup   = is_followup(query) and len(history) > 0
 
-    # fetch more chunks for MCQ generation
-    top_k  = 10 if mcq_count else 5
-    chunks = retrieve(query, topic=topic, top_k=top_k)
-
-    if not chunks:
-        context = "No relevant context found."
+    if followup:
+        # ── Follow-up: skip RAG, use history only ──────────────────────────
+        user_message = build_followup_prompt(query)
     else:
-        context = format_context(chunks)
-
-    # build the new user message with context injected
-    user_message = build_context_prompt(query, context)
-
-    # inject MCQ instructions if needed
-    if mcq_count:
-        user_message += f"\n\nIMPORTANT: Generate exactly {mcq_count} MCQs from the context above. \
-Format each MCQ as:\n**Q1.** Question\n- A) option\n- B) option\n- C) option\n- D) option\
-\n\nAfter all MCQs, provide the answer key in a markdown table with columns: | Q# | Answer | Explanation |"
-
-    # build full messages list: system + history + current user message
+        # ── New question: retrieve from Pinecone ───────────────────────────
+        top_k  = 10 if mcq_count else 5
+        chunks = retrieve(query, topic=topic, top_k=top_k)
+
+        if not chunks:
+            context = "No relevant context found in the notes."
+        else:
+            context = format_context(chunks)
+
+        user_message = build_context_prompt(query, context)
+
+        # inject MCQ instructions if needed
+        if mcq_count:
+            user_message += (
+                f"\n\nIMPORTANT: Generate exactly {mcq_count} MCQs from the context above. "
+                "Format each MCQ as:\n**Q1.** Question\n- A) option\n- B) option\n"
+                "- C) option\n- D) option\n\n"
+                "After all MCQs, provide the answer key in a markdown table "
+                "with columns: | Q# | Answer | Explanation |"
+            )
+
+    # ── Build messages: system + trimmed history + current user msg ────────
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    messages += trim_history(history[:-1])   # history WITHOUT the last user msg (we already have it)
+    messages += trim_history(history[:-1])   # history minus the last user msg
     messages.append({"role": "user", "content": user_message})
 
     response = client.chat.completions.create(
@@ -108,18 +133,16 @@ Format each MCQ as:\n**Q1.** Question\n- A) option\n- B) option\n- C) option\n-
 
 # ── Quick Test ────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
-    # simulate a follow-up conversation
     q1 = "Explain Process Registers in simple words"
     r1 = run_chain(q1, topic="ch-01-updated", history=[])
-    print("=== Response 1 ===")
-    print(r1)
+    print("=== Response 1 ===\n", r1)
 
     fake_history = [
         {"role": "user",      "content": q1},
         {"role": "assistant", "content": r1},
+        {"role": "user",      "content": "Translate the above into Roman Urdu"},
     ]
 
-    q2 = "Ab upar wali cheez ko Roman Urdu mein translate karo"
-    r2 = run_chain(q2, topic="ch-01-updated", history=fake_history + [{"role": "user", "content": q2}])
-    print("\n=== Response 2 (follow-up) ===")
-    print(r2)
+    q2 = "Translate the above into Roman Urdu"
+    r2 = run_chain(q2, topic="ch-01-updated", history=fake_history)
+    print("\n=== Response 2 (follow-up) ===\n", r2)