Update pipeline.py

pipeline.py

@@ -24,6 +24,7 @@ from classification_chain import get_classification_chain
 from refusal_chain import get_refusal_chain
 from tailor_chain import get_tailor_chain
 from cleaner_chain import get_cleaner_chain
+from contextualize_chain import get_contextualize_chain  # New Import for ContextualizeChain
 
 from langchain.llms.base import LLM
 
@@ -101,6 +102,7 @@ classification_chain = get_classification_chain()
 refusal_chain = get_refusal_chain()
 tailor_chain = get_tailor_chain()
 cleaner_chain = get_cleaner_chain()
+contextualize_chain = get_contextualize_chain()  # New Chain for Contextualizing User Queries
 
 ###############################################################################
 # 5) Build vectorstores & RAG
@@ -136,8 +138,11 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
     user_query = inputs["input"]
     chat_history = inputs.get("chat_history", [])
 
-    # 1) Classification
-    class_result = classification_chain.invoke({"query": user_query, "chat_history": chat_history})
+    # 1) Contextualize the query (ensure it's relevant to the conversation history)
+    contextualized_query = contextualize_chain.invoke({"user_query": user_query, "chat_history": chat_history})
+
+    # 2) Classification (using the contextualized query)
+    class_result = classification_chain.invoke({"query": contextualized_query, "chat_history": chat_history})
     classification = class_result.get("text", "").strip()
 
     if classification == "OutOfScope":
@@ -147,18 +152,18 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
 
     if classification == "Wellness":
         rag_result = wellness_rag_chain.invoke({
-            "query": user_query,
+            "query": contextualized_query,
             "chat_history": chat_history  # Pass history here
         })
         csv_answer = rag_result["result"].strip()
-        web_answer = do_web_search(user_query) if not csv_answer else ""
+        web_answer = do_web_search(contextualized_query) if not csv_answer else ""
         final_merged = cleaner_chain.merge(kb=csv_answer, web=web_answer, chat_history=chat_history)
         final_answer = tailor_chain.run({"response": final_merged, "chat_history": chat_history}).strip()
         return {"answer": final_answer}
 
     if classification == "Brand":
         rag_result = brand_rag_chain.invoke({
-            "query": user_query,
+            "query": contextualized_query,
             "chat_history": chat_history  # Pass history here
         })
         csv_answer = rag_result["result"].strip()