app.py

"""
app.py — Enterprise RAG System — Gradio 5 entry point.

Wires together all pipeline modules and renders the three-panel UI.
Uses gr.State() for per-session document isolation.
"""

import os
import logging
import gradio as gr

from src.ingestion import extract_text_from_pdf, validate_pdf
from src.chunking import chunk_text, chunk_statistics
from src.embeddings import build_faiss_index, embed_texts
from src.retrieval import retrieve_relevant_chunks
from src.generation import generate_answer
from src.evaluation import run_evaluation
from src.observability import trace_rag_query, get_observability_status
from src.metrics import record_query_metrics, get_metrics_summary
from src.utils import format_retrieved_chunks

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("enterprise-rag.app")


# ─────────────────────────────────────────────────────────────────────────────
# PIPELINE FUNCTIONS
# ─────────────────────────────────────────────────────────────────────────────

def process_pdf(pdf_file, chunk_size: int, chunk_overlap: int):
    """
    Full ingestion pipeline: PDF → text → chunks → embeddings → FAISS index.

    Gradio 5 passes uploaded files as a file path string directly.
    """
    if pdf_file is None:
        return (
            "⚠️ Please upload a PDF file first.",
            None,
            None,
            "",
        )

    try:
        # Gradio 5: pdf_file is a filepath string
        file_path = pdf_file if isinstance(pdf_file, str) else pdf_file.name

        with open(file_path, "rb") as f:
            file_bytes = f.read()

        # Validate size
        valid, size_msg = validate_pdf(file_bytes)
        if not valid:
            return f"❌ {size_msg}", None, None, ""

        # Step 1 — Extract text
        extraction = extract_text_from_pdf(file_bytes)
        if not extraction["success"]:
            return f"❌ {extraction['error']}", None, None, ""

        doc_text = extraction["text"]
        page_count = extraction["page_count"]

        # Step 2 — Chunk
        chunks = chunk_text(doc_text, int(chunk_size), int(chunk_overlap))
        if not chunks:
            return "❌ No chunks created. Document may be too short or empty.", None, None, ""

        stats = chunk_statistics(chunks)
        chunk_texts = [c["text"] for c in chunks]

        # Step 3 — Embed
        embeddings = embed_texts(chunk_texts)

        # Step 4 — Build FAISS index
        faiss_index = build_faiss_index(embeddings)

        status = (
            f"✅ Document processed successfully!\n\n"
            f"📄 Pages: {page_count}\n"
            f"📦 Chunks: {stats['count']}\n"
            f"📊 Avg chunk: {stats['avg_tokens']} tokens\n"
            f"🔢 Total tokens: {stats['total_tokens']}\n\n"
            f"Ready to answer questions."
        )

        doc_info = (
            f"**Loaded:** `{os.path.basename(file_path)}`  \n"
            f"Pages: {page_count} | Chunks: {stats['count']} | "
            f"Avg chunk: {stats['avg_tokens']} tokens"
        )

        logger.info(
            f"PDF ready: {stats['count']} chunks, "
            f"{stats['total_tokens']} total tokens"
        )
        return status, chunks, faiss_index, doc_info

    except Exception as e:
        logger.error(f"PDF processing failed: {e}")
        return f"❌ Error: {str(e)}", None, None, ""


def answer_question(query: str, chunks_state, index_state, top_k: int):
    """
    Full RAG query pipeline:
    embed query → retrieve → generate → evaluate → trace → display.
    """
    if not query or not query.strip():
        return "⚠️ Please enter a question.", "", "", "", ""

    if chunks_state is None or index_state is None:
        return (
            "⚠️ No document loaded. Please upload a PDF first.",
            "", "", "", "",
        )

    # Normalize chunks to plain text strings
    chunk_texts = [
        c["text"] if isinstance(c, dict) else c
        for c in chunks_state
    ]

    try:
        # Step 1 — Retrieve
        retrieval = retrieve_relevant_chunks(
            query=query,
            chunks=chunk_texts,
            faiss_index=index_state,
            top_k=int(top_k),
        )
        retrieved_chunks = retrieval["retrieved_chunks"]
        scores = retrieval["scores"]
        is_relevant = retrieval["is_relevant"]

        # Step 2 — Generate
        generation = generate_answer(
            query=query,
            context_chunks=retrieved_chunks,
            scores=scores,
            is_relevant=is_relevant,
        )
        answer = generation["answer"]
        prompt_tokens = generation["prompt_tokens"]
        response_tokens = generation["response_tokens"]
        gen_latency = generation["generation_latency_ms"]
        model_used = generation.get("model_used", "unknown")
        fallback_used = generation["fallback_used"]

        # Step 3 — Evaluate
        eval_scores = run_evaluation(
            query=query,
            answer=answer,
            context_chunks=retrieved_chunks,
            retrieval_scores=scores,
        )

        # Step 4 — Record metrics
        record_query_metrics(
            retrieval_latency_ms=retrieval["retrieval_latency_ms"],
            generation_latency_ms=gen_latency,
            prompt_tokens=prompt_tokens,
            response_tokens=response_tokens,
            eval_scores=eval_scores,
            fallback_used=fallback_used,
        )

        # Step 5 — Trace
        trace_rag_query(
            query=query,
            answer=answer,
            retrieved_chunks=retrieved_chunks,
            retrieval_scores=scores,
            eval_scores=eval_scores,
            retrieval_latency_ms=retrieval["retrieval_latency_ms"],
            generation_latency_ms=gen_latency,
            prompt_tokens=prompt_tokens,
            response_tokens=response_tokens,
            model_used=model_used,
            fallback_used=fallback_used,
        )

        # ── Format right-panel outputs ─────────────────────────────────────

        warning_text = f"\n\n⚠️ {retrieval['warning']}" if retrieval.get("warning") else ""
        chunks_display = format_retrieved_chunks(retrieved_chunks, scores) + warning_text

        metrics_display = get_metrics_summary()

        eval_display = (
            f"**Answer Quality Scores**\n\n"
            f"- Faithfulness: `{eval_scores['faithfulness']:.3f}` — grounded in context\n"
            f"- Answer Relevance: `{eval_scores['answer_relevance']:.3f}` — answers the question\n"
            f"- Context Precision: `{eval_scores['context_precision']:.3f}` — retrieval quality\n"
            f"- **Overall: `{eval_scores['overall']:.3f}`** {eval_scores['quality_label']}\n\n"
            f"{eval_scores.get('note', '')}"
        )

        obs_display = (
            f"{get_observability_status()}\n\n"
            f"**Last trace**\n"
            f"- Model: `{model_used}`\n"
            f"- Retrieval: `{retrieval['retrieval_latency_ms']:.0f}ms`\n"
            f"- Generation: `{gen_latency:.0f}ms`\n"
            f"- Total tokens: `{prompt_tokens + response_tokens}`\n"
            f"- Fallback used: `{'Yes' if fallback_used else 'No'}`"
        )

        return answer, chunks_display, metrics_display, eval_display, obs_display

    except Exception as e:
        logger.error(f"Query pipeline error: {e}")
        return f"❌ Pipeline error: {str(e)}", "", "", "", ""


# ─────────────────────────────────────────────────────────────────────────────
# GRADIO 5 UI
# ─────────────────────────────────────────────────────────────────────────────

with gr.Blocks(
    title="Enterprise RAG System",
    theme=gr.themes.Soft(primary_hue="blue"),
) as demo:

    # Per-session state — each user gets isolated chunks and FAISS index
    chunks_state = gr.State(None)
    index_state = gr.State(None)

    gr.Markdown(
        "# 🏢 Enterprise Knowledge Retrieval System\n"
        "**RAG pipeline · Groq LLM · FAISS · Evaluation · Observability**"
    )

    with gr.Row():

        # ── LEFT: Document Upload ─────────────────────────────────────────
        with gr.Column(scale=1, min_width=260):
            gr.Markdown("### 📁 Document Upload")

            pdf_input = gr.File(
                label="Upload PDF",
                file_types=[".pdf"],
            )

            with gr.Accordion("⚙️ Chunking Settings", open=False):
                chunk_size_slider = gr.Slider(
                    minimum=128,
                    maximum=1024,
                    value=512,
                    step=64,
                    label="Chunk Size (tokens)",
                    info="Larger = more context per chunk",
                )
                chunk_overlap_slider = gr.Slider(
                    minimum=0,
                    maximum=256,
                    value=64,
                    step=32,
                    label="Chunk Overlap (tokens)",
                    info="Prevents answer loss at boundaries",
                )
                top_k_slider = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=5,
                    step=1,
                    label="Top-K Retrieval",
                    info="Chunks returned per query",
                )

            process_btn = gr.Button("📥 Process Document", variant="primary")

            doc_status = gr.Textbox(
                label="Status",
                lines=6,
                interactive=False,
                value="No document loaded.",
            )

            doc_info_md = gr.Markdown("")

        # ── CENTER: Query & Answer ────────────────────────────────────────
        with gr.Column(scale=2, min_width=380):
            gr.Markdown("### 💬 Ask Questions")

            query_input = gr.Textbox(
                label="Your Question",
                placeholder="Ask anything about the uploaded document...",
                lines=3,
            )

            ask_btn = gr.Button("🔍 Get Answer", variant="primary", size="lg")

            answer_output = gr.Markdown(
                value="*Upload a document and ask a question to get started.*"
            )

            gr.Markdown(
                "---\n"
                "**Example questions after uploading:**\n"
                "- What are the main topics covered?\n"
                "- Summarize the key findings.\n"
                "- What risks or challenges are mentioned?\n"
                "- What are the specific numbers or statistics?"
            )

        # ── RIGHT: Observability Panel ────────────────────────────────────
        with gr.Column(scale=1, min_width=280):
            gr.Markdown("### 📊 Observability")

            with gr.Tabs():
                with gr.Tab("📄 Chunks"):
                    chunks_output = gr.Markdown(
                        value="*Retrieved context appears here after a query.*"
                    )

                with gr.Tab("📈 Metrics"):
                    metrics_output = gr.Markdown(
                        value="*Metrics appear after the first query.*"
                    )

                with gr.Tab("🧪 Evaluation"):
                    eval_output = gr.Markdown(
                        value="*Evaluation scores appear after a query.*"
                    )

                with gr.Tab("🔭 Traces"):
                    obs_output = gr.Markdown(
                        value=get_observability_status()
                    )

    # ── Event handlers ────────────────────────────────────────────────────

    process_btn.click(
        fn=process_pdf,
        inputs=[pdf_input, chunk_size_slider, chunk_overlap_slider],
        outputs=[doc_status, chunks_state, index_state, doc_info_md],
    )

    ask_btn.click(
        fn=answer_question,
        inputs=[query_input, chunks_state, index_state, top_k_slider],
        outputs=[answer_output, chunks_output, metrics_output, eval_output, obs_output],
    )

    query_input.submit(
        fn=answer_question,
        inputs=[query_input, chunks_state, index_state, top_k_slider],
        outputs=[answer_output, chunks_output, metrics_output, eval_output, obs_output],
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )