test_ui.py

import requests
import urllib3
import streamlit as st

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

API_URL = "https://idnameraj-rag-vs.hf.space"

st.set_page_config(page_title="LoomChat RAG", layout="centered")

st.title("LoomChat RAG")
st.caption("Upload documents and ask questions -- powered by Ollama + LanceDB")

# ── Status ──────────────────────────────────────────────────────────────────
@st.cache_data(ttl=30)
def fetch_health():
    try:
        r = requests.get(API_URL, timeout=10, verify=False)
        return r.json()
    except Exception:
        return None

info = fetch_health()
if info:
    c1, c2, c3 = st.columns(3)
    ollama_status = info.get("ollama", "unknown")
    c1.metric("Ollama", ollama_status)
    c2.metric("Model", info.get("ollama_model", "-"))
    c3.metric("Chunks in DB", info.get("chunks_in_db", 0))
else:
    st.error("API unreachable")
    st.stop()

st.divider()

# ── Upload ──────────────────────────────────────────────────────────────────
with st.expander("Upload a document", expanded=not bool(st.session_state.get("messages"))):
    uploaded_file = st.file_uploader(
        "Choose a file (PDF, TXT, CSV, DOCX)",
        type=["pdf", "txt", "csv", "docx"],
    )
    if uploaded_file:
        st.caption(f"{uploaded_file.name} - {uploaded_file.size / 1024:.0f} KB")
        if st.button("Upload"):
            with st.spinner("Processing..."):
                resp = requests.post(
                    f"{API_URL}/upload",
                    files={"file": (uploaded_file.name, uploaded_file, "application/octet-stream")},
                    timeout=120,
                    verify=False,
                )
            if resp.status_code == 200:
                data = resp.json()
                st.success(f"{data['chunks_stored']} chunks stored from {data['filename']}")
                fetch_health.clear()
                st.rerun()
            else:
                try:
                    detail = resp.json().get("detail", resp.text)
                except Exception:
                    detail = resp.text
                st.error(f"Upload failed: {detail}")

st.divider()

# ── Chat ────────────────────────────────────────────────────────────────────
if "messages" not in st.session_state:
    st.session_state.messages = []

# Render chat history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        if msg["role"] == "user":
            st.write(msg["content"])
        else:
            st.markdown(f"**Answer:** {msg['answer']}")
            if msg.get("sources"):
                with st.expander(f"View {len(msg['sources'])} source chunks"):
                    for i, src in enumerate(msg["sources"], 1):
                        st.markdown(f"**Chunk {i}**")
                        st.code(src, language=None)

# Chat input
query = st.chat_input("Ask a question about your documents...")

if query:
    # Show user message
    st.session_state.messages.append({"role": "user", "content": query})
    with st.chat_message("user"):
        st.write(query)

    # Get and show response
    with st.chat_message("assistant"):
        with st.spinner("Searching & generating..."):
            try:
                resp = requests.post(
                    f"{API_URL}/query",
                    json={"query": query},
                    timeout=180,
                    verify=False,
                )
            except requests.exceptions.Timeout:
                st.error("Request timed out. The model may be loading.")
                st.session_state.messages.append(
                    {"role": "assistant", "answer": "Request timed out.", "sources": []}
                )
                st.stop()

        if resp.status_code == 200:
            data = resp.json()
            answer = data["answer"]
            sources = data.get("sources", [])

            st.markdown(f"**Answer:** {answer}")
            if sources:
                with st.expander(f"View {len(sources)} source chunks"):
                    for i, src in enumerate(sources, 1):
                        st.markdown(f"**Chunk {i}**")
                        st.code(src, language=None)

            st.session_state.messages.append(
                {"role": "assistant", "answer": answer, "sources": sources}
            )
        else:
            try:
                detail = resp.json().get("detail", resp.text)
            except Exception:
                detail = resp.text
            st.error(detail)
            st.session_state.messages.append(
                {"role": "assistant", "answer": f"Error: {detail}", "sources": []}
            )