stapp.py

import streamlit as st
from dotenv import load_dotenv
import os
import shutil
import time

# ── env ──────────────────────────────────────────────────────────────────────
load_dotenv()

# ── page config ──────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="QGenie  – RAG Assistant",
    page_icon="📚",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ── paths ─────────────────────────────────────────────────────────────────────
BASE_DIR      = os.path.dirname(os.path.abspath(__file__))
DOCUMENT_DIR  = os.path.join(BASE_DIR, "documents")
CHROMA_PATH   = os.path.join(BASE_DIR, "chroma_db")
os.makedirs(DOCUMENT_DIR, exist_ok=True)

# ── custom CSS ────────────────────────────────────────────────────────────────
st.markdown("""
<style>
/* ── global ── */
html, body, [class*="css"] {
    font-family: 'Inter', sans-serif;
}
.stApp { background: #080a0f; color: #e2e8f0; }

/* ── sidebar ── */
section[data-testid="stSidebar"] {
    background: #0d0f14 !important;
    border-right: 1px solid #1e2535;
}
section[data-testid="stSidebar"] * { color: #a0aec0 !important; }
section[data-testid="stSidebar"] h1,
section[data-testid="stSidebar"] h2,
section[data-testid="stSidebar"] h3 { color: #00ffe7 !important; }

/* ── headings / labels ── */
h1, h2, h3, .stMarkdown h1 { color: #e2e8f0 !important; }
.label-cyan {
    color: #00ffe7;
    font-size: 11px;
    font-weight: 700;
    letter-spacing: .2em;
    font-family: 'Courier New', monospace;
    text-transform: uppercase;
}

/* ── chat bubbles ── */
.user-bubble {
    background: #00ffe7;
    color: #0d0f14;
    padding: 10px 16px;
    border-radius: 14px 14px 4px 14px;
    max-width: 65%;
    margin-left: auto;
    font-size: 14px;
    line-height: 1.6;
}
.ai-bubble {
    background: #0d1828;
    color: #e2e8f0;
    padding: 12px 16px;
    border-radius: 14px 14px 14px 4px;
    max-width: 75%;
    border: 1px solid #1e3a5a;
    font-size: 14px;
    line-height: 1.7;
    white-space: pre-wrap;
}
.source-tag {
    font-size: 11px;
    color: #2d5a7a;
    font-family: 'Courier New', monospace;
    margin-top: 4px;
}
.chat-turn { margin-bottom: 20px; }

/* ── log panel ── */
.log-box {
    background: #060810;
    border: 1px solid #1e2535;
    border-radius: 10px;
    padding: 12px 14px;
    font-family: 'Courier New', monospace;
    font-size: 12px;
    color: #4a5568;
    max-height: 260px;
    overflow-y: auto;
}
.log-line { color: #00ffe7; margin: 2px 0; }

/* ── file badge ── */
.file-badge {
    background: #0d0f14;
    border: 1px solid #1e2535;
    border-radius: 6px;
    padding: 6px 12px;
    font-size: 13px;
    font-family: 'Courier New', monospace;
    color: #a0aec0;
    margin: 3px 0;
    display: flex;
    align-items: center;
    gap: 8px;
}

/* ── buttons ── */
.stButton > button {
    background: #00ffe7 !important;
    color: #0d0f14 !important;
    font-weight: 700 !important;
    border: none !important;
    border-radius: 8px !important;
    transition: background .2s !important;
}
.stButton > button:hover { background: #00d4c5 !important; }

/* secondary buttons */
.btn-secondary > button {
    background: transparent !important;
    color: #4a5568 !important;
    border: 1px solid #1e2535 !important;
}
.btn-secondary > button:hover {
    color: #a0aec0 !important;
    border-color: #4a5568 !important;
}

/* danger button */
.btn-danger > button {
    background: transparent !important;
    color: #fc8181 !important;
    border: 1px solid #4a1515 !important;
}
.btn-danger > button:hover { background: #1f0d0d !important; }

/* ── inputs ── */
.stTextInput input, .stChatInput textarea {
    background: #0d0f14 !important;
    color: #e2e8f0 !important;
    border: 1px solid #1e2535 !important;
    border-radius: 8px !important;
}
.stTextInput input:focus, .stChatInput textarea:focus {
    border-color: #00ffe7 !important;
    box-shadow: none !important;
}

/* ── expander ── */
.streamlit-expanderHeader { color: #a0aec0 !important; background: #0d0f14 !important; }
.streamlit-expanderContent { background: #060810 !important; border: 1px solid #1e2535; }

/* ── tabs ── */
.stTabs [data-baseweb="tab-list"] { background: transparent; gap: 8px; }
.stTabs [data-baseweb="tab"] {
    background: #0d0f14;
    color: #4a5568 !important;
    border: 1px solid #1e2535;
    border-radius: 8px 8px 0 0;
}
.stTabs [aria-selected="true"] {
    background: #0d1828 !important;
    color: #00ffe7 !important;
    border-bottom-color: #0d1828 !important;
}

/* ── alerts ── */
.stSuccess { background: #0d1f14; border: 1px solid #1a4731; color: #68d391 !important; }
.stWarning { background: #1f1a0d; border: 1px solid #4a3a15; color: #f6e05e !important; }
.stError   { background: #1f0d0d; border: 1px solid #4a1515; color: #fc8181 !important; }
.stInfo    { background: #0d1828; border: 1px solid #1e3a5a; color: #63b3ed !important; }

/* ── divider ── */
hr { border-color: #1e2535 !important; }
</style>
""", unsafe_allow_html=True)


# ── session state defaults ────────────────────────────────────────────────────
for key, default in {
    "chat_history": [],   # list of {"question": ..., "answer": ..., "sources": [...]}
    "vectordb": None,
    "retriever": None,
    "backend_logs": [],
    "db_built": False,
}.items():
    if key not in st.session_state:
        st.session_state[key] = default


# ── helper: log ───────────────────────────────────────────────────────────────
def push_log(msg: str):
    st.session_state.backend_logs.append(msg)


# ── helper: list documents ────────────────────────────────────────────────────
def list_docs():
    try:
        return [f for f in os.listdir(DOCUMENT_DIR) if f.endswith((".pdf", ".txt"))]
    except Exception:
        return []


# ── helper: clear chroma ──────────────────────────────────────────────────────
def clear_chroma():
    if os.path.exists(CHROMA_PATH):
        shutil.rmtree(CHROMA_PATH)
    st.session_state.vectordb  = None
    st.session_state.retriever = None
    st.session_state.db_built  = False


# ── RAG backend ───────────────────────────────────────────────────────────────
def build_vectorstore():
    from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader, TextLoader
    from langchain_text_splitters import RecursiveCharacterTextSplitter
    from langchain_community.vectorstores import Chroma
    from langchain_huggingface import HuggingFaceEmbeddings

    push_log("📂 Scanning documents folder...")
    pdf_loader = DirectoryLoader(DOCUMENT_DIR, glob="*.pdf", loader_cls=PyPDFLoader)
    txt_loader = DirectoryLoader(DOCUMENT_DIR, glob="*.txt", loader_cls=TextLoader)

    push_log("📄 Loading PDFs...")
    pdf_docs = list(pdf_loader.lazy_load())
    push_log("📝 Loading TXT files...")
    txt_docs = list(txt_loader.lazy_load())
    docs = pdf_docs + txt_docs
    push_log(f"✅ {len(pdf_docs)} PDF(s) + {len(txt_docs)} TXT(s) loaded")

    if not docs:
        push_log("⚠️ No documents found. Upload files first.")
        return False

    push_log("✂️ Splitting into chunks (size=300, overlap=30)...")
    splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=30)
    chunks = splitter.split_documents(docs)
    push_log(f"📦 {len(chunks)} chunks created")

    push_log("🧠 Loading HuggingFace embedding model...")
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    push_log("✅ Embedding model ready")

    push_log("🗄️ Building ChromaDB...")
    vectordb = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=CHROMA_PATH,
    )
    retriever = vectordb.as_retriever(
        search_type="mmr",
        search_kwargs={"k": 3, "lambda_mult": 0.7},
    )
    st.session_state.vectordb  = vectordb
    st.session_state.retriever = retriever
    st.session_state.db_built  = True
    push_log("🚀 Vector DB ready!")
    return True


def load_existing_vectorstore():
    """Load an already-persisted Chroma DB from disk."""
    if not (os.path.exists(CHROMA_PATH) and os.listdir(CHROMA_PATH)):
        return False
    from langchain_community.vectorstores import Chroma
    from langchain_huggingface import HuggingFaceEmbeddings
    push_log("📂 Loading existing vector store from disk...")
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vectordb = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings)
    retriever = vectordb.as_retriever(
        search_type="mmr",
        search_kwargs={"k": 3, "lambda_mult": 0.7},
    )
    st.session_state.vectordb  = vectordb
    st.session_state.retriever = retriever
    st.session_state.db_built  = True
    push_log("✅ Vector store loaded from disk!")
    return True


def get_llm():
    from langchain_groq import ChatGroq
    return ChatGroq(
        model="llama-3.1-8b-instant",
        temperature=0.7,
        groq_api_key=os.getenv("GROQ_API_KEY"),
    )


def corrective_rag(query: str, k: int = 3):
    retriever = st.session_state.retriever
    if retriever is None:
        return "No documents available. Please upload files and build the vector store first.", []

    # Override k
    retriever.search_kwargs["k"] = k

    push_log("🔍 Retrieving relevant chunks...")
    retrieved = retriever.invoke(query)
    context  = "\n".join(d.page_content for d in retrieved)
    sources  = list({d.metadata.get("source", "unknown") for d in retrieved})

    llm = get_llm()

    push_log("🤔 Checking relevance...")
    eval_prompt = (
        f"Query: {query}\n"
        f"Retrieved Context:\n{context}\n\n"
        "Are these documents relevant enough to answer the query?\n"
        "Respond strictly with YES or NO."
    )
    evaluation = llm.invoke(eval_prompt).content.strip()
    push_log(f"📊 Relevance: {evaluation}")

    if "NO" in evaluation.upper():
        push_log("✏️ Rewriting query for better retrieval...")
        rewrite_prompt = (
            f"The query '{query}' did not retrieve relevant documents.\n"
            "Rewrite it to improve retrieval quality. Only return the improved query."
        )
        improved = llm.invoke(rewrite_prompt).content.strip()
        push_log(f"🔄 Improved query: {improved[:60]}...")
        retrieved = retriever.invoke(improved)
        context  = "\n".join(d.page_content for d in retrieved)
        sources  = list({d.metadata.get("source", "unknown") for d in retrieved})

    push_log("💬 Generating final answer...")
    final_prompt = (
        "Answer the question using ONLY the context below.\n"
        f"Context:\n{context}\n\n"
        f"Question: {query}\n\n"
        "Also mention the sources used at the end."
    )
    answer = llm.invoke(final_prompt).content
    push_log("✅ Answer ready!")
    return answer, sources


# ── sidebar ───────────────────────────────────────────────────────────────────
with st.sidebar:
    st.markdown("## ⚙️ Settings")
    st.markdown("---")

    groq_key = os.getenv("GROQ_API_KEY")
    if groq_key:
        st.success("Groq API Key Loaded ✅")
    else:
        st.error("Groq API Key NOT found ❌")
        manual_key = st.text_input("Enter Groq API Key", type="password")
        if manual_key:
            os.environ["GROQ_API_KEY"] = manual_key
            st.success("Key set ✅")

    st.markdown("---")
    st.markdown("### 🔎 Retrieval Settings")
    k_value = st.slider("Top-K Results", 1, 6, 3)

    st.markdown("---")
    st.markdown("### 🗄️ Vector Store")
    db_status = "🟢 Ready" if st.session_state.db_built else "🔴 Not built"
    st.markdown(f"**Status:** {db_status}")

    if st.button("🚀 Build / Rebuild Vector DB"):
        st.session_state.backend_logs = []
        clear_chroma()
        with st.spinner("Building vector store…"):
            ok = build_vectorstore()
        if ok:
            st.success("Vector DB built!")
        else:
            st.error("No documents found.")

    if not st.session_state.db_built and os.path.exists(CHROMA_PATH) and os.listdir(CHROMA_PATH):
        if st.button("📂 Load Existing DB from Disk"):
            st.session_state.backend_logs = []
            with st.spinner("Loading…"):
                load_existing_vectorstore()
            st.success("Loaded!")

    st.markdown("---")
    st.caption("Built with ❤️ using LangChain + Groq")


# ── main area: tabs ───────────────────────────────────────────────────────────
st.markdown('<p class="label-cyan">Document Intelligence Platform</p>', unsafe_allow_html=True)
st.title("📚 QGenie ")

tab_chat, tab_upload, tab_history, tab_logs = st.tabs(
    ["💬 Chat", "📁 Upload Documents", "📜 Chat History", "🖥️ Backend Logs"]
)


# ════════════════════════════════════════════════════════
# TAB 1 – CHAT
# ════════════════════════════════════════════════════════
with tab_chat:
    col_chat, col_tips = st.columns([3, 1])

    with col_tips:
        st.markdown("#### 💡 Tips")
        st.markdown("""
- Ask specific questions  
- *"Summarize this document"*  
- *"What does it say about X?"*  
- Upload docs in the **Upload** tab first  
- Build the vector DB from the sidebar  
        """)

        if st.session_state.chat_history:
            if st.button("🗑 Clear Chat"):
                st.session_state.chat_history = []
                st.rerun()

            # Download chat history
            lines = []
            for i, item in enumerate(st.session_state.chat_history, 1):
                lines.append(f"--- Q&A #{i} ---")
                lines.append(f"Q: {item['question']}")
                lines.append(f"A: {item['answer']}")
                if item["sources"]:
                    lines.append(f"Sources: {', '.join(os.path.basename(s) for s in item['sources'])}")
                lines.append("")
            st.download_button(
                "⬇️ Download History",
                data="\n".join(lines),
                file_name="chat_history.txt",
                mime="text/plain",
            )

    with col_chat:
        # Render existing chat history
        for item in st.session_state.chat_history:
            sources_str = ", ".join(os.path.basename(s) for s in item["sources"]) if item["sources"] else ""
            st.markdown(
                f'<div class="chat-turn">'
                f'<div style="display:flex;justify-content:flex-end;margin-bottom:8px;">'
                f'<div class="user-bubble">{item["question"]}</div></div>'
                f'<div class="ai-bubble">{item["answer"]}</div>'
                + (f'<div class="source-tag">📂 {sources_str}</div>' if sources_str else "")
                + "</div>",
                unsafe_allow_html=True,
            )

        if not st.session_state.chat_history:
            st.markdown(
                '<div style="text-align:center;padding:60px 0;color:#2d3748;">'
                '<div style="font-size:40px;">💬</div>'
                '<div style="font-size:15px;font-weight:600;margin-top:8px;">No messages yet</div>'
                '<div style="font-size:13px;margin-top:4px;">Ask a question about your uploaded documents</div>'
                "</div>",
                unsafe_allow_html=True,
            )

        # Chat input (always at bottom)
        user_query = st.chat_input("Ask something about your documents…")

        if user_query:
            if not st.session_state.db_built:
                # Try auto-loading from disk first
                if not load_existing_vectorstore():
                    st.warning("⚠️ Please build the vector database first (sidebar → 🚀 Build Vector DB).")
                    st.stop()

            st.session_state.backend_logs = []
            with st.spinner("🤖 Thinking…"):
                answer, sources = corrective_rag(user_query, k=k_value)

            st.session_state.chat_history.append({
                "question": user_query,
                "answer": answer,
                "sources": sources,
            })
            st.rerun()


# ════════════════════════════════════════════════════════
# TAB 2 – UPLOAD
# ════════════════════════════════════════════════════════
with tab_upload:
    col_up, col_files = st.columns([3, 2])

    with col_up:
        st.markdown('<p class="label-cyan">Document Ingestion</p>', unsafe_allow_html=True)
        st.subheader("Upload Documents")
        st.markdown("Upload **PDF** or **TXT** files. They'll be indexed when you build the vector store.")

        uploaded = st.file_uploader(
            "Drop files here or click to browse",
            type=["pdf", "txt"],
            accept_multiple_files=True,
            label_visibility="collapsed",
        )

        col_btn1, col_btn2 = st.columns([2, 1])
        with col_btn1:
            save_btn = st.button("⬆️ Save & Index Files")
        with col_btn2:
            st.markdown('<div class="btn-danger">', unsafe_allow_html=True)
            clear_all_btn = st.button("🗑 Clear All Files")
            st.markdown("</div>", unsafe_allow_html=True)

        if save_btn and uploaded:
            saved = []
            for f in uploaded:
                dest = os.path.join(DOCUMENT_DIR, f.name)
                with open(dest, "wb") as fp:
                    fp.write(f.read())
                saved.append(f.name)
            clear_chroma()
            st.success(f"✅ Saved {len(saved)} file(s): {', '.join(saved)}\nVector store invalidated — rebuild from the sidebar.")
        elif save_btn:
            st.warning("⚠️ No files selected.")

        if clear_all_btn:
            for fname in list_docs():
                os.remove(os.path.join(DOCUMENT_DIR, fname))
            clear_chroma()
            st.success("All documents cleared.")
            st.rerun()

    with col_files:
        st.markdown('<p class="label-cyan">Indexed Files</p>', unsafe_allow_html=True)
        docs = list_docs()
        if not docs:
            st.markdown(
                '<div style="color:#2d3748;font-family:Courier New,monospace;font-size:13px;">'
                "No files indexed yet</div>",
                unsafe_allow_html=True,
            )
        else:
            for fname in docs:
                col_f, col_del = st.columns([5, 1])
                with col_f:
                    st.markdown(
                        f'<div class="file-badge">📄 {fname}</div>',
                        unsafe_allow_html=True,
                    )
                with col_del:
                    if st.button("✕", key=f"del_{fname}"):
                        os.remove(os.path.join(DOCUMENT_DIR, fname))
                        clear_chroma()
                        st.rerun()


# ════════════════════════════════════════════════════════
# TAB 3 – HISTORY
# ════════════════════════════════════════════════════════
with tab_history:
    st.markdown('<p class="label-cyan">Conversation History</p>', unsafe_allow_html=True)
    st.subheader("Chat History")

    if not st.session_state.chat_history:
        st.info("No conversation history yet. Start chatting in the **Chat** tab.")
    else:
        for i, item in enumerate(reversed(st.session_state.chat_history), 1):
            with st.expander(f"Q{len(st.session_state.chat_history) - i + 1}: {item['question'][:80]}…"):
                st.markdown(f"**Question:** {item['question']}")
                st.markdown("**Answer:**")
                st.markdown(item["answer"])
                if item["sources"]:
                    src_names = [os.path.basename(s) for s in item["sources"]]
                    st.markdown(f"📂 **Sources:** {', '.join(src_names)}")


# ════════════════════════════════════════════════════════
# TAB 4 – BACKEND LOGS
# ════════════════════════════════════════════════════════
with tab_logs:
    st.markdown('<p class="label-cyan">System Logs</p>', unsafe_allow_html=True)
    st.subheader("Backend Logs")

    col_l1, col_l2 = st.columns([3, 1])
    with col_l2:
        if st.button("🗑 Clear Logs"):
            st.session_state.backend_logs = []
            st.rerun()

    logs = st.session_state.backend_logs
    if not logs:
        st.markdown(
            '<div class="log-box"><span style="color:#2d3748;">— no log entries —</span></div>',
            unsafe_allow_html=True,
        )
    else:
        lines_html = "".join(f'<div class="log-line">› {line}</div>' for line in logs)
        st.markdown(f'<div class="log-box">{lines_html}</div>', unsafe_allow_html=True)