Update pipeline.py

pipeline.py

@@ -1,7 +1,5 @@
 import os
 import getpass
-from pydantic_ai import Agent  # Import the Agent from pydantic_ai
-from pydantic_ai.models.mistral import MistralModel  # Import the Mistral model
 import spacy  # Import spaCy for NER functionality
 import pandas as pd
 from typing import Optional
@@ -13,25 +11,12 @@ 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 agent using Pydantic AI
+# Mistral Client Setup
+from mistralai import Mistral  # Import the Mistral client
 
+# Initialize Mistral API client
 mistral_api_key = os.environ.get("MISTRAL_API_KEY")  # Ensure your Mistral API key is set
-mistral_model = MistralModel("mistral-moderation-latest", api_key=mistral_api_key)  # Use a Mistral model
-mistral_agent = Agent(mistral_model)
+client = Mistral(api_key=mistral_api_key)
 
 # Load spaCy model for NER and download the spaCy model if not already installed
 def install_spacy_model():
@@ -49,21 +34,21 @@ install_spacy_model()
 # Load the spaCy model globally
 nlp = spacy.load("en_core_web_sm")
 
-# Function to moderate text using Pydantic AI's Mistral moderation model
+# Function to moderate text using Mistral moderation API
 def moderate_text(query: str) -> str:
     """
-    Classifies the query as harmful or not using Mistral Moderation via Pydantic AI.
+    Classifies the query as harmful or not using Mistral Moderation via Mistral API.
     Returns "OutOfScope" if harmful, otherwise returns the original query.
     """
     # Use the moderation API to evaluate if the query is harmful
-    response = mistral_agent.model.classifiers.moderate_chat(
+    response = client.classifiers.moderate_chat(
         model="mistral-moderation-latest",
         inputs=[
             {"role": "user", "content": query},
         ],
     )
-    
-    # Assuming the response contains 'results' with category scores
+
+    # Extracting category scores from response
     categories = response['results'][0]['categories']
 
     # Check if harmful content is flagged in moderation categories
@@ -72,7 +57,9 @@ def moderate_text(query: str) -> str:
        categories.get("dangerous_and_criminal_content", False) or \
        categories.get("selfharm", False):
         return "OutOfScope"
+    
     return query
+
 # 3) build_or_load_vectorstore (no changes)
 def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
     if os.path.exists(store_dir):