Update pipeline.py

pipeline.py

@@ -13,11 +13,15 @@ from langchain.llms.base import LLM  # Import LLM
 
 # Mistral Client Setup
 from mistralai import Mistral  # Import the Mistral client
+from pydantic_ai import Agent  # Import Pydantic AI's Agent
 
 # Initialize Mistral API client
 mistral_api_key = os.environ.get("MISTRAL_API_KEY")  # Ensure your Mistral API key is set
 client = Mistral(api_key=mistral_api_key)
 
+# Initialize Pydantic AI Agent (for text validation)
+pydantic_agent = Agent('mistral:mistral-large-latest', result_type=str)
+
 # Load spaCy model for NER and download the spaCy model if not already installed
 def install_spacy_model():
     try:
@@ -40,6 +44,14 @@ def moderate_text(query: str) -> str:
     Classifies the query as harmful or not using Mistral Moderation via Mistral API.
     Returns "OutOfScope" if harmful, otherwise returns the original query.
     """
+    # Validate the text type using Pydantic AI's Agent
+    try:
+        # Use Pydantic AI agent to ensure correct text type
+        pydantic_agent.run_sync(query)
+    except Exception as e:
+        print(f"Error validating text with Pydantic AI: {e}")
+        return "Invalid text format."
+    
     # Use the moderation API to evaluate if the query is harmful
     response = client.classifiers.moderate_chat(
         model="mistral-moderation-latest",