CHANGED1

pipeline.py

@@ -2,7 +2,7 @@
 import os
 import getpass
 import pandas as pd
-from typing import Optional, List
+from typing import Optional, Dict, Any
 
 from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceEmbeddings
@@ -12,17 +12,16 @@ from langchain.chains import RetrievalQA
 from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMModel
 import litellm
 
-# Import your classification/refusal/tailor/cleaner chains
+# For classification/refusal/tailor/cleaner logic
 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
 
-# For RAG chain building
 from langchain.llms.base import LLM
 
 ###############################################################################
-# 1) Environment: set up keys
+# 1) Environment Setup
 ###############################################################################
 if not os.environ.get("GEMINI_API_KEY"):
     os.environ["GEMINI_API_KEY"] = getpass.getpass("Enter your Gemini API Key: ")
@@ -30,7 +29,7 @@ if not os.environ.get("GROQ_API_KEY"):
     os.environ["GROQ_API_KEY"] = getpass.getpass("Enter your GROQ API Key: ")
 
 ###############################################################################
-# 2) Build or Load VectorStore
+# 2) VectorStore Building/Loading
 ###############################################################################
 def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
     if os.path.exists(store_dir):
@@ -70,7 +69,6 @@ def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
 def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
     class GeminiLangChainLLM(LLM):
         def _call(self, prompt: str, stop: Optional[list] = None, **kwargs) -> str:
-            # We'll treat the entire prompt as 'user' content
             messages = [{"role": "user", "content": prompt}]
             return llm_model(messages, stop_sequences=stop)
 
@@ -89,7 +87,7 @@ def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
     return rag_chain
 
 ###############################################################################
-# 4) Initialize your sub-chains
+# 4) Init Sub-Chains
 ###############################################################################
 classification_chain = get_classification_chain()
 refusal_chain = get_refusal_chain()
@@ -97,7 +95,7 @@ tailor_chain = get_tailor_chain()
 cleaner_chain = get_cleaner_chain()
 
 ###############################################################################
-# 5) Build VectorStores & RAG Chains
+# 5) Build VectorStores & RAG
 ###############################################################################
 wellness_csv = "AIChatbot.csv"
 brand_csv = "BrandAI.csv"
@@ -126,22 +124,21 @@ def do_web_search(query: str) -> str:
 ###############################################################################
 # 6) Orchestrator: run_with_chain_context
 ###############################################################################
-def run_with_chain_context(query: str, chat_history: list) -> str:
+def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
     """
-    Like run_with_chain, but also references `chat_history`.
-    We'll do single-turn classification, but pass chat_history 
-    to the RAG chain if needed.
-
-    Example usage:
-        chat_history = []
-        question = "What is Self-Reflection?"
-        resp1 = run_with_chain_context(question, chat_history)
-        # then chat_history.extend([...]) with HumanMessage/AIMessage
+    This function is called by the RunnableWithMessageHistory in my_memory_logic.py
+    inputs: { "input": <user_query>, "chat_history": <list of messages> }
+    Returns: { "answer": <final response> }
     """
-    print("DEBUG: Starting run_with_chain_context...")
 
-    # 1) Classification (no multi-turn, just single-turn classification)
-    class_result = classification_chain.invoke({"query": query})
+    user_query = inputs["input"]  # The user's new question
+    # You can optionally use inputs.get("chat_history") if needed
+    chat_history = inputs.get("chat_history", [])
+
+    print("DEBUG: Starting run_with_chain_context...")
+    print(f"User query: {user_query}")
+    # 1) Classification
+    class_result = classification_chain.invoke({"query": user_query})
     classification = class_result.get("text", "").strip()
     print("DEBUG: Classification =>", classification)
 
@@ -149,41 +146,34 @@ def run_with_chain_context(query: str, chat_history: list) -> str:
     if classification == "OutOfScope":
         refusal_text = refusal_chain.run({})
         final_refusal = tailor_chain.run({"response": refusal_text})
-        return final_refusal.strip()
+        return {"answer": final_refusal.strip()}
 
-    # 3) If Wellness => call wellness_rag_chain with chat_history
+    # 3) If Wellness => wellness RAG => if insufficient => web => unify => tailor
     if classification == "Wellness":
-        # pass the conversation to .invoke(...) so it can see it if needed
-        rag_result = wellness_rag_chain.invoke({
-            "input": query,
-            "chat_history": chat_history  # pass the entire list of prior messages
-        })
+        # pass chat_history if your chain can use it
+        rag_result = wellness_rag_chain.invoke({"input": user_query, "chat_history": chat_history})
         csv_answer = rag_result["result"].strip()
         if not csv_answer:
-            web_answer = do_web_search(query)
+            web_answer = do_web_search(user_query)
         else:
             lower_ans = csv_answer.lower()
             if any(phrase in lower_ans for phrase in ["i do not know", "not sure", "no context", "cannot answer"]):
-                web_answer = do_web_search(query)
+                web_answer = do_web_search(user_query)
             else:
                 web_answer = ""
-
         final_merged = cleaner_chain.merge(kb=csv_answer, web=web_answer)
-        final_answer = tailor_chain.run({"response": final_merged})
-        return final_answer.strip()
+        final_answer = tailor_chain.run({"response": final_merged}).strip()
+        return {"answer": final_answer}
 
-    # 4) If Brand => brand_rag_chain with chat_history
+    # 4) If Brand => brand RAG => tailor => return
     if classification == "Brand":
-        rag_result = brand_rag_chain.invoke({
-            "input": query,
-            "chat_history": chat_history
-        })
+        rag_result = brand_rag_chain.invoke({"input": 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})
-        return final_answer.strip()
+        final_answer = tailor_chain.run({"response": final_merged}).strip()
+        return {"answer": final_answer}
 
-    # fallback => refusal
+    # 5) fallback => refusal
     refusal_text = refusal_chain.run({})
-    final_refusal = tailor_chain.run({"response": refusal_text})
-    return final_refusal.strip()
+    final_refusal = tailor_chain.run({"response": refusal_text}).strip()
+    return {"answer": final_refusal}