app.py

import streamlit as st
import os
import glob
from dotenv import load_dotenv
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.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_groq import ChatGroq
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"
# 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

# 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

# --- Caching Functions ---

@st.cache_resource(show_spinner="Loading and Processing Documents...")
def load_and_process_documents(docs_path: str):
    """
    Loads documents from the specified path, processes them, and splits into chunks.
    Uses UnstructuredFileLoader for potentially better table extraction.
    """
    documents = []
    doc_files = []
    for ext in ["*.pdf", "*.docx", "*.xlsx", "*.json", "*.txt", "*.md"]:
        doc_files.extend(glob.glob(os.path.join(docs_path, ext)))

    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 []

    for file_path in doc_files:
        try:
            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 {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
    try:
        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 vector store: {e}")
        return None

@st.cache_resource(show_spinner="Initializing 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(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}")
        return None

# --- 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

# --- Main Application Logic ---
def main():
    load_dotenv()
    groq_api_key = os.getenv("GROQ_API_KEY")

    # --- 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>
        /* General body style */
        body {
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
            background-color: #f0f2f6; /* Light gray background */
        }
        /* 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 {
            background-color: #1E88E5; /* Catchy blue */
            color: white;
            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 on hover */
        }
        /* Status messages */
        .stAlert { /* For st.info, st.success, st.warning, st.error */
            border-radius: 5px;
        }
        /* 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)

    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...")

    if not groq_api_key:
        status_placeholder.error("GROQ API Key not configured. Application cannot start.")
        st.stop()

    # --- 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:
    """

    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}

    The customer's query is: {question}

    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.

    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()