Added Callback type annotation

pipeline.py

@@ -27,6 +27,8 @@ from langchain.chains import RetrievalQA, LLMChain
 from langchain.prompts import PromptTemplate
 from langchain.docstore.document import Document
 
+from langchain.callbacks.base import BaseCallbacks  # Updated import
+from langchain.callbacks.manager import CallbackManager
 # Custom chain imports
 from classification_chain import get_classification_chain
 from refusal_chain import get_refusal_chain
@@ -188,33 +190,35 @@ CACHE_SIZE_LIMIT = 1000
 #     logger.error(f"Failed to initialize ChatGroq: {e}")
 #     raise RuntimeError("ChatGroq initialization failed.") from e
 
-# Define a dummy BaseCache class locally
-class BaseCache:
-    def lookup(self, key: str):
-        return None  # Always return None, meaning no cache hit
 
-    def update(self, key: str, value: str):
-        pass  # Do nothing on cache update
 
-# Define a no-op cache that does nothing
-class NoOpCache(BaseCache):
-    pass
+# Define a no-op callback handler
+class NoOpCallbacks(BaseCallbacks):
+    """No-op callback handler."""
+    def on_llm_start(self, *args, **kwargs): pass
+    def on_llm_end(self, *args, **kwargs): pass
+    def on_llm_error(self, *args, **kwargs): pass
+    def on_chain_start(self, *args, **kwargs): pass
+    def on_chain_end(self, *args, **kwargs): pass
+    def on_chain_error(self, *args, **kwargs): pass
 
-# Assign the no-op cache to ChatGroq
-ChatGroq.cache = NoOpCache()
+# Create a callback manager with no-op callbacks
+callback_manager = CallbackManager([NoOpCallbacks()])
 
-# Rebuild the ChatGroq model to finalize its definition
-ChatGroq.model_rebuild()
-
-# Initialize ChatGroq without using caching
-fallback_groq_api_key = os.environ.get("GROQ_API_KEY_FALLBACK", "GROQ_API_KEY")
+# Initialize ChatGroq with the callback manager
 try:
+    fallback_groq_api_key = os.environ.get("GROQ_API_KEY_FALLBACK", "GROQ_API_KEY")
     groq_fallback_llm = ChatGroq(
-        model=GROQ_MODELS["default"],  # Replace with your model name
+        model=GROQ_MODELS["default"],
         temperature=0.7,
         groq_api_key=fallback_groq_api_key,
-        max_tokens=2048
+        max_tokens=2048,
+        callback_manager=callback_manager  # Add the callback manager here
     )
+    
+    # Rebuild the model after initialization
+    ChatGroq.model_rebuild()
+    
 except Exception as e:
     logger.error(f"Failed to initialize ChatGroq: {e}")
     raise RuntimeError("ChatGroq initialization failed.") from e