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 = "llama3-8b-8192"): # UPDATED MODEL
    """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:
        # Available models (check Groq documentation for the latest):
        # "llama3-8b-8192" (good balance of speed and capability)
        # "llama3-70b-8192" (more powerful, potentially slower)
        # "gemma-7b-it"
        llm = ChatGroq(temperature=0, groq_api_key=api_key, model_name=model_name)
        st.sidebar.info(f"LLM Initialized: {model_name}") # Add info about which model is used
        return llm
    except Exception as e:
        st.error(f"Error initializing Groq LLM: {e}")
        return None

# --- RAG Chain Setup ---
# ... (get_rag_chain function remains the same) ...

# --- 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 (remains the same)
    st.markdown("""
    <style>
        # ... (CSS content remains the same) ...
    </style>
    """, unsafe_allow_html=True)

    st.title("📚 Internal Knowledge Base AI 💡")
    
    st.sidebar.header("System Settings") # Changed from System Status for clarity
    
    # Model selection in sidebar (New Feature)
    available_models = ["llama3-8b-8192", "llama3-70b-8192", "gemma-7b-it"] # Add more as Groq supports them
    selected_model = st.sidebar.selectbox(
        "Select LLM Model:",
        available_models,
        index=available_models.index("llama3-8b-8192") # Default selection
    )
    st.sidebar.markdown("---")
    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 ---
    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()
        
        # Pass the selected model to get_llm
        llm = get_llm(groq_api_key, model_name=selected_model) 
        if not llm:
            # Error is already shown by get_llm, but update status_placeholder too
            status_placeholder.error("Failed to initialize LLM. Application cannot proceed.")
            st.stop()
            
        end_time = time.time()
        # status_placeholder is updated by get_llm or on success below
        status_placeholder.success(f"Application Ready! (Loaded in {end_time - start_time:.2f}s)")
    
    retriever = vector_store.as_retriever(search_kwargs={"k": 5})

    # --- Query Input and Response ---
    # ... (rest of the main function remains the same, including prompt templates, query input, button, and response display logic) ...

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

    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})
            
            current_model_info = st.sidebar.empty() # Placeholder for current mode info

            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"])):
                active_prompt_template = ORDER_STATUS_PROMPT
                current_model_info.info("Mode: Order Status Query")
            else:
                active_prompt_template = GENERAL_QA_PROMPT
                current_model_info.info("Mode: General Query")

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

    st.markdown("---")
    st.subheader("Response:")
    response_area = st.container()
    # Ensure response_area is robust against empty messages or incorrect last role
    last_assistant_message = "Ask a question to see the answer here."
    if st.session_state.messages and st.session_state.messages[-1]['role'] == 'assistant':
        last_assistant_message = st.session_state.messages[-1]['content']
    
    response_area.markdown(f"<div class='response-area'>{last_assistant_message}</div>", unsafe_allow_html=True)


    st.sidebar.markdown("---")
    st.sidebar.markdown("Built with ❤️ using Streamlit & Langchain & Groq")


if __name__ == "__main__":
    main()