streamlit_app.py

from pathlib import Path
from tempfile import NamedTemporaryFile

import streamlit as st

from app.config import get_settings
from app.ingest import ingest_local_file
from app.network import clear_dead_local_proxy
from app.rag import answer_query
from app.utils import SUPPORTED_FILE_EXTENSIONS


get_settings.cache_clear()
cleared_proxy_vars = clear_dead_local_proxy()

st.set_page_config(
    page_title="Multilingual RAG",
    layout="wide",
)


def show_health() -> None:
    settings = get_settings()
    st.success("Application is running.")
    st.json(
        {
            "status": "ok",
            "collection": settings.CHROMA_COLLECTION,
            "persist_dir": settings.CHROMA_PERSIST_DIR,
            "embedding_model": settings.EMBEDDING_MODEL,
            "llm_provider": "gemini",
            "gemini_model": settings.GEMINI_MODEL,
            "gemini_key_configured": bool(settings.GEMINI_API_KEY),
        }
    )


def save_uploaded_file(uploaded_file) -> Path:
    suffix = Path(uploaded_file.name).suffix
    with NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
        temp_file.write(uploaded_file.getbuffer())
        return Path(temp_file.name)


st.title("Multilingual RAG")
st.caption("Upload PDF, TXT, or Markdown files, then ask questions across languages.")

with st.sidebar:
    st.header("Status")
    if st.button("Check Health", use_container_width=True):
        st.session_state.health_checked = True

    if st.session_state.get("health_checked"):
        show_health()

    st.divider()
    st.header("Settings")
    settings = get_settings()
    st.write(f"Collection: `{settings.CHROMA_COLLECTION}`")
    st.write(f"LLM: `gemini`")
    st.write(f"Top K default: `{settings.TOP_K}`")
    if cleared_proxy_vars:
        st.caption("Ignored dead local proxy settings for model downloads.")

tab_ingest, tab_query = st.tabs(["Ingest File", "Ask Question"])

with tab_ingest:
    st.subheader("Upload and Ingest")
    uploaded_file = st.file_uploader(
        "Choose a file",
        type=[extension.lstrip(".") for extension in SUPPORTED_FILE_EXTENSIONS],
    )

    if uploaded_file is not None:
        st.write(f"Selected: `{uploaded_file.name}`")

    if st.button("Ingest Uploaded File", type="primary", disabled=uploaded_file is None):
        temp_path: Path | None = None
        try:
            temp_path = save_uploaded_file(uploaded_file)
            with st.spinner("Embedding and storing chunks. First run may download the embedding model."):
                chunks_added = ingest_local_file(temp_path, uploaded_file.name)
            st.success(f"Ingestion complete. Added {chunks_added} chunks.")
        except Exception as exc:
            st.error(f"Ingestion failed: {exc}")
        finally:
            if temp_path is not None:
                temp_path.unlink(missing_ok=True)

with tab_query:
    st.subheader("Ask a Question")
    query = st.text_area(
        "Question",
        placeholder="Example: What does the Hindi document say about AI?",
        height=120,
    )
    top_k = st.slider("Number of chunks to retrieve", min_value=1, max_value=20, value=settings.TOP_K)

    if st.button("Get Answer", type="primary", disabled=not query.strip()):
        try:
            with st.spinner("Retrieving context and generating answer."):
                answer, sources = answer_query(query.strip(), top_k)

            st.markdown("### Answer")
            st.write(answer)

            st.markdown("### Sources")
            if not sources:
                st.info("No sources returned.")
            for index, source in enumerate(sources, start=1):
                metadata = source.metadata
                label = metadata.get("source", f"Source {index}")
                with st.expander(f"{index}. {label}"):
                    st.json(metadata)
                    st.write(source.content)
        except Exception as exc:
            st.error(f"Query failed: {exc}")