app.py

import streamlit as st
import os
import glob
from dotenv import load_dotenv
import time
import threading
from twilio.rest import Client  # Import Twilio client
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 ---
# --- Moved groq_api_key here ---
load_dotenv()
groq_api_key = os.getenv("GROQ_API_KEY")
# groq_api_key = ""
DOCS_DIR = "docs"
EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
CACHE_DIR = ".streamlit_cache"
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, 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:
"""
MONITOR_INTERVAL_SECONDS = 30  # Add the constant for the monitoring interval
APP_START_TIME = time.time()

# Twilio Configuration (Add your Twilio credentials here)
account_sid = os.getenv("TWILIO_ACCOUNT_SID")
auth_token = os.getenv("TWILIO_AUTH_TOKEN")
twilio_number = os.getenv("TWILIO_PHONE_NUMBER")  # Your Twilio phone number


# 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=False)  # Disable spinner during initial load
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:
            print(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=False)  # Disable spinner during initial load
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 an empty FAISS instance instead of None to prevent the AttributeError.
        embeddings = HuggingFaceEmbeddings(model_name=_embedding_model_name)  # Initialize embeddings
        vector_store = FAISS.from_documents([], embeddings)  # Changed from None to FAISS.from_documents
        return vector_store


@st.cache_resource(show_spinner=False)  # Disable spinner during initial load
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 ---
def get_rag_chain(llm, retriever, prompt_template):
    """Creates the Retrieval QA chain."""
    prompt = PromptTemplate.from_template(prompt_template)
    rag_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )
    return rag_chain


# --- Twilio Helper Functions ---
def send_twilio_message(client, conversation_sid, message):
    """Sends a message to a Twilio Conversation."""
    try:
        client.conversations.v1.conversations(conversation_sid).messages.create(
            author="Internal Knowledge Base AI",  # Or some identifier
            body=message,
        )
    except Exception as e:
        print(f"❌ Error sending Twilio message: {e}")


def fetch_latest_incoming_message(client, conversation_sid):
    """Fetches the latest incoming message from a conversation."""
    try:
        messages = client.conversations.v1.conversations(conversation_sid).messages.list(limit=1)
        for msg in messages:
            if msg.direction == "inbound":
                return {
                    "body": msg.body,
                    "author": msg.author,
                    "timestamp": msg.date_created
                }
        return None  # No incoming message found
    except Exception as e:
        print(f"❌ Error fetching latest message: {e}")
        return None


def poll_conversation(convo_sid, rag_chain, client):
    """Polls a single conversation for new messages and responds."""
    last_processed_timestamp = {}
    while True:
        try:
            latest_msg = fetch_latest_incoming_message(client, convo_sid)
            if latest_msg:
                msg_time = latest_msg["timestamp"]
                if convo_sid not in last_processed_timestamp or msg_time > last_processed_timestamp[convo_sid]:
                    last_processed_timestamp[convo_sid] = msg_time
                    question = latest_msg["body"]
                    sender = latest_msg["author"]
                    print(f"\n📥 New message from {sender} in {convo_sid}: {question}")
                    # context = "\n\n".join(retrieve_chunks(question, index, embed_model, text_chunks))
                    # answer = generate_answer_with_groq(question, context)
                    try:
                        if "order" in question.lower() and (
                                "status" in question.lower() or "track" in question.lower() or "update" in question.lower() or any(
                                name_part.lower() in question.lower() for name_part in ["customer", "client", "name"])):
                            answer = rag_chain.invoke(question, config={'prompt': ORDER_STATUS_PROMPT})
                        else:
                            answer = rag_chain.invoke(question, config={'prompt': GENERAL_QA_PROMPT})
                        send_twilio_message(client, convo_sid, answer)
                        print(f"📤 Replied to {sender}: {answer}")
                    except Exception as e:
                        print(f"❌ Error during RAG chain invocation: {e}")
                        answer = "Sorry, I encountered an error while processing your request."
                        send_twilio_message(client, convo_sid, answer)
            time.sleep(3)  # Reduced polling interval
        except Exception as e:
            print(f"❌ Error in convo {convo_sid} polling:", e)
            time.sleep(5)


def poll_new_conversations(client, vector_store, llm):
    """Polls for new conversations and starts monitoring them."""
    processed_convos = set()
    print("➡️ Monitoring for new WhatsApp conversations...")
    rag_chain = get_rag_chain(llm, vector_store.as_retriever(search_kwargs={"k": 5}), GENERAL_QA_PROMPT)
    while True:
        try:
            conversations = client.conversations.v1.conversations.list(limit=20)  # Adjust limit as needed
            for convo in conversations:
                convo_full = client.conversations.v1.conversations(convo.sid).fetch()
                if convo.sid not in processed_convos and convo_full.date_created > APP_START_TIME:
                    participants = client.conversations.v1.conversations(convo.sid).participants.list()
                    for p in participants:
                        address = p.messaging_binding.get("address", "") if p.messaging_binding else ""
                        if address.startswith("whatsapp:"):
                            print(f"🆕 New WhatsApp convo found: {convo.sid}")
                            processed_convos.add(convo.sid)
                            threading.Thread(target=poll_conversation, args=(convo.sid, rag_chain, client),
                                             daemon=True).start()
        except Exception as e:
            print("❌ Error polling conversations:", e)
        time.sleep(MONITOR_INTERVAL_SECONDS)  # Use the defined interval


def start_conversation_monitor(client, vector_store, llm):
    """Starts the conversation monitoring process."""
    # Launch the new conversation polling in a separate thread
    threading.Thread(target=poll_new_conversations, args=(client, vector_store, llm),
                     daemon=True).start()
    print("🟢 WhatsApp monitoring started.")  # Add a message


# --- Main Application Logic ---
def main():
    # --- 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>
        .reportview-container .main .block-container{
            padding-top: 2rem;
            padding-bottom: 2rem;
        }
        .st-emotion-cache-z5fcl4 {
            padding-top: 1rem;
        }
        .response-area {
            background-color: #f0f2f6;
            padding: 15px;
            border-radius: 5px;
            margin-top: 10px;
        }
    </style>
    """, unsafe_allow_html=True)

    st.title("📚 Internal Knowledge Base AI 💡")

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

    # --- Initialize session state ---
    if "app_initialized" not in st.session_state:
        st.session_state.app_initialized = False

    # --- Start Button ---
    if not st.session_state.app_initialized:
        if st.button("Start"):  # Create a start button
            st.session_state.app_initialized = True  # set the session state to true
            st.rerun()  # Rerun the app to trigger the knowledge base loading

    # --- Knowledge Base Loading and LLM Initialization ---
    if st.session_state.app_initialized:  # only run if the app has been initialized
        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="llama3-8b-8192")  # Hardcoded to use llama3-8b-8192
            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)")

        # --- Initialize Twilio Client and Start Monitoring ---
        if account_sid and auth_token and twilio_number:
            try:
                client = Client(account_sid, auth_token)
                # Start the conversation monitor in a separate thread
                start_conversation_monitor(client, vector_store, llm)
                st.success("🟢 Monitoring new WhatsApp conversations...")
                st.info("⏳ Waiting for new messages...")
            except Exception as e:
                status_placeholder.error(f"Failed to initialize Twilio: {e}.  Check your credentials and network.")
                st.stop()
        else:
            st.warning("Twilio credentials not fully configured.  WhatsApp monitoring is disabled.")

        # --- Query Input and Response ---
        st.markdown("---")
        st.subheader("Ask a question about our documents:")

        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, vector_store.as_retriever(search_kwargs={"k": 5}), 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()