Update app.py

app.py

@@ -1,298 +1,350 @@
 import streamlit as st
 import os
+import glob
 from dotenv import load_dotenv
-from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader, UnstructuredFileLoader
-from langchain_community.vectorstores import FAISS
-from langchain_huggingface import HuggingFaceEmbeddings
+import time
+
+from langchain_community.document_loaders import (
+    PyPDFLoader,
+    Docx2txtLoader,
+    UnstructuredExcelLoader,
+    JSONLoader,
+    UnstructuredFileLoader # Generic loader, good for tables
+)
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.chains import RetrievalQA
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
 from langchain_groq import ChatGroq
-import time
-import glob
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain.schema.runnable import RunnablePassthrough
+from langchain.schema.output_parser import StrOutputParser
 
 # --- Configuration ---
 DOCS_DIR = "docs"
-CHUNK_SIZE = 1500
-CHUNK_OVERLAP = 200
-EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" # Good default, consider others for specific needs
-CACHE_DIR_FAISS = "faiss_index_cache" # Directory to cache FAISS index
+# Using a local sentence transformer model for embeddings
+EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+CACHE_DIR = ".streamlit_cache" # For potential disk-based caching if needed beyond Streamlit's default
 
-# --- Helper Functions ---
+# Create docs and cache directory if they don't exist
+if not os.path.exists(DOCS_DIR):
+    os.makedirs(DOCS_DIR)
+if not os.path.exists(CACHE_DIR):
+    os.makedirs(CACHE_DIR)
+
+# --- Helper Function for Document Loading ---
+def get_loader(file_path):
+    """Detects file type and returns appropriate Langchain loader."""
+    _, ext = os.path.splitext(file_path)
+    ext = ext.lower()
+    # Prioritize UnstructuredFileLoader for robust table and content extraction
+    # UnstructuredFileLoader can handle many types, but we can specify if needed
+    if ext in ['.pdf', '.docx', '.doc', '.xlsx', '.xls', '.json', '.txt', '.md', '.html', '.xml', '.eml', '.msg']:
+        return UnstructuredFileLoader(file_path, mode="elements", strategy="fast") # "elements" is good for tables
+    # Fallback or specific loaders if UnstructuredFileLoader has issues with a particular file
+    # elif ext == ".pdf":
+    #     return PyPDFLoader(file_path) # Basic PDF loader
+    # elif ext in [".docx", ".doc"]:
+    #     return Docx2txtLoader(file_path) # Basic DOCX loader
+    # elif ext in [".xlsx", ".xls"]:
+    #     return UnstructuredExcelLoader(file_path, mode="elements") # Unstructured for Excel
+    # elif ext == ".json":
+    #     return JSONLoader(file_path, jq_schema='.[]', text_content=False) # Adjust jq_schema as needed
+    else:
+        st.warning(f"Unsupported file type: {ext}. Skipping {os.path.basename(file_path)}")
+        return None
 
-def get_api_key():
-    """Loads GROQ API key from .env file or environment variables."""
-    load_dotenv()
-    groq_api_key = os.getenv("GROQ_API_KEY")
-    if not groq_api_key:
-        st.error("GROQ_API_KEY not found. Please set it in your environment variables or a .env file.")
-        st.stop()
-    return groq_api_key
+# --- Caching Functions ---
 
-@st.cache_resource(show_spinner="Loading and processing documents...")
-def load_and_process_documents(_docs_dir: str):
+@st.cache_resource(show_spinner="Loading and Processing Documents...")
+def load_and_process_documents(docs_path: str):
     """
-    Loads documents from the specified directory, processes them,
-    creates embeddings, and stores them in a FAISS vector store.
-    Caches the FAISS index to disk for faster subsequent loads.
+    Loads documents from the specified path, processes them, and splits into chunks.
+    Uses UnstructuredFileLoader for potentially better table extraction.
     """
-    if not os.path.exists(_docs_dir) or not os.listdir(_docs_dir):
-        st.warning(f"The '{_docs_dir}' directory is empty or does not exist. Please add your documents.")
-        return None
+    documents = []
+    doc_files = []
+    for ext in ["*.pdf", "*.docx", "*.xlsx", "*.json", "*.txt", "*.md"]:
+        doc_files.extend(glob.glob(os.path.join(docs_path, ext)))
 
-    st.write(f"Scanning for documents in '{_docs_dir}'...")
+    if not doc_files:
+        st.error(f"No documents found in the '{docs_path}' directory. Please add some documents.")
+        st.info("Supported formats: .pdf, .docx, .xlsx, .json, .txt, .md")
+        return []
 
-    # Using UnstructuredFileLoader for broader file type support including tables
-    # We'll use glob to find all files and pass them to UnstructuredFileLoader
-    all_files = []
-    supported_extensions = ["*.pdf", "*.docx", "*.doc", "*.xlsx", "*.xls", "*.json", "*.txt", "*.md", "*.html", "*.csv", "*.pptx"] # Add more if needed
-    for ext in supported_extensions:
-        all_files.extend(glob.glob(os.path.join(_docs_dir, ext)))
-
-    if not all_files:
-        st.warning(f"No supported documents found in '{_docs_dir}'. Supported types: {', '.join(supported_extensions)}")
-        return None
-
-    st.write(f"Found {len(all_files)} files to process: {', '.join([os.path.basename(f) for f in all_files])}")
-
-    docs = []
-    progress_bar = st.progress(0, text="Loading documents...")
-    for i, file_path in enumerate(all_files):
+    for file_path in doc_files:
         try:
-            st.write(f"Processing: {os.path.basename(file_path)}")
-            # UnstructuredFileLoader is good for various formats and attempts to handle tables.
-            # For complex tables in PDFs, more specialized parsers might be needed if Unstructured is insufficient.
-            # Common arguments for UnstructuredFileLoader for better table extraction:
-            # strategy="hi_res" (for PDFs with complex layouts, may require `detectron2` installation)
-            # mode="elements" or "paged"
-            # pdf_infer_table_structure=True (if using unstructured[pdf])
-            loader = UnstructuredFileLoader(file_path, mode="elements", strategy="fast") # Start with "fast", try "hi_res" if table extraction is poor
-            loaded_docs = loader.load()
-            docs.extend(loaded_docs)
+            st.write(f"Processing: {os.path.basename(file_path)}...") # Show progress
+            loader = get_loader(file_path)
+            if loader:
+                loaded_docs = loader.load()
+                # Add source metadata to each document for better traceability
+                for doc in loaded_docs:
+                    doc.metadata["source"] = os.path.basename(file_path)
+                documents.extend(loaded_docs)
         except Exception as e:
-            st.error(f"Error loading file {os.path.basename(file_path)}: {e}")
-        progress_bar.progress((i + 1) / len(all_files), text=f"Loaded {os.path.basename(file_path)}")
-
-    if not docs:
-        st.warning("No documents were successfully loaded or processed.")
-        return None
-
-    progress_bar.progress(1.0, text="Documents loaded. Splitting into chunks...")
-    time.sleep(0.5) # For UX
-
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
-    texts = text_splitter.split_documents(docs)
-
-    if not texts:
-        st.warning("Document splitting resulted in no text chunks. Check document content and splitter settings.")
+            st.error(f"Error loading {os.path.basename(file_path)}: {e}")
+            st.warning(f"Skipping file {os.path.basename(file_path)} due to error.")
+    
+    if not documents:
+        st.error("No documents were successfully loaded or processed.")
+        return []
+
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunked_documents = text_splitter.split_documents(documents)
+    
+    if not chunked_documents:
+        st.error("Document processing resulted in no text chunks. Check document content and parsing.")
+        return []
+        
+    st.success(f"Successfully loaded and processed {len(doc_files)} documents into {len(chunked_documents)} chunks.")
+    return chunked_documents
+
+@st.cache_resource(show_spinner="Creating Vector Store (Embeddings)...")
+def create_vector_store(_documents, _embedding_model_name: str):
+    """Creates a FAISS vector store from the given documents and embedding model."""
+    if not _documents:
+        st.warning("Cannot create vector store: No documents processed.")
         return None
-
-    st.write(f"Split documents into {len(texts)} chunks.")
-    progress_bar.progress(0, text="Generating embeddings and creating vector store... (This may take a while)")
-
-    # Initialize embeddings
     try:
-        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
-    except Exception as e:
-        st.error(f"Failed to load embedding model '{EMBEDDING_MODEL_NAME}': {e}")
-        st.error("Please ensure you have an internet connection and the model name is correct.")
-        st.stop()
-
-
-    # Create FAISS vector store and cache it
-    if os.path.exists(CACHE_DIR_FAISS) and os.listdir(CACHE_DIR_FAISS):
-        try:
-            st.write(f"Loading cached FAISS index from {CACHE_DIR_FAISS}...")
-            vector_store = FAISS.load_local(CACHE_DIR_FAISS, embeddings, allow_dangerous_deserialization=True) # Required for FAISS with HuggingFaceEmbeddings
-            st.write("FAISS index loaded from cache.")
-            progress_bar.progress(1.0, text="Vector store ready.")
-            return vector_store
-        except Exception as e:
-            st.warning(f"Could not load FAISS index from cache: {e}. Rebuilding index.")
-
-    try:
-        vector_store = FAISS.from_documents(texts, embeddings)
-        if not os.path.exists(CACHE_DIR_FAISS):
-            os.makedirs(CACHE_DIR_FAISS)
-        vector_store.save_local(CACHE_DIR_FAISS)
-        st.write(f"FAISS index created and saved to {CACHE_DIR_FAISS}.")
-        progress_bar.progress(1.0, text="Vector store ready.")
+        embeddings = HuggingFaceEmbeddings(model_name=_embedding_model_name)
+        vector_store = FAISS.from_documents(_documents, embedding=embeddings)
+        st.success("Vector Store created successfully!")
         return vector_store
     except Exception as e:
-        st.error(f"Error creating FAISS vector store: {e}")
+        st.error(f"Error creating vector store: {e}")
         return None
 
-
 @st.cache_resource(show_spinner="Initializing LLM...")
-def get_llm(_api_key: str):
-    """Initializes and returns the ChatGroq LLM."""
+def get_llm(api_key: str, model_name: str = "mixtral-8x7b-32768"): # "llama3-70b-8192" is another option
+    """Initializes the Groq LLM."""
+    if not api_key:
+        st.error("GROQ_API_KEY not found! Please set it in your environment variables or a .env file.")
+        return None
     try:
-        llm = ChatGroq(
-            groq_api_key=_api_key,
-            model_name="mixtral-8x7b-32768", # Or "llama3-70b-8192", "llama3-8b-8192", "gemma-7b-it"
-            temperature=0.2, # Adjust for creativity vs. factuality
-            # max_tokens=1024, # Optional: set max tokens
-        )
+        llm = ChatGroq(temperature=0, groq_api_key=api_key, model_name=model_name)
         return llm
     except Exception as e:
-        st.error(f"Error initializing GROQ LLM: {e}")
-        st.stop()
+        st.error(f"Error initializing Groq LLM: {e}")
+        return None
 
-# --- Streamlit App UI ---
+# --- RAG Chain Setup ---
+def get_rag_chain(llm, retriever, prompt_template_str):
+    """Creates a RAG chain with the given LLM, retriever, and prompt template."""
+    prompt = PromptTemplate(
+        template=prompt_template_str,
+        input_variables=["context", "question"]
+    )
+    
+    rag_chain = (
+        {"context": retriever, "question": RunnablePassthrough()}
+        | prompt
+        | llm
+        | StrOutputParser()
+    )
+    return rag_chain
 
-st.set_page_config(page_title="SmartQuery RAG", layout="wide", initial_sidebar_state="expanded")
+# --- Main Application Logic ---
+def main():
+    load_dotenv()
+    groq_api_key = os.getenv("GROQ_API_KEY")
 
-# --- Styling (Optional - for a "catchy" look) ---
-st.markdown("""
+    # --- UI Setup ---
+    st.set_page_config(page_title="Internal Knowledge Base AI", layout="wide", initial_sidebar_state="expanded")
+
+    # Custom CSS for a "catchy and elegant" design
+    st.markdown("""
     <style>
-        .stApp {
-            background-color: #f0f2f6; /* Light grey background */
+        /* General body style */
+        body {
+            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+            background-color: #f0f2f6; /* Light gray background */
         }
-        .stTextInput > div > div > input {
-            background-color: #ffffff;
+        /* Main content area */
+        .main .block-container {
+            padding-top: 2rem;
+            padding-bottom: 2rem;
+            padding-left: 3rem;
+            padding-right: 3rem;
+            background-color: #ffffff; /* White content background */
             border-radius: 10px;
+            box-shadow: 0 4px 12px rgba(0,0,0,0.1); /* Subtle shadow */
+        }
+        /* Title style */
+        h1 {
+            color: #1E88E5; /* Catchy blue */
+            text-align: center;
+            font-weight: 600;
+        }
+        /* Sidebar style */
+        .stSidebar {
+            background-color: #E3F2FD; /* Light blue sidebar */
+            padding: 10px;
+        }
+        .stSidebar .sidebar-content {
+             background-color: #E3F2FD;
+        }
+        /* Input box style */
+        .stTextInput > div > div > input {
+            background-color: #f8f9fa;
+            border-radius: 5px;
+            border: 1px solid #ced4da;
         }
+        /* Button style */
         .stButton > button {
-            border-radius: 10px;
-            background-color: #1E88E5; /* Blue */
+            background-color: #1E88E5; /* Catchy blue */
             color: white;
-            font-weight: bold;
+            border-radius: 5px;
+            padding: 0.5rem 1rem;
+            font-weight: 500;
+            border: none;
             transition: background-color 0.3s ease;
         }
         .stButton > button:hover {
-            background-color: #1565C0; /* Darker Blue */
-        }
-        .stSpinner > div > svg { /* Spinner color */
-            fill: #1E88E5;
+            background-color: #1565C0; /* Darker blue on hover */
         }
-        .loader-text {
-            font-size: 1.2em;
-            color: #333;
+        /* Status messages */
+        .stAlert { /* For st.info, st.success, st.warning, st.error */
+            border-radius: 5px;
         }
-        .ready-text {
-            font-size: 1.2em;
-            color: green;
-            font-weight: bold;
-        }
-        .response-container {
-            background-color: #ffffff;
-            padding: 20px;
-            border-radius: 10px;
-            box-shadow: 0 4px 8px rgba(0,0,0,0.1);
-            margin-top: 20px;
-        }
-        .response-header {
-            font-size: 1.5em;
-            color: #1E88E5;
-            margin-bottom: 10px;
+        /* Response area */
+        .response-area {
+            background-color: #f8f9fa;
+            padding: 1rem;
+            border-radius: 5px;
+            border: 1px solid #e0e0e0;
+            margin-top: 1rem;
+            min-height: 100px;
         }
     </style>
-""", unsafe_allow_html=True)
+    """, unsafe_allow_html=True)
 
+    st.title("📚 Internal Knowledge Base AI 💡")
+    
+    # Sidebar for status and information
+    st.sidebar.header("System Status")
+    status_placeholder = st.sidebar.empty()
+    status_placeholder.info("Initializing...")
 
-# --- Main Application Logic ---
-st.title("📄 SmartQuery RAG Assistant")
-st.markdown("Ask questions about your documents (Customer Orders, Company Policies, Financial Data, Products, Return Policies, etc.)")
+    if not groq_api_key:
+        status_placeholder.error("GROQ API Key not configured. Application cannot start.")
+        st.stop()
 
-# Load API Key
-groq_api_key = get_api_key()
+    # --- Knowledge Base Loading ---
+    # This will be cached after the first run
+    with st.spinner("Knowledge Base is loading... Please wait."):
+        start_time = time.time()
+        processed_documents = load_and_process_documents(DOCS_DIR)
+        if not processed_documents:
+            status_placeholder.error("Failed to load or process documents. Check logs and `docs` folder.")
+            st.stop()
+
+        vector_store = create_vector_store(processed_documents, EMBEDDING_MODEL_NAME)
+        if not vector_store:
+            status_placeholder.error("Failed to create vector store. Application cannot proceed.")
+            st.stop()
+        
+        llm = get_llm(groq_api_key)
+        if not llm:
+            status_placeholder.error("Failed to initialize LLM. Application cannot proceed.")
+            st.stop()
+            
+        end_time = time.time()
+        status_placeholder.success(f"Application Ready! (Loaded in {end_time - start_time:.2f}s)")
+    
+    retriever = vector_store.as_retriever(search_kwargs={"k": 5}) # Retrieve top 5 relevant chunks
+
+    # --- Query Input and Response ---
+    st.markdown("---")
+    st.subheader("Ask a question about our documents:")
+
+    # Prompt templates
+    GENERAL_QA_PROMPT = """
+    You are an AI assistant for our internal knowledge base.
+    Your goal is to provide accurate and concise answers based ONLY on the provided context.
+    Do not make up information. If the answer is not found in the context, state that clearly.
+    Ensure your answers are directly supported by the text.
+    Accuracy is paramount.
+
+    Context:
+    {context}
+
+    Question: {question}
+
+    Answer:
+    """
 
-# Sidebar for status and controls
-st.sidebar.header("Knowledge Base Status")
-status_placeholder = st.sidebar.empty()
-status_placeholder.markdown("<p class='loader-text'>Knowledge Base is loading...</p>", unsafe_allow_html=True)
+    ORDER_STATUS_PROMPT = """
+    You are an AI assistant helping with customer order inquiries.
+    Based ONLY on the following retrieved information from our order system and policies:
+    {context}
 
-# Create docs directory if it doesn't exist
-if not os.path.exists(DOCS_DIR):
-    os.makedirs(DOCS_DIR)
-    st.sidebar.info(f"'{DOCS_DIR}' directory created. Please add your documents there and refresh.")
+    The customer's query is: {question}
 
-# Load and process documents, then initialize RAG
-vector_store = load_and_process_documents(DOCS_DIR)
+    Please perform the following steps:
+    1. Carefully analyze the context for any order details (Order ID, Customer Name, Status, Items, Dates, etc.).
+    2. If an order matching the query (or related to a name in the query) is found in the context:
+       - Address the customer by their name if available in the order details (e.g., "Hello [Customer Name],").
+       - Provide ALL available information about their order, including Order ID, status, items, dates, and any other relevant details found in the context.
+       - Be comprehensive and clear.
+    3. If no specific order details are found in the context that match the query, or if the context is insufficient, politely state that you couldn't find the specific order information in the provided documents and suggest they contact support for further assistance.
+    4. Do NOT invent or infer any information not explicitly present in the context.
 
-if vector_store:
-    llm = get_llm(groq_api_key)
-    retriever = vector_store.as_retriever(
-        search_type="similarity", # "mmr" (Maximal Marginal Relevance) is another option
-        search_kwargs={"k": 5}  # Retrieve top 5 relevant chunks
-    )
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type="stuff", # Options: "stuff", "map_reduce", "refine", "map_rerank"
-        retriever=retriever,
-        return_source_documents=True # Set to True to see which documents were retrieved
-    )
-    status_placeholder.markdown("<p class='ready-text'>✅ Application is Ready. Ask your questions!</p>", unsafe_allow_html=True)
-    st.sidebar.success("Knowledge Base Loaded Successfully!")
-    if st.sidebar.button("🔄 Clear Cache & Reload Documents"):
-        # Clear specific caches or the entire cache
-        st.cache_resource.clear() # Clears all @st.cache_resource
-        # Could also selectively clear st.cache_data if used.
-        # Manually delete FAISS cache directory
-        if os.path.exists(CACHE_DIR_FAISS):
-            import shutil
-            shutil.rmtree(CACHE_DIR_FAISS)
-            st.sidebar.info(f"Cache '{CACHE_DIR_FAISS}' cleared.")
-        st.rerun()
-
-else:
-    status_placeholder.error("⚠️ Knowledge Base could not be loaded. Check messages above and ensure documents are in the 'docs' folder.")
-    st.stop()
-
-
-# --- User Interaction ---
-st.markdown("---")
-query = st.text_input("Enter your question:", placeholder="e.g., What is the return policy for electronics?")
-
-if st.button("Submit Query", type="primary"):
-    if not query:
-        st.warning("Please enter a question.")
-    else:
-        with st.spinner("🧠 Thinking... Fetching answer..."):
-            try:
-                start_time = time.time()
-                response = qa_chain.invoke({"query": query})
-                end_time = time.time()
-
-                st.markdown("<div class='response-container'>", unsafe_allow_html=True)
-                st.markdown("<p class='response-header'>💡 Answer:</p>", unsafe_allow_html=True)
-                st.write(response["result"])
-                st.markdown("</div>", unsafe_allow_html=True)
-
-                st.info(f"Response generated in {end_time - start_time:.2f} seconds.")
-
-                with st.expander("📚 Show Retrieved Sources"):
-                    if "source_documents" in response and response["source_documents"]:
-                        for i, doc in enumerate(response["source_documents"]):
-                            st.markdown(f"**Source {i+1} (from: {doc.metadata.get('source', 'N/A').split('/')[-1]})**")
-                            st.caption(doc.page_content[:500] + "..." if doc.page_content else "N/A") # Display first 500 chars
-                            st.markdown("---")
-                    else:
-                        st.write("No specific source documents were identified for this query.")
-
-            except Exception as e:
-                st.error(f"An error occurred while processing your query: {e}")
-
-# --- Suggestions for Improvement (as per prompt request) ---
-st.sidebar.markdown("---")
-st.sidebar.subheader("💡 Suggestions & Notes:")
-st.sidebar.markdown("""
-- **Table Data:** `UnstructuredFileLoader` attempts to parse tables. For PDFs with very complex tables, if accuracy is insufficient, consider:
-    - Pre-processing PDFs with tools like `Camelot` or `Tabula-py` to extract tables into CSV/Markdown, then load those.
-    - Exploring `unstructured` with `strategy="hi_res"` (may require `detectron2` and `brew install poppler` or similar for your OS). This is more computationally intensive.
-    - Fine-tuning embedding models or using models specialized for tabular data if table queries are critical.
-- **Accuracy:** "100% accuracy" is an ideal. RAG systems are powerful but can make mistakes. Improve by:
-    - Better chunking strategies.
-    - More advanced retrieval (e.g., HyDE, re-ranking).
-    - Prompt engineering for the QA chain.
-    - Using more powerful (and potentially slower/costlier) LLMs if available via GROQ.
-    - Regularly evaluating and curating the document set.
-- **Performance:** The current FAISS caching helps significantly. For very large datasets, explore more scalable vector DBs.
-- **UI/UX:** Added some basic styling. For more "catchy" UI, explore Streamlit Components or more elaborate CSS.
-- **Error Handling:** Added basic error checks. Robust applications need more comprehensive error management.
-- **Scalability:** For many concurrent users on Hugging Face, resource limits (CPU, RAM) for the free tier might be a bottleneck, especially during embedding.
-- **Embedding Model:** `all-MiniLM-L6-v2` is efficient. For higher accuracy with more complex content, consider models like `sentence-transformers/all-mpnet-base-v2` or domain-specific embeddings.
-- **Deployment:** Ensure `GROQ_API_KEY` is set as a secret in Hugging Face Spaces.
-""")
-
-st.markdown("---")
-st.markdown("<p style='text-align: center; color: grey;'>Powered by Streamlit, Langchain & GROQ</p>", unsafe_allow_html=True)
+    Answer:
+    """
+
+    # Use session state to store conversation history if desired, or just last query/response
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    query = st.text_input("Enter your question:", key="query_input", placeholder="e.g., 'What is the return policy?' or 'Status of order for John Doe?'")
+
+    if st.button("Submit", key="submit_button"):
+        if query:
+            st.session_state.messages.append({"role": "user", "content": query})
+            
+            # Determine prompt based on query type (simple keyword check)
+            # A more sophisticated intent detection could be used here (e.g., another LLM call, classifier)
+            if "order" in query.lower() and ("status" in query.lower() or "track" in query.lower() or "update" in query.lower() or any(name_part.lower() in query.lower() for name_part in ["customer", "client", "name"])): # Basic check for order status
+                active_prompt_template = ORDER_STATUS_PROMPT
+                st.sidebar.info("Using: Order Status Query Mode")
+            else:
+                active_prompt_template = GENERAL_QA_PROMPT
+                st.sidebar.info("Using: General Query Mode")
+
+            rag_chain = get_rag_chain(llm, retriever, active_prompt_template)
+
+            with st.spinner("Thinking..."):
+                try:
+                    response = rag_chain.invoke(query)
+                    st.session_state.messages.append({"role": "assistant", "content": response})
+                except Exception as e:
+                    st.error(f"Error during RAG chain invocation: {e}")
+                    response = "Sorry, I encountered an error while processing your request."
+                    st.session_state.messages.append({"role": "assistant", "content": response})
+        else:
+            st.warning("Please enter a question.")
+
+    # Display chat messages
+    st.markdown("---")
+    st.subheader("Response:")
+    response_area = st.container()
+    response_area.add_rows([ # Create a container with fixed height and scroll
+            st.markdown(f"<div class='response-area'>{st.session_state.messages[-1]['content'] if st.session_state.messages and st.session_state.messages[-1]['role'] == 'assistant' else 'Ask a question to see the answer here.'}</div>", unsafe_allow_html=True)
+    ])
+
+
+    # Optional: Display retrieved context for debugging or transparency
+    # if st.sidebar.checkbox("Show Retrieved Context (for debugging)"):
+    #     if query and vector_store: # Check if query and vector_store exist
+    #         docs = retriever.get_relevant_documents(query)
+    #         st.sidebar.subheader("Retrieved Context:")
+    #         for i, doc in enumerate(docs):
+    #             st.sidebar.text_area(f"Chunk {i+1} (Source: {doc.metadata.get('source', 'N/A')})", doc.page_content, height=150)
+
+    st.sidebar.markdown("---")
+    st.sidebar.markdown("Built with ❤️ using Streamlit & Langchain & Groq")
+
+
+if __name__ == "__main__":
+    main()