update

app.py

@@ -5,6 +5,8 @@ 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
+import streamlit as st
+from rag_sec.document_retriver import DocumentRetriever
 
 # Blockchain Logger
 blockchain_logger = BlockchainLogger()
@@ -65,29 +67,44 @@ if st.button("Validate Blockchain Integrity"):
 
 # Query System
 st.subheader("Query Files")
-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"):
-    if query:
-        # Process the query
-        result = system.process_query(query)
+# Initialize DocumentRetriever
+retriever = DocumentRetriever()
 
-        # Display the results
-        st.write("Query Status:", result.get("status"))
-        st.write("Query Response:", result.get("response"))
+@st.cache(allow_output_mutation=True)
+def load_retriever():
+    """Load documents into the retriever."""
+    retriever.load_documents()
+    return retriever
 
-        if "retrieved_documents" in result:
-            st.write("Retrieved Documents:")
-            for doc in result["retrieved_documents"]:
-                st.markdown(f"- {doc}")
+# Load the retriever and documents
+st.write("Loading documents...")
+retriever = load_retriever()
+st.write("Documents successfully loaded!")
 
-        if "blockchain_details" in result:
-            st.write("Blockchain Details:")
-            st.json(result["blockchain_details"])
+# Streamlit UI
+st.title("Document Search App")
+st.subheader("Enter a query to search for related documents")
 
-        if result.get("status") == "rejected":
-            st.error(f"Query Blocked: {result.get('message')}")
+# Query Input
+query = st.text_input("Enter your query (e.g., 'sports news', 'machine learning')")
+
+if st.button("Search"):
+    if query:
+        # Retrieve documents
+        results = retriever.retrieve(query)
+        if results == ["Document retrieval is not initialized."]:
+            st.error("Document retrieval is not initialized. Please reload the app.")
+        elif not results:
+            st.warning("No relevant documents found for your query.")
+        else:
+            st.success(f"Found {len(results)} relevant document(s).")
+            for idx, doc in enumerate(results, start=1):
+                st.write(f"### Document {idx}")
+                st.write(doc[:500])  # Display first 500 characters of each document
     else:
         st.warning("Please enter a query to search.")
+
+# Debugging Section
+if st.checkbox("Show Debug Information"):
+    st.write(f"Total documents loaded: {len(retriever.documents)}")

rag_sec/document_retriver.py

@@ -1,23 +1,17 @@
-import faiss
-from sklearn.feature_extraction.text import TfidfVectorizer
-import numpy as np
 from sklearn.datasets import fetch_20newsgroups
 
 class DocumentRetriever:
     def __init__(self):
         self.documents = []
 
-    def load_documents(self):
-        """Load 20 Newsgroups dataset."""
+    def load_documents(self, subset_size=500):
+        """Load a subset of 20 Newsgroups dataset."""
         newsgroups_data = fetch_20newsgroups(subset='all')
-        self.documents = newsgroups_data.data
-        if not self.documents:
-            print("No documents loaded!")
+        self.documents = newsgroups_data.data[:subset_size]  # Load only the first `subset_size` documents
+        print(f"Loaded {len(self.documents)} documents.")
 
     def retrieve(self, query):
         """Retrieve documents related to the query."""
         if not self.documents:
             return ["Document retrieval is not initialized."]
-        # Simple keyword match (can replace with advanced semantic similarity later)
         return [doc for doc in self.documents if query.lower() in doc.lower()]
-

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:
@@ -189,9 +189,27 @@ def main():
 
     return system
 
-
 if __name__ == "__main__":
-    main()
+    retriever = DocumentRetriever()
+    retriever.load_documents()
+
+    # Test queries
+    queries = [
+        "sports news",
+        "political debates",
+        "machine learning",
+        "space exploration"
+    ]
+
+    for query in queries:
+        print(f"\nQuery: {query}")
+        results = retriever.retrieve(query)
+        for idx, doc in enumerate(results, start=1):
+            print(f"\nResult {idx}:\n{doc[:500]}...\n")  # Show first 500 characters of each document
+
+
+# if __name__ == "__main__":
+#     main()
 
     # home_dir = Path(os.getenv("HOME", "/"))
     # data_dir = home_dir / "data-sets/aclImdb/train"