app.py

# app.py
import os
from typing import Dict, List

import streamlit as st
from langchain_community.llms import Ollama

from modules import (
    DataIngestor,
    RAGSettings,
    build_context_prompt,
    build_vector_store,
    list_collections,
)
from modules.config import ensure_directory
from modules.qdrant_utils import collection_exists


def run_app(settings: RAGSettings | None = None) -> None:
    """Render the Streamlit UI for ingesting and chatting with document collections."""
    st.set_page_config(page_title="Chat with Documents", layout="wide")

    cfg = settings or RAGSettings()
    ensure_directory(cfg.data_path)
    ensure_directory(cfg.log_dir)

    st.session_state.setdefault("messages", [])
    st.session_state.setdefault("selected_collection", cfg.collection_name)
    st.session_state.setdefault("active_collection", None)

    def refresh_collections() -> List[str]:
        """Fetch collection names from Qdrant sorted alphabetically."""
        collections = list_collections(cfg)
        collections.sort()
        return collections

    def write_files_to_directory(target_dir: str, files) -> List[str]:
        """Persist uploaded files to disk and return their filenames."""
        ensure_directory(target_dir)
        saved = []
        for uploaded_file in files:
            dest = os.path.join(target_dir, uploaded_file.name)
            with open(dest, "wb") as f:
                f.write(uploaded_file.getbuffer())
            saved.append(uploaded_file.name)
        return saved

    def create_progress_tracker():
        """Initialise progress state holders for ingestion feedback."""
        progress_bar = st.progress(0.0)
        stage_progress: Dict[str, float] = {"load": 0.0, "split": 0.0, "upsert": 0.0}
        stage_weights: Dict[str, float] = {"load": 0.1, "split": 0.4, "upsert": 0.5}
        stage_labels: Dict[str, str] = {
            "load": "Loading documents",
            "split": "Chunking documents",
            "upsert": "Writing to Qdrant",
        }
        return progress_bar, stage_progress, stage_weights, stage_labels

    def progress_callback_factory(progress_bar, stage_progress, stage_weights, stage_labels, status_box):
        """Return a callback that updates the Streamlit progress widgets."""
        def _callback(stage: str, current: int, total: int) -> None:
            total = max(total, 1)
            stage_progress[stage] = min(current / total, 1.0)
            cumulative = sum(stage_progress[s] * stage_weights[s] for s in stage_progress)
            progress_bar.progress(min(cumulative, 1.0))
            label = stage_labels.get(stage, stage.title())
            status_box.info(f"{label}: {current}/{total}")

        return _callback

    with st.sidebar:
        st.title("Menu")
        choice = st.radio("Choose page", ("Ingest Documents", "Chat with LLM"))
        if st.button("Refresh collections"):
            st.session_state.available_collections = refresh_collections()

    collections_cache = st.session_state.get("available_collections")
    if collections_cache is None:
        collections_cache = refresh_collections()
        st.session_state.available_collections = collections_cache

    if choice == "Ingest Documents":
        st.header("Document Ingestion")
        existing_collections = collections_cache
        default_collection = st.session_state.get("selected_collection", cfg.collection_name)

        if existing_collections:
            mode = st.radio("Collection mode", ("Use existing collection", "Create new collection"))
            if mode == "Use existing collection":
                try:
                    default_index = existing_collections.index(default_collection)
                except ValueError:
                    default_index = 0
                collection_name = st.selectbox("Select a collection", existing_collections, index=default_index)
            else:
                collection_name = st.text_input("New collection name", value=default_collection)
        else:
            st.info("No Qdrant collections found yet. Ingest to create one.")
            collection_name = st.text_input("Collection name", value=default_collection)

        uploaded_files = st.file_uploader("Upload PDF files", type="pdf", accept_multiple_files=True)
        chunk_size = st.number_input(
            "Chunk size",
            min_value=100,
            max_value=2000,
            value=cfg.chunk_size_default,
        )
        chunk_overlap = st.number_input(
            "Chunk overlap",
            min_value=0,
            max_value=500,
            value=cfg.chunk_overlap_default,
        )

        if st.button("Process"):
            if not collection_name:
                st.warning("Please provide a target collection name.")
            elif not uploaded_files:
                st.warning("Please upload at least one PDF file.")
            else:
                target_dir = os.path.join(cfg.data_path, collection_name)
                saved_files = write_files_to_directory(target_dir, uploaded_files)

                progress_bar, stage_progress, stage_weights, stage_labels = create_progress_tracker()
                status_box = st.empty()
                status_box.info(f"Saved {len(saved_files)} file(s) to '{target_dir}'.")

                update_progress = progress_callback_factory(
                    progress_bar, stage_progress, stage_weights, stage_labels, status_box
                )

                try:
                    collection_settings = cfg.with_collection(collection_name)
                    ingestor = DataIngestor(
                        settings=collection_settings,
                        data_path=target_dir,
                        chunk_size=chunk_size,
                        chunk_overlap=chunk_overlap,
                    )
                    written = ingestor.ingest_documents(file_format="pdf", progress_callback=update_progress)
                    progress_bar.progress(1.0)

                    if written is None:
                        status_box.warning("No new documents found to ingest.")
                    else:
                        status_box.success(f"Ingestion complete. Added or updated {written} chunks in '{collection_name}'.")
                        st.session_state.selected_collection = collection_name
                        st.session_state.available_collections = refresh_collections()
                except Exception as exc:
                    status_box.error(f"Ingestion error: {exc}")
                finally:
                    progress_bar.empty()

    elif choice == "Chat with LLM":
        st.header("Chat with your Documents")
        collections = collections_cache

        if not collections:
            st.info("No Qdrant collections available. Ingest documents to create one.")
        else:
            default_collection = st.session_state.get("selected_collection", collections[0])
            try:
                default_index = collections.index(default_collection)
            except ValueError:
                default_index = 0
            selected_collection = st.selectbox("Select a collection", collections, index=default_index)
            st.session_state.selected_collection = selected_collection

            if st.session_state.active_collection != selected_collection:
                st.session_state.messages = []
                st.session_state.active_collection = selected_collection

            if not collection_exists(selected_collection, cfg):
                st.warning(f"Collection '{selected_collection}' not found on Qdrant. Refresh collections or ingest data.")
            else:
                search_k = st.number_input("Retriever results (k)", min_value=1, max_value=10, value=cfg.retrieve_k, step=1)

                vector_db = build_vector_store(settings=cfg, collection_name=selected_collection)
                retriever = vector_db.as_retriever(search_kwargs={"k": search_k})
                llm = Ollama(model=cfg.gen_model, temperature=0.1)

                for msg in st.session_state.messages:
                    with st.chat_message(msg["role"]):
                        st.markdown(msg["content"])

                if prompt_text := st.chat_input("Ask a question…"):
                    st.session_state.messages.append({"role": "user", "content": prompt_text})
                    with st.chat_message("assistant"):
                        docs = retriever.invoke(prompt_text)
                        context = "\n\n".join([doc.page_content for doc in docs])

                        prompt = build_context_prompt()
                        full_prompt = prompt.format(context=context, question=prompt_text)

                        answer = llm.invoke(full_prompt)
                        st.markdown(answer)

                        with st.expander("Retrieved chunks"):
                            if not docs:
                                st.markdown("_No relevant chunks found._")
                            for doc in docs:
                                metadata = doc.metadata or {}
                                source = metadata.get("source", "N/A")
                                page = metadata.get("page")
                                page_text = f" (page {page})" if page is not None else ""
                                st.markdown(f"**Source:** {source}{page_text}")
                                st.markdown(f"> {doc.page_content}")
                                st.markdown("---")

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


if __name__ == "__main__":
    run_app()