CREATED PIPELINE RUNNABLE

pipeline.py

@@ -1,9 +1,15 @@
 # pipeline.py
+
 import os
 import getpass
 import pandas as pd
 from typing import Optional, Dict, Any
 
+# (Optional) from langchain.schema import RunnableConfig
+# If you have the latest "langchain_core", use from langchain_core.runnables.base import Runnable
+# or from langchain.runnables.base import Runnable (depending on your version)
+from langchain.runnables.base import Runnable
+
 from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
@@ -12,7 +18,7 @@ from langchain.chains import RetrievalQA
 from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMModel
 import litellm
 
-# For classification/refusal/tailor/cleaner logic
+# 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
@@ -21,7 +27,7 @@ from cleaner_chain import get_cleaner_chain
 from langchain.llms.base import LLM
 
 ###############################################################################
-# 1) Environment Setup
+# 1) Environment keys
 ###############################################################################
 if not os.environ.get("GEMINI_API_KEY"):
     os.environ["GEMINI_API_KEY"] = getpass.getpass("Enter your Gemini API Key: ")
@@ -29,11 +35,11 @@ if not os.environ.get("GROQ_API_KEY"):
     os.environ["GROQ_API_KEY"] = getpass.getpass("Enter your GROQ API Key: ")
 
 ###############################################################################
-# 2) VectorStore Building/Loading
+# 2) Build or load VectorStore
 ###############################################################################
 def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
     if os.path.exists(store_dir):
-        print(f"DEBUG: Found existing FAISS store at '{store_dir}'. Loading...")
+        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
@@ -64,7 +70,7 @@ def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
         return vectorstore
 
 ###############################################################################
-# 3) Build RAG chain for Gemini
+# 3) Build RAG chain
 ###############################################################################
 def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
     class GeminiLangChainLLM(LLM):
@@ -87,7 +93,7 @@ def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
     return rag_chain
 
 ###############################################################################
-# 4) Init Sub-Chains
+# 4) Initialize sub-chains
 ###############################################################################
 classification_chain = get_classification_chain()
 refusal_chain = get_refusal_chain()
@@ -95,15 +101,15 @@ tailor_chain = get_tailor_chain()
 cleaner_chain = get_cleaner_chain()
 
 ###############################################################################
-# 5) Build VectorStores & RAG
+# 5) Build vectorstores & RAG
 ###############################################################################
+gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
+
 wellness_csv = "AIChatbot.csv"
 brand_csv = "BrandAI.csv"
 wellness_store_dir = "faiss_wellness_store"
 brand_store_dir = "faiss_brand_store"
 
-gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
-
 wellness_vectorstore = build_or_load_vectorstore(wellness_csv, wellness_store_dir)
 brand_vectorstore = build_or_load_vectorstore(brand_csv, brand_store_dir)
 
@@ -122,35 +128,27 @@ def do_web_search(query: str) -> str:
     return response
 
 ###############################################################################
-# 6) Orchestrator: run_with_chain_context
+# 6) Orchestrator function: returns a dict => {"answer": "..."}
 ###############################################################################
 def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
     """
-    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> }
+    Called by the Runnable. 
+    inputs: { "input": <user_query>, "chat_history": <list of messages> (optional) }
+    Output: { "answer": <final string> }
     """
-
-    user_query = inputs["input"]  # The user's new question
-    # You can optionally use inputs.get("chat_history") if needed
+    user_query = inputs["input"]
     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)
 
-    # 2) If OutOfScope => refusal => tailor => return
     if classification == "OutOfScope":
         refusal_text = refusal_chain.run({})
         final_refusal = tailor_chain.run({"response": refusal_text})
         return {"answer": final_refusal.strip()}
 
-    # 3) If Wellness => wellness RAG => if insufficient => web => unify => tailor
     if classification == "Wellness":
-        # 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:
@@ -161,11 +159,11 @@ 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}
 
-    # 4) If Brand => brand RAG => tailor => return
     if classification == "Brand":
         rag_result = brand_rag_chain.invoke({"input": user_query, "chat_history": chat_history})
         csv_answer = rag_result["result"].strip()
@@ -173,7 +171,24 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
         final_answer = tailor_chain.run({"response": final_merged}).strip()
         return {"answer": final_answer}
 
-    # 5) fallback => refusal
+    # fallback
     refusal_text = refusal_chain.run({})
     final_refusal = tailor_chain.run({"response": refusal_text}).strip()
     return {"answer": final_refusal}
+
+
+###############################################################################
+# 7) Build a "Runnable" wrapper so .with_listeners() works
+###############################################################################
+from langchain.runnables.base import Runnable
+
+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()