Update app.py

app.py

@@ -1,357 +1,12 @@
-# main_metamask.py
+# app.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 advanced_rag import AdvancedRAG
 from metamask_component import metamask_connector
-
-
-class BlockchainEnabledRAG:
-    def __init__(self, 
-                 llm_model_name="deepseek-ai/DeepSeek-V3-0324",
-                 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
-        }
-
+from voice_component import voice_input_component
 
 # Helper function to initialize session state
 def initialize_session_state():
@@ -364,6 +19,10 @@ def initialize_session_state():
         st.session_state.temp_dir = None
     if "metamask_connected" not in st.session_state:
         st.session_state.metamask_connected = False
+    if "retrieval_method" not in st.session_state:
+        st.session_state.retrieval_method = "enhanced"
+    if "voice_transcript" not in st.session_state:
+        st.session_state.voice_transcript = ""
 
 # Helper function to clean up temporary files
 def cleanup_temp_files():
@@ -375,13 +34,22 @@ def cleanup_temp_files():
         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.set_page_config(
+        page_title="Advanced RAG with MetaMask and Voice", 
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
     
-    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")
+    st.title("🚀 Advanced RAG System with Blockchain Verification and Voice Input")
+    st.markdown("""
+    This application allows you to:
+    - Upload and process PDF documents
+    - Verify document authenticity on blockchain using MetaMask
+    - Ask questions using voice or text input
+    - Choose between direct retrieval or enhanced LLM-powered answers
+    """)
     
     # Initialize session state
     initialize_session_state()
@@ -425,12 +93,14 @@ def main():
             st.warning("No GPU detected. Running in CPU mode.")
         
         # Model selection
+        st.subheader("Model Selection")
         llm_model = st.selectbox(
             "LLM Model",
             options=[
-                "deepseek-ai/DeepSeek-V3-0324",
                 "mistralai/Mistral-7B-Instruct-v0.2",
-                "google/flan-t5-base",
+                "google/gemma-7b-it",
+                "google/flan-t5-xl",
+                "Salesforce/xgen-7b-8k-inst",
                 "tiiuae/falcon-7b-instruct"
             ],
             index=0
@@ -449,21 +119,20 @@ def main():
         use_gpu = st.checkbox("Use GPU Acceleration", value=gpu_available)
         
         # Blockchain configuration
-        st.header("🔗 Blockchain Configuration")
+        st.subheader("🔗 Blockchain Configuration")
         use_blockchain = st.checkbox("Enable Blockchain Verification", value=True)
         
         if use_blockchain:
-            contract_address = st.text_input("Contract Address", 
-                                            value="0x0000000000000000000000000000000000000000")
+            # Hardcoded contract address - replace with your deployed contract
+            contract_address = os.environ.get("CONTRACT_ADDRESS", "0x123abc...")  # Your pre-deployed contract
+            
+            st.info(f"Using pre-deployed contract: {contract_address[:10]}...")
             
             # 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"):
@@ -476,7 +145,7 @@ def main():
                 if use_blockchain and not contract_address:
                     st.error("Contract address is required for blockchain integration")
                 else:
-                    st.session_state.rag = BlockchainEnabledRAG(
+                    st.session_state.rag = AdvancedRAG(
                         llm_model_name=llm_model,
                         embedding_model_name=embedding_model,
                         chunk_size=chunk_size,
@@ -503,7 +172,7 @@ def main():
         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(
+                    st.session_state.rag = AdvancedRAG(
                         llm_model_name=llm_model,
                         embedding_model_name=embedding_model,
                         chunk_size=chunk_size,
@@ -531,6 +200,20 @@ def main():
                         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'}")
     
+    # Retrieval Method Selection
+    st.header("🔍 Retrieval Method")
+    retrieval_cols = st.columns(2)
+    
+    with retrieval_cols[0]:
+        if st.button("📄 Direct Retrieval", help="Get raw document chunks without LLM processing", use_container_width=True):
+            st.session_state.retrieval_method = "direct"
+            st.info("Using Direct Retrieval: Raw document passages will be returned without LLM processing")
+            
+    with retrieval_cols[1]:
+        if st.button("🧠 Enhanced Retrieval", help="Process results through LLM for comprehensive answers", use_container_width=True):
+            st.session_state.retrieval_method = "enhanced"
+            st.info("Using Enhanced Retrieval: Documents will be processed by LLM to generate comprehensive answers")
+    
     # Blockchain verification info
     if st.session_state.rag and st.session_state.rag.use_blockchain:
         if st.session_state.metamask_connected:
@@ -539,89 +222,141 @@ def main():
             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.header("💬 Chat")
+    
+    # Chat container
+    chat_container = st.container(height=400, border=True)
+    
+    with chat_container:
+        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")
+                        
+                        if "method" in message["content"]:
+                            method_name = "Direct Retrieval" if message["content"]["method"] == "direct" else "Enhanced Retrieval"
+                            st.caption(f"Method: {method_name}")
+                        
+                        # 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"])
+    
+    # Voice Input Section
+    st.header("🎤 Voice Input")
+    st.markdown("You can ask questions using your voice or type them below.")
+    
+    # Voice input component
+    voice_transcript = voice_input_component()
+    
+    # Update session state with voice transcript if not empty
+    if voice_transcript and voice_transcript.strip():
+        st.session_state.voice_transcript = voice_transcript.strip()
+        st.success(f"Voice input received: {voice_transcript}")
+    
+    # Chat input - show the voice transcript in the text input
+    prompt_placeholder = "Ask a question about your PDFs..."
+    if st.session_state.voice_transcript:
+        prompt_placeholder = st.session_state.voice_transcript
     
     # Chat input
-    if prompt := st.chat_input("Ask a question about your PDFs..."):
+    prompt = st.chat_input(prompt_placeholder)
+    
+    # Process either voice input or text input
+    if prompt or st.session_state.voice_transcript:
+        # Prioritize text input over voice input
+        if prompt:
+            user_input = prompt
+        else:
+            user_input = st.session_state.voice_transcript
+            # Clear voice transcript after using it
+            st.session_state.voice_transcript = ""
+            # Rerun to clear the voice input display
+            st.rerun()
+        
         # Add user message to chat
-        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.session_state.messages.append({"role": "user", "content": user_input})
         
         # Display user message
-        with st.chat_message("user"):
-            st.markdown(prompt)
+        with chat_container:
+            with st.chat_message("user"):
+                st.markdown(user_input)
         
         # 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})
+            with chat_container:
+                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")
+            with chat_container:
+                with st.chat_message("assistant"):
+                    # Get retrieval method
+                    method = st.session_state.retrieval_method
                     
-                    # 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]}...")
+                    # Get response using specified method
+                    response = st.session_state.rag.ask(user_input, method=method)
+                    st.session_state.messages.append({"role": "assistant", "content": response})
                     
-                    # 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]}...")
+                    if isinstance(response, dict):
+                        st.markdown(response["answer"])
+                        
+                        if "query_time" in response:
+                            st.caption(f"Response time: {response['query_time']:.2f} seconds")
+                        
+                        if "method" in response:
+                            method_name = "Direct Retrieval" if response["method"] == "direct" else "Enhanced Retrieval"
+                            st.caption(f"Method: {method_name}")
+                        
+                        # 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']}**")
                                     
-                                st.text(source["content"])
-                                st.divider()
-                else:
-                    st.markdown(response)
+                                    # 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})
+            with chat_container:
+                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
@@ -629,5 +364,4 @@ if __name__ == "__main__":
     # Register cleanup function
     atexit.register(cleanup_temp_files)
     
-    main()
-    
+    main()