Update pipeline.py

pipeline.py

@@ -29,6 +29,7 @@ from cleaner_chain import get_cleaner_chain
 
 from langchain.llms.base import LLM
 
+
 ###############################################################################
 # 1) Environment keys
 ###############################################################################
@@ -137,7 +138,7 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
     user_query = inputs["input"]
     chat_history = inputs.get("chat_history", [])
 
-    # Classification step
+    # 1) Classification
     class_result = classification_chain.invoke({"query": user_query})
     classification = class_result.get("text", "").strip()
 
@@ -147,7 +148,7 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
         return {"answer": final_refusal.strip()}
 
     if classification == "Wellness":
-        rag_result = wellness_rag_chain.invoke({"input": user_query, "chat_history": chat_history})
+        rag_result = wellness_rag_chain.invoke({"query": user_query, "chat_history": chat_history})
         csv_answer = rag_result["result"].strip()
         if not csv_answer:
             web_answer = do_web_search(user_query)
@@ -162,7 +163,7 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
         return {"answer": final_answer}
 
     if classification == "Brand":
-        rag_result = brand_rag_chain.invoke({"input": user_query, "chat_history": chat_history})
+        rag_result = brand_rag_chain.invoke({"query": user_query, "chat_history": chat_history})
         csv_answer = rag_result["result"].strip()
         final_merged = cleaner_chain.merge(kb=csv_answer, web="")
         final_answer = tailor_chain.run({"response": final_merged}).strip()