chatbot/app.py

"""
KALRO Maize Research Chatbot — Streamlit UI
Run: streamlit run chatbot/app.py
"""

import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from dotenv import load_dotenv
load_dotenv(Path(__file__).parent.parent / ".env")  # no-op on HF Spaces (secrets via env vars)

import streamlit as st
from chatbot.chat import Chatbot
from chatbot.retriever import Retriever

# ── page config ────────────────────────────────────────────────────────────────

st.set_page_config(
    page_title="KALRO Maize Research Assistant",
    page_icon="🌽",
    layout="wide",
)

# ── shared retriever (loaded once per server, shared across all user sessions) ─

@st.cache_resource(show_spinner="Building knowledge base — first run only, ~30s...")
def _load_retriever() -> Retriever:
    db_path = Path(__file__).parent / "db" / "chroma.sqlite3"
    if not db_path.exists():
        embed_script = Path(__file__).parent / "pipeline" / "embed.py"
        subprocess.run([sys.executable, str(embed_script)], check=True)
    return Retriever()

# ── session state ──────────────────────────────────────────────────────────────

if "chatbot" not in st.session_state:
    st.session_state.chatbot = Chatbot(retriever=_load_retriever())

if "messages" not in st.session_state:
    st.session_state.messages = []     # list of {role, content, meta}


# ── sidebar ────────────────────────────────────────────────────────────────────

with st.sidebar:
    st.markdown("## 🌽 KALRO Maize Research Assistant")
    st.markdown(
        "Ask questions about maize production, varieties, diseases, "
        "soil & water management, and extension in Kenya."
    )
    st.divider()

    if st.button("New conversation", use_container_width=True):
        st.session_state.chatbot.reset()
        st.session_state.messages = []
        st.rerun()

    st.divider()
    st.markdown("**Knowledge base**")
    st.markdown("- 16 research papers & manuals\n- Wiki: 163 structured chunks\n- PDFs: 1,013 raw text chunks")
    st.markdown("**Model:** Claude Sonnet 4.6")
    st.markdown("**Embedding:** all-MiniLM-L6-v2")

    # Show sources for the last answer
    if st.session_state.messages:
        last = st.session_state.messages[-1]
        if last["role"] == "assistant" and last.get("meta"):
            meta = last["meta"]
            st.divider()
            st.markdown("**Sources used (last answer)**")
            pdf_flag = meta.get("pdf_fallback_used", False)
            if pdf_flag:
                st.warning("PDF fallback triggered — answer draws from raw document text")
            for chunk in meta.get("chunks_used", []):
                tag = "📄 Wiki" if chunk.layer == "wiki" else "📑 PDF"
                dist = f"{chunk.distance:.2f}"
                st.markdown(f"{tag} `{dist}` {chunk.citation()}")


# ── main chat area ─────────────────────────────────────────────────────────────

st.markdown("## 🌽 KALRO Maize Research Assistant")
st.caption("Answers grounded in KALRO research papers, field trials, and training manuals.")

# Render conversation history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])
        if msg["role"] == "assistant" and msg.get("meta"):
            meta = msg["meta"]
            with st.expander("Sources", expanded=False):
                for chunk in meta.get("chunks_used", []):
                    tag = "📄 Wiki" if chunk.layer == "wiki" else "📑 PDF"
                    st.markdown(f"- {tag} **{chunk.citation()}** (dist: {chunk.distance:.2f})")

# Chat input
if prompt := st.chat_input("Ask about Kenya maize research..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("Searching knowledge base..."):
            response = st.session_state.chatbot.ask(prompt)

        st.markdown(response.answer)

        with st.expander("Sources", expanded=False):
            if response.pdf_fallback_used:
                st.warning("PDF fallback triggered — wiki confidence was low; answer also draws from raw document text.")
            for chunk in response.chunks_used:
                tag = "📄 Wiki" if chunk.layer == "wiki" else "📑 PDF"
                st.markdown(f"- {tag} **{chunk.citation()}** (dist: {chunk.distance:.2f})")

    st.session_state.messages.append({
        "role": "assistant",
        "content": response.answer,
        "meta": {
            "pdf_fallback_used": response.pdf_fallback_used,
            "chunks_used": response.chunks_used,
        },
    })