update

app.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from PIL import Image
 from rag_sec.document_search_system import DocumentSearchSystem
 from chainguard.blockchain_logger import BlockchainLogger
+from rag_sec.document_search_system import main
 
 # Blockchain Logger
 blockchain_logger = BlockchainLogger()
@@ -64,14 +65,29 @@ if st.button("Validate Blockchain Integrity"):
 
 # Query System
 st.subheader("Query Files")
-query = st.text_input("Enter your query (e.g., 'Good comedy')")
+system = main()  # Initialize system with Neo4j and load documents
+
+# Query Input
+query = st.text_input("Enter your query", placeholder="E.g., 'Good comedy'")
 if st.button("Search"):
-    result = system.process_query(query)
-    st.write("Query Status:", result.get("status"))
-    st.write("Query Response:", result.get("response"))
-    if "retrieved_documents" in result:
-        st.write("Retrieved Documents:", result["retrieved_documents"])
-    if "blockchain_details" in result:
-        st.write("Blockchain Details:", result["blockchain_details"])
-    if result.get("status") == "rejected":
-        st.error(f"Query Blocked: {result.get('message')}")
+    if query:
+        # Process the query
+        result = system.process_query(query)
+
+        # Display the results
+        st.write("Query Status:", result.get("status"))
+        st.write("Query Response:", result.get("response"))
+
+        if "retrieved_documents" in result:
+            st.write("Retrieved Documents:")
+            for doc in result["retrieved_documents"]:
+                st.markdown(f"- {doc}")
+
+        if "blockchain_details" in result:
+            st.write("Blockchain Details:")
+            st.json(result["blockchain_details"])
+
+        if result.get("status") == "rejected":
+            st.error(f"Query Blocked: {result.get('message')}")
+    else:
+        st.warning("Please enter a query to search.")

rag_sec/document_search_system.py

@@ -7,10 +7,10 @@ import sys
 from os import path
 
 sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))
-from .bad_query_detector import BadQueryDetector
-from .query_transformer import QueryTransformer
-from .document_retriver import DocumentRetriever
-from .senamtic_response_generator import SemanticResponseGenerator
+from bad_query_detector import BadQueryDetector
+from query_transformer import QueryTransformer
+from document_retriver import DocumentRetriever
+from senamtic_response_generator import SemanticResponseGenerator
 
 
 class DataTransformer:
@@ -171,12 +171,11 @@ class DocumentSearchSystem:
         return self.data_transformer.validate_blockchain()
 
 
-if __name__ == "__main__":
-
+def main():
+    # Path to the dataset directory
     home_dir = Path(os.getenv("HOME", "/"))
     data_dir = home_dir / "data-sets/aclImdb/train"
 
-
     # Initialize system with Neo4j credentials
     system = DocumentSearchSystem(
         neo4j_uri="neo4j+s://0ca71b10.databases.neo4j.io",
@@ -184,21 +183,42 @@ if __name__ == "__main__":
         neo4j_password="HwGDOxyGS1-79nLeTiX5bx5ohoFSpvHCmTv8IRgt-lY"
     )
 
+    # Load documents into the retriever
     system.retriever.load_documents(data_dir)
-    # Perform a normal query
-    normal_query = "Good comedy ."
-    print("\nNormal Query Result:")
-    result = system.process_query(normal_query)
-    print("Status:", result["status"])
-    print("Response:", result["response"])
-    print("Retrieved Documents:", result["retrieved_documents"])
-    print("Blockchain Details:", result["blockchain_details"])
-
-    # Perform a malicious query
-    malicious_query = "DROP TABLE users; SELECT * FROM sensitive_data;"
-    print("\nMalicious Query Result:")
-    result = system.process_query(malicious_query)
-    print("Status:", result["status"])
-    print("Message:", result.get("message"))
+    print("Documents successfully loaded.")
+
+    return system
+
+
+if __name__ == "__main__":
+    main()
+
+    # home_dir = Path(os.getenv("HOME", "/"))
+    # data_dir = home_dir / "data-sets/aclImdb/train"
+    #
+    #
+    # # Initialize system with Neo4j credentials
+    # system = DocumentSearchSystem(
+    #     neo4j_uri="neo4j+s://0ca71b10.databases.neo4j.io",
+    #     neo4j_user="neo4j",
+    #     neo4j_password="HwGDOxyGS1-79nLeTiX5bx5ohoFSpvHCmTv8IRgt-lY"
+    # )
+    #
+    # system.retriever.load_documents(data_dir)
+    # # Perform a normal query
+    # normal_query = "Good comedy ."
+    # print("\nNormal Query Result:")
+    # result = system.process_query(normal_query)
+    # print("Status:", result["status"])
+    # print("Response:", result["response"])
+    # print("Retrieved Documents:", result["retrieved_documents"])
+    # print("Blockchain Details:", result["blockchain_details"])
+    #
+    # # Perform a malicious query
+    # malicious_query = "DROP TABLE users; SELECT * FROM sensitive_data;"
+    # print("\nMalicious Query Result:")
+    # result = system.process_query(malicious_query)
+    # print("Status:", result["status"])
+    # print("Message:", result.get("message"))
 
 

rag_sec/requirements.txt

@@ -4,3 +4,4 @@ numpy
 scikit-learn
 faiss-cpu
 pandas
+transformers