backend/gradio_ui/documents_tab.py

# documents_tab.py
"""Document upload, list, delete, and selection."""

from __future__ import annotations

from typing import List, Tuple

import gradio as gr

from core.bootstrap import get_db, get_ingestion_service
from gradio_ui.state import SearchMode, footer_hint


def _doc_table_rows() -> List[List[str]]:
    db = get_db()
    return [
        [
            doc.get("document_id", ""),
            doc.get("document_name", ""),
            str(doc.get("chunk_count", 0)),
        ]
        for doc in db.list_documents()
    ]


def _doc_choices() -> List[Tuple[str, str]]:
    db = get_db()
    return [
        (doc.get("document_name") or doc.get("document_id", ""), doc.get("document_id", ""))
        for doc in db.list_documents()
        if doc.get("document_id")
    ]


def _choice_label(doc_id: str, name: str, chunks: str) -> str:
    return f"{name} · {chunks} chunks"


def _choices_with_meta(rows: List[List[str]]) -> List[Tuple[str, str]]:
    return [
        (_choice_label(doc_id, name, chunks), doc_id)
        for doc_id, name, chunks in rows
        if doc_id
    ]


def _merge_selection(
    selected: List[str],
    doc_id: str,
    search_mode: SearchMode,
) -> List[str]:
    selected = list(selected or [])
    if not doc_id:
        return selected
    if search_mode == "single":
        return [doc_id]
    if doc_id not in selected:
        selected.append(doc_id)
    return selected


def upload_document(
    file_path: str | None,
    selected: List[str],
    search_mode: SearchMode,
) -> Tuple[str, List[List[str]], gr.update, gr.update, List[str], str]:
    rows = _doc_table_rows()
    choices = _choices_with_meta(rows)
    if not file_path:
        return (
            "",
            rows,
            gr.update(choices=choices, value=selected),
            gr.update(choices=[c[1] for c in _doc_choices()]),
            selected,
            footer_hint(search_mode, selected, len(rows), "documents"),
        )

    with open(file_path, "rb") as handle:
        content = handle.read()

    filename = file_path.rsplit("\\", 1)[-1].rsplit("/", 1)[-1]
    ext = filename.lower().rsplit(".", 1)[-1]
    ingestion = get_ingestion_service()

    if ext == "pdf":
        result = ingestion.ingest_pdf(content, filename)
    elif ext in ("png", "jpg", "jpeg", "tiff", "bmp", "webp"):
        result = ingestion.ingest_image(content, filename)
    else:
        return (
            "Unsupported file type. Use PDF or image files.",
            rows,
            gr.update(choices=choices, value=selected),
            gr.update(choices=[c[1] for c in _doc_choices()]),
            selected,
            footer_hint(search_mode, selected, len(rows), "documents"),
        )

    doc_id = result.get("document_id", "")
    rows = _doc_table_rows()
    choices = _choices_with_meta(rows)
    new_selected = _merge_selection(selected, doc_id, search_mode)
    msg = f"Indexed **{filename}** ({result.get('chunks_indexed', 0)} chunks)."
    return (
        msg,
        rows,
        gr.update(choices=choices, value=new_selected),
        gr.update(choices=[c[1] for c in _doc_choices()]),
        new_selected,
        footer_hint(search_mode, new_selected, len(rows), "documents"),
    )


def delete_document(
    doc_id: str,
    selected: List[str],
    search_mode: SearchMode,
) -> Tuple[str, List[List[str]], gr.update, List[str], str]:
    rows = _doc_table_rows()
    choices = _choices_with_meta(rows)
    if not doc_id:
        return (
            "Select a document to delete.",
            rows,
            gr.update(choices=choices, value=selected),
            selected,
            footer_hint(search_mode, selected, len(rows), "documents"),
        )

    get_db().delete_document(doc_id)
    new_selected = [d for d in (selected or []) if d != doc_id]
    rows = _doc_table_rows()
    choices = _choices_with_meta(rows)
    return (
        "Deleted document.",
        rows,
        gr.update(choices=choices, value=new_selected),
        new_selected,
        footer_hint(search_mode, new_selected, len(rows), "documents"),
    )


def on_doc_selection(selected: List[str], search_mode: SearchMode) -> Tuple[List[str], str, gr.update]:
    if search_mode == "single" and selected and len(selected) > 1:
        selected = [selected[-1]]
    rows = _doc_table_rows()
    choices = _choices_with_meta(rows)
    return (
        selected,
        footer_hint(search_mode, selected, len(_doc_choices())),
        gr.update(choices=choices, value=selected),
    )


def on_search_mode_change(search_mode: SearchMode, selected: List[str]) -> Tuple[List[str], str]:
    if search_mode == "single" and selected and len(selected) > 1:
        selected = [selected[0]]
    return selected, footer_hint(search_mode, selected, len(_doc_choices()))


def doc_select_update(selected: List[str] | None = None) -> gr.update:
    rows = _doc_table_rows()
    return gr.update(choices=_choices_with_meta(rows), value=selected or [])


def build_documents_panel() -> dict:
    with gr.Column(elem_classes=["fs-modal-card", "fs-docs-card"]):
        with gr.Row(elem_classes=["fs-modal-header-row"]):
            gr.HTML(
                """
                <div class="fs-modal-header">
                  <h2>Documents</h2>
                  <p>Upload PDFs or images, then select which to use for RAG</p>
                </div>
                """,
                padding=False,
            )
            close_btn = gr.Button("✕", elem_classes=["fs-modal-close"], scale=0, min_width=44)

        upload_status = gr.Markdown("", elem_classes=["fs-upload-status"])
        file_upload = gr.File(
            label="Drag & drop PDF or images · PDF, PNG, JPG, TIFF",
            file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".webp"],
            type="filepath",
            elem_classes=["fs-upload-zone"],
            file_count="single",
            height=140,
        )
        upload_btn = gr.Button("Upload & index", elem_classes=["fs-gold-btn"])
        doc_select = gr.CheckboxGroup(
            label="Select documents for RAG",
            info="Hybrid RAG: pick any subset · Single Document: one at a time",
            choices=[],
            value=[],
            elem_classes=["fs-doc-select"],
        )
        done_btn = gr.Button("Done — back to chat", elem_classes=["fs-gold-btn", "fs-done-btn"])
        doc_table = gr.Dataframe(
            headers=["Document ID", "Name", "Chunks"],
            datatype=["str", "str", "str"],
            interactive=False,
            visible=False,
        )
        with gr.Row(elem_classes=["fs-doc-delete-row"]):
            delete_dropdown = gr.Dropdown(label="Delete document", choices=[], interactive=True)
            delete_btn = gr.Button("Delete", variant="stop", scale=0, min_width=100)

    return {
        "upload_status": upload_status,
        "file_upload": file_upload,
        "upload_btn": upload_btn,
        "doc_table": doc_table,
        "doc_select": doc_select,
        "close_btn": close_btn,
        "done_btn": done_btn,
        "delete_dropdown": delete_dropdown,
        "delete_btn": delete_btn,
        "doc_choices": _doc_choices,
        "doc_table_rows": _doc_table_rows,
        "upload_document": upload_document,
        "delete_document": delete_document,
        "on_doc_selection": on_doc_selection,
        "on_search_mode_change": on_search_mode_change,
        "doc_select_update": doc_select_update,
    }