Update pipeline.py

pipeline.py

@@ -13,7 +13,24 @@ from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMMod
 import subprocess  # Import subprocess to run shell commands
 from langchain.llms.base import LLM  # Import LLM
 
+from classification_chain import get_classification_chain
+from refusal_chain import get_refusal_chain
+from tailor_chain import get_tailor_chain
+from cleaner_chain import get_cleaner_chain, CleanerChain
+
+
+# 1) Environment: set up keys if missing
+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: ")
+if not os.environ.get("MISTRAL_API_KEY"):
+    os.environ["MISTRAL_API_KEY"] = getpass.getpass("Enter your Mistral API Key: ")
+
+# Initialize Mistral client
+mistral_client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
 # Initialize Mistral agent using Pydantic AI
+
 mistral_api_key = os.environ.get("MISTRAL_API_KEY")  # Ensure your Mistral API key is set
 mistral_model = MistralModel("mistral-large-latest", api_key=mistral_api_key)  # Use a Mistral model
 mistral_agent = Agent(mistral_model)