app.py

"""
app.py  –  Industrial Document Q&A System
Streamlit-based UI that showcases Endee vector database for semantic
search and RAG over industrial safety / gas detection documentation.
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))

import streamlit as st
from indexer import DocumentIndexer
from retriever import Retriever, RAGAnswerGenerator


# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="IndustrialDocSearch | Powered by Endee",
    page_icon="🔬",
    layout="wide",
)

# ── CSS ───────────────────────────────────────────────────────────────────────
st.markdown("""
<style>
.doc-card {
    background: #f8f9fa;
    border-left: 4px solid #1f77b4;
    padding: 1rem 1.2rem;
    border-radius: 6px;
    margin-bottom: 1rem;
}
.doc-title { font-size: 1.05rem; font-weight: 600; color: #1a1a2e; }
.doc-meta  { font-size: 0.8rem;  color: #6c757d; margin: 0.2rem 0 0.5rem; }
.score-badge {
    background: #1f77b4; color: white;
    padding: 2px 10px; border-radius: 12px; font-size: 0.78rem;
}
.answer-box {
    background: #e8f4fd;
    border: 1px solid #90caf9;
    border-radius: 8px;
    padding: 1.2rem;
    margin-bottom: 1.5rem;
}
</style>
""", unsafe_allow_html=True)


# ── Sidebar ───────────────────────────────────────────────────────────────────
with st.sidebar:
    st.image("https://endee.io/favicon.ico", width=40)
    st.title("⚙️ Settings")
    endee_host = st.text_input("Endee Host", value="http://localhost:8080")
    auth_token = st.text_input("Auth Token (optional)", type="password", value="")
    top_k = st.slider("Results to retrieve (top-k)", 1, 10, 5)
    gemini_key = st.text_input("Google Gemini API Key (optional)", type="password",
                               help="Enables AI-generated answers. Free tier available.")

    if gemini_key:
        os.environ["GOOGLE_API_KEY"] = gemini_key

    st.markdown("---")
    st.subheader("🗄️ Index Management")
    if st.button("🔄 Re-index Documents", use_container_width=True):
        with st.spinner("Indexing documents into Endee…"):
            try:
                indexer = DocumentIndexer(host=endee_host, auth_token=auth_token)
                indexer.setup(recreate=True)
                st.success("✅ Documents indexed successfully!")
                st.cache_resource.clear()
            except Exception as e:
                st.error(f"Indexing failed: {e}")

    st.markdown("---")
    st.caption("Built with [Endee](https://endee.io) vector database")


# ── Cached resources ──────────────────────────────────────────────────────────
@st.cache_resource
def get_retriever(host, token):
    return Retriever(host=host, auth_token=token)

@st.cache_resource
def get_rag():
    return RAGAnswerGenerator()


# ── Main page ─────────────────────────────────────────────────────────────────
st.title("🔬 IndustrialDocSearch")
st.markdown(
    "**Semantic Search & RAG Q&A** over industrial safety and gas detection documents, "
    "powered by the [Endee](https://github.com/endee-io/endee) vector database."
)

# Category filter
CATEGORIES = [
    "All", "Safety", "Maintenance", "Calibration", "Configuration",
    "Installation", "Inspection", "Monitoring", "Integration",
    "Technical Reference", "Compliance", "Emergency",
]
col_q, col_cat = st.columns([3, 1])
with col_q:
    query = st.text_input(
        "🔍 Ask a question or search for a topic",
        placeholder="e.g. How do I calibrate an H2S sensor?",
        label_visibility="collapsed",
    )
with col_cat:
    category = st.selectbox("Category", CATEGORIES, label_visibility="collapsed")

search_btn = st.button("Search", type="primary", use_container_width=False)

# ── Search logic ──────────────────────────────────────────────────────────────
if search_btn and query.strip():
    retriever = get_retriever(endee_host, auth_token)
    rag = get_rag()

    cat_filter = None if category == "All" else category

    with st.spinner("Querying Endee vector database…"):
        try:
            results = retriever.search(query, top_k=top_k, category_filter=cat_filter)
        except Exception as e:
            st.error(f"Search failed: {e}")
            st.info(
                "Make sure the Endee server is running and documents are indexed. "
                "Use the sidebar **Re-index Documents** button to populate the database."
            )
            st.stop()

    if not results:
        st.warning("No results found. Try a different query or re-index the documents.")
        st.stop()

    # ── RAG Answer ────────────────────────────────────────────────────────────
    st.subheader("💡 Answer")
    with st.spinner("Generating answer…"):
        answer, is_llm = rag.generate(query, results[:3])

    st.markdown(f'<div class="answer-box">{answer}</div>', unsafe_allow_html=True)
    if is_llm:
        st.caption("🤖 AI-generated answer grounded in retrieved documents (Gemini)")
    else:
        st.caption("📄 Extractive answer from top result (set Gemini API key for AI answers)")

    # ── Retrieved Documents ───────────────────────────────────────────────────
    st.subheader(f"📚 Top {len(results)} Retrieved Documents")
    for doc in results:
        score_str = f"{doc['score']:.4f}" if doc["score"] is not None else "N/A"
        st.markdown(
            f'<div class="doc-card">'
            f'<div class="doc-title">{doc["title"]}</div>'
            f'<div class="doc-meta">📂 {doc["category"]} &nbsp;|&nbsp; '
            f'<span class="score-badge">Similarity: {score_str}</span></div>'
            f'<p style="font-size:0.92rem;color:#333;margin:0">{doc["content"][:400]}…</p>'
            f'</div>',
            unsafe_allow_html=True,
        )

elif search_btn:
    st.warning("Please enter a search query.")

# ── Example queries ───────────────────────────────────────────────────────────
st.markdown("---")
st.subheader("💬 Example Queries")
examples = [
    "What are the alarm thresholds for oxygen deficiency?",
    "How often should a dew point analyzer be calibrated?",
    "Where should gas sensors be placed for H2S detection?",
    "What should I do in case of a gas leak?",
    "Difference between catalytic bead and infrared sensors?",
    "OSHA requirements for confined space entry",
]
cols = st.columns(3)
for i, ex in enumerate(examples):
    if cols[i % 3].button(ex, key=f"ex_{i}", use_container_width=True):
        st.session_state["_example"] = ex
        st.rerun()

# Pick up example query selection
if "_example" in st.session_state:
    st.query_params["q"] = st.session_state.pop("_example")


# ── System status ─────────────────────────────────────────────────────────────
st.markdown("---")
with st.expander("🔧 System Architecture"):
    st.markdown("""
    ```
    User Query
        │
        ▼
    Sentence Transformer (all-MiniLM-L6-v2)
        │  Produces 384-dim cosine-normalized vector
        ▼
    Endee Vector Database  ← Documents indexed here at startup
        │  HNSW approximate nearest-neighbor search
        ▼
    Top-K Results (with metadata)
        │
        ▼
    (Optional) Google Gemini 1.5 Flash
        │  RAG prompt: context + question → grounded answer
        ▼
    Streamlit UI
    ```
    - **Endee** handles vector storage, indexing (HNSW), and k-NN retrieval
    - **sentence-transformers** generates free, local embeddings (no API key needed)
    - **Gemini 1.5 Flash** (optional, free tier) generates natural language answers
    """)