Add Dockerfile and final setup for Hugging Face deployment

app.py

@@ -1,100 +1,209 @@
+# app.py
+
 import os
+import time
+import uuid
 from flask import Flask, request, render_template, session, jsonify, Response, stream_with_context
 from werkzeug.utils import secure_filename
 from rag_processor import create_rag_chain
-import time
+
+# ============================ ADDITIONS START ============================
+from gtts import gTTS
+import io
+import re # <-- Import the regular expression module
+# ============================ ADDITIONS END ==============================
+
+# Document Loaders
+from langchain_community.document_loaders import (
+    TextLoader,
+    UnstructuredPDFLoader,
+    Docx2txtLoader,
+    UnstructuredImageLoader,
+)
+
+# Text Splitter, Embeddings, Retrievers
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.retrievers import EnsembleRetriever
+from langchain_community.retrievers import BM25Retriever
+from langchain_community.chat_message_histories import ChatMessageHistory
 
 # --- Basic Flask App Setup ---
 app = Flask(__name__)
-# A secret key is needed for session management
-app.config['SECRET_KEY'] = os.urandom(24) 
-# Configure the upload folder
+app.config['SECRET_KEY'] = os.urandom(24)
 app.config['UPLOAD_FOLDER'] = 'uploads'
-# Ensure the upload folder exists
 os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
 
-# In-memory storage for RAG chains to avoid re-creating them on every request.
-# In a production scenario, you might want a more persistent cache like Redis.
+# --- In-memory Storage & Global Model Loading ---
 rag_chains = {}
+message_histories = {}
 
-@app.route('/', methods=['GET'])
-def index():
+# Load the embedding model once when the application starts for efficiency.
+print("Loading embedding model...")
+EMBEDDING_MODEL = HuggingFaceEmbeddings(model_name="all-miniLM-L6-v2")
+print("Embedding model loaded successfully.")
+
+# A dictionary to map file extensions to their corresponding loader classes
+LOADER_MAPPING = {
+    ".txt": TextLoader,
+    ".pdf": UnstructuredPDFLoader,
+    ".docx": Docx2txtLoader,
+    ".jpeg": UnstructuredImageLoader, ".jpg": UnstructuredImageLoader, ".png": UnstructuredImageLoader,
+}
+
+def get_session_history(session_id: str) -> ChatMessageHistory:
     """
-    Renders the main page.
+    Retrieves the chat history for a given session ID. If it doesn't exist,
+    a new history object is created.
     """
+    if session_id not in message_histories:
+        message_histories[session_id] = ChatMessageHistory()
+    return message_histories[session_id]
+
+@app.route('/', methods=['GET'])
+def index():
+    """Renders the main page."""
     return render_template('index.html')
 
 @app.route('/upload', methods=['POST'])
-def upload_file():
-    """
-    Handles file uploads and processing.
-    """
-    if 'file' not in request.files:
-        return jsonify({'status': 'error', 'message': 'No file part in the request.'}), 400
-    
-    file = request.files['file']
+def upload_files():
+    """Handles multiple file uploads, processing, and RAG chain creation."""
+    files = request.files.getlist('file')
+
+    if not files or all(f.filename == '' for f in files):
+        return jsonify({'status': 'error', 'message': 'No selected files.'}), 400
+
+    all_docs = []
+    all_filenames = []
     
-    if file.filename == '':
-        return jsonify({'status': 'error', 'message': 'No selected file.'}), 400
-        
-    if file:
-        filename = secure_filename(file.filename)
-        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-        file.save(filepath)
+    try:
+        print(f"Processing {len(files)} files...")
         
-        try:
-            # --- RAG Chain Creation ---
-            print(f"Creating RAG chain for {filename}...")
-            # Simulate a delay for demonstration purposes of the loading animation
-            time.sleep(2) 
-            rag_chains[filename] = create_rag_chain(filepath)
-            print("RAG chain created successfully.")
+        for file in files:
+            if file and file.filename:
+                filename = secure_filename(file.filename)
+                all_filenames.append(filename)
+                filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+                file.save(filepath)
 
-            # Store the filename in the user's session
-            session['filename'] = filename
-            
-            return jsonify({'status': 'success', 'filename': filename})
+                file_extension = os.path.splitext(filename)[1].lower()
+                if file_extension not in LOADER_MAPPING:
+                    print(f"Skipping unsupported file type: {filename}")
+                    continue
+                
+                loader_class = LOADER_MAPPING[file_extension]
+                loader_kwargs = {}
+                if file_extension in [".jpeg", ".jpg", ".png"]:
+                    loader_kwargs['mode'] = 'single'
+                if file_extension == ".pdf":
+                    loader_kwargs['languages'] = ['eng']
 
-        except Exception as e:
-            print(f"Error creating RAG chain: {e}")
-            if os.path.exists(filepath):
-                os.remove(filepath)
-            return jsonify({'status': 'error', 'message': f'Failed to process file: {str(e)}'}), 500
+                loader = loader_class(filepath, **loader_kwargs)
+                all_docs.extend(loader.load())
 
-    return jsonify({'status': 'error', 'message': 'An unexpected error occurred.'}), 500
+        if not all_docs:
+             return jsonify({'status': 'error', 'message': 'No processable files were uploaded.'}), 400
 
+        # --- Process all documents together ---
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        splits = text_splitter.split_documents(all_docs)
+
+        print("Creating vector store for all documents...")
+        vectorstore = FAISS.from_documents(documents=splits, embedding=EMBEDDING_MODEL)
+
+        bm25_retriever = BM25Retriever.from_documents(splits)
+        bm25_retriever.k = 5
+        faiss_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
+        ensemble_retriever = EnsembleRetriever(
+            retrievers=[bm25_retriever, faiss_retriever],
+            weights=[0.5, 0.5]
+        )
+        
+        session_id = str(uuid.uuid4())
+        rag_chains[session_id] = create_rag_chain(ensemble_retriever, get_session_history)
+        print(f"RAG chain created for session {session_id} with {len(all_filenames)} documents.")
+
+        session['session_id'] = session_id
+        
+        display_filenames = ", ".join(all_filenames)
+        return jsonify({'status': 'success', 'filename': display_filenames})
+
+    except Exception as e:
+        print(f"Error creating RAG chain: {e}")
+        return jsonify({'status': 'error', 'message': f'Failed to process files: {str(e)}'}), 500
 
 @app.route('/chat', methods=['POST'])
 def chat():
-    """
-    Handles chat messages from the user and streams the response.
-    """
+    """Handles chat messages and streams the response with memory."""
     data = request.get_json()
     question = data.get('question')
-    filename = session.get('filename')
+    session_id = session.get('session_id')
 
-    if not question:
-        return jsonify({'status': 'error', 'message': 'Question is missing.'}), 400
+    if not all([question, session_id]):
+        return jsonify({'status': 'error', 'message': 'Missing data in request.'}), 400
         
-    if not filename or filename not in rag_chains:
-        return jsonify({'status': 'error', 'message': 'File not uploaded or processed yet.'}), 400
+    if session_id not in rag_chains:
+        return jsonify({'status': 'error', 'message': 'Session not found. Please upload documents again.'}), 400
 
     try:
-        rag_chain = rag_chains[filename]
+        rag_chain = rag_chains[session_id]
+        config = {"configurable": {"session_id": session_id}}
         
         def generate():
             """A generator function to stream the response."""
-            for chunk in rag_chain.stream(question):
+            for chunk in rag_chain.stream({"question": question, "config": config}):
                 yield chunk
         
-        # Use stream_with_context to ensure the generator has access to the request context
         return Response(stream_with_context(generate()), mimetype='text/plain')
 
     except Exception as e:
         print(f"Error during chat invocation: {e}")
-        # This error won't be sent as a stream, handle appropriately
         return Response("An error occurred while getting the answer.", status=500, mimetype='text/plain')
 
-if __name__ == '__main__':
-    app.run(debug=True, port=5001)
+# ============================ ADDITIONS START ============================
 
+def clean_markdown_for_tts(text: str) -> str:
+    """Removes markdown formatting for cleaner text-to-speech output."""
+    # Remove bold (**text**) and italics (*text* or _text_)
+    text = re.sub(r'\*(\*?)(.*?)\1\*', r'\2', text)
+    text = re.sub(r'\_(.*?)\_', r'\1', text)
+    # Remove inline code (`code`)
+    text = re.sub(r'`(.*?)`', r'\1', text)
+    # Remove headings (e.g., #, ##, ###)
+    text = re.sub(r'^\s*#{1,6}\s+', '', text, flags=re.MULTILINE)
+    # Remove list item markers (*, -, 1.)
+    text = re.sub(r'^\s*[\*\-]\s+', '', text, flags=re.MULTILINE)
+    text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE)
+    # Remove blockquotes (>)
+    text = re.sub(r'^\s*>\s?', '', text, flags=re.MULTILINE)
+    # Replace multiple newlines with a single space
+    text = re.sub(r'\n+', ' ', text)
+    return text.strip()
+
+@app.route('/tts', methods=['POST'])
+def text_to_speech():
+    """Generates audio from text and returns it as an MP3 stream."""
+    data = request.get_json()
+    text = data.get('text')
+
+    if not text:
+        return jsonify({'status': 'error', 'message': 'No text provided.'}), 400
+
+    try:
+        # --- FIX IS HERE: Clean the text before sending to gTTS ---
+        clean_text = clean_markdown_for_tts(text)
+        
+        tts = gTTS(clean_text, lang='en')
+        mp3_fp = io.BytesIO()
+        tts.write_to_fp(mp3_fp)
+        mp3_fp.seek(0)
+        return Response(mp3_fp, mimetype='audio/mpeg')
+    except Exception as e:
+        print(f"Error in TTS generation: {e}")
+        return jsonify({'status': 'error', 'message': 'Failed to generate audio.'}), 500
+# ============================ ADDITIONS END ==============================
+
+
+if __name__ == '__main__':
+    app.run(debug=True, port=5001)

rag_processor.py

@@ -1,55 +1,34 @@
+# rag_processor.py
+
 import os
 from dotenv import load_dotenv
-
-# Document Loaders
-from langchain_community.document_loaders import (
-    TextLoader,
-    PyPDFLoader,
-    Docx2txtLoader,
-    UnstructuredImageLoader,
-)
-
-# Text Splitter
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-
-# Embeddings
-from langchain_huggingface import HuggingFaceEmbeddings
-
-# Vector Stores
-from langchain_community.vectorstores import FAISS
+from operator import itemgetter # <--- ADD THIS IMPORT
 
 # LLM
 from langchain_groq import ChatGroq
 
 # Prompting
-from langchain.prompts import PromptTemplate
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 
 # Chains
 from langchain_core.runnables import RunnableParallel, RunnablePassthrough
 from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables.history import RunnableWithMessageHistory
 
-# A dictionary to map file extensions to their corresponding loader classes
-LOADER_MAPPING = {
-    ".txt": TextLoader,
-    ".pdf": PyPDFLoader,
-    ".docx": Docx2txtLoader,
-    ".jpeg": UnstructuredImageLoader,
-    ".jpg": UnstructuredImageLoader,
-    ".png": UnstructuredImageLoader,
-}
-
-def create_rag_chain(filepath):
+def create_rag_chain(retriever, get_session_history_func):
     """
-    Creates a Retrieval-Augmented Generation (RAG) chain from a given file path.
+    Creates an advanced Retrieval-Augmented Generation (RAG) chain with hybrid search,
+    query rewriting, answer refinement, and conversational memory.
 
     Args:
-        filepath (str): The path to the document file.
+        retriever: A configured LangChain retriever object.
+        get_session_history_func: A function to get the chat history for a session.
 
     Returns:
-        A LangChain runnable object representing the RAG chain.
+        A LangChain runnable object representing the RAG chain with memory.
         
     Raises:
-        ValueError: If the file extension is not supported.
+        ValueError: If the GROQ_API_KEY is missing.
     """
     # Load environment variables from .env file
     load_dotenv()
@@ -57,75 +36,92 @@ def create_rag_chain(filepath):
     if not api_key:
         raise ValueError("GROQ_API_KEY not found in environment variables.")
 
-    # --- 1. Load Document ---
-    print("Loading document...")
-    file_extension = "." + filepath.rsplit(".", 1)[-1].lower()
+    # --- 1. Initialize the LLM ---
+    # Updated model_name to a standard, high-performance Groq model
+    llm = ChatGroq(model_name="llama-3.1-8b-instant", api_key=api_key, temperature=1)
+
+    # --- 2. Create Query Rewriting Chain 🧠 ---
+    print("\nSetting up query rewriting chain...")
+    rewrite_template = """You are an expert at rewriting user questions for a vector database.
+You are here to help the user with their document.
+Based on the chat history, reformulate the follow-up question to be a standalone question.
+This new query should be optimized to find the most relevant documents in a knowledge base.
+Do NOT answer the question, only provide the rewritten, optimized question.
+
+Chat History:
+{chat_history}
+
+Follow-up Question: {question}
+Standalone Question:"""
+    rewrite_prompt = ChatPromptTemplate.from_messages([
+        ("system", rewrite_template),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "Based on our conversation, reformulate this question to be a standalone query: {question}")
+    ])
+    query_rewriter = rewrite_prompt | llm | StrOutputParser()
     
-    if file_extension in LOADER_MAPPING:
-        loader_class = LOADER_MAPPING[file_extension]
-        # For image loaders, mode="single" can be useful to treat the image as one document
-        if file_extension in [".jpeg", ".jpg", ".png"]:
-             loader = loader_class(filepath, mode="single")
-        else:
-             loader = loader_class(filepath)
-        docs = loader.load()
-    else:
-        raise ValueError(f"Unsupported file type: '{file_extension}'")
+    # --- 3. Create Main RAG Chain with Memory ---
+    print("\nSetting up main RAG chain...")
+    rag_template = """You are an expert assistant named `Cognichat`.Whenver user ask you about who you are , simply say you are `Cognichat`.
+    You are developed by Ritesh and Alish.
+    Your job is to provide accurate and helpful answers based ONLY on the provided context.
+If the information is not in the context, clearly state that you don't know the answer.
+Provide a clear and concise answer.
+
+Context:
+{context}"""
+    rag_prompt = ChatPromptTemplate.from_messages([
+        ("system", rag_template),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{question}"),
+    ])
+
+    # ============================ FIX IS HERE ============================
+    # Parallel process to fetch context and correctly pass through question and history.
+    # We use itemgetter to select the specific keys from the input dictionary.
+    setup_and_retrieval = RunnableParallel({
+        "context": query_rewriter | retriever,
+        "question": itemgetter("question"),
+        "chat_history": itemgetter("chat_history"),
+    })
+    # =====================================================================
     
-    print(f"Document loaded successfully. Number of pages/docs: {len(docs)}")
-
-    # --- 2. Split Text ---
-    print("\nSplitting document into chunks...")
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    splits = text_splitter.split_documents(docs)
-    print(f"{len(splits)} chunks created.")
-
-    # --- 3. Create Embeddings ---
-    print("\nInitializing Hugging Face embeddings model...")
-    model_name = "BAAI/bge-base-en-v1.5"
-    model_kwargs = {'device': 'cpu'}
-    encode_kwargs = {'normalize_embeddings': True}
-    embeddings = HuggingFaceEmbeddings(
-        model_name=model_name,
-        model_kwargs=model_kwargs,
-        encode_kwargs=encode_kwargs
-    )
-    print("Embeddings model loaded.")
-
-    # --- 4. Create Vector Store ---
-    print("\nCreating FAISS vector store from document chunks...")
-    vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
-    print("Vector store created successfully.")
-
-    # --- 5. Create Retriever ---
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-
-    # --- 6. Define the Prompt Template ---
-    template = """
-You are an expert program. Your job is to provide accurate and helpful answers based ONLY on the provided context.
-If the information is not in the context, say that you don't know the answer.
-Keep your more ellaborated and  explain in a clear way.
-
-Context: {context} 
-
-Question: {question} 
-
-Answer:
-"""
-    prompt = PromptTemplate.from_template(template)
-
-    # --- 7. Initialize the LLM ---
-    llm = ChatGroq(model_name="llama-3.1-8b-instant", api_key=api_key, temperature=0)
-
-    # --- 8. Create the RAG Chain ---
-    rag_chain = (
-        RunnableParallel(
-            context=retriever,
-            question=RunnablePassthrough()
-        )
-        | prompt
+    # The initial RAG chain
+    conversational_rag_chain = (
+        setup_and_retrieval
+        | rag_prompt
         | llm
         | StrOutputParser()
     )
+
+    # Wrap the chain with memory management
+    chain_with_memory = RunnableWithMessageHistory(
+        conversational_rag_chain,
+        get_session_history_func,
+        input_messages_key="question",
+        history_messages_key="chat_history",
+    )
+
+    # --- 4. Create Answer Refinement Chain ✨ ---
+    print("\nSetting up answer refinement chain...")
+    refine_template = """You are an expert at editing and refining content.
+Your task is to take a given answer and improve its clarity, structure, and readability.
+Use formatting such as bold text, bullet points, or numbered lists where it enhances the explanation.
+Do not add any new information that wasn't in the original answer.
+
+Original Answer:
+{answer}
+
+Refined Answer:"""
+    refine_prompt = ChatPromptTemplate.from_template(refine_template)
+    refinement_chain = refine_prompt | llm | StrOutputParser()
+
+    # --- 5. Combine Everything into the Final Chain ---
+    # The final chain passes the output of the memory-enabled chain to the refinement chain
+    # Note: We need to adapt the input for the refinement chain
+    final_chain = (
+        lambda input_dict: {"answer": chain_with_memory.invoke(input_dict, config=input_dict.get('config'))}
+    ) | refinement_chain
     
-    return rag_chain
+    print("\nFinalizing the complete chain with memory...")
+    return final_chain

templates/index.html

@@ -82,6 +82,22 @@
             100% { transform: rotate(360deg); }
         }
         
+        /* Typing Indicator Animation */
+        .typing-indicator span {
+            height: 10px;
+            width: 10px;
+            background-color: #9E9E9E;
+            border-radius: 50%;
+            display: inline-block;
+            animation: bounce 1.4s infinite ease-in-out both;
+        }
+        .typing-indicator span:nth-child(1) { animation-delay: -0.32s; }
+        .typing-indicator span:nth-child(2) { animation-delay: -0.16s; }
+        @keyframes bounce {
+            0%, 80%, 100% { transform: scale(0); }
+            40% { transform: scale(1.0); }
+        }
+
         /* Markdown Styling */
         .markdown-content p { margin-bottom: 0.75rem; line-height: 1.75; }
         .markdown-content ul, .markdown-content ol { margin-left: 1.5rem; margin-bottom: 0.75rem; }
@@ -93,12 +109,22 @@
         .markdown-content pre .copy-code-btn { position: absolute; top: 0.5rem; right: 0.5rem; background-color: #e8eaed; border: 1px solid #dadce0; color: #5f6368; padding: 0.3rem 0.6rem; border-radius: 0.25rem; cursor: pointer; opacity: 0; transition: opacity 0.2s; font-size: 0.8em;}
         .dark .markdown-content pre .copy-code-btn { background-color: #3c4043; border-color: #5f6368; color: #e8eaed; }
         .markdown-content pre:hover .copy-code-btn { opacity: 1; }
+        
+        /* Spinner for the TTS button */
+        .tts-button-loader {
+            width: 16px;
+            height: 16px;
+            border: 2px solid currentColor; /* Use button's text color */
+            border-radius: 50%;
+            display: inline-block;
+            box-sizing: border-box;
+            animation: rotation 0.8s linear infinite;
+            border-bottom-color: transparent; /* Makes it a half circle spinner */
+        }
     </style>
 </head>
-<body class="w-screen h-screen dark"> <!-- Default to dark mode -->
-
+<body class="w-screen h-screen dark">
     <main id="main-content" class="h-full flex flex-col transition-opacity duration-500">
-        <!-- Chat Area -->
         <div id="chat-container" class="hidden flex-1 flex flex-col w-full mx-auto overflow-hidden">
             <header class="text-center p-4 border-b border-[var(--card-border)] flex-shrink-0">
                 <h1 class="text-xl font-medium">Chat with your Docs</h1>
@@ -106,36 +132,34 @@
             </header>
             <div id="chat-window" class="flex-1 overflow-y-auto p-4 md:p-6 lg:p-10">
                 <div id="chat-content" class="max-w-4xl mx-auto space-y-8">
-                    <!-- Chat messages will be appended here -->
                 </div>
             </div>
             <div class="p-4 flex-shrink-0 bg-[var(--background)] border-t border-[var(--card-border)]">
                 <form id="chat-form" class="max-w-4xl mx-auto bg-[var(--card)] rounded-full p-2 flex items-center shadow-sm border border-transparent focus-within:border-[var(--primary)] transition-colors">
-                    <input type="text" id="chat-input" placeholder="Ask a question about your document..." class="flex-grow bg-transparent focus-outline-none px-4 text-sm" autocomplete="off">
-                    <button type="submit" id="chat-submit-btn" class="bg-[var(--primary)] hover:bg-[var(--primary-hover)] text-white p-2 rounded-full transition-all duration-200 disabled:opacity-50 disabled:cursor-not-allowed disabled:bg-gray-500" title="butn">
+                    <input type="text" id="chat-input" placeholder="Ask a question about your documents..." class="flex-grow bg-transparent focus:outline-none px-4 text-sm" autocomplete="off">
+                    <button type="submit" id="chat-submit-btn" class="bg-[var(--primary)] hover:bg-[var(--primary-hover)] text-white p-2 rounded-full transition-all duration-200 disabled:opacity-50 disabled:cursor-not-allowed disabled:bg-gray-500" title="Send">
                         <svg class="w-5 h-5" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M3.49941 11.5556L11.555 3.5L12.4438 4.38889L6.27721 10.5556H21.9994V11.5556H6.27721L12.4438 17.7222L11.555 18.6111L3.49941 10.5556V11.5556Z" transform="rotate(180, 12.7497, 11.0556)" fill="currentColor"></path></svg>
                     </button>
                 </form>
             </div>
         </div>
 
-        <!-- Upload Area -->
         <div id="upload-container" class="flex-1 flex flex-col items-center justify-center p-8 transition-opacity duration-300">
             <div class="text-center">
                 <h1 class="text-5xl font-medium mb-4">Upload docs to chat</h1>
                 <div id="drop-zone" class="w-full max-w-lg text-center border-2 border-dashed border-[var(--card-border)] rounded-2xl p-10 transition-all duration-300 cursor-pointer bg-[var(--card)] hover:border-[var(--primary)]">
-                    <input id="file-upload" type="file" class="hidden" accept=".pdf,.txt,.docx,.jpg,.jpeg,.png" placeholder="upload your documents here">
+                    <input id="file-upload" type="file" class="hidden" accept=".pdf,.txt,.docx,.jpg,.jpeg,.png" multiple title="input">
                     <svg class="mx-auto h-12 w-12 text-gray-400" fill="none" viewBox="0 0 24 24" stroke-width="1.5" stroke="currentColor" ><path stroke-linecap="round" stroke-linejoin="round" d="M12 16.5V9.75m0 0l3-3m-3 3l-3 3M6.75 19.5a4.5 4.5 0 01-1.41-8.775 5.25 5.25 0 0110.233-2.33 3 3 0 013.758 3.848A3.752 3.752 0 0118 19.5H6.75z"></path></svg>
-                    <p class="mt-4 text-sm font-medium">Drag & drop a file or click to upload</p>
+                    <p class="mt-4 text-sm font-medium">Drag & drop files or click to upload</p>
                     <p id="file-name" class="mt-2 text-xs text-gray-500"></p>
                 </div>
             </div>
         </div>
         
-        <!-- Loading Overlay -->
-        <div id="loading-overlay" class="hidden fixed inset-0 bg-[var(--background)] bg-opacity-80 backdrop-blur-sm flex flex-col items-center justify-center z-50">
+        <div id="loading-overlay" class="hidden fixed inset-0 bg-[var(--background)] bg-opacity-80 backdrop-blur-sm flex flex-col items-center justify-center z-50 text-center p-4">
             <div class="loader"></div>
-            <p id="loading-text" class="mt-6 text-sm">Processing...</p>
+            <p id="loading-text" class="mt-6 text-sm font-medium"></p>
+            <p id="loading-subtext" class="mt-2 text-xs text-gray-500 dark:text-gray-400"></p>
         </div>
     </main>
 
@@ -148,6 +172,7 @@
             const fileNameSpan = document.getElementById('file-name');
             const loadingOverlay = document.getElementById('loading-overlay');
             const loadingText = document.getElementById('loading-text');
+            const loadingSubtext = document.getElementById('loading-subtext');
 
             const chatForm = document.getElementById('chat-form');
             const chatInput = document.getElementById('chat-input');
@@ -156,8 +181,6 @@
             const chatContent = document.getElementById('chat-content');
             const chatFilename = document.getElementById('chat-filename');
 
-            let selectedFile = null;
-
             // --- File Upload Logic ---
             dropZone.addEventListener('click', () => fileUploadInput.click());
             
@@ -171,67 +194,51 @@
 
             dropZone.addEventListener('drop', (e) => {
                 const files = e.dataTransfer.files;
-                if (files.length > 0) handleFile(files[0]);
+                if (files.length > 0) handleFiles(files);
             });
             
             fileUploadInput.addEventListener('change', (e) => {
-                if (e.target.files.length > 0) handleFile(e.target.files[0]);
+                if (e.target.files.length > 0) handleFiles(e.target.files);
             });
 
             function preventDefaults(e) { e.preventDefault(); e.stopPropagation(); }
             
-            async function handleFile(file) {
-                selectedFile = file;
-                fileNameSpan.textContent = `Selected: ${file.name}`;
-                await uploadAndProcessFile();
-            }
-
-            async function uploadAndProcessFile() {
-                if (!selectedFile) return;
-
+            async function handleFiles(files) {
                 const formData = new FormData();
-                formData.append('file', selectedFile);
+                let fileNames = [];
+                for (const file of files) {
+                    formData.append('file', file);
+                    fileNames.push(file.name);
+                }
+
+                fileNameSpan.textContent = `Selected: ${fileNames.join(', ')}`;
+                await uploadAndProcessFiles(formData, fileNames);
+            }
 
+            async function uploadAndProcessFiles(formData, fileNames) {
                 loadingOverlay.classList.remove('hidden');
+                loadingText.textContent = `Processing ${fileNames.length} document(s)...`;
+                loadingSubtext.textContent = "For large documents or OCR, setup may take a few minutes to build the knowledge base.";
                 
-                const loadingSteps = [
-                    `Uploading ${selectedFile.name}...`,
-                    "Parsing document...",
-                    "Extracting text...",
-                    "Creating embeddings (this may take a moment)...",
-                    "Building knowledge base..."
-                ];
-                let stepIndex = 0;
-                loadingText.textContent = loadingSteps[stepIndex];
-
-                const stepInterval = setInterval(() => {
-                    stepIndex++;
-                    if (stepIndex < loadingSteps.length) {
-                        loadingText.textContent = loadingSteps[stepIndex];
-                    } else {
-                        loadingText.textContent = "Finalizing...";
-                    }
-                }, 1500);
-
                 try {
                     const response = await fetch('/upload', { method: 'POST', body: formData });
                     const result = await response.json();
 
                     if (!response.ok) throw new Error(result.message || 'Unknown error occurred.');
                     
-                    chatFilename.textContent = `Chatting with ${result.filename}`;
+                    chatFilename.textContent = `Chatting with: ${result.filename}`;
                     uploadContainer.classList.add('hidden');
                     chatContainer.classList.remove('hidden');
-                    appendMessage("I've analyzed your document. What would you like to know?", "bot");
+                    appendMessage("I've analyzed your documents. What would you like to know?", "bot");
 
                 } catch (error) {
                     console.error('Upload error:', error);
                     alert(`Error: ${error.message}`);
                 } finally {
-                    clearInterval(stepInterval);
                     loadingOverlay.classList.add('hidden');
+                    loadingSubtext.textContent = ''; 
                     fileNameSpan.textContent = '';
-                    selectedFile = null;
+                    fileUploadInput.value = ''; // Reset file input
                 }
             }
             
@@ -246,8 +253,9 @@
                 chatInput.disabled = true;
                 chatSubmitBtn.disabled = true;
 
-                const botMessageContainer = appendMessage('', 'bot');
-                const contentDiv = botMessageContainer.querySelector('.markdown-content');
+                const typingIndicator = showTypingIndicator();
+                let botMessageContainer = null;
+                let contentDiv = null;
 
                 try {
                     const response = await fetch('/chat', {
@@ -258,6 +266,10 @@
 
                     if (!response.ok) throw new Error(`Server error: ${response.statusText}`);
 
+                    typingIndicator.remove();
+                    botMessageContainer = appendMessage('', 'bot');
+                    contentDiv = botMessageContainer.querySelector('.markdown-content');
+
                     const reader = response.body.getReader();
                     const decoder = new TextDecoder();
                     let fullResponse = '';
@@ -271,10 +283,17 @@
                         scrollToBottom();
                     }
                     contentDiv.querySelectorAll('pre').forEach(addCopyButton);
+                    
+                    addTextToSpeechControls(botMessageContainer, fullResponse); 
 
                 } catch (error) {
                     console.error('Chat error:', error);
-                    contentDiv.innerHTML = `<p class="text-red-500">Error: ${error.message}</p>`;
+                    if (typingIndicator) typingIndicator.remove();
+                    if (contentDiv) {
+                        contentDiv.innerHTML = `<p class="text-red-500">Error: ${error.message}</p>`;
+                    } else {
+                        appendMessage(`Error: ${error.message}`, 'bot');
+                    }
                 } finally {
                     chatInput.disabled = false;
                     chatSubmitBtn.disabled = false;
@@ -283,6 +302,7 @@
             });
 
             // --- UI Helper Functions ---
+            
             function appendMessage(text, sender) {
                 const messageWrapper = document.createElement('div');
                 messageWrapper.className = `flex items-start gap-4`;
@@ -301,9 +321,13 @@
                 const contentDiv = document.createElement('div');
                 contentDiv.className = 'text-base markdown-content';
                 contentDiv.innerHTML = marked.parse(text);
-                
+
+                const controlsContainer = document.createElement('div');
+                controlsContainer.className = 'tts-controls mt-2';
+
                 messageBubble.appendChild(senderName);
                 messageBubble.appendChild(contentDiv);
+                messageBubble.appendChild(controlsContainer);
                 messageWrapper.innerHTML = iconSVG;
                 messageWrapper.appendChild(messageBubble);
                 
@@ -312,6 +336,35 @@
 
                 return messageBubble;
             }
+
+            function showTypingIndicator() {
+                const indicatorWrapper = document.createElement('div');
+                indicatorWrapper.className = `flex items-start gap-4`;
+                indicatorWrapper.id = 'typing-indicator';
+                
+                const iconSVG = `<div class="bg-gray-200 dark:bg-gray-700 rounded-full flex-shrink-0 mt-1 text-xl flex items-center justify-center w-10 h-10">✨</div>`;
+                
+                const messageBubble = document.createElement('div');
+                messageBubble.className = 'flex-1 pt-1';
+                
+                const senderName = document.createElement('p');
+                senderName.className = 'font-medium text-sm mb-1';
+                senderName.textContent = 'CogniChat is thinking...';
+                
+                const indicator = document.createElement('div');
+                indicator.className = 'typing-indicator';
+                indicator.innerHTML = '<span></span><span></span><span></span>';
+                
+                messageBubble.appendChild(senderName);
+                messageBubble.appendChild(indicator);
+                indicatorWrapper.innerHTML = iconSVG;
+                indicatorWrapper.appendChild(messageBubble);
+                
+                chatContent.appendChild(indicatorWrapper);
+                scrollToBottom();
+
+                return indicatorWrapper;
+            }
             
             function scrollToBottom() {
                 chatWindow.scrollTo({
@@ -334,9 +387,100 @@
                     });
                 });
             }
+
+            // ============================ MODIFICATIONS START ==============================
+            let currentAudio = null;
+            let currentPlayingButton = null;
+            
+            const playIconSVG = `<svg class="w-5 h-5" fill="currentColor" viewBox="0 0 24 24"><path d="M8 5v14l11-7z"/></svg>`;
+            const pauseIconSVG = `<svg class="w-5 h-5" fill="currentColor" viewBox="0 0 24 24"><path d="M6 19h4V5H6v14zm8-14v14h4V5h-4z"/></svg>`;
+
+            function addTextToSpeechControls(messageBubble, text) {
+                const ttsControls = messageBubble.querySelector('.tts-controls');
+                if (text.trim().length > 0) {
+                    const speakButton = document.createElement('button');
+                    // --- STYLING CHANGE HERE: Brighter blue color for better visibility ---
+                    speakButton.className = 'speak-btn px-3 py-1.5 bg-blue-600 text-white rounded-full text-sm font-medium hover:bg-blue-700 transition-colors flex items-center gap-2 disabled:opacity-50 disabled:cursor-not-allowed';
+                    speakButton.title = 'Listen to this message';
+                    // --- EMOJI ADDED ---
+                    speakButton.innerHTML = `🔊 ${playIconSVG} <span>Listen</span>`;
+                    ttsControls.appendChild(speakButton);
+                    speakButton.addEventListener('click', () => handleTTS(text, speakButton));
+                }
+            }
+
+            // --- BUG FIX: Reworked the entire function for correct pause/resume/stop logic ---
+            async function handleTTS(text, button) {
+                // Case 1: The clicked button is already playing or paused.
+                if (button === currentPlayingButton) {
+                    if (currentAudio && !currentAudio.paused) { // If playing, pause it.
+                        currentAudio.pause();
+                        button.innerHTML = `🔊 ${playIconSVG} <span>Listen</span>`;
+                    } else if (currentAudio && currentAudio.paused) { // If paused, resume it.
+                        currentAudio.play();
+                        button.innerHTML = `🔊 ${pauseIconSVG} <span>Pause</span>`;
+                    }
+                    return;
+                }
+
+                // Case 2: A new button is clicked. Stop any other audio.
+                if (currentAudio) {
+                    currentAudio.pause();
+                }
+                resetAllSpeakButtons();
+                
+                currentPlayingButton = button;
+                button.innerHTML = `<div class="tts-button-loader"></div> <span>Loading...</span>`;
+                button.disabled = true;
+
+                try {
+                    const response = await fetch('/tts', {
+                        method: 'POST',
+                        headers: { 'Content-Type': 'application/json' },
+                        body: JSON.stringify({ text: text })
+                    });
+                    if (!response.ok) throw new Error('Failed to generate audio.');
+                    
+                    const blob = await response.blob();
+                    const audioUrl = URL.createObjectURL(blob);
+                    currentAudio = new Audio(audioUrl);
+                    
+                    currentAudio.play().catch(e => { throw e; });
+                    button.innerHTML = `🔊 ${pauseIconSVG} <span>Pause</span>`;
+
+                    currentAudio.onended = () => {
+                        button.innerHTML = `🔊 ${playIconSVG} <span>Listen</span>`;
+                        currentAudio = null;
+                        currentPlayingButton = null;
+                    };
+                    
+                    currentAudio.onerror = (e) => {
+                        console.error('Audio playback error:', e);
+                        throw new Error('Could not play the generated audio.');
+                    };
+
+                } catch (error) {
+                    console.error('TTS Error:', error);
+                    alert('Failed to play audio. Please try again.');
+                    resetAllSpeakButtons(); // Reset state on error
+                } finally {
+                    button.disabled = false;
+                }
+            }
+
+            function resetAllSpeakButtons() {
+                document.querySelectorAll('.speak-btn').forEach(btn => {
+                    btn.innerHTML = `🔊 ${playIconSVG} <span>Listen</span>`;
+                    btn.disabled = false;
+                });
+                if (currentAudio) {
+                    currentAudio.pause();
+                    currentAudio = null;
+                }
+                currentPlayingButton = null;
+            }
+            // ============================ MODIFICATIONS END ==============================
         });
     </script>
 </body>
-</html>
-
-
+</html>