fix/Update in the ui

app.py

@@ -1,6 +1,9 @@
+#!/usr/bin/python
+ # -*- coding: utf-8 -*-
 import os
 import uuid
-from flask import Flask, request, render_template, session, jsonify, Response
+from flask import Flask, request, render_template, session, jsonify, \
+    Response, stream_with_context # Added stream_with_context
 from werkzeug.utils import secure_filename
 from rag_processor import create_rag_chain
 from typing import Sequence, Any, List
@@ -9,12 +12,15 @@ import re
 import io
 from gtts import gTTS
 from langchain_core.documents import Document
-from langchain_community.document_loaders import TextLoader, Docx2txtLoader
+from langchain_community.document_loaders import TextLoader, \
+    Docx2txtLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever
-from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
+from langchain.retrievers import EnsembleRetriever, \
+    ContextualCompressionRetriever
+from langchain.retrievers.document_compressors.base import \
+    BaseDocumentCompressor
 from langchain_community.retrievers import BM25Retriever
 from langchain_community.chat_message_histories import ChatMessageHistory
 from langchain.storage import InMemoryStore
@@ -23,15 +29,15 @@ from sentence_transformers.cross_encoder import CrossEncoder
 app = Flask(__name__)
 app.config['SECRET_KEY'] = os.urandom(24)
 
-# --- FIX: Use STRING keys for the dictionary ---
 # Maps temperature strings (from the form) to the mode labels
 TEMPERATURE_LABELS = {
-    "0.2": "Precise",
-    "0.4": "Confident",
-    "0.6": "Balanced",
-    "0.8": "Flexible",
-    "1.0": "Creative"
-}
+    '0.2': 'Precise',
+    '0.4': 'Confident',
+    '0.6': 'Balanced',
+    '0.8': 'Flexible',
+    '1.0': 'Creative',
+    }
+
 
 class LocalReranker(BaseDocumentCompressor):
     model: Any
@@ -40,153 +46,196 @@ class LocalReranker(BaseDocumentCompressor):
     class Config:
         arbitrary_types_allowed = True
 
-    def compress_documents(
-        self, documents: Sequence[Document], query: str, callbacks=None
-    ) -> Sequence[Document]:
+    def compress_documents(self, documents: Sequence[Document], query: str,
+                           callbacks=None) -> Sequence[Document]:
         if not documents:
             return []
         pairs = [[query, doc.page_content] for doc in documents]
         scores = self.model.predict(pairs, show_progress_bar=False)
         doc_scores = list(zip(documents, scores))
-        sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1], reverse=True)
+        sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1],
+                                   reverse=True)
         top_docs = []
-        for doc, score in sorted_doc_scores[: self.top_n]:
-            doc.metadata["rerank_score"] = float(score)
+        for (doc, score) in sorted_doc_scores[:self.top_n]:
+            doc.metadata['rerank_score'] = float(score)
             top_docs.append(doc)
         return top_docs
 
+
 def create_optimized_parent_child_chunks(all_docs):
     if not all_docs:
-        print("❌ CHUNKING: No input documents provided!")
-        return [], [], []
-
-    parent_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=900, chunk_overlap=200, separators=["\n\n", "\n", ". ", "! ", "? ", "; ", ", ", " ", ""]
-    )
-    child_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=350, chunk_overlap=80, separators=["\n", ". ", "! ", "? ", "; ", ", ", " ", ""]
-    )
+        print ('❌ CHUNKING: No input documents provided!')
+        return ([], [], [])
+
+    parent_splitter = RecursiveCharacterTextSplitter(chunk_size=900,
+            chunk_overlap=200, separators=['\n\n', '\n', '. ', '! ',
+            '? ', '; ', ', ', ' ', ''])
+    child_splitter = RecursiveCharacterTextSplitter(chunk_size=350,
+            chunk_overlap=80, separators=['\n', '. ', '! ', '? ', '; ',
+            ', ', ' ', ''])
     parent_docs = parent_splitter.split_documents(all_docs)
     doc_ids = [str(uuid.uuid4()) for _ in parent_docs]
     child_docs = []
 
-    for i, parent_doc in enumerate(parent_docs):
+    for (i, parent_doc) in enumerate(parent_docs):
         parent_id = doc_ids[i]
         children = child_splitter.split_documents([parent_doc])
-        for j, child in enumerate(children):
-            child.metadata.update({
-                "doc_id": parent_id, "chunk_index": j, "total_chunks": len(children),
-                "is_first_chunk": j == 0, "is_last_chunk": j == len(children) - 1,
-            })
+        for (j, child) in enumerate(children):
+            child.metadata.update({'doc_id': parent_id,
+                                  'chunk_index': j,
+                                  'total_chunks': len(children),
+                                  'is_first_chunk': j == 0,
+                                  'is_last_chunk': j == len(children)
+                                  - 1})
             if len(children) > 1:
-                if j == 0: child.page_content = "[Beginning] " + child.page_content
-                elif j == len(children) - 1: child.page_content = "[Continues...] " + child.page_content
+                if j == 0:
+                    child.page_content = '[Beginning] ' \
+                        + child.page_content
+                elif j == len(children) - 1:
+                    child.page_content = '[Continues...] ' \
+                        + child.page_content
             child_docs.append(child)
-    
-    print(f"✅ CHUNKING: Created {len(parent_docs)} parent and {len(child_docs)} child chunks.")
-    return parent_docs, child_docs, doc_ids
 
-def get_context_aware_parents(docs: List[Document], store: InMemoryStore) -> List[Document]:
-    if not docs: return []
-    parent_scores, child_content_by_parent = {}, {}
+    print (f"✅ CHUNKING: Created {len(parent_docs)} parent and {len(child_docs)} child chunks."
+           )
+    return (parent_docs, child_docs, doc_ids)
+
+
+def get_context_aware_parents(docs: List[Document],
+                              store: InMemoryStore) -> List[Document]:
+    if not docs:
+        return []
+    (parent_scores, child_content_by_parent) = ({}, {})
     for doc in docs:
-        parent_id = doc.metadata.get("doc_id")
+        parent_id = doc.metadata.get('doc_id')
         if parent_id:
-            parent_scores[parent_id] = parent_scores.get(parent_id, 0) + 1
-            if parent_id not in child_content_by_parent: child_content_by_parent[parent_id] = []
+            parent_scores[parent_id] = parent_scores.get(parent_id, 0) \
+                + 1
+            if parent_id not in child_content_by_parent:
+                child_content_by_parent[parent_id] = []
             child_content_by_parent[parent_id].append(doc.page_content)
-    
+
     parent_ids = list(parent_scores.keys())
     parents = store.mget(parent_ids)
     enhanced_parents = []
-    
-    for i, parent in enumerate(parents):
+
+    for (i, parent) in enumerate(parents):
         if parent is not None:
             parent_id = parent_ids[i]
             if parent_id in child_content_by_parent:
-                child_excerpts = "\n".join(child_content_by_parent[parent_id][:3])
-                enhanced_content = f"{parent.page_content}\n\nRelevant excerpts:\n{child_excerpts}"
-                enhanced_parent = Document(
-                    page_content=enhanced_content,
-                    metadata={**parent.metadata, "child_relevance_score": parent_scores[parent_id], "matching_children": len(child_content_by_parent[parent_id])}
-                )
+                child_excerpts = '\n'.join(child_content_by_parent[parent_id][:3])
+                enhanced_content = \
+                    f"{parent.page_content}\n\nRelevant excerpts:\n{child_excerpts}"
+                enhanced_parent = \
+                    Document(page_content=enhanced_content,
+                             metadata={**parent.metadata,
+                             'child_relevance_score': parent_scores[parent_id],
+                             'matching_children': len(child_content_by_parent[parent_id])})
                 enhanced_parents.append(enhanced_parent)
         else:
-            print(f"❌ PARENT_FETCH: Parent {parent_ids[i]} not found in store!")
+            print (f"❌ PARENT_FETCH: Parent {parent_ids[i]} not found in store!"
+                   )
 
-    enhanced_parents.sort(key=lambda p: p.metadata.get("child_relevance_score", 0), reverse=True)
+    enhanced_parents.sort(key=lambda p: p.metadata.get('child_relevance_score',
+                         0), reverse=True)
     return enhanced_parents
 
-is_hf_spaces = bool(os.getenv("SPACE_ID") or os.getenv("SPACES_ZERO_GPU"))
+
+is_hf_spaces = bool(os.getenv('SPACE_ID') or os.getenv('SPACES_ZERO_GPU'
+                     ))
 app.config['UPLOAD_FOLDER'] = '/tmp/uploads' if is_hf_spaces else 'uploads'
 
 try:
     os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
-    print(f"📁 Upload folder ready: {app.config['UPLOAD_FOLDER']}")
+    print (f"📁 Upload folder ready: {app.config['UPLOAD_FOLDER']}")
 except Exception as e:
-    print(f"❌ Failed to create upload folder, falling back to /tmp: {e}")
+    print (f"❌ Failed to create upload folder, falling back to /tmp: {e}"
+           )
     app.config['UPLOAD_FOLDER'] = '/tmp/uploads'
     os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
 
 session_data = {}
 message_histories = {}
 
-print("🔄 Loading embedding model...")
+print ('🔄 Loading embedding model...')
 try:
-    EMBEDDING_MODEL = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2",
-        model_kwargs={'device': 'cpu'},
-        encode_kwargs={'normalize_embeddings': True}
-    )
-    print("✅ Embedding model loaded.")
-except Exception as e: print(f"❌ FATAL: Could not load embedding model. Error: {e}"); raise e
-
-print("🔄 Loading reranker model...")
+    EMBEDDING_MODEL = \
+        HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2'
+                              , model_kwargs={'device': 'cpu'},
+                              encode_kwargs={'normalize_embeddings': True})
+    print ('✅ Embedding model loaded.')
+except Exception as e:
+    print (f"❌ FATAL: Could not load embedding model. Error: {e}")
+    raise e
+
+print ('🔄 Loading reranker model...')
 try:
-    RERANKER_MODEL = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", device='cpu')
-    print("✅ Reranker model loaded.")
-except Exception as e: print(f"❌ FATAL: Could not load reranker model. Error: {e}"); raise e
+    RERANKER_MODEL = \
+        CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2',
+                     device='cpu')
+    print ('✅ Reranker model loaded.')
+except Exception as e:
+    print (f"❌ FATAL: Could not load reranker model. Error: {e}")
+    raise e
+
 
 def load_pdf_with_fallback(filepath):
     try:
         docs = []
         with fitz.open(filepath) as pdf_doc:
-            for page_num, page in enumerate(pdf_doc):
+            for (page_num, page) in enumerate(pdf_doc):
                 text = page.get_text()
-                if text.strip(): docs.append(Document(page_content=text, metadata={"source": os.path.basename(filepath), "page": page_num + 1}))
+                if text.strip():
+                    docs.append(Document(page_content=text,
+                                metadata={'source': os.path.basename(filepath),
+                                'page': page_num + 1}))
         if docs:
-            print(f"✅ Loaded PDF: {os.path.basename(filepath)} - {len(docs)} pages")
+            print (f"✅ Loaded PDF: {os.path.basename(filepath)} - {len(docs)} pages"
+                   )
             return docs
-        else: raise ValueError("No text content found in PDF.")
-    except Exception as e: print(f"❌ PyMuPDF failed for {filepath}: {e}"); raise
+        else:
+            raise ValueError('No text content found in PDF.')
+    except Exception as e:
+        print (f"❌ PyMuPDF failed for {filepath}: {e}")
+        raise
+
+
+LOADER_MAPPING = {'.txt': TextLoader, '.pdf': load_pdf_with_fallback,
+                  '.docx': Docx2txtLoader}
 
-LOADER_MAPPING = {".txt": TextLoader, ".pdf": load_pdf_with_fallback, ".docx": Docx2txtLoader}
 
 def get_session_history(session_id: str) -> ChatMessageHistory:
-    if session_id not in message_histories: message_histories[session_id] = ChatMessageHistory()
+    if session_id not in message_histories:
+        message_histories[session_id] = ChatMessageHistory()
     return message_histories[session_id]
 
+
 @app.route('/health', methods=['GET'])
-def health_check(): return jsonify({'status': 'healthy'}), 200
+def health_check():
+    return (jsonify({'status': 'healthy'}), 200)
+
 
 @app.route('/', methods=['GET'])
-def index(): return render_template('index.html')
+def index():
+    return render_template('index.html')
+
 
 @app.route('/upload', methods=['POST'])
 def upload_files():
     files = request.files.getlist('file')
-    
-    # Get temperature as a string for the dictionary key
+
     temperature_str = request.form.get('temperature', '0.2')
-    temperature = float(temperature_str) # Convert to float for the LLM
-    model_name = request.form.get('model_name', 'moonshotai/kimi-k2-instruct')
-    print(f"⚙️ UPLOAD: Model: {model_name}, Temp: {temperature}")
+    temperature = float(temperature_str)
+    model_name = request.form.get('model_name',
+                                  'moonshotai/kimi-k2-instruct')
+    print (f"⚙️ UPLOAD: Model: {model_name}, Temp: {temperature}")
 
     if not files or all(f.filename == '' for f in files):
-        return jsonify({'status': 'error', 'message': 'No selected files.'}), 400
+        return (jsonify({'status': 'error',
+                'message': 'No selected files.'}), 400)
 
-    all_docs, processed_files, failed_files = [], [], []
-    print(f"📁 Processing {len(files)} file(s)...")
+    (all_docs, processed_files, failed_files) = ([], [], [])
+    print (f"📁 Processing {len(files)} file(s)...")
     for file in files:
         if file and file.filename:
             filename = secure_filename(file.filename)
@@ -194,128 +243,217 @@ def upload_files():
             try:
                 file.save(filepath)
                 file_ext = os.path.splitext(filename)[1].lower()
-                if file_ext not in LOADER_MAPPING: raise ValueError("Unsupported file format.")
+                if file_ext not in LOADER_MAPPING:
+                    raise ValueError('Unsupported file format.')
                 loader_func = LOADER_MAPPING[file_ext]
-                docs = loader_func(filepath) if file_ext == ".pdf" else loader_func(filepath).load()
-                if not docs: raise ValueError("No content extracted.")
+                docs = loader_func(filepath) if file_ext == '.pdf' \
+                    else loader_func(filepath).load()
+                if not docs:
+                    raise ValueError('No content extracted.')
                 all_docs.extend(docs)
                 processed_files.append(filename)
             except Exception as e:
-                print(f"✗ Error processing {filename}: {e}")
+                print (f"✗ Error processing {filename}: {e}")
                 failed_files.append(f"{filename} ({e})")
 
     if not all_docs:
-        return jsonify({'status': 'error', 'message': f"Failed to process all files. Reasons: {', '.join(failed_files)}"}), 400
+        return (jsonify({'status': 'error',
+                'message': f"Failed to process all files. Reasons: {', '.join(failed_files)}"
+                }), 400)
 
-    print(f"✅ UPLOAD: Processed {len(processed_files)} files.")
+    print (f"✅ UPLOAD: Processed {len(processed_files)} files.")
     try:
-        print("🔄 Starting RAG pipeline setup...")
-        parent_docs, child_docs, doc_ids = create_optimized_parent_child_chunks(all_docs)
-        if not child_docs: raise ValueError("No child documents created during chunking.")
-        
-        vectorstore = FAISS.from_documents(child_docs, EMBEDDING_MODEL)
-        store = InMemoryStore(); store.mset(list(zip(doc_ids, parent_docs)))
-        print(f"✅ Indexed {len(child_docs)} document chunks.")
+        print ('🔄 Starting RAG pipeline setup...')
+        (parent_docs, child_docs, doc_ids) = \
+            create_optimized_parent_child_chunks(all_docs)
+        if not child_docs:
+            raise ValueError('No child documents created during chunking.')
 
-        bm25_retriever = BM25Retriever.from_documents(child_docs); bm25_retriever.k = 12
-        faiss_retriever = vectorstore.as_retriever(search_kwargs={"k": 12})
-        ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, faiss_retriever], weights=[0.6, 0.4])
+        vectorstore = FAISS.from_documents(child_docs, EMBEDDING_MODEL)
+        store = InMemoryStore()
+        store.mset(list(zip(doc_ids, parent_docs)))
+        print (f"✅ Indexed {len(child_docs)} document chunks.")
+
+        bm25_retriever = BM25Retriever.from_documents(child_docs)
+        bm25_retriever.k = 12
+        faiss_retriever = vectorstore.as_retriever(search_kwargs={'k': 12})
+        ensemble_retriever = \
+            EnsembleRetriever(retrievers=[bm25_retriever,
+                               faiss_retriever], weights=[0.6, 0.4])
         reranker = LocalReranker(model=RERANKER_MODEL, top_n=5)
-        def get_parents(docs: List[Document]) -> List[Document]: return get_context_aware_parents(docs, store)
-        compression_retriever = ContextualCompressionRetriever(base_compressor=reranker, base_retriever=ensemble_retriever)
+
+        def get_parents(docs: List[Document]) -> List[Document]:
+            return get_context_aware_parents(docs, store)
+
+        compression_retriever = \
+            ContextualCompressionRetriever(base_compressor=reranker,
+                base_retriever=ensemble_retriever)
         final_retriever = compression_retriever | get_parents
 
         session_id = str(uuid.uuid4())
-        rag_chain, api_key_manager = create_rag_chain(
-            retriever=final_retriever, get_session_history_func=get_session_history,
-            model_name=model_name, temperature=temperature
-        )
-        # Store the float temperature
-        session_data[session_id] = {'chain': rag_chain, 'model_name': model_name, 'temperature': temperature, 'api_key_manager': api_key_manager}
-        
+        (rag_chain, api_key_manager) = \
+            create_rag_chain(retriever=final_retriever,
+                             get_session_history_func=get_session_history,
+                             model_name=model_name,
+                             temperature=temperature)
+
+        session_data[session_id] = {'chain': rag_chain,
+                                    'model_name': model_name,
+                                    'temperature': temperature,
+                                    'api_key_manager': api_key_manager}
+
         success_msg = f"Processed: {', '.join(processed_files)}"
-        if failed_files: success_msg += f". Failed: {', '.join(failed_files)}"
-        
-        # Get the mode label using the STRING key
-        mode_label = TEMPERATURE_LABELS.get(temperature_str, temperature_str)
-        
-        print(f"✅ UPLOAD COMPLETE: Session {session_id} is ready.")
-        # Return all info needed by the frontend
+        if failed_files:
+            success_msg += f". Failed: {', '.join(failed_files)}"
+
+        mode_label = TEMPERATURE_LABELS.get(temperature_str,
+                temperature_str)
+
+        print (f"✅ UPLOAD COMPLETE: Session {session_id} is ready.")
+
         return jsonify({
-            'status': 'success', 
-            'filename': success_msg, 
+            'status': 'success',
+            'filename': success_msg,
             'session_id': session_id,
             'model_name': model_name,
-            'mode': mode_label
-        })
+            'mode': mode_label,
+            })
     except Exception as e:
-        import traceback; traceback.print_exc()
-        return jsonify({'status': 'error', 'message': f'RAG setup failed: {e}'}), 500
+        import traceback
+        traceback.print_exc()
+        return (jsonify({'status': 'error',
+                'message': f'RAG setup failed: {e}'}), 500)
 
-@app.route('/chat', methods=['POST'])
+# --- CORRECTED: Added 'GET' to methods and handle request args ---
+@app.route('/chat', methods=['POST', 'GET'])
 def chat():
-    data = request.get_json()
-    question, session_id = data.get('question'), data.get('session_id') or session.get('session_id')
+    # Handle GET request (used by EventSource)
+    if request.method == 'GET':
+        question = request.args.get('question')
+        session_id = request.args.get('session_id')
+        print(f"Received GET request for chat: session={session_id}, question={question[:50]}...") # Log GET request
+    # Handle POST request (if you ever need it again)
+    elif request.method == 'POST':
+        data = request.get_json()
+        question = data.get('question')
+        session_id = data.get('session_id') or session.get('session_id')
+        print(f"Received POST request for chat: session={session_id}, question={question[:50]}...") # Log POST request
+    else:
+        return (jsonify({'status': 'error', 'message': 'Method not allowed'}), 405)
+
+    # --- Validation ---
+    if not question:
+        error_msg = "Error: No question provided."
+        print(f"❌ CHAT Validation Error: {error_msg}")
+        if request.method == 'GET':
+            # For SSE, need to yield an error event, not return plain text
+            def error_stream():
+                yield f'data: {{"error": "{error_msg}"}}\n\n'
+            return Response(stream_with_context(error_stream()), mimetype='text/event-stream', status=400)
+        return jsonify({'status': 'error','message': error_msg}), 400
 
-    if not question: return jsonify({'status': 'error', 'message': 'No question provided.'}), 400
     if not session_id or session_id not in session_data:
-        print(f"❌ CHAT: Invalid session {session_id}.")
-        return jsonify({'status': 'error', 'message': 'Invalid session. Please upload documents first.'}), 400
-
+        error_msg = "Error: Invalid session. Please upload documents first."
+        print(f"❌ CHAT Validation Error: Invalid session {session_id}.")
+        if request.method == 'GET':
+            def error_stream():
+                yield f'data: {{"error": "{error_msg}"}}\n\n'
+            return Response(stream_with_context(error_stream()), mimetype='text/event-stream', status=400)
+        return jsonify({'status': 'error', 'message': error_msg }), 400
+
+    # --- Process Request ---
     try:
         session_info = session_data[session_id]
         rag_chain = session_info['chain']
-        
-        # --- START: BUGFIX & FEATURE UPDATE ---
-        
-        # 1. Get model name from session
         model_name = session_info['model_name']
-        
-        # 2. Get temperature (float) and convert to string for lookup
         temperature_float = session_info['temperature']
         temperature_str = str(temperature_float)
-        
-        # 3. Get the correct mode label
         mode_label = TEMPERATURE_LABELS.get(temperature_str, temperature_str)
-        
-        # --- END: BUGFIX & FEATURE UPDATE ---
-        
-        print(f"💬 CHAT: Invoking chain for session {session_id}...")
-        answer = rag_chain.invoke({"question": question}, config={"configurable": {"session_id": session_id}})
-        print(f"✅ CHAT: Answer generated.")
-        
-        # Return all info needed by the frontend
-        return jsonify({
-            'answer': answer,
-            'model_name': model_name,
-            'mode': mode_label
-        })
-        
+
+        print (f"💬 CHAT: Streaming response for session {session_id} (Model: {model_name}, Temp: {temperature_float})...")
+
+        def generate_chunks():
+            full_response = ''
+            try:
+                stream_iterator = rag_chain.stream({'question': question},
+                        config={'configurable': {'session_id': session_id}})
+
+                for chunk in stream_iterator:
+                    if isinstance(chunk, str): # Ensure it's a string chunk
+                        full_response += chunk
+                        token_escaped = chunk.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
+                        model_name_escaped = model_name.replace('"', '\\"')
+                        mode_label_escaped = mode_label.replace('"', '\\"')
+                        yield f'data: {{"token": "{token_escaped}", "model_name": "{model_name_escaped}", "mode": "{mode_label_escaped}"}}\n\n'
+                    else:
+                        # Handle potential other types if stream yields non-strings
+                        print(f"⚠️ Received non-string chunk: {type(chunk)}")
+
+
+                print ('✅ CHAT: Streaming finished successfully.')
+                # Optionally update session history or store full response if needed later
+                # get_session_history(session_id).add_ai_message(full_response)
+
+            except Exception as e:
+                print(f"❌ CHAT Error during streaming generation: {e}")
+                import traceback
+                traceback.print_exc()
+                error_msg = f"Error during response generation: {str(e)}".replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
+                yield f'data: {{"error": "{error_msg}"}}\n\n'
+
+        # Return the streaming response
+        return Response(stream_with_context(generate_chunks()), mimetype='text/event-stream')
+
     except Exception as e:
-        import traceback; traceback.print_exc()
-        return jsonify({'status': 'error', 'message': f'Error during chat: {e}'}), 500
+        # Catch errors during setup before streaming starts
+        print(f"❌ CHAT Setup Error: {e}")
+        import traceback
+        traceback.print_exc()
+        error_msg = f"Error setting up chat stream: {str(e)}"
+        if request.method == 'GET':
+             def error_stream():
+                yield f'data: {{"error": "{error_msg.replace("\"", "\\\"").replace("n", "\\n")}"}}\n\n'
+             return Response(stream_with_context(error_stream()), mimetype='text/event-stream', status=500)
+        return (jsonify({'status': 'error', 'message': error_msg}), 500)
+
 
 def clean_markdown_for_tts(text: str) -> str:
-    text = re.sub(r'\*(\*?)(.*?)\1\*', r'\2', text); text = re.sub(r'\_(.*?)\_', r'\1', text)
-    text = re.sub(r'`(.*?)`', r'\1', text); text = re.sub(r'^\s*#{1,6}\s+', '', text, flags=re.MULTILINE)
-    text = re.sub(r'^\s*[\*\-]\s+', '', text, flags=re.MULTILINE); text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE)
-    text = re.sub(r'^\s*>\s?', '', text, flags=re.MULTILINE); text = re.sub(r'^\s*[-*_]{3,}\s*$', '', text, flags=re.MULTILINE)
-    text = re.sub(r'\n+', ' ', text)
+    # --- Simplified cleaning for TTS ---
+    text = re.sub(r'\[.*?\]\(.*?\)', '', text) # Remove links
+    text = re.sub(r'[`*_#]', '', text) # Remove formatting chars
+    text = re.sub(r'^\s*[\-\*\+]\s+', '', text, flags=re.MULTILINE) # Remove list markers
+    text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE) # Remove numbered list markers
+    text = re.sub(r'^\s*>\s?', '', text, flags=re.MULTILINE) # Remove blockquote markers
+    text = re.sub(r'\n+', ' ', text) # Replace newlines with spaces
+    text = re.sub(r'\s{2,}', ' ', text) # Collapse multiple spaces
     return text.strip()
 
+
 @app.route('/tts', methods=['POST'])
 def text_to_speech():
-    data = request.get_json(); text = data.get('text')
-    if not text: return jsonify({'status': 'error', 'message': 'No text provided.'}), 400
+    data = request.get_json()
+    text = data.get('text')
+    if not text:
+        return (jsonify({'status': 'error',
+                'message': 'No text provided.'}), 400)
     try:
-        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)
+        clean_text = clean_markdown_for_tts(text)
+        if not clean_text:
+             return (jsonify({'status': 'error', 'message': 'No speakable text found.'}), 400)
+
+        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"❌ TTS Error: {e}")
-        return jsonify({'status': 'error', 'message': 'Failed to generate audio.'}), 500
+        print (f"❌ TTS Error: {e}")
+        return (jsonify({'status': 'error',
+                'message': 'Failed to generate audio.'}), 500)
+
 
 if __name__ == '__main__':
-    port = int(os.environ.get("PORT", 7860))
-    print(f"🚀 Starting Flask app on port {port}")
-    app.run(host="0.0.0.0", port=port, debug=False, threaded=False)
+    port = int(os.environ.get('PORT', 7860))
+    print (f"🚀 Starting Flask app on port {port}")
+    # Use threaded=True for better handling of concurrent requests during streaming
+    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)

rag_processor.py

@@ -231,7 +231,7 @@ RESPONSE GUIDELINES:
 - Add blank lines between sections for readability
 
 **Source Citation:**
-- Always cite information using: [Source: filename, Page: X]
+- Always cite information using: [Source: filename, Page: X] and cite it at the end of the entire answer only
 - Place citations at the end of your final answer only 
 - Do not cite sources within the body of your answer
 - Multiple sources: [Source: doc1.pdf, Page: 3; doc2.pdf, Page: 7]

templates/index.html

@@ -1,500 +1,821 @@
 <!DOCTYPE html>
-<html lang="en">
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>CogniChat - Chat with your Documents</title>
-    <script src="https://cdn.tailwindcss.com"></script>
-    <link rel="preconnect" href="https://fonts.googleapis.com">
-    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
-    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=Google+Sans:wght@400;500;700&family=Roboto:wght@400;500&display=swap" rel="stylesheet">
-    <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
-    <style>
-        :root {
-            --background: #f0f4f9;
-            --foreground: #1f1f1f;
-            --primary: #1a73e8;
-            --primary-hover: #1867cf;
-            --card: #ffffff;
-            --card-border: #dadce0;
-            --input-bg: #e8f0fe;
-            --user-bubble: #d9e7ff;
-            --bot-bubble: #f1f3f4;
-            --select-bg: #ffffff;
-            --select-border: #dadce0;
-            --select-text: #1f1f1f;
-        }
-
-        .dark {
-            --background: #111827;
-            --foreground: #e5e7eb;
-            --primary: #3b82f6;
-            --primary-hover: #60a5fa;
-            --card: #1f2937;
-            --card-border: #4b5563;
-            --input-bg: #374151;
-            --user-bubble: #374151;
-            --bot-bubble: #374151;
-            --select-bg: #374151;
-            --select-border: #6b7280;
-            --select-text: #f3f4f6;
-        }
-        
-        body {
-            font-family: 'Inter', 'Google Sans', 'Roboto', sans-serif;
-            background-color: var(--background);
-            color: var(--foreground);
-            overflow: hidden;
-        }
-        
-        #chat-window::-webkit-scrollbar { width: 8px; }
-        #chat-window::-webkit-scrollbar-track { background: transparent; }
-        #chat-window::-webkit-scrollbar-thumb { background-color: #4b5563; border-radius: 20px; }
-        .dark #chat-window::-webkit-scrollbar-thumb { background-color: #5f6368; }
-
-        .drop-zone--over {
-            border-color: var(--primary);
-            box-shadow: 0 0 20px rgba(59, 130, 246, 0.4);
-        }
-        
-        .loader {
-            width: 48px;
-            height: 48px;
-            border: 3px solid var(--card-border);
-            border-radius: 50%;
-            display: inline-block;
-            position: relative;
-            box-sizing: border-box;
-            animation: rotation 1s linear infinite;
-        }
-        .loader::after {
-            content: '';  
-            box-sizing: border-box;
+ <html lang="en">
+ <head>
+     <meta charset="UTF-8">
+     <meta name="viewport" content="width=device-width, initial-scale=1.0">
+     <title>CogniChat - Chat with your Documents</title>
+     <script src="https://cdn.tailwindcss.com"></script>
+     <link rel="preconnect" href="https://fonts.googleapis.com">
+     <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+     <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=Google+Sans:wght@400;500;700&family=Roboto:wght@400;500&display=swap" rel="stylesheet">
+     <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
+     <style>
+         :root {
+             --background: #f0f4f9;
+             --foreground: #1f1f1f;
+             --primary: #1a73e8;
+             --primary-hover: #1867cf;
+             --card: #ffffff;
+             --card-border: #dadce0;
+             --input-bg: #e8f0fe;
+             --user-bubble: #d9e7ff;
+             --bot-bubble: #f1f3f4;
+             --select-bg: #ffffff;
+             --select-border: #dadce0;
+             --select-text: #1f1f1f;
+         }
+
+         .dark {
+             --background: #111827; /* Darker background */
+             --foreground: #e5e7eb;
+             --primary: #3b82f6; /* Adjusted primary blue */
+             --primary-hover: #60a5fa; /* Lighter hover blue */
+             --card: #1f2937;    /* Dark card background */
+             --card-border: #4b5563; /* Greyer border */
+             --input-bg: #374151;    /* Darker input background */
+             --user-bubble: #374151; /* Darker user bubble */
+             --bot-bubble: #374151;  /* Darker bot bubble */
+             --select-bg: #374151;
+             --select-border: #6b7280;
+             --select-text: #f3f4f6;
+             --code-bg: #2d2d2d;     /* Specific background for code blocks */
+             --code-text: #d4d4d4;   /* Light grey text for code */
+             --copy-btn-bg: #4a4a4a;
+             --copy-btn-hover-bg: #5a5a5a;
+             --copy-btn-text: #e0e0e0;
+         }
+
+         body {
+             font-family: 'Inter', 'Google Sans', 'Roboto', sans-serif;
+             background-color: var(--background);
+             color: var(--foreground);
+             overflow: hidden; /* Prevent body scroll */
+         }
+
+         #chat-window {
+             scroll-behavior: smooth; /* Ensure smooth programatic scroll */
+         }
+         #chat-window::-webkit-scrollbar { width: 8px; }
+         #chat-window::-webkit-scrollbar-track { background: transparent; }
+         #chat-window::-webkit-scrollbar-thumb { background-color: #4b5563; border-radius: 20px; }
+         .dark #chat-window::-webkit-scrollbar-thumb { background-color: #5f6368; }
+
+         .drop-zone--over {
+             border-color: var(--primary);
+             box-shadow: 0 0 20px rgba(59, 130, 246, 0.4);
+         }
+
+         .loader {
+             width: 48px;
+             height: 48px;
+             border: 3px solid var(--card-border);
+             border-radius: 50%;
+             display: inline-block;
+             position: relative;
+             box-sizing: border-box;
+             animation: rotation 1s linear infinite;
+         }
+         .loader::after {
+             content: '';
+             box-sizing: border-box;
+             position: absolute;
+             left: 50%;
+             top: 50%;
+             transform: translate(-50%, -50%);
+             width: 56px;
+             height: 56px;
+             border-radius: 50%;
+             border: 3px solid;
+             border-color: var(--primary) transparent;
+         }
+
+         @keyframes rotation {
+             0% { transform: rotate(0deg); }
+             100% { transform: rotate(360deg); }
+         }
+
+         /* --- Updated Typing Indicator --- */
+         .typing-indicator {
+             display: inline-flex; /* Changed to inline-flex */
+             align-items: center;
+             padding: 8px 0; /* Add some vertical padding */
+         }
+         .typing-indicator span {
+             height: 8px; /* Slightly smaller dots */
+             width: 8px;
+             margin: 0 2px;
+             background-color: #9E9E9E;
+             border-radius: 50%;
+             opacity: 0; /* Start invisible */
+             animation: typing-pulse 1.4s infinite ease-in-out;
+         }
+         .typing-indicator span:nth-child(1) { animation-delay: 0s; }
+         .typing-indicator span:nth-child(2) { animation-delay: 0.2s; }
+         .typing-indicator span:nth-child(3) { animation-delay: 0.4s; }
+
+         @keyframes typing-pulse {
+             0%, 100% { opacity: 0; transform: scale(0.7); }
+             50% { opacity: 1; transform: scale(1); }
+         }
+        /* --- End Typing Indicator --- */
+
+        /* --- Updated Markdown Styling --- */
+         .markdown-content { /* Base styles for the content area */
+             line-height: 1.75;
+         }
+         .markdown-content p { margin-bottom: 1rem; }
+         .markdown-content h1, .markdown-content h2, .markdown-content h3,
+         .markdown-content h4, .markdown-content h5, .markdown-content h6 {
+             font-weight: 600;
+             margin-top: 1.5rem;
+             margin-bottom: 0.75rem;
+             line-height: 1.3;
+         }
+         .markdown-content h1 { font-size: 1.5em; border-bottom: 1px solid var(--card-border); padding-bottom: 0.3rem;}
+         .markdown-content h2 { font-size: 1.25em; }
+         .markdown-content h3 { font-size: 1.1em; }
+         .markdown-content ul, .markdown-content ol { padding-left: 1.75rem; margin-bottom: 1rem; }
+         .markdown-content li { margin-bottom: 0.5rem; }
+         .markdown-content a { color: var(--primary); text-decoration: none; font-weight: 500; }
+         .markdown-content a:hover { text-decoration: underline; }
+         .markdown-content strong, .markdown-content b { font-weight: 600; } /* Ensure bold works */
+         .markdown-content blockquote {
+             border-left: 4px solid var(--card-border);
+             padding-left: 1rem;
+             margin-left: 0;
+             margin-bottom: 1rem;
+             color: #a0aec0; /* Lighter text for quotes */
+         }
+        /* --- Code Block Styling --- */
+         .markdown-content pre {
+             position: relative;
+             background-color: var(--code-bg);
+             border: 1px solid var(--card-border);
+             border-radius: 0.5rem;
+             margin-bottom: 1rem;
+             font-size: 0.9em;
+             color: var(--code-text);
+             overflow: hidden; /* Hide horizontal overflow until hovered/focused */
+         }
+         .markdown-content pre code {
+             display: block;
+             padding: 1rem;
+             overflow-x: auto; /* Enable horizontal scroll on the code itself */
+             background: none !important; /* Override potential highlight.js background */
+             font-family: Consolas, Monaco, 'Andale Mono', 'Ubuntu Mono', monospace;
+             white-space: pre; /* Ensure whitespace is preserved */
+         }
+          /* --- Copy Button Styling --- */
+         .markdown-content pre .copy-code-btn {
             position: absolute;
-            left: 50%;
-            top: 50%;
-            transform: translate(-50%, -50%);
-            width: 56px;
-            height: 56px;
-            border-radius: 50%;
-            border: 3px solid;
-            border-color: var(--primary) transparent;
-        }
-
-        @keyframes rotation {
-            0% { transform: rotate(0deg); }
-            100% { transform: rotate(360deg); }
-        }
-        
-        .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-content p { margin-bottom: 1rem; line-height: 1.75; }
-        .markdown-content h1, .markdown-content h2, .markdown-content h3 { font-weight: 600; margin-top: 1.5rem; margin-bottom: 0.75rem; line-height: 1.3; }
-        .markdown-content h1 { font-size: 1.5em; border-bottom: 1px solid var(--card-border); padding-bottom: 0.3rem;}
-        .markdown-content h2 { font-size: 1.25em; }
-        .markdown-content h3 { font-size: 1.1em; }
-        .markdown-content ul, .markdown-content ol { padding-left: 1.75rem; margin-bottom: 1rem; }
-        .markdown-content li { margin-bottom: 0.5rem; }
-        .markdown-content a { color: var(--primary); text-decoration: none; font-weight: 500; }
-        .markdown-content pre { position: relative; background-color: #2e2f32; border: 1px solid var(--card-border); border-radius: 0.5rem; margin-bottom: 1rem; font-size: 0.9em;}
-        .markdown-content pre code { background: none; padding: 1rem; display: block; overflow-x: auto; }
-        .markdown-content pre .copy-code-btn { position: absolute; top: 0.5rem; right: 0.5rem; background-color: #3c4043; border: 1px solid #5f6368; color: #e8eaed; padding: 0.3rem 0.6rem; border-radius: 0.25rem; cursor: pointer; opacity: 0; transition: opacity 0.2s; font-size: 0.8em;}
-        .markdown-content pre:hover .copy-code-btn { opacity: 1; }
-        
-        .tts-button-loader {
-            width: 16px;
-            height: 16px;
-            border: 2px solid currentColor;
-            border-radius: 50%;
-            display: inline-block;
-            box-sizing: border-box;
-            animation: rotation 0.8s linear infinite;
-            border-bottom-color: transparent;
-        }
-
-        .select-wrapper {
-            position: relative;
-        }
-        .select-wrapper select {
-            background-color: var(--select-bg);
-            border: 1px solid var(--select-border);
-            color: var(--select-text);
-            padding: 0.75rem 2.5rem 0.75rem 1rem;
-            border-radius: 0.75rem;
-            font-size: 0.875rem;
-            width: 100%;
-            appearance: none;
-            -webkit-appearance: none;
-            transition: all 0.2s ease-in-out;
+            top: 0.5rem;
+            right: 0.5rem;
+            background-color: var(--copy-btn-bg);
+            border: 1px solid var(--card-border);
+            color: var(--copy-btn-text);
+            padding: 0.3rem 0.6rem;
+            border-radius: 0.25rem;
             cursor: pointer;
-            background-image: url("data:image/svg+xml,%3csvg xmlns='http://www.w3.org/2000/svg' fill='none' viewBox='0 0 20 20'%3e%3cpath stroke='%239ca3af' stroke-linecap='round' stroke-linejoin='round' stroke-width='1.5' d='M6 8l4 4 4-4'/%3e%3c/svg%3e");
-            background-position: right 0.75rem center;
-            background-repeat: no-repeat;
-            background-size: 1.25em 1.25em;
-        }
-    </style>
-</head>
-<body class="w-screen h-screen dark">
-    <main id="main-content" class="h-full flex flex-col transition-opacity duration-500">
-        <div id="chat-container" class="hidden flex-1 flex flex-col w-full mx-auto overflow-hidden">
-            
-            <!-- --- CORRECT HEADER (Center/Right Layout) --- -->
-            <header class="p-4 border-b border-[var(--card-border)] flex-shrink-0 flex justify-between items-center w-full">
-                <div class="w-1/4"></div>
-                <div class="w-1/2 text-center">
-                    <h1 class="text-xl font-medium tracking-wide">CogniChat</h1>
-                    <p id="chat-filename" class="text-xs text-gray-400 mt-1 truncate"></p>
-                </div>
-                <div id="chat-session-info" class="w-1/4 text-right text-xs">
-                    <!-- This will be populated by JavaScript -->
-                </div>
-            </header>
-            <!-- --- END 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"></div>
-            </div>
-            <div class="p-4 flex-shrink-0 bg-opacity-50 backdrop-blur-md 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-lg border border-[var(--card-border)] focus-within:ring-2 focus-within:ring-[var(--primary)] transition-all">
-                    <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.5 rounded-full transition-all duration-200 disabled:opacity-50 disabled:cursor-not-allowed" title="Send">
-                        <svg class="w-5 h-5" viewBox="0 0 20 20" fill="currentColor"><path fill-rule="evenodd" d="M10 18a8 8 0 100-16 8 8 0 000 16zm3.707-8.707l-3-3a1 1 0 00-1.414 1.414L10.586 9H7a1 1 0 100 2h3.586l-1.293 1.293a1 1 0 101.414 1.414l3-3a1 1 0 000-1.414z" clip-rule="evenodd"></path></svg>
-                    </button>
-                </form>
-            </div>
-        </div>
-
-        <div id="upload-container" class="flex-1 flex flex-col items-center justify-center p-8 transition-opacity duration-300">
-            <div class="text-center max-w-xl w-full">
-                <h1 class="text-5xl font-bold mb-3 tracking-tight">CogniChat</h1>
-                <p class="text-lg text-gray-400 mb-8">Upload your documents to start a conversation.</p>
-                <div class="mb-8 p-5 bg-[var(--card)] rounded-2xl border border-[var(--card-border)] shadow-lg">
-                    <div class="flex flex-col sm:flex-row items-center gap-6">
-                        <div class="w-full sm:w-1/2">
-                            <div class="flex items-center gap-2 mb-2">
-                                <svg class="w-5 h-5 text-gray-400" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 20 20" fill="currentColor"><path d="M7 3a1 1 0 000 2h6a1 1 0 100-2H7zM4 7a1 1 0 011-1h10a1 1 0 110 2H5a1 1 0 01-1-1zM2 11a2 2 0 012-2h12a2 2 0 012 2v4a2 2 0 01-2 2H4a2 2 0 01-2-2v-4z" /></svg>
-                                <label for="model-select" class="block text-sm font-medium text-gray-300">Model</label>
-                            </div>
-                            <div class="select-wrapper">
-                                <select id="model-select" name="model_name">
-                                    <option value="moonshotai/kimi-k2-instruct" selected>Kimi Instruct</option>
-                                    <option value="openai/gpt-oss-20b">GPT OSS 20b</option>
-                                    <option value="llama-3.3-70b-versatile">Llama 3.3 70b</option>
-                                    <option value="llama-3.1-8b-instant">Llama 3.1 8b Instant</option>
-                                </select>
-                            </div>
-                        </div>
-                        <div class="w-full sm:w-1/2">
+            opacity: 0; /* Initially hidden */
+            transition: opacity 0.2s, background-color 0.2s;
+            font-size: 0.8em;
+            display: flex; /* For icon alignment */
+            align-items: center;
+            gap: 0.25rem;
+         }
+         .markdown-content pre .copy-code-btn:hover {
+             background-color: var(--copy-btn-hover-bg);
+         }
+         .markdown-content pre:hover .copy-code-btn {
+             opacity: 1; /* Show on hover */
+         }
+        /* --- Inline Code Styling --- */
+         .markdown-content code:not(pre code) {
+             background-color: rgba(110, 118, 129, 0.4);
+             padding: 0.2em 0.4em;
+             margin: 0 0.1em; /* Add slight horizontal margin */
+             font-size: 85%;
+             border-radius: 6px;
+             font-family: Consolas, Monaco, 'Andale Mono', 'Ubuntu Mono', monospace;
+         }
+        /* --- End Markdown Styling --- */
+
+         .tts-button-loader {
+             width: 16px;
+             height: 16px;
+             border: 2px solid currentColor;
+             border-radius: 50%;
+             display: inline-block;
+             box-sizing: border-box;
+             animation: rotation 0.8s linear infinite;
+             border-bottom-color: transparent;
+         }
+
+         .select-wrapper {
+             position: relative;
+         }
+         .select-wrapper select {
+             background-color: var(--select-bg);
+             border: 1px solid var(--select-border);
+             color: var(--select-text);
+             padding: 0.75rem 2.5rem 0.75rem 1rem;
+             border-radius: 0.75rem;
+             font-size: 0.875rem;
+             width: 100%;
+             appearance: none;
+             -webkit-appearance: none;
+             transition: all 0.2s ease-in-out;
+             cursor: pointer;
+             background-image: url("data:image/svg+xml,%3csvg xmlns='http://www.w3.org/2000/svg' fill='none' viewBox='0 0 20 20'%3e%3cpath stroke='%239ca3af' stroke-linecap='round' stroke-linejoin='round' stroke-width='1.5' d='M6 8l4 4 4-4'/%3e%3c/svg%3e");
+             background-position: right 0.75rem center;
+             background-repeat: no-repeat;
+             background-size: 1.25em 1.25em;
+         }
+         /* Add styles for prose class if needed, tailwind typography plugin is good for this */
+         /* Example: .prose dark:prose-invert ... */
+     </style>
+ </head>
+ <body class="w-screen h-screen dark">
+     <main id="main-content" class="h-full flex flex-col transition-opacity duration-500">
+         <div id="chat-container" class="hidden flex-1 flex flex-col w-full mx-auto overflow-hidden">
+
+             <header class="p-4 border-b border-[var(--card-border)] flex-shrink-0 flex justify-between items-center w-full">
+                 <div class="w-1/4"></div> <div class="w-1/2 text-center">
+                     <h1 class="text-xl font-medium tracking-wide">CogniChat ✨</h1>
+                     <p id="chat-filename" class="text-xs text-gray-400 mt-1 truncate"></p>
+                 </div>
+                 <div id="chat-session-info" class="w-1/4 text-right text-xs space-y-1 pr-4">
+                     </div>
+             </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"></div>
+             </div>
+             <div class="p-4 flex-shrink-0 bg-opacity-50 backdrop-blur-md 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-lg border border-[var(--card-border)] focus-within:ring-2 focus-within:ring-[var(--primary)] transition-all">
+                     <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.5 rounded-full transition-all duration-200 disabled:opacity-50 disabled:cursor-not-allowed" title="Send">
+                         <svg class="w-5 h-5" viewBox="0 0 20 20" fill="currentColor"><path fill-rule="evenodd" d="M10 18a8 8 0 100-16 8 8 0 000 16zm3.707-8.707l-3-3a1 1 0 00-1.414 1.414L10.586 9H7a1 1 0 100 2h3.586l-1.293 1.293a1 1 0 101.414 1.414l3-3a1 1 0 000-1.414z" clip-rule="evenodd"></path></svg>
+                     </button>
+                 </form>
+             </div>
+         </div>
+
+         <div id="upload-container" class="flex-1 flex flex-col items-center justify-center p-8 transition-opacity duration-300">
+             <div class="text-center max-w-xl w-full">
+                 <h1 class="text-5xl font-bold mb-3 tracking-tight">CogniChat ✨</h1>
+                 <p class="text-lg text-gray-400 mb-8">Upload your documents to start a conversation.</p>
+                 <div class="mb-8 p-5 bg-[var(--card)] rounded-2xl border border-[var(--card-border)] shadow-lg">
+                     <div class="flex flex-col sm:flex-row items-center gap-6">
+                         <div class="w-full sm:w-1/2">
                              <div class="flex items-center gap-2 mb-2">
-                                <svg class="w-5 h-5 text-gray-400" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 20 20" fill="currentColor"><path fill-rule="evenodd" d="M5.5 16a3.5 3.5 0 100-7 3.5 3.5 0 000 7zM12 5.5a3.5 3.5 0 11-7 0 3.5 3.5 0 017 0zM14.5 16a3.5 3.5 0 100-7 3.5 3.5 0 000 7z" clip-rule="evenodd" /></svg>
-                                <label for="temperature-select" class="block text-sm font-medium text-gray-300">Mode</label>
-                            </div>
+                                 <svg class="w-5 h-5 text-gray-400" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 20 20" fill="currentColor"><path d="M7 3a1 1 0 000 2h6a1 1 0 100-2H7zM4 7a1 1 0 011-1h10a1 1 0 110 2H5a1 1 0 01-1-1zM2 11a2 2 0 012-2h12a2 2 0 012 2v4a2 2 0 01-2 2H4a2 2 0 01-2-2v-4z" /></svg>
+                                 <label for="model-select" class="block text-sm font-medium text-gray-300">Model</label>
+                             </div>
                              <div class="select-wrapper">
-                                <select id="temperature-select" name="temperature">
-                                    <option value="0.2" selected>0.2 - Precise</option>
-                                    <option value="0.4">0.4 - Confident</option>
-                                    <option value="0.6">0.6 - Balanced</option>
-                                    <option value="0.8">0.8 - Flexible</option>
-                                    <option value="1.0">1.0 - Creative</option>
-                                </select>
-                            </div>
-                        </div>
-                    </div>
-                    <p class="text-xs text-gray-500 mt-4 text-center">Higher creativity modes may reduce factual accuracy.</p>
-                </div>
-                <div id="drop-zone" class="w-full text-center border-2 border-dashed border-[var(--card-border)] rounded-2xl p-10 transition-all duration-300 cursor-pointer hover:bg-[var(--card)] hover:border-[var(--primary)]">
-                    <div class="flex flex-col items-center justify-center pointer-events-none">
-                        <svg class="mx-auto h-12 w-12 text-gray-500" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="1.5" 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 text-gray-400">Drag & drop files or <span class="text-[var(--primary)] font-semibold">click to upload</span></p>
-                        <p class="text-xs text-gray-400 mt-1">Supports PDF, DOCX, TXT</p>
-                        <p id="file-name" class="mt-2 text-xs text-gray-500"></p>
-                    </div>
-                    <input id="file-upload" type="file" class="hidden" accept=".pdf,.txt,.docx" multiple>
-                </div>
-            </div>
-        </div>
-        
-        <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 class="loader"></div>
-            <p id="loading-text" class="mt-6 text-sm font-medium"></p>
-            <p id="loading-subtext" class="mt-2 text-xs text-gray-400"></p>
-        </div>
-    </main>
-
-    <script>
-        document.addEventListener('DOMContentLoaded', () => {
-            const uploadContainer = document.getElementById('upload-container');
-            const chatContainer = document.getElementById('chat-container');
-            const dropZone = document.getElementById('drop-zone');
-            const fileUploadInput = document.getElementById('file-upload');
-            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');
-            const chatSubmitBtn = document.getElementById('chat-submit-btn');
-            const chatWindow = document.getElementById('chat-window');
-            const chatContent = document.getElementById('chat-content');
-            const modelSelect = document.getElementById('model-select');
-            const temperatureSelect = document.getElementById('temperature-select');
-            const chatFilename = document.getElementById('chat-filename');
-            const chatSessionInfo = document.getElementById('chat-session-info');
-
-            let sessionId = sessionStorage.getItem('cognichat_session_id');
-
-            dropZone.addEventListener('click', () => fileUploadInput.click());
-            
-            ['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => {
-                dropZone.addEventListener(eventName, e => {e.preventDefault(); e.stopPropagation();});
-            });
-            ['dragenter', 'dragover'].forEach(eventName => dropZone.addEventListener(eventName, () => dropZone.classList.add('drop-zone--over')));
-            ['dragleave', 'drop'].forEach(eventName => dropZone.addEventListener(eventName, () => dropZone.classList.remove('drop-zone--over')));
-            
-            dropZone.addEventListener('drop', (e) => {
-                if (e.dataTransfer.files.length > 0) handleFiles(e.dataTransfer.files);
-            });
-            fileUploadInput.addEventListener('change', (e) => {
-                if (e.target.files.length > 0) handleFiles(e.target.files);
-            });
-
-            async function handleFiles(files) {
-                const formData = new FormData();
-                let fileNames = Array.from(files).map(f => f.name);
-                for (const file of files) { formData.append('file', file); }
-
-                formData.append('model_name', modelSelect.value);
-                formData.append('temperature', temperatureSelect.value);
-
-                fileNameSpan.textContent = `Selected: ${fileNames.join(', ')}`;
-                await uploadAndProcessFiles(formData);
-            }
-
-            async function uploadAndProcessFiles(formData) {
-                loadingOverlay.classList.remove('hidden');
-                loadingText.textContent = `Processing document(s)...`;
-                loadingSubtext.textContent = "Creating a knowledge base may take a minute. Please hold on tight!";
-                
-                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.');
-                    
-                    sessionId = result.session_id;
-                    sessionStorage.setItem('cognichat_session_id', sessionId);
-                    
-                    chatFilename.innerHTML = `Chatting with: <strong>${result.filename}</strong>`;
-                    chatFilename.title = result.filename;
-
-                    chatSessionInfo.innerHTML = `
-                        <span class="text-gray-500 dark:text-gray-500 italic block hover:text-gray-300 transition-colors cursor-pointer" onclick="location.reload()">
-                            Refresh to change settings
-                        </span>`;
-                    
-                    const modelOption = modelSelect.querySelector(`option[value="${result.model_name}"]`);
-                    const simpleModelName = modelOption ? modelOption.textContent : result.model_name;
-
-                    const modelInfo = {
-                        model: result.model_name,
-                        mode: result.mode,
-                        simpleModelName: simpleModelName
-                    };
-
-                    uploadContainer.classList.add('hidden');
-                    chatContainer.classList.remove('hidden');
-                    appendMessage("I've analyzed your documents. What would you like to know?", "bot", modelInfo);
-
-                } catch (error) {
-                    console.error('Upload error:', error);
-                    alert(`Error: ${error.message}`);
-                } finally {
-                    loadingOverlay.classList.add('hidden');
-                    fileNameSpan.textContent = '';
-                    fileUploadInput.value = '';
-                }
-            }
-            
-            chatForm.addEventListener('submit', async (e) => {
-                e.preventDefault();
-                const question = chatInput.value.trim();
-                if (!question) return;
-
-                appendMessage(question, 'user');
-                chatInput.value = '';
-                chatInput.disabled = true;
-                chatSubmitBtn.disabled = true;
-
-                const typingIndicator = showTypingIndicator();
-                
-                try {
-                    const response = await fetch('/chat', {
-                        method: 'POST',
-                        headers: { 'Content-Type': 'application/json' },
-                        body: JSON.stringify({ question, session_id: sessionId }),
-                    });
-                    if (!response.ok) throw new Error(`Server error: ${response.statusText}`);
-
-                    const result = await response.json();
-                    
-                    const modelOption = modelSelect.querySelector(`option[value="${result.model_name}"]`);
-                    const simpleModelName = modelOption ? modelOption.textContent : result.model_name;
-                    const modelInfo = {
-                        model: result.model_name,
-                        mode: result.mode,
-                        simpleModelName: simpleModelName
-                    };
-                    
-                    typingIndicator.remove();
-                    const botMessageContainer = appendMessage('', 'bot', modelInfo); 
-                    const contentDiv = botMessageContainer.querySelector('.markdown-content');
-                    contentDiv.innerHTML = marked.parse(result.answer);
-                    contentDiv.querySelectorAll('pre').forEach(addCopyButton);
-                    scrollToBottom();
-                    addTextToSpeechControls(botMessageContainer, result.answer); 
-                } catch (error) {
-                    console.error('Chat error:', error);
-                    typingIndicator.remove();
-                    appendMessage(`Error: ${error.message}`, 'bot');
-                } finally {
-                    chatInput.disabled = false;
-                    chatSubmitBtn.disabled = false;
-                    chatInput.focus();
-                }
-            });
-
-            // --- FINAL, CORRECT appendMessage function ---
-            function appendMessage(text, sender, modelInfo = null) {
-                const messageWrapper = document.createElement('div');
-                const iconSVG = sender === 'user' 
-                    ? `<div class="bg-blue-200 dark:bg-gray-700 p-2.5 rounded-full flex-shrink-0 mt-1"><svg class="w-5 h-5 text-blue-700 dark:text-blue-300" viewBox="0 0 24 24"><path fill="currentColor" d="M12 12c2.21 0 4-1.79 4-4s-1.79-4-4-4-4 1.79-4 4 1.79 4 4 4zm0 2c-2.67 0-8 1.34-8 4v2h16v-2c0-2.66-5.33-4-8-4z"></path></svg></div>` 
-                    : `<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>`;
-                
-                let senderHTML;
-                if (sender === 'user') {
-                    senderHTML = '<p class="font-medium text-sm mb-1">You</p>';
-                } else {
-                    let modelInfoHTML = '';
-                    if (modelInfo && modelInfo.simpleModelName) {
-                        modelInfoHTML = `
-                            <span class="ml-2 text-xs font-normal text-gray-400">
-                                (Model: ${modelInfo.simpleModelName} • Mode: ${modelInfo.mode})
-                            </span>
-                        `;
-                    }
-                    senderHTML = `<div class="font-medium text-sm mb-1 flex items-center">CogniChat ${modelInfoHTML}</div>`;
-                }
-
-                messageWrapper.className = `flex items-start gap-4`;
-                messageWrapper.innerHTML = `
-                    ${iconSVG}
-                    <div class="flex-1 pt-1">
-                        ${senderHTML}
-                        <div class="text-base markdown-content">${text ? marked.parse(text) : ''}</div>
-                        <div class="tts-controls mt-2"></div>
-                    </div>
-                `;
-                chatContent.appendChild(messageWrapper);
-                scrollToBottom();
-                return messageWrapper.querySelector('.flex-1');
-            }
-
-            function showTypingIndicator() {
-                const indicator = document.createElement('div');
-                indicator.id = 'typing-indicator';
-                indicator.className = `flex items-start gap-4`;
-                indicator.innerHTML = `<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><div class="flex-1 pt-1"><p class="font-medium text-sm mb-1">CogniChat is thinking...</p><div class="typing-indicator"><span></span><span></span><span></span></div></div>`;
-                chatContent.appendChild(indicator);
-                scrollToBottom();
-                return indicator;
-            }
-            
-            function scrollToBottom() { chatWindow.scrollTo({ top: chatWindow.scrollHeight, behavior: 'smooth' }); }
-            
-            function addCopyButton(pre) {
-                const button = document.createElement('button');
-                button.className = 'copy-code-btn';
-                button.textContent = 'Copy';
-                pre.appendChild(button);
-                button.addEventListener('click', () => {
-                    navigator.clipboard.writeText(pre.querySelector('code').innerText)
-                        .then(() => {
-                            button.textContent = 'Copied!';
-                            setTimeout(() => button.textContent = 'Copy', 2000);
-                        });
-                });
-            }
-
-            // (TTS functions remain unchanged)
-            let currentAudio, currentPlayingButton;
-            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) {
-                if (!text.trim()) return;
-                const speakButton = document.createElement('button');
-                speakButton.className = 'speak-btn mt-2 px-3 py-1.5 bg-blue-700 text-white rounded-full text-sm font-medium hover:bg-blue-800 transition-colors flex items-center gap-2 disabled:opacity-50';
-                speakButton.title = 'Listen to this message';
-                speakButton.innerHTML = `${playIconSVG} <span>Listen</span>`;
-                messageBubble.querySelector('.tts-controls').appendChild(speakButton);
-                speakButton.addEventListener('click', () => handleTTS(text, speakButton));
-            }
-
-            async function handleTTS(text, button) {
-                if (button === currentPlayingButton) {
-                    if (currentAudio && !currentAudio.paused) {
-                        currentAudio.pause();
-                        button.innerHTML = `${playIconSVG} <span>Listen</span>`;
-                    } else if (currentAudio && currentAudio.paused) {
-                        currentAudio.play();
-                        button.innerHTML = `${pauseIconSVG} <span>Pause</span>`;
-                    }
-                    return;
-                }
-                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 }) });
-                    if (!response.ok) throw new Error('Failed to generate audio.');
-                    const blob = await response.blob();
-                    currentAudio = new Audio(URL.createObjectURL(blob));
-                    currentAudio.play();
-                    button.innerHTML = `${pauseIconSVG} <span>Pause</span>`;
-                    currentAudio.onended = resetAllSpeakButtons;
-                } catch (error) {
-                    console.error('TTS Error:', error);
-                    resetAllSpeakButtons();
-                } 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;
-            }
-        });
-    </script>
-</body>
-</html>
+                                 <select id="model-select" name="model_name">
+                                     <option value="moonshotai/kimi-k2-instruct" selected>Kimi Instruct</option>
+                                     <option value="openai/gpt-oss-20b">GPT OSS 20b</option>
+                                     <option value="llama-3.3-70b-versatile">Llama 3.3 70b</option>
+                                     <option value="llama-3.1-8b-instant">Llama 3.1 8b Instant</option>
+                                 </select>
+                             </div>
+                         </div>
+                         <div class="w-full sm:w-1/2">
+                              <div class="flex items-center gap-2 mb-2">
+                                 <svg class="w-5 h-5 text-gray-400" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 20 20" fill="currentColor"><path fill-rule="evenodd" d="M5.5 16a3.5 3.5 0 100-7 3.5 3.5 0 000 7zM12 5.5a3.5 3.5 0 11-7 0 3.5 3.5 0 017 0zM14.5 16a3.5 3.5 0 100-7 3.5 3.5 0 000 7z" clip-rule="evenodd" /></svg>
+                                 <label for="temperature-select" class="block text-sm font-medium text-gray-300">Mode</label>
+                             </div>
+                              <div class="select-wrapper">
+                                 <select id="temperature-select" name="temperature">
+                                     <option value="0.2" selected>0.2 - Precise</option>
+                                     <option value="0.4">0.4 - Confident</option>
+                                     <option value="0.6">0.6 - Balanced</option>
+                                     <option value="0.8">0.8 - Flexible</option>
+                                     <option value="1.0">1.0 - Creative</option>
+                                 </select>
+                             </div>
+                         </div>
+                     </div>
+                     <p class="text-xs text-gray-500 mt-4 text-center">Higher creativity modes may reduce factual accuracy.</p>
+                 </div>
+                 <div id="drop-zone" class="w-full text-center border-2 border-dashed border-[var(--card-border)] rounded-2xl p-10 transition-all duration-300 cursor-pointer hover:bg-[var(--card)] hover:border-[var(--primary)]">
+                     <div class="flex flex-col items-center justify-center pointer-events-none">
+                         <svg class="mx-auto h-12 w-12 text-gray-500" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="1.5" 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 text-gray-400">Drag & drop files or <span class="text-[var(--primary)] font-semibold">click to upload</span></p>
+                         <p class="text-xs text-gray-400 mt-1">Supports PDF, DOCX, TXT</p>
+                         <p id="file-name" class="mt-2 text-xs text-gray-500"></p>
+                     </div>
+                     <input id="file-upload" type="file" class="hidden" accept=".pdf,.txt,.docx" multiple>
+                 </div>
+             </div>
+         </div>
+
+         <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 class="loader"></div>
+             <p id="loading-text" class="mt-6 text-sm font-medium"></p>
+             <p id="loading-subtext" class="mt-2 text-xs text-gray-400"></p>
+         </div>
+     </main>
+
+     <script>
+         document.addEventListener('DOMContentLoaded', () => {
+             const uploadContainer = document.getElementById('upload-container');
+             const chatContainer = document.getElementById('chat-container');
+             const dropZone = document.getElementById('drop-zone');
+             const fileUploadInput = document.getElementById('file-upload');
+             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');
+             const chatSubmitBtn = document.getElementById('chat-submit-btn');
+             const chatWindow = document.getElementById('chat-window');
+             const chatContent = document.getElementById('chat-content');
+             const modelSelect = document.getElementById('model-select');
+             const temperatureSelect = document.getElementById('temperature-select');
+             const chatFilename = document.getElementById('chat-filename');
+             const chatSessionInfo = document.getElementById('chat-session-info');
+
+             let sessionId = sessionStorage.getItem('cognichat_session_id');
+             let currentModelInfo = JSON.parse(sessionStorage.getItem('cognichat_model_info')); // Load model info
+
+             // --- Initialize Marked.js options ---
+             marked.setOptions({
+                 breaks: true, // Convert single line breaks to <br>
+                 gfm: true,    // Enable GitHub Flavored Markdown
+                 // Optional: Add highlight.js integration if library is loaded
+                 // highlight: function(code, lang) {
+                 //   const language = hljs.getLanguage(lang) ? lang : 'plaintext';
+                 //   return hljs.highlight(code, { language }).value;
+                 // }
+             });
+
+             // --- Restore Chat State if Session Exists ---
+             if (sessionId && currentModelInfo) {
+                 console.log("Restoring session:", sessionId);
+                 // Fetch chat history if needed or just show the chat interface
+                 // For now, just show the chat interface and header info
+                 uploadContainer.classList.add('hidden');
+                 chatContainer.classList.remove('hidden');
+                 // You might want to fetch and display previous messages here later
+                 chatFilename.innerHTML = `Chatting with: <strong class="font-semibold">${sessionStorage.getItem('cognichat_filename') || 'documents'}</strong>`;
+                 chatFilename.title = sessionStorage.getItem('cognichat_filename') || 'documents';
+                 chatSessionInfo.innerHTML = `
+                     <p>Model: ${currentModelInfo.simpleModelName}</p>
+                     <p>Mode: ${currentModelInfo.mode}</p>
+                     <button class="mt-1 text-xs text-blue-400 hover:text-blue-300 focus:outline-none" onclick="sessionStorage.clear(); location.reload();">New Chat</button>`;
+             }
+
+
+             // --- File Upload Logic ---
+             dropZone.addEventListener('click', () => fileUploadInput.click());
+
+             ['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => {
+                 dropZone.addEventListener(eventName, e => {e.preventDefault(); e.stopPropagation();}, false);
+                 document.body.addEventListener(eventName, e => {e.preventDefault(); e.stopPropagation();}, false);
+             });
+             ['dragenter', 'dragover'].forEach(eventName => dropZone.addEventListener(eventName, () => dropZone.classList.add('drop-zone--over')));
+             ['dragleave', 'drop'].forEach(eventName => dropZone.addEventListener(eventName, () => dropZone.classList.remove('drop-zone--over')));
+
+             dropZone.addEventListener('drop', (e) => {
+                 if (e.dataTransfer.files.length > 0) handleFiles(e.dataTransfer.files);
+             });
+             fileUploadInput.addEventListener('change', (e) => {
+                 if (e.target.files.length > 0) handleFiles(e.target.files);
+             });
+
+             async function handleFiles(files) {
+                 const formData = new FormData();
+                 let fileNames = Array.from(files).map(f => f.name);
+                 for (const file of files) { formData.append('file', file); }
+
+                 formData.append('model_name', modelSelect.value);
+                 formData.append('temperature', temperatureSelect.value);
+
+                 fileNameSpan.textContent = `Selected: ${fileNames.join(', ')}`;
+                 await uploadAndProcessFiles(formData);
+             }
+
+             async function uploadAndProcessFiles(formData) {
+                 loadingOverlay.classList.remove('hidden');
+                 loadingText.textContent = `Processing document(s)...`;
+                 loadingSubtext.textContent = "Creating a knowledge base... this might take a minute 🧠";
+                 chatContent.innerHTML = ''; // Clear previous chat content on new upload
+
+                 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 during upload.');
+
+                     sessionId = result.session_id;
+                     sessionStorage.setItem('cognichat_session_id', sessionId);
+
+                     const modelOption = modelSelect.querySelector(`option[value="${result.model_name}"]`);
+                     const simpleModelName = modelOption ? modelOption.textContent : result.model_name; // Adjust if needed
+
+                     currentModelInfo = {
+                         model: result.model_name,
+                         mode: result.mode,
+                         simpleModelName: simpleModelName // Use the derived simpler name
+                     };
+                     sessionStorage.setItem('cognichat_model_info', JSON.stringify(currentModelInfo)); // Store model info
+                     sessionStorage.setItem('cognichat_filename', result.filename); // Store filename
+
+                     chatFilename.innerHTML = `Chatting with: <strong class="font-semibold">${result.filename}</strong>`;
+                     chatFilename.title = result.filename;
+
+                     chatSessionInfo.innerHTML = `
+                         <p>Model: ${currentModelInfo.simpleModelName}</p>
+                         <p>Mode: ${currentModelInfo.mode}</p>
+                         <button class="mt-1 text-xs text-blue-400 hover:text-blue-300 focus:outline-none" onclick="sessionStorage.clear(); location.reload();">New Chat</button>`;
+
+                     uploadContainer.classList.add('hidden');
+                     chatContainer.classList.remove('hidden');
+                     appendMessage("Hello! 👋 I've analyzed your documents. What would you like to know?", "bot", currentModelInfo);
+
+                 } catch (error) {
+                     console.error('Upload error:', error);
+                     alert(`Error processing files: ${error.message}`);
+                     sessionStorage.clear(); // Clear session if upload fails
+                 } finally {
+                     loadingOverlay.classList.add('hidden');
+                     fileNameSpan.textContent = '';
+                     fileUploadInput.value = '';
+                 }
+             }
+
+             // --- Chat Logic (Using Server-Sent Events - UPDATED FOR STREAMING & INDICATOR) ---
+             chatForm.addEventListener('submit', async (e) => {
+                 e.preventDefault();
+                 const question = chatInput.value.trim();
+                 if (!question || !sessionId) {
+                     console.warn("Submit ignored: No question or session ID.");
+                     return;
+                 }
+
+                 appendMessage(question, 'user');
+                 chatInput.value = '';
+                 chatInput.disabled = true;
+                 chatSubmitBtn.disabled = true;
+
+                 let botMessageContainer;
+                 let contentDiv;
+                 let fullResponse = '';
+                 let eventSource = null;
+                 let inactivityTimeout = null;
+                 let streamClosedCleanly = false; // Flag to check if stream ended normally vs error
+                 let typingIndicatorElement = null; // Store indicator element
+
+                 // Function to finalize chat (called on error, timeout, or successful completion)
+                 function finalizeChat(isError = false) {
+                     console.log(`Finalizing chat. Was error: ${isError}, Stream ended cleanly: ${streamClosedCleanly}`);
+                     if (eventSource) {
+                         eventSource.close();
+                         eventSource = null;
+                         console.log("SSE connection explicitly closed in finalizeChat.");
+                     }
+                     if (inactivityTimeout) {
+                         clearTimeout(inactivityTimeout);
+                         inactivityTimeout = null;
+                     }
+                     // Remove indicator if it's still there
+                     if (typingIndicatorElement && typingIndicatorElement.parentNode) {
+                         typingIndicatorElement.parentNode.removeChild(typingIndicatorElement);
+                         typingIndicatorElement = null;
+                     }
+
+
+                     if (botMessageContainer && contentDiv) {
+                        const hasErrorMsg = contentDiv.innerHTML.includes('⚠️');
+                        // Ensure final render, apply copy buttons and TTS ONLY if response wasn't an error
+                        if (!hasErrorMsg && fullResponse) {
+                            // Re-parse the complete response to ensure correct final Markdown
+                            contentDiv.innerHTML = marked.parse(fullResponse);
+                            // Apply final touches like copy buttons and TTS
+                            contentDiv.querySelectorAll('pre').forEach(addCopyButton);
+                            addTextToSpeechControls(botMessageContainer, fullResponse);
+                            // Optional: Final highlighting if using highlight.js
+                            // contentDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
+                        }
+                         scrollToBottom(true); // Ensure scrolled to the end
+                     }
+
+                     // Always re-enable input fields
+                     chatInput.disabled = false;
+                     chatSubmitBtn.disabled = false;
+                     chatInput.focus();
+                 }
+
+                 try {
+                     // Create the bot message container *before* starting the stream
+                     botMessageContainer = appendMessage('', 'bot', currentModelInfo); // Append empty bot message
+                     contentDiv = botMessageContainer.querySelector('.markdown-content');
+
+                     // Show typing indicator *inside* the contentDiv
+                     typingIndicatorElement = showTypingIndicator();
+                     if (contentDiv) {
+                         contentDiv.appendChild(typingIndicatorElement);
+                         scrollToBottom(true); // Scroll to show indicator
+                     } else {
+                         console.error("Could not find contentDiv to append typing indicator.");
+                     }
+
+
+                     // Establish SSE connection via GET request
+                     const chatUrl = `/chat?question=${encodeURIComponent(question)}&session_id=${encodeURIComponent(sessionId)}`;
+                     console.log("Connecting to SSE:", chatUrl);
+                     eventSource = new EventSource(chatUrl);
+
+                     eventSource.onopen = () => {
+                         console.log("SSE Connection opened.");
+                         // Remove indicator when connection opens and stream is about to start
+                         if (typingIndicatorElement && typingIndicatorElement.parentNode) {
+                            typingIndicatorElement.parentNode.removeChild(typingIndicatorElement);
+                            typingIndicatorElement = null;
+                         }
+                         streamClosedCleanly = false; // Reset flag on new connection
+                     };
+
+                     eventSource.onmessage = (event) => {
+                         // Remove indicator on first message just in case onopen didn't fire reliably
+                          if (typingIndicatorElement && typingIndicatorElement.parentNode) {
+                            typingIndicatorElement.parentNode.removeChild(typingIndicatorElement);
+                            typingIndicatorElement = null;
+                         }
+
+                         // Reset inactivity timeout on each message
+                         if (inactivityTimeout) clearTimeout(inactivityTimeout);
+                         inactivityTimeout = setTimeout(() => {
+                             console.log("Inactivity timeout triggered after message.");
+                             streamClosedCleanly = true; // Assume normal end
+                             finalizeChat(false);
+                         }, 5000); // 5 seconds of inactivity
+
+                         let data;
+                         try {
+                              data = JSON.parse(event.data);
+                         } catch (parseError){
+                             console.error("Failed to parse SSE data:", event.data, parseError);
+                             contentDiv.innerHTML += `<p class="text-red-400 text-sm">Error receiving data chunk.</p>`;
+                             return;
+                         }
+
+                         if (data.error) {
+                             console.error('SSE Error from server:', data.error);
+                             contentDiv.innerHTML = `<p class="text-red-500 font-semibold">⚠️ Server Error: ${data.error}</p>`;
+                             streamClosedCleanly = false;
+                             finalizeChat(true); // Pass true for error
+                             return;
+                         }
+
+                         if (data.token !== undefined && data.token !== null) {
+                             fullResponse += data.token;
+                             // Update content by parsing the accumulated response
+                             contentDiv.innerHTML = marked.parse(fullResponse);
+                             scrollToBottom(); // Scroll smoothly as content arrives
+                         }
+                     };
+
+                     eventSource.onerror = (error) => {
+                         console.error('SSE connection error event:', error);
+                         // Remove indicator on error
+                         if (typingIndicatorElement && typingIndicatorElement.parentNode) {
+                             typingIndicatorElement.parentNode.removeChild(typingIndicatorElement);
+                             typingIndicatorElement = null;
+                         }
+                         // Don't show generic error if we received data and the stream likely just closed normally
+                         if (!fullResponse && !streamClosedCleanly) { // Only show error if nothing received AND not already cleanly closed
+                             const errorMsg = "⚠️ Connection error. Please try again.";
+                             if (contentDiv) {
+                                 contentDiv.innerHTML = `<p class="text-red-500 font-semibold">${errorMsg}</p>`;
+                             } else {
+                                 // Fallback if container wasn't created somehow
+                                 appendMessage(errorMsg, 'bot', currentModelInfo); // Pass model info here too
+                             }
+                              streamClosedCleanly = false;
+                         } else if (!streamClosedCleanly) {
+                              // If we received data, assume it's a normal closure misinterpreted as error
+                              console.log("SSE connection closed (likely normal end detected by onerror).");
+                              streamClosedCleanly = true; // Mark as clean closure NOW
+                         } else {
+                              console.log("SSE onerror event after stream already marked cleanly closed.")
+                         }
+                         finalizeChat(!streamClosedCleanly); // Finalize, indicate error if not clean
+                     };
+
+                 } catch (error) {
+                     // For setup errors before SSE starts
+                     console.error('Chat setup error:', error);
+                      // Remove indicator on setup error
+                     if (typingIndicatorElement && typingIndicatorElement.parentNode) {
+                         typingIndicatorElement.parentNode.removeChild(typingIndicatorElement);
+                         typingIndicatorElement = null;
+                     }
+                      if (botMessageContainer && contentDiv) {
+                          contentDiv.innerHTML = `<p class="text-red-500 font-semibold">⚠️ Error starting chat: ${error.message}</p>`;
+                     } else {
+                         appendMessage(`Error starting chat: ${error.message}`, 'bot', currentModelInfo); // Pass model info
+                     }
+                      finalizeChat(true);
+                 }
+             });
+
+
+             // --- UI Helper Functions ---
+
+             function appendMessage(text, sender, modelInfo = null) {
+                 const messageWrapper = document.createElement('div');
+                 const iconSVG = sender === 'user'
+                     ? `<div class="bg-blue-200 dark:bg-gray-700 p-2.5 rounded-full flex-shrink-0 mt-1 self-start"><svg class="w-5 h-5 text-blue-700 dark:text-blue-300" viewBox="0 0 24 24"><path fill="currentColor" d="M12 12c2.21 0 4-1.79 4-4s-1.79-4-4-4-4 1.79-4 4 1.79 4 4 4zm0 2c-2.67 0-8 1.34-8 4v2h16v-2c0-2.66-5.33-4-8-4z"></path></svg></div>`
+                     : `<div class="bg-gray-200 dark:bg-gray-700 rounded-full flex-shrink-0 mt-1 self-start text-xl flex items-center justify-center w-10 h-10">✨</div>`;
+
+                 let senderHTML;
+                 if (sender === 'user') {
+                     senderHTML = '<p class="font-medium text-sm mb-1">You</p>';
+                 } else {
+                     let modelInfoHTML = '';
+                     const displayInfo = modelInfo || currentModelInfo;
+                     if (displayInfo && displayInfo.simpleModelName) {
+                         modelInfoHTML = `
+                             <span class="ml-2 text-xs font-normal text-gray-400">
+                                 (Model: ${displayInfo.simpleModelName} | Mode: ${displayInfo.mode})
+                             </span>
+                         `;
+                     }
+                     senderHTML = `<div class="font-medium text-sm mb-1 flex items-center">CogniChat ${modelInfoHTML}</div>`;
+                 }
+
+                 messageWrapper.className = `flex items-start gap-3`;
+                 // Ensure markdown-content div exists even if text is empty for the indicator
+                 messageWrapper.innerHTML = `
+                     ${iconSVG}
+                     <div class="flex-1 pt-1 min-w-0"> ${senderHTML}
+                         <div class="text-base markdown-content prose dark:prose-invert max-w-none">${text ? marked.parse(text) : ''}</div>
+                         <div class="tts-controls mt-2"></div>
+                     </div>
+                 `;
+                 chatContent.appendChild(messageWrapper);
+                 // Force scroll only when adding user message or initial bot message with content
+                 if (sender === 'user' || text) {
+                     scrollToBottom(true);
+                 }
+                 // Return the container that holds the sender name and content div
+                 return messageWrapper.querySelector('.flex-1');
+             }
+
+             // --- UPDATED showTypingIndicator ---
+             function showTypingIndicator() {
+                 const indicator = document.createElement('div');
+                 indicator.className = 'typing-indicator'; // Use the main class
+                 indicator.innerHTML = '<span></span><span></span><span></span>';
+                 // Don't append here, just return the element
+                 return indicator;
+             }
+             // --- End UPDATED showTypingIndicator ---
+
+
+            function scrollToBottom(force = false) {
+                 const isNearBottom = chatWindow.scrollHeight - chatWindow.clientHeight <= chatWindow.scrollTop + 150; // Threshold
+
+                 if (force || isNearBottom) {
+                     requestAnimationFrame(() => { // Use rAF for smoother render loop
+                         chatWindow.scrollTo({
+                             top: chatWindow.scrollHeight,
+                             behavior: 'smooth'
+                         });
+                     });
+                 }
+             }
+
+             function addCopyButton(pre) {
+                 if (pre.querySelector('.copy-code-btn')) return;
+
+                 const button = document.createElement('button');
+                 // Updated classes for better styling
+                 button.className = 'copy-code-btn absolute top-2 right-2 p-1 rounded bg-[var(--copy-btn-bg)] text-[var(--copy-btn-text)] hover:bg-[var(--copy-btn-hover-bg)] transition-opacity duration-200 flex items-center gap-1 text-xs';
+                 button.innerHTML = `<svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"></path></svg> Copy`;
+                 pre.style.position = 'relative'; // Ensure parent is relative for absolute positioning
+                 pre.appendChild(button);
+
+                 button.addEventListener('click', () => {
+                     const code = pre.querySelector('code')?.innerText || '';
+                     navigator.clipboard.writeText(code)
+                         .then(() => {
+                             button.textContent = 'Copied!';
+                             setTimeout(() => button.innerHTML = `<svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"></path></svg> Copy`, 1500);
+                         })
+                         .catch(err => {
+                             console.error('Failed to copy code: ', err);
+                             button.textContent = 'Error';
+                             setTimeout(() => button.innerHTML = `<svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"></path></svg> Copy`, 1500);
+                         });
+                 });
+             }
+
+             // --- TTS Functions ---
+             let currentAudio = null;
+             let currentPlayingButton = null;
+             const playIconSVG = `<svg class="w-4 h-4" fill="currentColor" viewBox="0 0 20 20"><path d="M6.3 2.841A1.5 1.5 0 004 4.11V15.89a1.5 1.5 0 002.3 1.269l9.344-5.89a1.5 1.5 0 000-2.538L6.3 2.84z"></path></svg>`;
+             const pauseIconSVG = `<svg class="w-4 h-4" fill="currentColor" viewBox="0 0 20 20"><path d="M5.75 4.75a.75.75 0 00-.75.75v9.5c0 .414.336.75.75.75h1.5a.75.75 0 00.75-.75v-9.5a.75.75 0 00-.75-.75h-1.5zm6.5 0a.75.75 0 00-.75.75v9.5c0 .414.336.75.75.75h1.5a.75.75 0 00.75-.75v-9.5a.75.75 0 00-.75-.75h-1.5z"></path></svg>`;
+
+             function addTextToSpeechControls(messageBubble, text) {
+                 if (!text || !text.trim()) return;
+                 // Ensure messageBubble is the correct container (the one returned by appendMessage)
+                 const ttsControls = messageBubble.querySelector('.tts-controls');
+                 if (!ttsControls || ttsControls.querySelector('.speak-btn')) return;
+
+                 const speakButton = document.createElement('button');
+                 speakButton.className = 'speak-btn mt-2 px-3 py-1.5 bg-blue-700 text-white rounded-full text-xs font-medium hover:bg-blue-800 transition-colors flex items-center gap-1.5 disabled:opacity-50 disabled:cursor-not-allowed';
+                 speakButton.title = 'Listen to this message';
+                 speakButton.innerHTML = `${playIconSVG} <span>Listen</span>`;
+                 ttsControls.appendChild(speakButton);
+                 speakButton.addEventListener('click', () => handleTTS(text, speakButton));
+             }
+
+             async function handleTTS(text, button) {
+                 if (!text || !text.trim()) return;
+
+                 if (button === currentPlayingButton) {
+                     if (currentAudio && !currentAudio.paused) {
+                         currentAudio.pause();
+                         button.innerHTML = `${playIconSVG} <span>Listen</span>`;
+                     } else if (currentAudio && currentAudio.paused) {
+                         currentAudio.play().catch(e => {console.error("Audio resume error:", e); resetAllSpeakButtons();});
+                         button.innerHTML = `${pauseIconSVG} <span>Pause</span>`;
+                     }
+                     return;
+                 }
+
+                 resetAllSpeakButtons();
+                 currentPlayingButton = button;
+                 button.innerHTML = `<div class="tts-button-loader mr-1"></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(`TTS generation failed (${response.status})`);
+                     const blob = await response.blob();
+                     if (!blob || blob.size === 0) throw new Error("Received empty audio blob.");
+
+                     const audioUrl = URL.createObjectURL(blob);
+                     currentAudio = new Audio(audioUrl);
+
+                     await currentAudio.play();
+                     button.innerHTML = `${pauseIconSVG} <span>Pause</span>`;
+                     button.disabled = false;
+
+                     currentAudio.onended = () => {
+                         if (button === currentPlayingButton) resetAllSpeakButtons();
+                     };
+                     currentAudio.onerror = (e) => {
+                         console.error('Audio object error:', e);
+                         alert('Error playing audio.');
+                         resetAllSpeakButtons();
+                     };
+
+                 } catch (error) {
+                     console.error('TTS Handling Error:', error);
+                     alert(`Failed to play audio: ${error.message}`);
+                     resetAllSpeakButtons();
+                 }
+             }
+
+             function resetAllSpeakButtons() {
+                 document.querySelectorAll('.speak-btn').forEach(resetSpecificButton);
+                 if (currentAudio) {
+                     currentAudio.pause();
+                     currentAudio.onended = null;
+                     currentAudio.onerror = null;
+                      // It's generally safer to let the browser manage object URLs unless memory becomes an issue
+                     // If needed: URL.revokeObjectURL(currentAudio.src);
+                     currentAudio = null;
+                 }
+                 currentPlayingButton = null;
+             }
+
+             function resetSpecificButton(btn){
+                  btn.innerHTML = `${playIconSVG} <span>Listen</span>`;
+                  btn.disabled = false;
+                  btn.removeAttribute('data-state');
+             }
+             // --- End of TTS Functions ---
+
+             // Optional: Initialize highlight.js after DOM load if using
+             // if (typeof hljs !== 'undefined') {
+             //    document.addEventListener('DOMContentLoaded', (event) => {
+             //        document.querySelectorAll('pre code').forEach((block) => {
+             //            hljs.highlightElement(block);
+             //        });
+             //    });
+             // }
+
+         });
+     </script>
+ </body>
+ </html>