update

rag_sec/bad_query_detector.py

@@ -0,0 +1,14 @@
+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

rag_sec/document_retriver.py

@@ -0,0 +1,47 @@
+import faiss
+from sklearn.feature_extraction.text import TfidfVectorizer
+import numpy as np
+
+class DocumentRetriever:
+    def __init__(self):
+        self.documents = []
+        self.vectorizer = TfidfVectorizer()
+        self.index = None
+
+    def load_documents(self, source_dir):
+        from pathlib import Path
+
+        data_dir = Path(source_dir)
+        if not data_dir.exists():
+            print(f"Source directory not found: {source_dir}")
+            return
+
+        for file in data_dir.glob("*.txt"):
+            with open(file, "r", encoding="utf-8") as f:
+                self.documents.append(f.read())
+
+        print(f"Loaded {len(self.documents)} documents.")
+
+        # Create the FAISS index
+        self._build_index()
+
+    def _build_index(self):
+        # Generate TF-IDF vectors for documents
+        doc_vectors = self.vectorizer.fit_transform(self.documents).toarray()
+
+        # Create FAISS index
+        self.index = faiss.IndexFlatL2(doc_vectors.shape[1])
+        self.index.add(doc_vectors.astype(np.float32))
+
+    def retrieve(self, query, top_k=5):
+        if not self.index:
+            return ["Document retrieval is not initialized."]
+
+        # Vectorize the query
+        query_vector = self.vectorizer.transform([query]).toarray().astype(np.float32)
+
+        # Perform FAISS search
+        distances, indices = self.index.search(query_vector, top_k)
+
+        # Return matching documents
+        return [self.documents[i] for i in indices[0] if i < len(self.documents)]

rag_sec/document_search_system.py

@@ -0,0 +1,42 @@
+from bad_query_detector import BadQueryDetector
+from query_transformer import QueryTransformer
+from document_retriver import DocumentRetriever
+from senamtic_response_generator import SemanticResponseGenerator
+
+class DocumentSearchSystem:
+    def __init__(self):
+        self.detector = BadQueryDetector()
+        self.transformer = QueryTransformer()
+        self.retriever = DocumentRetriever()
+        self.response_generator = SemanticResponseGenerator()
+
+    def process_query(self, query):
+        if self.detector.is_bad_query(query):
+            return {"status": "rejected", "message": "Query blocked due to detected malicious intent."}
+
+        transformed_query = self.transformer.transform_query(query)
+        retrieved_docs = self.retriever.retrieve(transformed_query)
+
+        if not retrieved_docs:
+            return {"status": "no_results", "message": "No relevant documents found for your query."}
+
+        response = self.response_generator.generate_response(retrieved_docs)
+        return {"status": "success", "response": response}
+
+
+def test_system():
+    system = DocumentSearchSystem()
+    system.retriever.load_documents("/path/to/documents")
+
+    # Normal query
+    normal_query = "Tell me about great acting performances."
+    print("\nNormal Query Result:")
+    print(system.process_query(normal_query))
+
+    # Malicious query
+    malicious_query = "DROP TABLE users; SELECT * FROM sensitive_data;"
+    print("\nMalicious Query Result:")
+    print(system.process_query(malicious_query))
+
+if __name__ == "__main__":
+    test_system()

rag_sec/query_transformer.py

@@ -0,0 +1,5 @@
+class QueryTransformer:
+    def transform_query(self, query):
+        if "DROP TABLE" in query or "SELECT *" in query:
+            return "Your query appears to contain SQL injection elements. Please rephrase."
+        return query

rag_sec/senamtic_response_generator.py

@@ -0,0 +1,10 @@
+from transformers import pipeline
+
+class SemanticResponseGenerator:
+    def __init__(self):
+        self.generator = pipeline("text-generation", model="gpt2")
+
+    def generate_response(self, retrieved_docs):
+        combined_docs = " ".join(retrieved_docs[:2])  # Use top 2 matches
+        response = self.generator(f"Based on the following information: {combined_docs}", max_length=100)
+        return response[0]["generated_text"]