feat/used the parentdocument retriever

app.py

@@ -3,28 +3,20 @@
 import os
 import time
 import uuid
-from flask import Flask, request, render_template, session, jsonify, Response, stream_with_context
+from flask import Flask, request, render_template, session, jsonify, Response
 from werkzeug.utils import secure_filename
 from rag_processor import create_rag_chain
-from typing import Sequence, Any
-from langchain_experimental.text_splitter import SemanticChunker
-
-from gtts import gTTS
-import io
+from typing import Sequence, Any, List
+import fitz
 import re
+import io
+from gtts import gTTS
 
-
+from langchain_core.documents import Document
 from langchain_community.document_loaders import (
     TextLoader,
-    PyPDFLoader,
     Docx2txtLoader,
 )
-
-# Additional imports for robust PDF handling
-from langchain_core.documents import Document
-import fitz 
-
-# Text Splitter, Embeddings, Retrievers
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
@@ -32,8 +24,8 @@ from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriev
 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
 from sentence_transformers.cross_encoder import CrossEncoder
-import numpy as np
 
 
 app = Flask(__name__)
@@ -42,7 +34,7 @@ app.config['SECRET_KEY'] = os.urandom(24)
 
 class LocalReranker(BaseDocumentCompressor):
     model: Any
-    top_n: int = 5
+    top_n: int = 3
 
     class Config:
         arbitrary_types_allowed = True
@@ -62,8 +54,11 @@ class LocalReranker(BaseDocumentCompressor):
         doc_scores = list(zip(documents, scores))
         sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1], reverse=True)
         
-        return [doc for doc, score in sorted_doc_scores[:self.top_n]]
-
+        top_docs = []
+        for doc, score in sorted_doc_scores[:self.top_n]:
+            doc.metadata['rerank_score'] = float(score)
+            top_docs.append(doc)
+        return top_docs
 
 is_hf_spaces = bool(os.getenv("SPACE_ID") or os.getenv("SPACES_ZERO_GPU"))
 if is_hf_spaces:
@@ -71,451 +66,225 @@ if is_hf_spaces:
 else:
     app.config['UPLOAD_FOLDER'] = 'uploads'
 
-
 try:
     os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
     print(f"✓ Upload folder ready: {app.config['UPLOAD_FOLDER']}")
 except Exception as e:
     print(f"✗ Failed to create upload folder {app.config['UPLOAD_FOLDER']}: {e}")
-    # Fallback to /tmp if the configured path fails
     app.config['UPLOAD_FOLDER'] = '/tmp/uploads'
     os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
     print(f"✓ Using fallback upload folder: {app.config['UPLOAD_FOLDER']}")
 
-
-
-
-
 rag_chains = {}
 message_histories = {}
 
 print("Loading embedding model...")
-
-
-cache_base = os.path.expanduser("~/.cache") if os.path.expanduser("~") != "~" else "/tmp/hf_cache"
-os.environ.setdefault('HF_HOME', f'{cache_base}/huggingface')
-os.environ.setdefault('HF_HUB_CACHE', f'{cache_base}/huggingface/hub')
-os.environ.setdefault('TRANSFORMERS_CACHE', f'{cache_base}/transformers')
-os.environ.setdefault('SENTENCE_TRANSFORMERS_HOME', f'{cache_base}/sentence_transformers')
-
-
-cache_dirs = [
-    os.environ['HF_HOME'],
-    os.environ['HF_HUB_CACHE'],
-    os.environ['TRANSFORMERS_CACHE'],
-    os.environ['SENTENCE_TRANSFORMERS_HOME']
-]
-
-for cache_dir in cache_dirs:
-    try:
-        os.makedirs(cache_dir, mode=0o777, exist_ok=True)
-        # Set permissions explicitly after creation
-        os.chmod(cache_dir, 0o777)
-        print(f"Cache directory ready: {cache_dir}")
-    except PermissionError as pe:
-        print(f"Permission error for {cache_dir}: {pe}")
-        # Try to use a fallback directory in /tmp
-        fallback_dir = cache_dir.replace('/app/.cache', '/tmp/hf_cache')
-        try:
-            os.makedirs(fallback_dir, mode=0o777, exist_ok=True)
-            print(f"Using fallback cache directory: {fallback_dir}")
-            # Update environment variable to point to fallback
-            env_var_map = {
-                '/app/.cache/huggingface': 'HF_HOME',
-                '/app/.cache/huggingface/hub': 'HF_HUB_CACHE',
-                '/app/.cache/transformers': 'TRANSFORMERS_CACHE',
-                '/app/.cache/sentence_transformers': 'SENTENCE_TRANSFORMERS_HOME'
-            }
-            if cache_dir in env_var_map:
-                os.environ[env_var_map[cache_dir]] = fallback_dir
-        except Exception as fe:
-            print(f"Failed to create fallback directory {fallback_dir}: {fe}")
-    except Exception as e:
-        print(f"Warning: Could not create {cache_dir}: {e}")
-
-
-
-# Try loading embedding model with error handling and fallbacks
 try:
-    print("Attempting to load embedding model...")
     EMBEDDING_MODEL = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-miniLM-L6-v2",
-        model_kwargs={'device': 'cpu'},  # Ensure CPU usage
-        cache_folder=os.environ.get('SENTENCE_TRANSFORMERS_HOME')
+        model_name="sentence-transformers/all-MiniLM-L6-v2",
+        model_kwargs={'device': 'cpu'}
     )
-    print("Embedding model loaded successfully.")
+    print("✓ Embedding model loaded successfully.")
 except Exception as e:
-    print(f"Error loading embedding model with full name: {e}")
-    print("Trying shortened model name...")
-    try:
-        EMBEDDING_MODEL = HuggingFaceEmbeddings(
-            model_name="all-miniLM-L6-v2",
-            model_kwargs={'device': 'cpu'},
-            cache_folder=os.environ.get('SENTENCE_TRANSFORMERS_HOME')
-        )
-        print("Embedding model with shortened name loaded successfully.")
-    except Exception as e2:
-        print(f"Failed to load embedding model with shortened name: {e2}")
-        print("Trying without cache folder specification...")
-        try:
-            EMBEDDING_MODEL = HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-miniLM-L6-v2",
-                model_kwargs={'device': 'cpu'}
-            )
-            print("Embedding model loaded without cache folder specification.")
-        except Exception as e3:
-            print(f"Final attempt failed: {e3}")
-            # Use a simpler fallback model or raise the error
-            raise Exception(f"Could not load any embedding model. Last error: {e3}")
-        
-
+    print(f"✗ FATAL: Could not load embedding model. Error: {e}")
+    raise e
 
 print("Loading local re-ranking model...")
-RERANKER_MODEL = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", device='cpu')
-print("Re-ranking model loaded successfully.")
-
+try:
+    RERANKER_MODEL = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", device='cpu')
+    print("✓ Re-ranking model loaded successfully.")
+except Exception as e:
+    print(f"✗ FATAL: Could not load reranker model. Error: {e}")
+    raise e
 
 def load_pdf_with_fallback(filepath):
-    """
-    Tries to load PDF using multiple methods for better reliability.
-    Returns a list of Document objects.
-    """
-    # Method 1: Try PyPDFLoader (primary method)
     try:
-        loader = PyPDFLoader(filepath)
-        docs = loader.load()
-        if docs and any(doc.page_content.strip() for doc in docs):
-            print(f"✓ Successfully loaded PDF using PyPDFLoader: {filepath}")
-            return docs
-    except Exception as e:
-        print(f"PyPDFLoader failed for {filepath}: {e}")
-    
-    # Method 2: Try PyMuPDF (fitz) as fallback
-    try:
-        import fitz
         docs = []
-        pdf_doc = fitz.open(filepath)
-        
-        for page_num in range(pdf_doc.page_count):
-            page = pdf_doc[page_num]
-            text = page.get_text()
-            
-            if text.strip():  # Only add pages with content
-                doc = Document(
-                    page_content=text,
-                    metadata={
-                        "source": filepath,
-                        "page": page_num + 1,
-                        "total_pages": pdf_doc.page_count
-                    }
-                )
-                docs.append(doc)
-        
-        pdf_doc.close()
-        
-        if docs:
-            print(f"✓ Successfully loaded PDF using PyMuPDF: {filepath}")
-            return docs
-            
-    except Exception as e:
-        print(f"PyMuPDF failed for {filepath}: {e}")
-    
-    # Method 3: Try to extract text using pdfplumber (if available)
-    try:
-        import pdfplumber
-        docs = []
-        
-        with pdfplumber.open(filepath) as pdf:
-            for page_num, page in enumerate(pdf.pages):
-                text = page.extract_text()
-                if text and text.strip():
-                    doc = Document(
+        with fitz.open(filepath) as 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": filepath,
+                            "source": os.path.basename(filepath),
                             "page": page_num + 1,
-                            "total_pages": len(pdf.pages)
                         }
-                    )
-                    docs.append(doc)
-        
+                    ))
         if docs:
-            print(f"✓ Successfully loaded PDF using pdfplumber: {filepath}")
+            print(f"✓ Successfully loaded PDF with PyMuPDF: {filepath}")
             return docs
-            
-    except ImportError:
-        print("pdfplumber not available, skipping method 3")
+        else:
+            raise ValueError("No text content found in PDF.")
     except Exception as e:
-        print(f"pdfplumber failed for {filepath}: {e}")
-    
-    # If all methods fail, raise the original error
-    raise Exception(f"Could not load PDF {filepath} using any available method. The file may be corrupted, password-protected, or in an unsupported format.")
+        print(f"✗ PyMuPDF failed for {filepath}: {e}")
+        raise
 
-# A dictionary to map file extensions to their corresponding loader classes
 LOADER_MAPPING = {
     ".txt": TextLoader,
-    ".pdf": load_pdf_with_fallback,  # Use our custom PDF loader
+    ".pdf": load_pdf_with_fallback,
     ".docx": Docx2txtLoader,
 }
 
 def get_session_history(session_id: str) -> ChatMessageHistory:
-    """
-    Retrieves the chat history for a given session ID. If it doesn't exist,
-    a new history object is created.
-    """
     if session_id not in message_histories:
         message_histories[session_id] = ChatMessageHistory()
     return message_histories[session_id]
 
 @app.route('/health', methods=['GET'])
 def health_check():
-    """Health check endpoint for monitoring."""
-    return jsonify({'status': 'healthy', 'message': 'CogniChat is running'}), 200
+    return jsonify({'status': 'healthy'}), 200
 
 @app.route('/', methods=['GET'])
 def index():
-    """Renders the main page."""
     return render_template('index.html')
 
 @app.route('/upload', methods=['POST'])
 def upload_files():
-    """Handles multiple file uploads, processing, and RAG chain creation."""
     files = request.files.getlist('file')
-
     if not files or all(f.filename == '' for f in files):
         return jsonify({'status': 'error', 'message': 'No selected files.'}), 400
 
     all_docs = []
-    all_filenames = []
-    
-    try:
-        print(f"Processing {len(files)} files...")
-        processed_files = []
-        failed_files = []
-        
-        for file in files:
-            if file and file.filename:
-                filename = secure_filename(file.filename)
-                all_filenames.append(filename)
-                filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+    processed_files, failed_files = [], []
+
+    for file in files:
+        if file and file.filename:
+            filename = secure_filename(file.filename)
+            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+            try:
+                file.save(filepath)
+                file_ext = os.path.splitext(filename)[1].lower()
+                if file_ext not in LOADER_MAPPING:
+                    raise ValueError("Unsupported file format.")
                 
-                try:
-                    file.save(filepath)
-                    print(f"✓ Successfully saved file: {filepath}")
-                except Exception as save_error:
-                    print(f"✗ Failed to save file {filename}: {save_error}")
-                    failed_files.append(f"{filename} (failed to save: {save_error})")
-                    continue
-
-                file_extension = os.path.splitext(filename)[1].lower()
-                if file_extension not in LOADER_MAPPING:
-                    print(f"Skipping unsupported file type: {filename}")
-                    failed_files.append(f"{filename} (unsupported format)")
-                    continue
+                loader_func = LOADER_MAPPING[file_ext]
+                docs = loader_func(filepath) if file_ext == ".pdf" else loader_func(filepath).load()
                 
-                # Add file-specific error handling
-                try:
-                    loader_or_function = LOADER_MAPPING[file_extension]
-                    
-                    # Handle custom PDF loader function vs regular loader classes
-                    if file_extension == ".pdf":
-                        file_docs = loader_or_function(filepath)  # Direct function call
-                    else:
-                        loader = loader_or_function(filepath)  # Create loader instance
-                        file_docs = loader.load()
-                    
-                    if file_docs:  # Only add if documents were successfully loaded
-                        all_docs.extend(file_docs)
-                        processed_files.append(filename)
-                        print(f"✓ Successfully processed: {filename} ({len(file_docs)} pages/sections)")
-                    else:
-                        failed_files.append(f"{filename} (no content extracted)")
-                        print(f"⚠ Warning: No content extracted from {filename}")
-                        
-                except Exception as file_error:
-                    error_msg = str(file_error)
-                    print(f"✗ Error processing {filename}: {error_msg}")
-                    
-                    # Provide specific error messages for common PDF issues
-                    if "endstream" in error_msg.lower():
-                        failed_files.append(f"{filename} (corrupted PDF - missing endstream marker)")
-                    elif "pdf" in error_msg.lower() or "pypdf" in error_msg.lower():
-                        failed_files.append(f"{filename} (PDF parsing error)")
-                    elif "permission" in error_msg.lower():
-                        failed_files.append(f"{filename} (password protected or permission denied)")
-                    elif "encoding" in error_msg.lower():
-                        failed_files.append(f"{filename} (text encoding issue)")
-                    else:
-                        failed_files.append(f"{filename} ({error_msg[:50]}...)" if len(error_msg) > 50 else f"{filename} ({error_msg})")
-                    
-                    continue
-
-        # Check if any documents were successfully processed
-        if not all_docs:
-            if failed_files:
-                error_msg = f"Failed to process all uploaded files:\n" + "\n".join(failed_files[:5])
-                if len(failed_files) > 5:
-                    error_msg += f"\n...and {len(failed_files) - 5} more files"
-                return jsonify({'status': 'error', 'message': error_msg}), 400
-            else:
-                return jsonify({'status': 'error', 'message': 'No processable files were uploaded.'}), 400
-
-        # --- Process all documents together ---
-        print(f"Successfully processed {len(processed_files)} files, creating knowledge base...")
-        text_splitter = SemanticChunker(EMBEDDING_MODEL, breakpoint_threshold_type="percentile")
-        splits = text_splitter.split_documents(all_docs)
-
-        print("Creating vector store for all documents...")
-        vectorstore = FAISS.from_documents(documents=splits, embedding=EMBEDDING_MODEL)
-
-        bm25_retriever = BM25Retriever.from_documents(splits)
-        bm25_retriever.k = 7
-        faiss_retriever = vectorstore.as_retriever(search_kwargs={"k": 7})
+                if not docs:
+                    raise ValueError("No content extracted.")
+
+                all_docs.extend(docs)
+                processed_files.append(filename)
+                print(f"✓ Successfully processed: {filename}")
+            except Exception as e:
+                error_msg = str(e)
+                print(f"✗ Error processing {filename}: {error_msg}")
+                failed_files.append(f"{filename} ({error_msg})")
+
+    if not all_docs:
+        error_summary = "Failed to process all files."
+        if failed_files:
+            error_summary += " Reasons: " + ", ".join(failed_files)
+        return jsonify({'status': 'error', 'message': error_summary}), 400
+
+    try:
+        print("Starting RAG pipeline setup...")
+
+        vectorstore = FAISS.from_texts([""], EMBEDDING_MODEL) 
+        store = InMemoryStore()
+        
+        parent_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
+        child_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=50)
+
+        parent_docs = parent_splitter.split_documents(all_docs)
+        doc_ids = [str(uuid.uuid4()) for _ in parent_docs]
+        
+        child_docs = []
+        for i, doc in enumerate(parent_docs):
+            _id = doc_ids[i]
+            sub_docs = child_splitter.split_documents([doc])
+            for child in sub_docs:
+                child.metadata["doc_id"] = _id
+            child_docs.extend(sub_docs)
+        
+        store.mset(list(zip(doc_ids, parent_docs)))
+        vectorstore.add_documents(child_docs)
+        print(f"✓ Stored {len(parent_docs)} parent docs and indexed {len(child_docs)} child docs.")
+
+        bm25_retriever = BM25Retriever.from_documents(child_docs)
+        bm25_retriever.k = 8
+
+        faiss_retriever = vectorstore.as_retriever(search_kwargs={"k": 8})
+
         ensemble_retriever = EnsembleRetriever(
             retrievers=[bm25_retriever, faiss_retriever],
-            weights=[0.5, 0.5]
+            weights=[0.4, 0.6]
         )
-        reranker = LocalReranker(model=RERANKER_MODEL, top_n=5)
+        print("✓ Created Hybrid Retriever for child documents.")
+
+        reranker = LocalReranker(model=RERANKER_MODEL, top_n=4)
+        
+        def get_parents(docs: List[Document]) -> List[Document]:
+            parent_ids = {d.metadata["doc_id"] for d in docs}
+            return store.mget(list(parent_ids))
 
         compression_retriever = ContextualCompressionRetriever(
-        base_compressor=reranker, 
-        base_retriever=ensemble_retriever
-                                        )
-    
+            base_compressor=reranker, base_retriever=ensemble_retriever
+        )
         
-        session_id = str(uuid.uuid4())
-        print(f"Creating RAG chain for session {session_id}...")
+        final_retriever = compression_retriever | get_parents
         
-        try:
-            rag_chain = create_rag_chain(compression_retriever, get_session_history)
-            rag_chains[session_id] = rag_chain
-            print(f"✓ RAG chain created successfully for session {session_id} with {len(processed_files)} documents.")
-        except Exception as rag_error:
-            print(f"✗ Failed to create RAG chain: {rag_error}")
-            import traceback
-            traceback.print_exc()
-            raise rag_error
+        print("✓ Final retriever chain created: (Hybrid -> Rerank) -> Parent Fetch")
 
+        session_id = str(uuid.uuid4())
+        rag_chain = create_rag_chain(final_retriever, get_session_history)
+        rag_chains[session_id] = rag_chain
         session['session_id'] = session_id
-        print(f"✓ Session stored in Flask session: {session_id}")
-        
-        # Prepare success message with file processing summary
+
         success_msg = f"Successfully processed: {', '.join(processed_files)}"
         if failed_files:
-            success_msg += f"\n\nFailed to process {len(failed_files)} file(s):\n" + "\n".join(failed_files[:3])
-            if len(failed_files) > 3:
-                success_msg += f"\n...and {len(failed_files) - 3} more"
-        
+            success_msg += f"\nFailed to process: {', '.join(failed_files)}"
+
         return jsonify({
             'status': 'success', 
             'filename': success_msg,
-            'processed_count': len(processed_files),
-            'failed_count': len(failed_files),
-            'session_id': session_id  # Include session_id in response for client to store
+            'session_id': session_id
         })
 
     except Exception as e:
-        print(f"Error creating RAG chain: {e}")
-        return jsonify({'status': 'error', 'message': f'Failed to process files: {str(e)}'}), 500
+        import traceback
+        traceback.print_exc()
+        return jsonify({'status': 'error', 'message': f'Failed during RAG setup: {e}'}), 500
 
 @app.route('/chat', methods=['POST'])
 def chat():
-    """Handles chat messages and streams the response with memory."""
     data = request.get_json()
     question = data.get('question')
-    
-    # Try to get session_id from multiple sources (Flask session or request body)
     session_id = session.get('session_id') or data.get('session_id')
 
-    print(f"\n=== Chat Request Debug ===")
-    print(f"Raw request data: {data}")
-    print(f"Question: '{question}' (type: {type(question)})")
-    print(f"Session ID from Flask session: {session.get('session_id')}")
-    print(f"Session ID from request body: {data.get('session_id')}")
-    print(f"Final session ID used: {session_id}")
-    print(f"Available RAG chains: {list(rag_chains.keys())}")
-    print(f"Total RAG chains: {len(rag_chains)}")
-    
-    # Check each condition individually with detailed logging
-    if not question:
-        print("✗ FAILURE: No question provided")
-        return jsonify({'status': 'error', 'message': 'No question provided.'}), 400
-    else:
-        print("✓ Question is valid")
-        
-    if not session_id:
-        print("✗ FAILURE: No session ID in Flask session")
-        return jsonify({'status': 'error', 'message': 'No session found. Please upload documents first.'}), 400
-    else:
-        print(f"✓ Session ID found: {session_id}")
-        
-    if session_id not in rag_chains:
-        print(f"✗ FAILURE: Session ID {session_id} not found in RAG chains")
-        print(f"Available sessions: {list(rag_chains.keys())}")
-        return jsonify({'status': 'error', 'message': 'Session not found. Please upload documents again.'}), 400
-    else:
-        print(f"✓ RAG chain found for session: {session_id}")
+    if not question or not session_id or session_id not in rag_chains:
+        return jsonify({'status': 'error', 'message': 'Invalid session or no question provided.'}), 400
 
     try:
         rag_chain = rag_chains[session_id]
         config = {"configurable": {"session_id": session_id}}
         
-        def generate():
-            """A generator function to stream the response."""
-            # Create the input dictionary with question and config
-            input_data = {
-                "question": question,
-                "config": config
-            }
-            
-            try:
-                # For streaming, we need to handle the chain differently
-                # The final chain doesn't support streaming, so we'll invoke it directly
-                response = rag_chain.invoke(input_data)
-                yield response
-            except Exception as stream_error:
-                print(f"Streaming error: {stream_error}")
-                # Fallback to direct invocation
-                try:
-                    response = rag_chain.invoke(input_data)
-                    yield response
-                except Exception as invoke_error:
-                    print(f"Invoke error: {invoke_error}")
-                    yield "I apologize, but I'm having trouble processing your question. Please try again or upload your documents again."
+        # Invoke the chain, which will return a string
+        answer_string = rag_chain.invoke({"question": question}, config=config)
         
-        return Response(stream_with_context(generate()), mimetype='text/plain')
+        # --- THIS IS THE FIX ---
+        # Directly use the returned string in the JSON response.
+        return jsonify({'answer': answer_string})
 
     except Exception as e:
         print(f"Error during chat invocation: {e}")
-        return Response("An error occurred while getting the answer.", status=500, mimetype='text/plain')
-
+        return jsonify({'status': 'error', 'message': 'An error occurred while getting the answer.'}), 500
 
 def clean_markdown_for_tts(text: str) -> str:
-    """Removes markdown formatting for cleaner text-to-speech output."""
-    # Remove bold (**text**) and italics (*text* or _text_)
     text = re.sub(r'\*(\*?)(.*?)\1\*', r'\2', text)
     text = re.sub(r'\_(.*?)\_', r'\1', text)
-    # Remove inline code (`code`)
     text = re.sub(r'`(.*?)`', r'\1', text)
-    # Remove headings (e.g., #, ##, ###)
     text = re.sub(r'^\s*#{1,6}\s+', '', text, flags=re.MULTILINE)
-    # Remove list item markers (*, -, 1.)
     text = re.sub(r'^\s*[\*\-]\s+', '', text, flags=re.MULTILINE)
     text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE)
-    # Remove blockquotes (>)
     text = re.sub(r'^\s*>\s?', '', text, flags=re.MULTILINE)
-    
-    # --- NEW ADDITION HERE ---
-    # Remove horizontal rules (---, ***, ===)
     text = re.sub(r'^\s*[-*_]{3,}\s*$', '', text, flags=re.MULTILINE)
-    
-    # Replace multiple newlines with a single space
     text = re.sub(r'\n+', ' ', text)
     return text.strip()
 
 @app.route('/tts', methods=['POST'])
 def text_to_speech():
-    """Generates audio from text and returns it as an MP3 stream."""
     data = request.get_json()
     text = data.get('text')
 
@@ -523,9 +292,7 @@ def text_to_speech():
         return jsonify({'status': 'error', 'message': 'No text provided.'}), 400
 
     try:
-        # --- FIX IS HERE: Clean the text before sending to gTTS ---
         clean_text = clean_markdown_for_tts(text)
-        
         tts = gTTS(clean_text, lang='en')
         mp3_fp = io.BytesIO()
         tts.write_to_fp(mp3_fp)
@@ -535,79 +302,6 @@ def text_to_speech():
         print(f"Error in TTS generation: {e}")
         return jsonify({'status': 'error', 'message': 'Failed to generate audio.'}), 500
 
-
-
-@app.route('/debug', methods=['GET'])
-def debug_info():
-    """Debug endpoint to check configuration."""
-    api_key = os.getenv("GROQ_API_KEY")
-    is_hf_spaces = bool(os.getenv("SPACE_ID") or os.getenv("SPACES_ZERO_GPU"))
-    
-    return jsonify({
-        'environment': 'Hugging Face Spaces' if is_hf_spaces else 'Local/Other',
-        'upload_folder': app.config['UPLOAD_FOLDER'],
-        'upload_folder_writable': os.access(app.config['UPLOAD_FOLDER'], os.W_OK),
-        'groq_api_key_set': bool(api_key and api_key != "your_groq_api_key_here"),
-        'groq_api_key_placeholder': api_key == "your_groq_api_key_here" if api_key else False,
-        'sessions_count': len(rag_chains),
-        'session_ids': list(rag_chains.keys()),
-        'flask_session_id': session.get('session_id'),
-        'flask_session_keys': list(session.keys()) if session else [],
-        'embedding_model_loaded': 'EMBEDDING_MODEL' in globals(),
-        'space_id': os.getenv("SPACE_ID"),
-        'hf_spaces_detected': is_hf_spaces
-    })
-
-@app.route('/test-session', methods=['GET', 'POST'])
-def test_session():
-    """Test endpoint to debug Flask session issues."""
-    if request.method == 'POST':
-        # Set a test session value
-        session['test_key'] = 'test_value'
-        session['timestamp'] = str(time.time())
-        return jsonify({
-            'action': 'session_set',
-            'test_key': session.get('test_key'),
-            'timestamp': session.get('timestamp'),
-            'session_keys': list(session.keys())
-        })
-    else:
-        # Read session values
-        return jsonify({
-            'action': 'session_read',
-            'test_key': session.get('test_key'),
-            'timestamp': session.get('timestamp'), 
-            'session_id': session.get('session_id'),
-            'session_keys': list(session.keys()),
-            'has_session_data': bool(session)
-        })
-
 if __name__ == '__main__':
-    print(f"\n=== Application Startup Complete at {time.strftime('%Y-%m-%d %H:%M:%S')} ===")
-    
-    # Check environment
-    is_hf_spaces = bool(os.getenv("SPACE_ID") or os.getenv("SPACES_ZERO_GPU"))
-    api_key = os.getenv('GROQ_API_KEY')
-    
-    print(f"Environment: {'Hugging Face Spaces' if is_hf_spaces else 'Local Development'}")
-    print(f"Upload folder: {app.config['UPLOAD_FOLDER']}")
-    
-    if is_hf_spaces:
-        print(f"Space ID: {os.getenv('SPACE_ID', 'Not set')}")
-        print(f"GROQ API Key: {'Set via HF Secrets' if api_key and api_key != 'your_groq_api_key_here' else 'NOT SET - Add to Space Secrets'}")
-    else:
-        print(f"GROQ API Key: {'Set' if api_key and api_key != 'your_groq_api_key_here' else 'NOT SET'}")
-    
-    print(f"Embedding model: {'Loaded' if 'EMBEDDING_MODEL' in globals() else 'NOT LOADED'}")
-    
-    if not api_key or api_key == 'your_groq_api_key_here':
-        print("\n⚠️  WARNING: GROQ API KEY NOT CONFIGURED!")
-        if is_hf_spaces:
-            print("   → Go to your Space Settings > Repository Secrets")
-            print("   → Add GROQ_API_KEY as a secret")
-        else:
-            print("   → Update .env file with your GROQ API key")
-    
-    # Use port 7860 for Hugging Face Spaces, fallback to 5001 for local development
     port = int(os.environ.get("PORT", 7860))
-    app.run(host="0.0.0.0", port=port, debug=False)
+    app.run(host="0.0.0.0", port=port, debug=False)

rag_processor.py

@@ -35,9 +35,11 @@ def create_rag_chain(retriever, get_session_history_func):
     
     # Get API key from environment (works for both HF Spaces secrets and local .env)
     api_key = os.getenv("GROQ_API_KEY")
+    print("key loaded")
     
     if not api_key or api_key == "your_groq_api_key_here":
         error_msg = "GROQ_API_KEY not found or not configured properly.\n"
+        print("Not found key")
         
         # Detect if running in Hugging Face Spaces
         if os.getenv("SPACE_ID") or os.getenv("SPACES_ZERO_GPU"):
@@ -56,7 +58,7 @@ def create_rag_chain(retriever, get_session_history_func):
 
     # --- 1. Initialize the LLM ---
     # Updated model_name to a standard, high-performance Groq model
-    llm = ChatGroq(model_name="llama-3.1-8b-instant", api_key=api_key, temperature=1)
+    llm = ChatGroq(model_name="llama-3.1-8b-instant", api_key=api_key, temperature=0.1)
 
     # --- 2. Create Query Rewriting Chain 🧠 ---
     print("\nSetting up query rewriting chain...")
@@ -80,14 +82,12 @@ Standalone Question:"""
     
     # --- 3. Create Main RAG Chain with Memory ---
     print("\nSetting up main RAG chain...")
-    rag_template = """You are an expert assistant named `Cognichat`, developed by Ritesh and Alish.
-Your primary job is to provide accurate, comprehensive answers based ONLY on the provided context document.
-Act as an expert historian analyzing a text. Synthesize information from all available context snippets to form a complete answer.
-Your answer must be grounded in the text. Quote key phrases where necessary but do not simply copy-paste.
-Structure your answers clearly using markdown, such as bold text and bullet points, for readability.
-
-**Crucially, do not state that the information is missing from the context.** Assume the provided text contains the necessary information and do your best to construct the answer from it. If the context is truly insufficient, answer the question as best you can with the information available without mentioning the context's limitations.
-
+    rag_template = """You are CogniChat, an expert document analysis assistant.
+IMPORTANT RULES:
+1. ONLY use information from the provided context
+2. If the answer isn't in the context, say "I cannot find this information in the uploaded documents"
+3. Cite specific sections when possible (e.g., "According to Section 2.3...")
+4. Be concise but comprehensive
 Context:
 {context}"""
     rag_prompt = ChatPromptTemplate.from_messages([
@@ -100,10 +100,11 @@ Context:
     # Parallel process to fetch context and correctly pass through question and history.
     # We use itemgetter to select the specific keys from the input dictionary.
     setup_and_retrieval = RunnableParallel({
-        "context": query_rewriter | retriever,
-        "question": itemgetter("question"),
-        "chat_history": itemgetter("chat_history"),
-    })
+    "context": RunnablePassthrough.assign(
+        rewritten_question=query_rewriter
+    ) | (lambda x: x["rewritten_question"]) | retriever,
+    "question": itemgetter("question"),
+    "chat_history": itemgetter("chat_history")})
     # =====================================================================
     
     # The initial RAG chain

templates/index.html

@@ -98,17 +98,122 @@
             40% { transform: scale(1.0); }
         }
 
-        /* Markdown Styling */
-        .markdown-content p { margin-bottom: 0.75rem; line-height: 1.75; }
-        .markdown-content ul, .markdown-content ol { margin-left: 1.5rem; margin-bottom: 0.75rem; }
-        .markdown-content code { background-color: rgba(0,0,0,0.05); padding: 0.2rem 0.4rem; border-radius: 0.25rem; font-family: 'Roboto Mono', monospace; font-size: 0.9em; }
-        .dark .markdown-content code { background-color: rgba(255,255,255,0.1); }
-        .markdown-content pre { position: relative; background-color: #f8f9fa; border: 1px solid var(--card-border); border-radius: 0.5rem; margin-bottom: 1rem; }
-        .dark .markdown-content pre { background-color: #2e2f32; }
-        .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: #e8eaed; border: 1px solid #dadce0; color: #5f6368; padding: 0.3rem 0.6rem; border-radius: 0.25rem; cursor: pointer; opacity: 0; transition: opacity 0.2s; font-size: 0.8em;}
-        .dark .markdown-content pre .copy-code-btn { background-color: #3c4043; border-color: #5f6368; color: #e8eaed; }
-        .markdown-content pre:hover .copy-code-btn { opacity: 1; }
+        /* Enhanced Markdown Styling for better readability and aesthetics */
+        .markdown-content p { 
+            margin-bottom: 1rem; 
+            line-height: 1.75; 
+        }
+        .markdown-content h1, .markdown-content h2, .markdown-content h3, .markdown-content h4 {
+            font-family: 'Google Sans', sans-serif;
+            font-weight: 700;
+            margin-top: 1.75rem;
+            margin-bottom: 1rem;
+            line-height: 1.3;
+        }
+        .markdown-content h1 { font-size: 1.75em; border-bottom: 1px solid var(--card-border); padding-bottom: 0.5rem; }
+        .markdown-content h2 { font-size: 1.5em; }
+        .markdown-content h3 { font-size: 1.25em; }
+        .markdown-content h4 { font-size: 1.1em; }
+        .markdown-content ul, .markdown-content ol { 
+            padding-left: 1.75rem; 
+            margin-bottom: 1rem; 
+        }
+        .markdown-content li {
+            margin-bottom: 0.5rem;
+        }
+        .dark .markdown-content ul > li::marker { color: var(--primary); }
+        .markdown-content ul > li::marker { color: var(--primary); }
+        .markdown-content a {
+            color: var(--primary);
+            text-decoration: none;
+            font-weight: 500;
+            border-bottom: 1px solid transparent;
+            transition: all 0.2s ease-in-out;
+        }
+        .markdown-content a:hover {
+            border-bottom-color: var(--primary-hover);
+        }
+        .markdown-content blockquote {
+            margin: 1.5rem 0;
+            padding-left: 1.5rem;
+            border-left: 4px solid var(--card-border);
+            color: #6c757d;
+            font-style: italic;
+        }
+        .dark .markdown-content blockquote {
+            color: #adb5bd;
+        }
+        .markdown-content hr {
+            border: none;
+            border-top: 1px solid var(--card-border);
+            margin: 2rem 0;
+        }
+        .markdown-content table {
+            width: 100%;
+            border-collapse: collapse;
+            margin: 1.5rem 0;
+            font-size: 0.9em;
+            box-shadow: 0 1px 3px rgba(0,0,0,0.05);
+            border-radius: 8px;
+            overflow: hidden;
+        }
+        .markdown-content th, .markdown-content td {
+            border: 1px solid var(--card-border);
+            padding: 0.75rem 1rem;
+            text-align: left;
+        }
+        .markdown-content th {
+            background-color: var(--bot-bubble);
+            font-weight: 500;
+        }
+        .markdown-content code { 
+            background-color: rgba(0,0,0,0.05); 
+            padding: 0.2rem 0.4rem; 
+            border-radius: 0.25rem; 
+            font-family: 'Roboto Mono', monospace; 
+            font-size: 0.9em; 
+        }
+        .dark .markdown-content code { 
+            background-color: rgba(255,255,255,0.1); 
+        }
+        .markdown-content pre { 
+            position: relative; 
+            background-color: #f8f9fa; 
+            border: 1px solid var(--card-border); 
+            border-radius: 0.5rem; 
+            margin-bottom: 1rem; 
+        }
+        .dark .markdown-content pre { 
+            background-color: #2e2f32; 
+        }
+        .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: #e8eaed; 
+            border: 1px solid #dadce0; 
+            color: #5f6368; 
+            padding: 0.3rem 0.6rem; 
+            border-radius: 0.25rem; 
+            cursor: pointer; 
+            opacity: 0; 
+            transition: opacity 0.2s; 
+            font-size: 0.8em;
+        }
+        .dark .markdown-content pre .copy-code-btn { 
+            background-color: #3c4043; 
+            border-color: #5f6368; 
+            color: #e8eaed; 
+        }
+        .markdown-content pre:hover .copy-code-btn { 
+            opacity: 1; 
+        }
         
         /* Spinner for the TTS button */
         .tts-button-loader {
@@ -181,7 +286,6 @@
             const chatContent = document.getElementById('chat-content');
             const chatFilename = document.getElementById('chat-filename');
 
-            // Persist the active Flask session between requests
             let sessionId = null;
             const storedSessionId = sessionStorage.getItem('cognichat_session_id');
             if (storedSessionId) {
@@ -257,7 +361,7 @@
                     loadingOverlay.classList.add('hidden');
                     loadingSubtext.textContent = ''; 
                     fileNameSpan.textContent = '';
-                    fileUploadInput.value = ''; // Reset file input
+                    fileUploadInput.value = '';
                 }
             }
             
@@ -289,26 +393,28 @@
                     });
 
                     if (!response.ok) throw new Error(`Server error: ${response.statusText}`);
+                    
+                    // ============================ MODIFICATION START ==============================
+                    // Parse the JSON response instead of reading a stream
+                    const result = await response.json();
+                    const answer = result.answer; // Extract the 'answer' field
+
+                    if (!answer) {
+                        throw new Error("Received an empty or invalid response from the server.");
+                    }
 
                     typingIndicator.remove();
                     botMessageContainer = appendMessage('', 'bot');
                     contentDiv = botMessageContainer.querySelector('.markdown-content');
 
-                    const reader = response.body.getReader();
-                    const decoder = new TextDecoder();
-                    let fullResponse = '';
-
-                    while (true) {
-                        const { value, done } = await reader.read();
-                        if (done) break;
-                        
-                        fullResponse += decoder.decode(value, { stream: true });
-                        contentDiv.innerHTML = marked.parse(fullResponse);
-                        scrollToBottom();
-                    }
+                    // Use the extracted answer for rendering
+                    contentDiv.innerHTML = marked.parse(answer);
                     contentDiv.querySelectorAll('pre').forEach(addCopyButton);
+                    scrollToBottom(); // Scroll after content is added
                     
-                    addTextToSpeechControls(botMessageContainer, fullResponse); 
+                    // Use the extracted answer for TTS
+                    addTextToSpeechControls(botMessageContainer, answer);
+                    // ============================ MODIFICATION END ==============================
 
                 } catch (error) {
                     console.error('Chat error:', error);
@@ -344,7 +450,10 @@
 
                 const contentDiv = document.createElement('div');
                 contentDiv.className = 'text-base markdown-content';
-                contentDiv.innerHTML = marked.parse(text);
+                // Only parse if text is not empty
+                if (text) {
+                    contentDiv.innerHTML = marked.parse(text);
+                }
 
                 const controlsContainer = document.createElement('div');
                 controlsContainer.className = 'tts-controls mt-2';
@@ -412,7 +521,7 @@
                 });
             }
 
-            // ============================ MODIFICATIONS START ==============================
+            // --- Text-to-Speech Logic ---
             let currentAudio = null;
             let currentPlayingButton = null;
             
@@ -424,7 +533,6 @@
                 const ttsControls = messageBubble.querySelector('.tts-controls');
                 if (text.trim().length > 0) {
                     const speakButton = document.createElement('button');
-                    // STYLING CHANGE HERE: Replaced theme variables with specific dark blue colors.
                     speakButton.className = 'speak-btn px-4 py-2 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 disabled:cursor-not-allowed';
                     speakButton.title = 'Listen to this message';
                     speakButton.setAttribute('data-state', 'play');
@@ -435,24 +543,19 @@
             }
 
             async function handleTTS(text, button) {
-                // BUG FIX: Reworked the logic to correctly handle pause/resume.
-                
-                // Case 1: The clicked button is already active (playing or paused).
                 if (button === currentPlayingButton) {
-                    if (currentAudio && !currentAudio.paused) { // If it's playing, pause it.
+                    if (currentAudio && !currentAudio.paused) {
                         currentAudio.pause();
                         button.setAttribute('data-state', 'paused');
                         button.innerHTML = `${playIconSVG} <span>Play</span>`;
-                    } else if (currentAudio && currentAudio.paused) { // If it's paused, resume it.
+                    } else if (currentAudio && currentAudio.paused) {
                         currentAudio.play();
                         button.setAttribute('data-state', 'playing');
                         button.innerHTML = `${pauseIconSVG} <span>Pause</span>`;
                     }
-                    return; // Stop the function here.
+                    return;
                 }
 
-                // Case 2: A new button is clicked (or no audio is active).
-                // Stop any other audio that might be playing.
                 resetAllSpeakButtons();
                 
                 currentPlayingButton = button;
@@ -488,7 +591,7 @@
                     button.setAttribute('data-state', 'error');
                     button.innerHTML = `${playIconSVG} <span>Error</span>`;
                     alert('Failed to play audio. Please try again.');
-                    resetAllSpeakButtons(); // Reset state on error
+                    resetAllSpeakButtons();
                 } finally {
                     button.disabled = false;
                 }
@@ -506,8 +609,7 @@
                 }
                 currentPlayingButton = null;
             }
-            // ============================ MODIFICATIONS END ==============================
         });
     </script>
 </body>
-</html><!DOCTYPE html>
+</html>