Update pipeline.py

pipeline.py

@@ -20,7 +20,7 @@ except LookupError:
     english_words = set(words.words())
 
 # LangChain / Groq / LLM imports
-from langchain_groq import ChatGroq
+# from langchain_groq import ChatGroq
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.chains import RetrievalQA, LLMChain
@@ -29,8 +29,10 @@ from langchain.docstore.document import Document
 
 # from langchain.callbacks.base import BaseCallbacks  # Updated import
 # from langchain.callbacks.manager import CallbackManager
-from langchain.callbacks import StdOutCallbackHandler
+# from langchain.callbacks import StdOutCallbackHandler
+
 # Custom chain imports
+from groq_client import GroqClient
 from classification_chain import get_classification_chain
 from refusal_chain import get_refusal_chain
 from tailor_chain import get_tailor_chain
@@ -41,7 +43,7 @@ from tailor_chain_wellnessBrand import get_tailor_chain_wellnessBrand
 from mistralai import Mistral
 
 # Google Gemini LLM
-from langchain_google_genai import ChatGoogleGenerativeAI
+# from langchain_google_genai import ChatGoogleGenerativeAI
 
 # Web search
 # from smolagents import DuckDuckGoSearchTool, ManagedAgent, HfApiModel, CodeAgent
@@ -191,44 +193,73 @@ CACHE_SIZE_LIMIT = 1000
 #     logger.error(f"Failed to initialize ChatGroq: {e}")
 #     raise RuntimeError("ChatGroq initialization failed.") from e
 
-
-from langchain_core.callbacks import CallbackManager
-from langchain_core.callbacks.base import BaseCallbackHandler
-from langchain_core.caches import BaseCache
-
-# Initialize ChatGroq with the callback manager
 try:
-    # First, ensure BaseCache is defined (if needed)
-    class NoCache(BaseCache):
-        """Simple no-op cache implementation."""
-        def __init__(self): 
-            pass
-            
-        def lookup(self, prompt, llm_string): 
-            return None
-            
-        def update(self, prompt, llm_string, return_val): 
-            pass
-            
-        def clear(self): 
-            pass
-    
-    # Call model_rebuild() before initializing
-    ChatGroq.model_rebuild()
-    
-    # Then initialize the model
+        # Get API key from environment variables
     fallback_groq_api_key = os.environ.get("GROQ_API_KEY_FALLBACK", os.environ.get("GROQ_API_KEY"))
-    groq_fallback_llm = ChatGroq(
+    
+    if not fallback_groq_api_key:
+        logger.warning("No Groq API key found for fallback LLM")
+        return None
+        
+    # Initialize the Groq client as fallback
+    groq_fallback_llm = GroqClient(
         model=GROQ_MODELS["default"],
         temperature=0.7,
         groq_api_key=fallback_groq_api_key,
-        max_tokens=2048,
-        callback_manager=CallbackManager([])  # Use CallbackManager instead of raw list
+        max_tokens=2048
     )
+    
+    # Verify connection by checking models (optional)
+    try:
+        groq_fallback_llm.list_models()
+        logger.info("Fallback Groq LLM initialized successfully")
+    except Exception as e:
+        logger.error(f"Fallback Groq LLM connection test failed: {e}")
+        return None
+        
+    
         
 except Exception as e:
-    logger.error(f"Failed to initialize ChatGroq: {e}")
-    raise RuntimeError("ChatGroq initialization failed.") from e
+        logger.error(f"Failed to initialize fallback Groq LLM: {e}")
+        return None
+
+# from langchain_core.callbacks import CallbackManager
+# from langchain_core.callbacks.base import BaseCallbackHandler
+# from langchain_core.caches import BaseCache
+
+# # Initialize ChatGroq with the callback manager
+# try:
+#     # First, ensure BaseCache is defined (if needed)
+#     class NoCache(BaseCache):
+#         """Simple no-op cache implementation."""
+#         def __init__(self): 
+#             pass
+            
+#         def lookup(self, prompt, llm_string): 
+#             return None
+            
+#         def update(self, prompt, llm_string, return_val): 
+#             pass
+            
+#         def clear(self): 
+#             pass
+    
+#     # Call model_rebuild() before initializing
+#     ChatGroq.model_rebuild()
+    
+#     # Then initialize the model
+#     fallback_groq_api_key = os.environ.get("GROQ_API_KEY_FALLBACK", os.environ.get("GROQ_API_KEY"))
+#     groq_fallback_llm = ChatGroq(
+#         model=GROQ_MODELS["default"],
+#         temperature=0.7,
+#         groq_api_key=fallback_groq_api_key,
+#         max_tokens=2048,
+#         callback_manager=CallbackManager([])  # Use CallbackManager instead of raw list
+#     )
+        
+# except Exception as e:
+#     logger.error(f"Failed to initialize ChatGroq: {e}")
+#     raise RuntimeError("ChatGroq initialization failed.") from e
 
 
 # -------------------------------------------------------