rag_sec/bad_query_detector.py

from transformers import pipeline

class BadQueryDetector:
    def __init__(self):
        self.detector = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")

    def is_bad_query(self, query):
        result = self.detector(query)[0]
        label = result["label"]
        score = result["score"]
        if label == "NEGATIVE" and score > 0.8:
            print(f"Detected malicious query with high confidence ({score:.4f}): {query}")
            return True
        return False