Update pipeline.py

pipeline.py

@@ -5,8 +5,14 @@ import getpass
 import pandas as pd
 from typing import Optional, Dict, Any
 
-# Correct import for Runnable
-from langchain.schema import Runnable
+# Conditional import for Runnable from available locations
+try:
+    from langchain_core.runnables.base import Runnable
+except ImportError:
+    try:
+        from langchain.runnables.base import Runnable
+    except ImportError:
+        raise ImportError("Cannot find Runnable class. Please upgrade LangChain or check your installation.")
 
 from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceEmbeddings
@@ -16,7 +22,6 @@ 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
@@ -129,15 +134,10 @@ def do_web_search(query: str) -> str:
 # 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", [])
 
-    # 1) Classification
+    # Classification step
     class_result = classification_chain.invoke({"query": user_query})
     classification = class_result.get("text", "").strip()
 
@@ -157,7 +157,6 @@ 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}
@@ -169,7 +168,6 @@ 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}
 
-    # fallback
     refusal_text = refusal_chain.run({})
     final_refusal = tailor_chain.run({"response": refusal_text}).strip()
     return {"answer": final_refusal}
@@ -177,14 +175,8 @@ def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
 ###############################################################################
 # 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()