Upload 2 files

main-metamask.py

@@ -0,0 +1,631 @@
+# main_metamask.py
+import os
+import tempfile
+import shutil
+import PyPDF2
+import streamlit as st
+import torch
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.llms import HuggingFaceHub
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+from langchain.docstore.document import Document
+from langchain.prompts import PromptTemplate
+import time
+import psutil
+import uuid
+import atexit
+from blockchain_utils_metamask import BlockchainManagerMetaMask
+from metamask_component import metamask_connector
+
+
+class BlockchainEnabledRAG:
+    def __init__(self, 
+                 llm_model_name="mistralai/Mistral-7B-Instruct-v0.2",
+                 embedding_model_name="sentence-transformers/all-MiniLM-L6-v2",
+                 chunk_size=1000,
+                 chunk_overlap=200,
+                 use_gpu=True,
+                 use_blockchain=False,
+                 contract_address=None):
+        """
+        Initialize the GPU-efficient RAG system with MetaMask blockchain integration.
+        
+        Args:
+            llm_model_name: The HuggingFace model for text generation
+            embedding_model_name: The HuggingFace model for embeddings
+            chunk_size: Size of document chunks
+            chunk_overlap: Overlap between chunks
+            use_gpu: Whether to use GPU acceleration
+            use_blockchain: Whether to enable blockchain verification
+            contract_address: Address of the deployed RAG Document Verifier contract
+        """
+        self.llm_model_name = llm_model_name
+        self.embedding_model_name = embedding_model_name
+        self.use_gpu = use_gpu and torch.cuda.is_available()
+        self.use_blockchain = use_blockchain
+        
+        # Device selection for embeddings
+        self.device = "cuda" if self.use_gpu else "cpu"
+        st.sidebar.info(f"Using device: {self.device}")
+        
+        # Initialize text splitter
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            length_function=len,
+        )
+        
+        # Initialize embeddings model
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name=embedding_model_name,
+            model_kwargs={"device": self.device}
+        )
+        
+        # Initialize LLM using HuggingFaceHub instead of Ollama
+        try:
+            # Use HF_TOKEN from environment variables
+            hf_token = os.environ.get("HF_TOKEN")
+            if not hf_token:
+                st.warning("No HuggingFace token found. Using model without authentication.")
+                
+            self.llm = HuggingFaceHub(
+                repo_id=llm_model_name,
+                huggingfacehub_api_token=hf_token,
+                model_kwargs={"temperature": 0.7, "max_length": 1024}
+            )
+        except Exception as e:
+            st.error(f"Error initializing LLM: {str(e)}")
+            st.info("Trying to initialize with default model...")
+            # Fallback to a smaller model
+            self.llm = HuggingFaceHub(
+                repo_id="google/flan-t5-small",
+                model_kwargs={"temperature": 0.7, "max_length": 512}
+            )
+        
+        # Initialize vector store
+        self.vector_store = None
+        self.documents_processed = 0
+        
+        # Monitoring stats
+        self.processing_times = {}
+        
+        # Initialize blockchain manager if enabled
+        self.blockchain = None
+        if use_blockchain:
+            try:
+                self.blockchain = BlockchainManagerMetaMask(
+                    contract_address=contract_address
+                )
+                st.sidebar.success("Blockchain manager initialized. Please connect MetaMask to continue.")
+            except Exception as e:
+                st.sidebar.error(f"Failed to initialize blockchain manager: {str(e)}")
+                self.use_blockchain = False
+
+    def update_blockchain_connection(self, metamask_info):
+        """Update blockchain connection with MetaMask info."""
+        if self.blockchain and metamask_info:
+            self.blockchain.update_connection(
+                is_connected=metamask_info.get("connected", False),
+                user_address=metamask_info.get("address"),
+                network_id=metamask_info.get("network_id")
+            )
+            return self.blockchain.is_connected
+        return False
+
+    def process_pdfs(self, pdf_files):
+        """Process PDF files, create a vector store, and verify documents on blockchain."""
+        all_docs = []
+        
+        with st.status("Processing PDF files...") as status:
+            # Create temporary directory for file storage
+            temp_dir = tempfile.mkdtemp()
+            st.session_state['temp_dir'] = temp_dir
+            
+            # Monitor processing time and memory usage
+            start_time = time.time()
+            
+            # Track memory before processing
+            mem_before = psutil.virtual_memory().used / (1024 * 1024 * 1024)  # GB
+            
+            # Process each PDF file
+            for i, pdf_file in enumerate(pdf_files):
+                try:
+                    file_start_time = time.time()
+                    
+                    # Save uploaded file to temp directory
+                    pdf_path = os.path.join(temp_dir, pdf_file.name)
+                    with open(pdf_path, "wb") as f:
+                        f.write(pdf_file.getbuffer())
+                    
+                    status.update(label=f"Processing {pdf_file.name} ({i+1}/{len(pdf_files)})...")
+                    
+                    # Extract text from PDF
+                    text = ""
+                    with open(pdf_path, "rb") as f:
+                        pdf = PyPDF2.PdfReader(f)
+                        for page_num in range(len(pdf.pages)):
+                            page = pdf.pages[page_num]
+                            page_text = page.extract_text()
+                            if page_text:
+                                text += page_text + "\n\n"
+                    
+                    # Create documents
+                    docs = [Document(page_content=text, metadata={"source": pdf_file.name})]
+                    
+                    # Split documents into chunks
+                    split_docs = self.text_splitter.split_documents(docs)
+                    
+                    all_docs.extend(split_docs)
+                    
+                    # Verify document on blockchain if enabled and connected
+                    if self.use_blockchain and self.blockchain and self.blockchain.is_connected:
+                        try:
+                            # Create a unique document ID
+                            document_id = f"{pdf_file.name}_{uuid.uuid4().hex[:8]}"
+                            
+                            # Verify document on blockchain
+                            status.update(label=f"Verifying {pdf_file.name} on blockchain...")
+                            verification = self.blockchain.verify_document(document_id, pdf_path)
+                            
+                            if verification.get('status'):  # Success
+                                st.sidebar.success(f"✅ {pdf_file.name} verified on blockchain")
+                                if 'tx_hash' in verification:
+                                    st.sidebar.info(f"Transaction: {verification['tx_hash'][:10]}...")
+                                
+                                # Add blockchain metadata to documents
+                                for doc in split_docs:
+                                    doc.metadata["blockchain"] = {
+                                        "verified": True,
+                                        "document_id": document_id,
+                                        "document_hash": verification.get("document_hash", ""),
+                                        "tx_hash": verification.get("tx_hash", ""),
+                                        "block_number": verification.get("block_number", 0)
+                                    }
+                            else:
+                                st.sidebar.warning(f"❌ Failed to verify {pdf_file.name} on blockchain")
+                                if 'error' in verification:
+                                    st.sidebar.error(f"Error: {verification['error']}")
+                        except Exception as e:
+                            st.sidebar.error(f"Blockchain verification error: {str(e)}")
+                    elif self.use_blockchain:
+                        st.sidebar.warning("MetaMask not connected. Document not verified on blockchain.")
+                    
+                    file_end_time = time.time()
+                    processing_time = file_end_time - file_start_time
+                    
+                    st.sidebar.success(f"Processed {pdf_file.name}: {len(split_docs)} chunks in {processing_time:.2f}s")
+                    self.processing_times[pdf_file.name] = {
+                        "chunks": len(split_docs),
+                        "time": processing_time
+                    }
+                    
+                except Exception as e:
+                    st.sidebar.error(f"Error processing {pdf_file.name}: {str(e)}")
+            
+            # Create vector store if we have documents
+            if all_docs:
+                status.update(label="Building vector index...")
+                try:
+                    # Record the time taken to build the index
+                    index_start_time = time.time()
+                    
+                    # Create the vector store using FAISS
+                    self.vector_store = FAISS.from_documents(all_docs, self.embeddings)
+                    
+                    index_end_time = time.time()
+                    index_time = index_end_time - index_start_time
+                    
+                    # Track memory after processing
+                    mem_after = psutil.virtual_memory().used / (1024 * 1024 * 1024)  # GB
+                    mem_used = mem_after - mem_before
+                    
+                    total_time = time.time() - start_time
+                    
+                    status.update(label=f"Completed processing {len(all_docs)} chunks in {total_time:.2f}s", state="complete")
+                    
+                    # Save performance metrics
+                    self.processing_times["index_building"] = index_time
+                    self.processing_times["total_time"] = total_time
+                    self.processing_times["memory_used_gb"] = mem_used
+                    self.documents_processed = len(all_docs)
+                    
+                    return True
+                except Exception as e:
+                    st.error(f"Error creating vector store: {str(e)}")
+                    status.update(label="Error creating vector store", state="error")
+                    return False
+            else:
+                status.update(label="No content extracted from PDFs", state="error")
+                return False
+
+    def ask(self, query):
+        """Ask a question and get an answer based on the PDFs with blockchain logging."""
+        if not self.vector_store:
+            return "Please upload and process PDF files first."
+            
+        try:
+            # Custom prompt
+            prompt_template = """
+            You are an AI assistant that provides accurate information based on PDF documents.
+            
+            Use the following context to answer the question. Be detailed and precise in your answer.
+            If the answer is not in the context, say "I don't have enough information to answer this question."
+            
+            Context:
+            {context}
+            
+            Question: {question}
+            
+            Answer:
+            """
+            PROMPT = PromptTemplate(
+                template=prompt_template, 
+                input_variables=["context", "question"]
+            )
+            
+            # Start timing the query
+            query_start_time = time.time()
+            
+            # Create QA chain
+            chain_type_kwargs = {"prompt": PROMPT}
+            qa = RetrievalQA.from_chain_type(
+                llm=self.llm,
+                chain_type="stuff",
+                retriever=self.vector_store.as_retriever(search_kwargs={"k": 4}),
+                chain_type_kwargs=chain_type_kwargs,
+                return_source_documents=True
+            )
+            
+            # Get answer
+            with st.status("Searching documents and generating answer..."):
+                response = qa({"query": query})
+                
+            answer = response["result"]
+            source_docs = response["source_documents"]
+            
+            # Calculate query time
+            query_time = time.time() - query_start_time
+            
+            # Format sources
+            sources = []
+            for i, doc in enumerate(source_docs):
+                # Extract blockchain verification info if available
+                blockchain_info = None
+                if "blockchain" in doc.metadata:
+                    blockchain_info = {
+                        "verified": doc.metadata["blockchain"]["verified"],
+                        "document_id": doc.metadata["blockchain"]["document_id"],
+                        "tx_hash": doc.metadata["blockchain"]["tx_hash"]
+                    }
+                
+                sources.append({
+                    "content": doc.page_content[:300] + "..." if len(doc.page_content) > 300 else doc.page_content,
+                    "source": doc.metadata.get("source", "Unknown"),
+                    "blockchain": blockchain_info
+                })
+            
+            # Log query to blockchain if enabled and connected
+            blockchain_log = None
+            if self.use_blockchain and self.blockchain and self.blockchain.is_connected:
+                try:
+                    with st.status("Logging query to blockchain..."):
+                        log_result = self.blockchain.log_query(query, answer)
+                        
+                        if log_result.get("status"):  # Success
+                            blockchain_log = {
+                                "logged": True,
+                                "query_id": log_result.get("query_id", ""),
+                                "tx_hash": log_result.get("tx_hash", ""),
+                                "block_number": log_result.get("block_number", 0)
+                            }
+                        else:
+                            st.error(f"Error logging to blockchain: {log_result.get('error', 'Unknown error')}")
+                except Exception as e:
+                    st.error(f"Error logging to blockchain: {str(e)}")
+                
+            return {
+                "answer": answer,
+                "sources": sources,
+                "query_time": query_time,
+                "blockchain_log": blockchain_log
+            }
+                
+        except Exception as e:
+            st.error(f"Error generating answer: {str(e)}")
+            return f"Error: {str(e)}"
+
+    def get_performance_metrics(self):
+        """Return performance metrics for the RAG system."""
+        if not self.processing_times:
+            return None
+            
+        return {
+            "documents_processed": self.documents_processed,
+            "index_building_time": self.processing_times.get("index_building", 0),
+            "total_processing_time": self.processing_times.get("total_time", 0),
+            "memory_used_gb": self.processing_times.get("memory_used_gb", 0),
+            "device": self.device,
+            "embedding_model": self.embedding_model_name,
+            "blockchain_enabled": self.use_blockchain,
+            "blockchain_connected": self.blockchain.is_connected if self.blockchain else False
+        }
+
+
+# Helper function to initialize session state
+def initialize_session_state():
+    """Initialize Streamlit session state variables."""
+    if "rag" not in st.session_state:
+        st.session_state.rag = None
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    if "temp_dir" not in st.session_state:
+        st.session_state.temp_dir = None
+    if "metamask_connected" not in st.session_state:
+        st.session_state.metamask_connected = False
+
+# Helper function to clean up temporary files
+def cleanup_temp_files():
+    """Clean up temporary files when application exits."""
+    if st.session_state.get('temp_dir') and os.path.exists(st.session_state.temp_dir):
+        try:
+            shutil.rmtree(st.session_state.temp_dir)
+            print(f"Cleaned up temporary directory: {st.session_state.temp_dir}")
+        except Exception as e:
+            print(f"Error cleaning up temporary directory: {e}")
+
+
+# Streamlit UI
+def main():
+    st.set_page_config(page_title="Blockchain-Enabled RAG System", layout="wide")
+    
+    st.title("🚀 GPU-Accelerated PDF Question Answering with MetaMask Blockchain Verification")
+    st.markdown("Upload PDFs, verify them on blockchain with MetaMask, and ask questions with audit log")
+    
+    # Initialize session state
+    initialize_session_state()
+    
+    # MetaMask Connection Section
+    st.header("🦊 MetaMask Connection")
+    st.markdown("Connect your MetaMask wallet to verify documents and log queries on the blockchain.")
+    
+    # Add MetaMask connector and get connection info
+    metamask_info = metamask_connector()
+    
+    # Display MetaMask connection status
+    if metamask_info and metamask_info.get("connected"):
+        st.success(f"✅ MetaMask Connected: {metamask_info.get('address')}")
+        st.info(f"Network: {metamask_info.get('network_name')}")
+        st.session_state.metamask_connected = True
+    else:
+        st.warning("⚠️ MetaMask not connected. Please connect your wallet to use blockchain features.")
+        st.session_state.metamask_connected = False
+    
+    # Update RAG system with MetaMask connection if needed
+    if st.session_state.rag and metamask_info:
+        is_connected = st.session_state.rag.update_blockchain_connection(metamask_info)
+        if is_connected:
+            st.success("RAG system updated with MetaMask connection")
+    
+    # Sidebar for configuration and file upload
+    with st.sidebar:
+        st.header("⚙️ Configuration")
+        
+        # GPU Detection
+        gpu_available = torch.cuda.is_available()
+        if gpu_available:
+            try:
+                gpu_info = torch.cuda.get_device_properties(0)
+                st.success(f"GPU detected: {gpu_info.name} ({gpu_info.total_memory / 1024**3:.1f} GB)")
+            except Exception as e:
+                st.warning(f"GPU detected but couldn't get properties: {str(e)}")
+                st.info("Running with limited GPU information")
+        else:
+            st.warning("No GPU detected. Running in CPU mode.")
+        
+        # Model selection
+        llm_model = st.selectbox(
+            "LLM Model",
+            options=[
+                "mistralai/Mistral-7B-Instruct-v0.2",
+                "google/flan-t5-base",
+                "tiiuae/falcon-7b-instruct"
+            ],
+            index=0
+        )
+        
+        embedding_model = st.selectbox(
+            "Embedding Model",
+            options=[
+                "sentence-transformers/all-mpnet-base-v2",
+                "sentence-transformers/all-MiniLM-L6-v2",
+                "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+            ],
+            index=1  # all-MiniLM-L6-v2 is smaller and faster
+        )
+        
+        use_gpu = st.checkbox("Use GPU Acceleration", value=gpu_available)
+        
+        # Blockchain configuration
+        st.header("🔗 Blockchain Configuration")
+        use_blockchain = st.checkbox("Enable Blockchain Verification", value=True)
+        
+        if use_blockchain:
+            contract_address = st.text_input("Contract Address", 
+                                            value="0x0000000000000000000000000000000000000000")
+            
+            # Display MetaMask connection status in sidebar
+            if metamask_info and metamask_info.get("connected"):
+                st.success(f"✅ MetaMask Connected: {metamask_info.get('address')[:10]}...")
+            else:
+                st.warning("⚠️ MetaMask not connected. Please connect your wallet above.")
+            
+            if not contract_address or contract_address == "0x0000000000000000000000000000000000000000":
+                st.error("Please deploy the contract and enter its address")
+        
+        # Advanced options
+        with st.expander("Advanced Options"):
+            chunk_size = st.slider("Chunk Size", 100, 2000, 1000)
+            chunk_overlap = st.slider("Chunk Overlap", 0, 500, 200)
+        
+        # Initialize button
+        if st.button("Initialize System"):
+            with st.spinner("Initializing RAG system..."):
+                if use_blockchain and not contract_address:
+                    st.error("Contract address is required for blockchain integration")
+                else:
+                    st.session_state.rag = BlockchainEnabledRAG(
+                        llm_model_name=llm_model,
+                        embedding_model_name=embedding_model,
+                        chunk_size=chunk_size,
+                        chunk_overlap=chunk_overlap,
+                        use_gpu=use_gpu and gpu_available,
+                        use_blockchain=use_blockchain,
+                        contract_address=contract_address if use_blockchain else None
+                    )
+                    
+                    # Update with current MetaMask connection if available
+                    if use_blockchain and metamask_info:
+                        st.session_state.rag.update_blockchain_connection(metamask_info)
+                    
+                    st.success(f"System initialized with {embedding_model} on {st.session_state.rag.device}")
+                    if use_blockchain:
+                        if metamask_info and metamask_info.get("connected"):
+                            st.success("Blockchain verification enabled with MetaMask")
+                        else:
+                            st.warning("Blockchain verification enabled but MetaMask not connected")
+        
+        st.header("📄 Upload Documents")
+        uploaded_files = st.file_uploader("Select PDFs", type="pdf", accept_multiple_files=True)
+        
+        if uploaded_files and st.button("Process PDFs"):
+            if not st.session_state.rag:
+                with st.spinner("Initializing RAG system..."):
+                    st.session_state.rag = BlockchainEnabledRAG(
+                        llm_model_name=llm_model,
+                        embedding_model_name=embedding_model,
+                        chunk_size=chunk_size,
+                        chunk_overlap=chunk_overlap,
+                        use_gpu=use_gpu and gpu_available,
+                        use_blockchain=use_blockchain,
+                        contract_address=contract_address if use_blockchain else None
+                    )
+                    
+                    # Update with current MetaMask connection if available
+                    if use_blockchain and metamask_info:
+                        st.session_state.rag.update_blockchain_connection(metamask_info)
+            
+            success = st.session_state.rag.process_pdfs(uploaded_files)
+            if success:
+                metrics = st.session_state.rag.get_performance_metrics()
+                if metrics:
+                    st.success("PDFs processed successfully!")
+                    with st.expander("💹 Performance Metrics"):
+                        st.markdown(f"**Documents processed:** {metrics['documents_processed']} chunks")
+                        st.markdown(f"**Index building time:** {metrics['index_building_time']:.2f} seconds")
+                        st.markdown(f"**Total processing time:** {metrics['total_processing_time']:.2f} seconds")
+                        st.markdown(f"**Memory used:** {metrics['memory_used_gb']:.2f} GB")
+                        st.markdown(f"**Device used:** {metrics['device']}")
+                        st.markdown(f"**Blockchain verification:** {'Enabled' if metrics['blockchain_enabled'] else 'Disabled'}")
+                        st.markdown(f"**Blockchain connected:** {'Yes' if metrics.get('blockchain_connected') else 'No'}")
+    
+    # Blockchain verification info
+    if st.session_state.rag and st.session_state.rag.use_blockchain:
+        if st.session_state.metamask_connected:
+            st.info("🔗 Blockchain verification is enabled with MetaMask. Documents are cryptographically verified and queries are logged with immutable audit trail.")
+        else:
+            st.warning("🔗 Blockchain verification is enabled but MetaMask is not connected. Please connect your MetaMask wallet to use blockchain features.")
+    
+    # Display chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            if message["role"] == "user":
+                st.markdown(message["content"])
+            else:
+                if isinstance(message["content"], dict):
+                    st.markdown(message["content"]["answer"])
+                    
+                    if "query_time" in message["content"]:
+                        st.caption(f"Response time: {message['content']['query_time']:.2f} seconds")
+                    
+                    # Display blockchain log if available
+                    if "blockchain_log" in message["content"] and message["content"]["blockchain_log"]:
+                        blockchain_log = message["content"]["blockchain_log"]
+                        st.success(f"✅ Query logged on blockchain | Transaction: {blockchain_log['tx_hash'][:10]}...")
+                    
+                    # Display sources in expander
+                    if "sources" in message["content"] and message["content"]["sources"]:
+                        with st.expander("📄 View Sources"):
+                            for i, source in enumerate(message["content"]["sources"]):
+                                st.markdown(f"**Source {i+1}: {source['source']}**")
+                                
+                                # Show blockchain verification if available
+                                if source.get("blockchain"):
+                                    st.success(f"✅ Verified on blockchain | TX: {source['blockchain']['tx_hash'][:10]}...")
+                                    
+                                st.text(source["content"])
+                                st.divider()
+                else:
+                    st.markdown(message["content"])
+    
+    # Chat input
+    if prompt := st.chat_input("Ask a question about your PDFs..."):
+        # Add user message to chat
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        
+        # Display user message
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        
+        # Check if system is initialized
+        if not st.session_state.rag:
+            with st.chat_message("assistant"):
+                message = "Please initialize the system and process PDFs first."
+                st.markdown(message)
+                st.session_state.messages.append({"role": "assistant", "content": message})
+        
+        # Get response if vector store is ready
+        elif st.session_state.rag.vector_store:
+            with st.chat_message("assistant"):
+                response = st.session_state.rag.ask(prompt)
+                st.session_state.messages.append({"role": "assistant", "content": response})
+                
+                if isinstance(response, dict):
+                    st.markdown(response["answer"])
+                    
+                    if "query_time" in response:
+                        st.caption(f"Response time: {response['query_time']:.2f} seconds")
+                    
+                    # Display blockchain log if available
+                    if "blockchain_log" in response and response["blockchain_log"]:
+                        blockchain_log = response["blockchain_log"]
+                        st.success(f"✅ Query logged on blockchain | Transaction: {blockchain_log['tx_hash'][:10]}...")
+                    
+                    # Display sources in expander
+                    if "sources" in response and response["sources"]:
+                        with st.expander("📄 View Sources"):
+                            for i, source in enumerate(response["sources"]):
+                                st.markdown(f"**Source {i+1}: {source['source']}**")
+                                
+                                # Show blockchain verification if available
+                                if source.get("blockchain"):
+                                    st.success(f"✅ Verified on blockchain | TX: {source['blockchain']['tx_hash'][:10]}...")
+                                    
+                                st.text(source["content"])
+                                st.divider()
+                else:
+                    st.markdown(response)
+        else:
+            with st.chat_message("assistant"):
+                message = "Please upload and process PDF files first."
+                st.markdown(message)
+                st.session_state.messages.append({"role": "assistant", "content": message})
+
+
+# Main entry point
+if __name__ == "__main__":
+    # Register cleanup function
+    atexit.register(cleanup_temp_files)
+    
+    main()

metamask-requirements.txt

@@ -0,0 +1,13 @@
+streamlit>=1.30.0
+langchain>=0.0.292
+langchain-huggingface>=0.0.3
+langchain-community>=0.0.3
+python-dotenv>=1.0.0
+transformers>=4.33.3
+faiss-cpu>=1.7.4
+sentence-transformers>=2.2.2
+PyPDF2>=3.0.1
+psutil>=5.9.5
+web3>=6.10.0
+uuid>=1.30
+huggingface-hub>=0.17.3