Update pipeline.py

pipeline.py

@@ -55,10 +55,15 @@ def moderate_text(query: str) -> str:
     Classifies the query as harmful or not using Mistral Moderation via Pydantic AI.
     Returns "OutOfScope" if harmful, otherwise returns the original query.
     """
-    # Use the agent to interact with the moderation model, here assuming predict() method
-    response = mistral_agent.predict({"inputs": [query]})  # Adjust if necessary
+    # Use the moderation API to evaluate if the query is harmful
+    response = mistral_agent.model.classifiers.moderate_chat(
+        model="mistral-moderation-latest",
+        inputs=[
+            {"role": "user", "content": query},
+        ],
+    )
     
-    # Assuming response contains a 'results' field similar to the previous example
+    # Assuming the response contains 'results' with category scores
     categories = response['results'][0]['categories']
 
     # Check if harmful content is flagged in moderation categories
@@ -68,7 +73,6 @@ def moderate_text(query: str) -> str:
        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):