Update pipeline.py

pipeline.py

@@ -67,33 +67,35 @@ def classify_query(query: str) -> str:
     classification = class_result.get("text", "").strip()
     return classification if classification != "OutOfScope" else "OutOfScope"
 
-# Function to moderate text using Mistral moderation API (sync version)
-def moderate_text(query: str) -> str:
+# Function to moderate text using Mistral moderation API (async version)
+async def moderate_text(query: str) -> str:
     try:
-        # Use Pydantic AI for text validation synchronously
-        pydantic_agent.run(query)  # This is a synchronous call
+        # Use Pydantic AI to validate the text
+        await pydantic_agent.run(query)  # Use async run for Pydantic validation
     except Exception as e:
         print(f"Error validating text: {e}")
         return "Invalid text format."
     
-    # Mistral moderation, no need for await as it's synchronous
-    response = client.classifiers.moderate_chat(
+    # Call the Mistral moderation API
+    response = await client.classifiers.moderate_chat(
         model="mistral-moderation-latest",
         inputs=[{"role": "user", "content": query}]
     )
     
-    # Extract moderation categories
-    categories = response['results'][0]['categories']
-    
-    # Check for harmful categories and return "OutOfScope" if any are found
-    if categories.get("violence_and_threats", False) or \
-       categories.get("hate_and_discrimination", False) or \
-       categories.get("dangerous_and_criminal_content", False) or \
-       categories.get("selfharm", False):
-        return "OutOfScope"
+    # Assuming the response is an object of type 'ClassificationResponse', 
+    # check if it has a 'results' attribute, and then access its categories
+    if hasattr(response, 'results') and response.results:
+        categories = response.results[0].categories
+        # Check if harmful categories are present
+        if categories.get("violence_and_threats", False) or \
+           categories.get("hate_and_discrimination", False) or \
+           categories.get("dangerous_and_criminal_content", False) or \
+           categories.get("selfharm", False):
+            return "OutOfScope"
     
     return query
 
+
 # Function to build or load the vector store from CSV data
 def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
     if os.path.exists(store_dir):