app.py

"""
app.py -- Gradio UI for the RAG-lite Question Answering System.

It contains the main application logic and UI components, including:
  - Confidence score display: every answer shows a confidence label
    (High / Medium / Low + percentage) derived from the generator.
  - Groq streaming: when Groq mode is selected, the answer textbox
    updates token-by-token instead of waiting for the full response.
  - Structured logging: all pipeline stages print [INFO] / [WARN] lines
    so the terminal is readable during demos and debugging.
  - Confidence is averaged across sub-questions for multi/compound inputs.
"""

import gradio as gr
from retriever import build_index, retrieve_top_chunks
from generator import generate_answer, fallback_answer, stream_groq_answer, merge_chunks
from utils import (
    clean_text,
    format_debug_info,
    format_context_display,
    format_multi_answer_output,
)
from question_parser import parse_questions, is_multi_question
from history_manager import (
    add_qa_to_history,
    get_context_previews,
    get_entry_by_preview,
    format_context_entry_display,
    format_full_history_summary,
)
import os
from dotenv import load_dotenv

load_dotenv()

# CONFIG
# Threshold is not used as a hard filter.
# Retrieval always returns the top TOP_K chunks by score.
# Scores are used only for confidence display and terminal logging.

TOP_K = 5

# UI THEME AND CSS  (module-level so launch() can access them)
CSS = """
    .answer-box textarea  { font-size: 1.05rem; font-weight: 500; }
    .context-box textarea { font-size: 0.88rem; color: #555; }
    .debug-box textarea   { font-family: monospace; font-size: 0.82rem; }
    .history-box textarea { font-family: monospace; font-size: 0.84rem; }
    .status-box textarea  { font-size: 0.9rem; }
    .conf-box textarea    { font-size: 0.9rem; font-weight: 600; }
    footer { display: none !important; }
"""

THEME = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="slate",
    neutral_hue="slate",
)


# STEP 1: INDEX TEXT

def index_text(text: str):
    """
    Chunk and embed the paragraph once. Result stored in gr.State.
    Returns (status_message, chunks, cleaned_text).
    """
    if not text or not text.strip():
        return "Please enter some text before indexing.", None, ""

    if len(text.strip().split()) < 10:
        return "Text is too short. Add at least a few sentences.", None, ""

    try:
        cleaned = clean_text(text)
        chunks  = build_index(cleaned)

        if not chunks:
            return "Could not extract content from this text.", None, ""

        word_count = len(cleaned.split())
        print(f"[INFO] Indexed {len(chunks)} chunk(s) from {word_count} words")
        msg = (
            f"Text indexed -- {len(chunks)} chunk(s) from {word_count} words. "
            f"Select a generation mode and ask questions below."
        )
        return msg, chunks, cleaned

    except Exception as e:
        print(f"[ERROR] Indexing failed: {e}")
        return f"Error while indexing: {str(e)}", None, ""


# CONFIDENCE LABEL FORMATTER

def _format_confidence(score: float) -> str:
    """Convert 0-1 float to a human-readable label for the UI textbox."""
    if score >= 0.75:
        return f"High  ({score:.0%})"
    elif score >= 0.45:
        return f"Medium  ({score:.0%})"
    elif score > 0.0:
        return f"Low  ({score:.0%})"
    return "N/A"


# CORE: SINGLE QUESTION PIPELINE

def _answer_single(
    question: str,
    chunks: list,
    mode: str,
    groq_api_key: str,
    show_debug: bool,
    active_context: str = "",
) -> dict:
    """
    Retrieval + generation for one question against the cached index.

    retrieve_top_chunks always returns the top_k chunks by score regardless of score value. The max
    score is used as a confidence signal but never blocks generation.

    active_context is the raw paragraph text passed to retrieval for
    pronoun resolution ("What are its types?" -> subject extracted from text).

    Returns dict: {question, answer, context, debug, confidence}
    """
    result = {
        "question":   question,
        "answer":     "",
        "context":    "",
        "debug":      "",
        "confidence": 0.0,
    }

    # Always retrieve -- threshold=0.0 means no filtering, just ranking
    top_chunks, top_scores = retrieve_top_chunks(
        question=question,
        chunks=chunks,
        top_k=TOP_K,
        threshold=0.0,
        context_text=active_context,
    )

    # If somehow nothing came back (empty index), handle gracefully
    if not top_chunks:
        print(f"[WARN] Index is empty for: {question[:60]}")
        result["answer"] = "No text has been indexed yet. Please index a paragraph first."
        return result

    # Log the top score so we can monitor retrieval quality in the terminal
    max_score = max(top_scores) if top_scores else 0.0
    if max_score < 0.20:
        print(f"[WARN] Low retrieval confidence ({max_score:.3f}) for: {question[:60]}")
    else:
        print(f"[INFO] Retrieval max score: {max_score:.3f} for: {question[:60]}")

    result["context"] = format_context_display(
        [c["text"] for c in top_chunks],
        [{"start_word": c["start_word"], "end_word": c["end_word"]}
         for c in top_chunks],
    )

    # generate_answer returns (answer, success, confidence)
    answer, success, confidence = generate_answer(
        context_chunks=top_chunks,
        question=question,
        mode=mode,
        groq_api_key=groq_api_key,
    )

    result["answer"]     = answer if success else fallback_answer(top_chunks[0]["text"])
    result["confidence"] = confidence if success else 0.0

    if show_debug:
        result["debug"] = format_debug_info(
            [c["text"] for c in top_chunks], top_scores
        )

    return result



# STEP 2A: SINGLE / COMPOUND ROUTING

def answer_question(
    question: str,
    chunks: list,
    active_context: str,
    search_history: list,
    mode: str,
    groq_api_key: str,
    show_debug: bool,
):
    """
    Route input through single, multi, or compound question paths.
    Returns (answer, context, debug, confidence_label, search_history).
    """
    if not chunks:
        msg = "No text indexed yet. Go to Step 1 and click Index Text first."
        return msg, "", "", "N/A", search_history

    if not question or not question.strip():
        return "Please enter a question.", "", "", "N/A", search_history

    if len(question.strip()) < 3:
        return "Question too short. Try something more specific.", "", "", "N/A", search_history

    search_history = search_history or []

    if is_multi_question(question):
        return _handle_multi_question(
            question, chunks, active_context,
            search_history, mode, groq_api_key, show_debug,
        )

    result = _answer_single(question, chunks, mode, groq_api_key, show_debug, active_context)
    search_history = add_qa_to_history(
        search_history, active_context,
        result["question"], result["answer"], mode,
    )
    return (
        result["answer"],
        result["context"],
        result["debug"],
        _format_confidence(result["confidence"]),
        search_history,
    )


# ------------------------------------------------------------------
# STEP 2B: MULTI-QUESTION HANDLER
# ------------------------------------------------------------------

def _handle_multi_question(
    raw_input: str,
    chunks: list,
    active_context: str,
    search_history: list,
    mode: str,
    groq_api_key: str,
    show_debug: bool,
):
    """
    Answer each parsed sub-question independently.
    Cached index reused for all -- no redundant embedding.
    Confidence is the average across all sub-questions.
    """
    questions          = parse_questions(raw_input)
    qa_pairs           = []
    context_parts      = []
    debug_parts        = []
    confidence_scores  = []

    for i, q in enumerate(questions, start=1):
        result = _answer_single(q, chunks, mode, groq_api_key, show_debug, active_context)
        qa_pairs.append({"question": result["question"], "answer": result["answer"]})
        confidence_scores.append(result["confidence"])

        if result["context"]:
            context_parts.append(f"Question {i}: {q}\n{result['context']}")
        if show_debug and result["debug"]:
            debug_parts.append(f"Question {i}: {q}\n{result['debug']}")

        search_history = add_qa_to_history(
            search_history, active_context,
            result["question"], result["answer"], mode,
        )

    avg_conf = (sum(confidence_scores) / len(confidence_scores)
                if confidence_scores else 0.0)

    return (
        format_multi_answer_output(qa_pairs),
        "\n\n".join(context_parts),
        "\n\n".join(debug_parts),
        _format_confidence(avg_conf),
        search_history,
    )


# ------------------------------------------------------------------
# GROQ STREAMING HANDLER  (yields partial answers to the UI)
# ------------------------------------------------------------------

def handle_ask_streaming(
    question, chunks, active_context, search_history,
    mode_label, groq_key, show_debug,
):
    """
    Generator function for Gradio streaming.

    For Groq mode:
      - Shows "Connecting to Groq API... (max 15 seconds)" immediately
        so the user knows something is happening.
      - Yields partial tokens as they stream in.
      - If no token arrives within 15 seconds, stops and shows a retry message.

    For local mode:
      - Runs synchronously and yields the final answer once.

    Always yields a 5-tuple: (answer, context, debug, confidence_label, history)
    """
    import time

    mode         = resolve_mode_string(mode_label)
    resolved_key = (groq_key or "").strip() or os.environ.get("GROQ_API_KEY", "")

    if not chunks:
        yield ("No text indexed yet. Go to Step 1 and click Index Text first.",
               "", "", "N/A", search_history or [])
        return

    if not question or len(question.strip()) < 3:
        yield ("Please enter a valid question.", "", "", "N/A", search_history or [])
        return

    search_history = search_history or []

    # Multi-question: process synchronously (streaming not needed per sub-question)
    if is_multi_question(question):
        result = _handle_multi_question(
            question, chunks, active_context,
            search_history, mode, resolved_key, show_debug,
        )
        yield result
        return

    # Retrieval -- always returns chunks, no threshold filter
    top_chunks, top_scores = retrieve_top_chunks(
        question=question,
        chunks=chunks,
        top_k=TOP_K,
        threshold=0.0,
        context_text=active_context,
    )

    if not top_chunks:
        yield ("No text has been indexed yet.", "", "", "N/A", search_history)
        return

    max_score = max(top_scores) if top_scores else 0.0
    if max_score < 0.20:
        print(f"[WARN] Low retrieval confidence ({max_score:.3f}) for: {question[:60]}")

    context_display = format_context_display(
        [c["text"] for c in top_chunks],
        [{"start_word": c["start_word"], "end_word": c["end_word"]}
         for c in top_chunks],
    )
    debug_out = ""
    if show_debug:
        debug_out = format_debug_info([c["text"] for c in top_chunks], top_scores)

    # ---- GROQ: stream with 15-second timeout guard ----
    if mode == "groq":
        from generator import merge_chunks as _merge, _score_groq_confidence

        merged      = _merge(top_chunks)
        final_answer = ""
        TIMEOUT_SEC  = 15

        # Show an immediate waiting message so the UI does not look frozen
        yield (
            "Connecting to Groq API... (max 15 seconds)",
            context_display, debug_out, "Connecting...", search_history
        )

        start_time   = time.time()
        got_first_token = False

        for partial in stream_groq_answer(merged, question, resolved_key):
            elapsed = time.time() - start_time

            # Hard timeout -- stop waiting after 15 seconds
            if elapsed > TIMEOUT_SEC and not got_first_token:
                final_answer = (
                    "Groq API is taking longer than expected. "
                    "Please click Get Answer again to retry."
                )
                yield (final_answer, context_display, debug_out, "N/A", search_history)
                return

            got_first_token = True
            final_answer    = partial
            yield (partial, context_display, debug_out, "Generating...", search_history)

        if not final_answer or final_answer.startswith("Groq API is taking"):
            return

        from generator import clean_answer
        final_answer   = clean_answer(final_answer)
        confidence     = _score_groq_confidence(final_answer)
        conf_label     = _format_confidence(confidence)
        search_history = add_qa_to_history(
            search_history, active_context, question, final_answer, mode
        )
        print(f"[INFO] Groq streaming complete | confidence: {confidence:.3f}")
        yield (final_answer, context_display, debug_out, conf_label, search_history)
        return

    # ---- LOCAL FLAN-T5: single synchronous call ----
    answer, success, confidence = generate_answer(
        context_chunks=top_chunks,
        question=question,
        mode=mode,
        groq_api_key=resolved_key,
    )
    final_answer = answer if success else fallback_answer(top_chunks[0]["text"])
    conf_label   = _format_confidence(confidence if success else 0.0)

    search_history = add_qa_to_history(
        search_history, active_context, question, final_answer, mode
    )
    yield (final_answer, context_display, debug_out, conf_label, search_history)


# ------------------------------------------------------------------
# HISTORY TAB HANDLERS
# ------------------------------------------------------------------

def refresh_history_tab(search_history: list):
    summary = format_full_history_summary(search_history)
    choices = get_context_previews(search_history)
    return summary, gr.update(choices=choices, value=None), ""


def load_context_detail(selected_preview: str, search_history: list) -> str:
    if not selected_preview:
        return ""
    entry = get_entry_by_preview(search_history, selected_preview)
    return format_context_entry_display(entry)


def clear_all_history(search_history: list):
    return [], "History cleared.", gr.update(choices=[], value=None), ""


# ------------------------------------------------------------------
# UI HELPERS
# ------------------------------------------------------------------

def toggle_groq_key_visibility(mode: str):
    """Show the API key field only when Groq mode is selected."""
    return gr.update(visible=("Groq" in mode))


def resolve_mode_string(mode_label: str) -> str:
    """Map the dropdown label to the internal mode string used by generator.py."""
    if "Groq" in mode_label:
        return "groq"
    return "local"


# ------------------------------------------------------------------
# GRADIO UI
# ------------------------------------------------------------------

def build_ui():

    with gr.Blocks(title="RAG-lite QA System") as demo:

        # Per-session state
        chunk_store    = gr.State(value=None)
        active_context = gr.State(value="")
        search_history = gr.State(value=[])

        with gr.Tabs():

            # ============================================================
            # TAB 1 -- Question Answering
            # ============================================================
            with gr.Tab("Question Answering"):

                gr.Markdown(
                    """
                    # RAG-lite Question Answering System

                    Context-only answering. No internet. No hallucinations.

                    Index a paragraph once, then ask questions using the local FLAN-T5
                    model or the Groq-hosted LLaMA 3 model. Groq answers stream
                    token-by-token in real time. Every answer shows a confidence score.
                    All Q&A is saved to the Search History tab grouped by paragraph.
                    """
                )

                # -- Step 1: Text input --
                gr.Markdown("### Step 1 -- Provide your text")
                text_input = gr.Textbox(
                    label="Input Paragraph / Text",
                    placeholder="Paste your paragraph or article excerpt here...",
                    lines=8,
                    max_lines=30,
                )
                index_btn    = gr.Button("Index Text", variant="secondary", size="lg")
                index_status = gr.Textbox(
                    label="Index Status",
                    interactive=False,
                    lines=1,
                    elem_classes=["status-box"],
                    placeholder="Status will appear here after indexing...",
                )

                gr.Markdown("---")

                # -- Step 2: Mode + question --
                gr.Markdown(
                    "### Step 2 -- Select generation mode and ask a question\n\n"
                    "Local (FLAN-T5): runs on device, no API key, loads in ~10 seconds.\n\n"
                    "Groq (LLaMA 3): calls the Groq API, streams answers in real time, "
                    "requires a free key from console.groq.com set in your .env file."
                )

                with gr.Row():
                    mode_selector = gr.Dropdown(
                        label="Generation Mode",
                        choices=["Local (FLAN-T5)", "Groq (LLaMA 3)"],
                        value="Local (FLAN-T5)",
                        interactive=True,
                        scale=2,
                    )
                    groq_key_input = gr.Textbox(
                        label="Groq API Key (optional override)",
                        placeholder="Leave blank to use key from .env file",
                        type="password",
                        visible=False,
                        scale=3,
                    )

                mode_selector.change(
                    fn=toggle_groq_key_visibility,
                    inputs=mode_selector,
                    outputs=groq_key_input,
                )

                gr.Markdown("---")

                with gr.Row():
                    with gr.Column(scale=3):
                        question_input = gr.Textbox(
                            label="Your Question(s)",
                            placeholder=(
                                "Single:   What is machine learning?\n"
                                "Multi:    What is ML? How does it work? What are its uses?\n"
                                "Compound: What is ML and how does it differ from AI?"
                            ),
                            lines=3,
                        )
                    with gr.Column(scale=1):
                        debug_checkbox = gr.Checkbox(label="Show Debug Info", value=False)
                        ask_btn = gr.Button("Get Answer", variant="primary", size="lg")

                gr.Markdown("---")

                answer_output = gr.Textbox(
                    label="Answer",
                    lines=6,
                    interactive=False,
                    elem_classes=["answer-box"],
                )
                confidence_output = gr.Textbox(
                    label="Answer Confidence",
                    lines=1,
                    interactive=False,
                    elem_classes=["conf-box"],
                    placeholder="Confidence score will appear here after answering...",
                )
                context_output = gr.Textbox(
                    label="Relevant Context Used",
                    lines=5,
                    interactive=False,
                    elem_classes=["context-box"],
                )
                debug_output = gr.Textbox(
                    label="Debug Info -- Retrieved Chunks and Scores",
                    lines=8,
                    interactive=False,
                    elem_classes=["debug-box"],
                    visible=False,
                )

                debug_checkbox.change(
                    fn=lambda v: gr.update(visible=v),
                    inputs=debug_checkbox,
                    outputs=debug_output,
                )

                gr.Examples(
                    examples=[
                        [
                            "The Amazon rainforest, often referred to as the lungs of the Earth, "
                            "produces about 20% of the world's oxygen. It spans over 5.5 million "
                            "square kilometers across nine countries in South America, with Brazil "
                            "containing about 60% of the rainforest. The Amazon is home to an "
                            "estimated 10% of all species on Earth, including over 40,000 plant "
                            "species, 1,300 bird species, and 3,000 types of fish. Deforestation "
                            "remains a major threat, with approximately 17% of the forest lost in "
                            "the last 50 years due to logging, agriculture, and urban development.",
                        ],
                        [
                            "Machine learning is a subset of artificial intelligence that gives "
                            "systems the ability to automatically learn and improve from experience "
                            "without being explicitly programmed. It focuses on developing computer "
                            "programs that can access data and use it to learn for themselves. The "
                            "process begins with observations or data, such as examples, direct "
                            "experience, or instruction, to look for patterns in data and make "
                            "better decisions in the future. The primary aim is to allow computers "
                            "to learn automatically without human intervention or assistance and "
                            "adjust actions accordingly.",
                        ],
                    ],
                    inputs=[text_input],
                    label="Load an example paragraph (then click Index Text)",
                )

                gr.Markdown(
                    """
                    ---
                    Text is indexed once per paragraph. Ask multiple questions with question
                    mark separators, or compound questions with conjunctions such as "and",
                    "but", "while". All answers are saved to the Search History tab with
                    the generation model and confidence score recorded.
                    """
                )

            # ============================================================
            # TAB 2 -- Search History
            # ============================================================
            with gr.Tab("Search History") as history_tab:

                gr.Markdown(
                    """
                    ## Search History

                    All Q&A activity from this session, organized by the paragraph each
                    question was asked on. Each entry shows the model that produced the
                    answer (Local FLAN-T5 or Groq LLaMA 3), enabling direct comparison
                    of outputs from the same question asked with different backends.

                    Select a paragraph from the dropdown to see all questions and answers
                    recorded for it.
                    """
                )

                history_summary = gr.Textbox(
                    label="Session Overview",
                    interactive=False,
                    lines=8,
                    elem_classes=["history-box"],
                    placeholder="Switch to this tab after asking questions to see your history...",
                )

                gr.Markdown("---")
                gr.Markdown("### Browse by Paragraph")

                with gr.Row():
                    context_dropdown = gr.Dropdown(
                        label="Select a Paragraph",
                        choices=[],
                        value=None,
                        interactive=True,
                        info="Each entry represents one paragraph you indexed",
                        scale=3,
                    )
                    clear_history_btn = gr.Button(
                        "Clear All History",
                        variant="stop",
                        size="sm",
                        scale=1,
                    )

                context_detail = gr.Textbox(
                    label="Q&A Detail for Selected Paragraph",
                    interactive=False,
                    lines=18,
                    elem_classes=["history-box"],
                    placeholder="Select a paragraph above to see its Q&A pairs here...",
                )

        # ------------------------------------------------------------------
        # EVENT WIRING
        # ------------------------------------------------------------------

        # Index text
        index_btn.click(
            fn=index_text,
            inputs=[text_input],
            outputs=[index_status, chunk_store, active_context],
        )

        # Ask button -- uses streaming generator so Groq updates live
        ask_btn.click(
            fn=handle_ask_streaming,
            inputs=[
                question_input, chunk_store, active_context, search_history,
                mode_selector, groq_key_input, debug_checkbox,
            ],
            outputs=[
                answer_output, context_output, debug_output,
                confidence_output, search_history,
            ],
        )

        # Enter key in question box also submits
        question_input.submit(
            fn=handle_ask_streaming,
            inputs=[
                question_input, chunk_store, active_context, search_history,
                mode_selector, groq_key_input, debug_checkbox,
            ],
            outputs=[
                answer_output, context_output, debug_output,
                confidence_output, search_history,
            ],
        )

        # History tab auto-refresh
        history_tab.select(
            fn=refresh_history_tab,
            inputs=[search_history],
            outputs=[history_summary, context_dropdown, context_detail],
        )

        context_dropdown.change(
            fn=load_context_detail,
            inputs=[context_dropdown, search_history],
            outputs=[context_detail],
        )

        clear_history_btn.click(
            fn=clear_all_history,
            inputs=[search_history],
            outputs=[search_history, history_summary, context_dropdown, context_detail],
        )

    return demo


# ------------------------------------------------------------------
# ENTRY POINT
# ------------------------------------------------------------------

if __name__ == "__main__":
    demo = build_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=THEME,
        css=CSS,
    )