Update pipeline.py

pipeline.py

@@ -5,10 +5,8 @@ import getpass
 import pandas as pd
 from typing import Optional, Dict, Any
 
-try:
-    from langchain.runnables.base import Runnable
-except ImportError:
-    from langchain_core.runnables.base import Runnable
+# Correct import for Runnable
+from langchain.schema import Runnable
 
 from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceEmbeddings
@@ -18,6 +16,7 @@ from langchain.chains import RetrievalQA
 from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMModel
 import litellm
 
+# Classification/Refusal/Tailor/Cleaner
 from classification_chain import get_classification_chain
 from refusal_chain import get_refusal_chain
 from tailor_chain import get_tailor_chain
@@ -25,27 +24,83 @@ from cleaner_chain import get_cleaner_chain
 
 from langchain.llms.base import LLM
 
-# Environment keys
+###############################################################################
+# 1) Environment keys
+###############################################################################
 if not os.environ.get("GEMINI_API_KEY"):
     os.environ["GEMINI_API_KEY"] = getpass.getpass("Enter your Gemini API Key: ")
 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
+###############################################################################
 def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
-    # ... [unchanged code for building/loading vectorstore] ...
-    # Use your previously provided implementation here.
-    # For brevity, not repeating this section.
-    pass
-
+    if os.path.exists(store_dir):
+        print(f"DEBUG: Found existing FAISS store at '{store_dir}'. Loading from disk.")
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1")
+        vectorstore = FAISS.load_local(store_dir, embeddings)
+        return vectorstore
+    else:
+        print(f"DEBUG: Building new store from CSV: {csv_path}")
+        df = pd.read_csv(csv_path)
+        df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
+        df.columns = df.columns.str.strip()
+
+        if "Answer" in df.columns:
+            df.rename(columns={"Answer": "Answers"}, inplace=True)
+        if "Question" not in df.columns and "Question " in df.columns:
+            df.rename(columns={"Question ": "Question"}, inplace=True)
+
+        if "Question" not in df.columns or "Answers" not in df.columns:
+            raise ValueError("CSV must have 'Question' and 'Answers' columns.")
+
+        docs = []
+        for _, row in df.iterrows():
+            q = str(row["Question"])
+            ans = str(row["Answers"])
+            doc = Document(page_content=ans, metadata={"question": q})
+            docs.append(doc)
+
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1")
+        vectorstore = FAISS.from_documents(docs, embedding=embeddings)
+        vectorstore.save_local(store_dir)
+        return vectorstore
+
+###############################################################################
+# 3) Build RAG chain
+###############################################################################
 def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
-    # ... [unchanged code for building a RAG chain] ...
-    pass
-
+    class GeminiLangChainLLM(LLM):
+        def _call(self, prompt: str, stop: Optional[list] = None, **kwargs) -> str:
+            messages = [{"role": "user", "content": prompt}]
+            return llm_model(messages, stop_sequences=stop)
+
+        @property
+        def _llm_type(self) -> str:
+            return "custom_gemini"
+
+    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
+    gemini_as_llm = GeminiLangChainLLM()
+    rag_chain = RetrievalQA.from_chain_type(
+        llm=gemini_as_llm,
+        chain_type="stuff",
+        retriever=retriever,
+        return_source_documents=True
+    )
+    return rag_chain
+
+###############################################################################
+# 4) Initialize sub-chains
+###############################################################################
 classification_chain = get_classification_chain()
 refusal_chain = get_refusal_chain()
 tailor_chain = get_tailor_chain()
 cleaner_chain = get_cleaner_chain()
 
+###############################################################################
+# 5) Build vectorstores & RAG
+###############################################################################
 gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
 
 wellness_csv = "AIChatbot.csv"
@@ -70,14 +125,21 @@ def do_web_search(query: str) -> str:
     response = manager_agent.run(search_query)
     return response
 
+###############################################################################
+# 6) Orchestrator function: returns a dict => {"answer": "..."}
+###############################################################################
 def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
+    """
+    Called by the Runnable. 
+    inputs: { "input": <user_query>, "chat_history": <list of messages> (optional) }
+    Output: { "answer": <final string> }
+    """
     user_query = inputs["input"]
     chat_history = inputs.get("chat_history", [])
 
-    print("DEBUG: Starting run_with_chain_context...")
+    # 1) Classification
     class_result = classification_chain.invoke({"query": user_query})
     classification = class_result.get("text", "").strip()
-    print("DEBUG: Classification =>", classification)
 
     if classification == "OutOfScope":
         refusal_text = refusal_chain.run({})
@@ -85,8 +147,7 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
         return {"answer": final_refusal.strip()}
 
     if classification == "Wellness":
-        # Use the correct key "query" instead of "input"
-        rag_result = wellness_rag_chain.invoke({"query": user_query, "chat_history": chat_history})
+        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(user_query)
@@ -96,24 +157,34 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
                 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}).strip()
         return {"answer": final_answer}
 
     if classification == "Brand":
-        rag_result = brand_rag_chain.invoke({"query": user_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}).strip()
         return {"answer": final_answer}
 
+    # fallback
     refusal_text = refusal_chain.run({})
     final_refusal = tailor_chain.run({"response": refusal_text}).strip()
     return {"answer": final_refusal}
 
-# Runnable wrapper for my_memory_logic.py
+###############################################################################
+# 7) Build a "Runnable" wrapper so .with_listeners() works
+###############################################################################
+
 class PipelineRunnable(Runnable[Dict[str, Any], Dict[str, str]]):
+    """
+    Wraps run_with_chain_context(...) in a Runnable 
+    so that RunnableWithMessageHistory can attach listeners.
+    """
     def invoke(self, input: Dict[str, Any], config: Optional[Any] = None) -> Dict[str, str]:
         return run_with_chain_context(input)
 
+# Export an instance of PipelineRunnable for use in my_memory_logic.py
 pipeline_runnable = PipelineRunnable()