server.py

"""
Document Summarizer — MCP SSE Server
======================================
Accepts files (PDF, scanned PDF, DOCX, DOC) uploaded as base64 and returns
summaries produced by a Map-Reduce → LLM pipeline identical to the Colab
notebook.

Transport : SSE (Server-Sent Events)
Protocol  : Model Context Protocol (MCP) 2024-11-05
Endpoints : GET  /sse          — SSE stream (clients connect here)
            POST /messages/    — JSON-RPC messages endpoint

Start:
    python server.py                  # listens on 0.0.0.0:8000
    MCP_PORT=9000 python server.py    # custom port
"""

from __future__ import annotations

import base64
import json
import logging
import asyncio
import os
import subprocess
import tempfile
from pathlib import Path
from typing import Optional
from functools import partial

import pytesseract
pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"

# ── Third-party ─────────────────────────────────────────────────────────────
from mcp.server.fastmcp import FastMCP, Context
from openai import OpenAI
import pdfplumber
import fitz          # PyMuPDF
from PIL import Image
from docx import Document as DocxDocument
from rouge_score import rouge_scorer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-7s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

# ═════════════════════════════════════════════════════════════════════════════
# FastMCP server instance
# ═════════════════════════════════════════════════════════════════════════════

mcp = FastMCP(
    "Document Summarizer",
    instructions=(
        "Summarizes PDF (text & scanned), DOCX, and DOC files using a "
        "Map-Reduce + LLM pipeline powered by OpenAI."
    ),
)


# ═════════════════════════════════════════════════════════════════════════════
# STAGE 0 — LLM client helper
# ═════════════════════════════════════════════════════════════════════════════

def _make_client(model: str) -> tuple[OpenAI, str]:
    key = os.getenv("OPENAI_API_KEY", "")
    if not key:
        raise ValueError("OPENAI_API_KEY environment variable is not set.")
    return OpenAI(api_key=key), model


def _llm_call(
    client: OpenAI,
    model: str,
    user_message: str,
    system_message: str = "You are a helpful assistant.",
    max_tokens: int = 1000,
) -> str:
    response = client.chat.completions.create(
        model=model,
        max_tokens=max_tokens,
        messages=[
            {"role": "system", "content": system_message},
            {"role": "user",   "content": user_message},
        ],
    )
    return response.choices[0].message.content


# ═════════════════════════════════════════════════════════════════════════════
# STAGE 1 — Text extraction  (all synchronous — run in executor)
# ═════════════════════════════════════════════════════════════════════════════

def _is_scanned_pdf(filepath: str) -> bool:
    with pdfplumber.open(filepath) as pdf:
        for page in pdf.pages[:3]:
            text = page.extract_text() or ""
            if len(text.strip()) > 50:
                return False
    return True


def _extract_word(filepath: str) -> str:
    doc = DocxDocument(filepath)
    return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip())


def _extract_pdf(filepath: str) -> str:
    pages = []
    with pdfplumber.open(filepath) as pdf:
        for i, page in enumerate(pdf.pages):
            text = page.extract_text() or ""
            if text.strip():
                pages.append(f"[Page {i + 1}]\n{text}")
    return "\n\n".join(pages)


def _extract_scanned_pdf(filepath: str) -> str:
    doc = fitz.open(filepath)
    pages = []
    for i, page in enumerate(doc):
        mat = fitz.Matrix(200 / 72, 200 / 72)
        pix = page.get_pixmap(matrix=mat)
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        text = pytesseract.image_to_string(img)
        if text.strip():
            pages.append(f"[Page {i + 1}]\n{text}")
    return "\n\n".join(pages)


def _extract_doc(filepath: str) -> str:
    for cmd in [["antiword", filepath], ["catdoc", "-w", filepath]]:
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            text = result.stdout.strip()
            if text:
                log.info(f"    .doc extracted via {cmd[0]}: {len(text):,} chars")
                return text
        except (FileNotFoundError, subprocess.CalledProcessError):
            continue

    log.warning("Falling back to raw OLE text scan for .doc file")
    import re
    raw = Path(filepath).read_bytes()
    decoded = raw.decode("latin-1", errors="replace")
    runs = re.findall(r"[\x20-\x7e]{4,}", decoded)
    text = "\n".join(runs)
    if text:
        return text
    raise RuntimeError(
        f"Could not extract text from {Path(filepath).name}. "
        "Ensure 'antiword' is listed in packages.txt on HF Spaces."
    )


def extract_text(filepath: str, tmp_dir: str) -> dict:
    path = Path(filepath)
    ext  = path.suffix.lower()
    if ext == ".doc":
        text, file_type = _extract_doc(filepath), "word_doc"
    elif ext == ".docx":
        text, file_type = _extract_word(filepath), "word"
    elif ext == ".pdf":
        if _is_scanned_pdf(filepath):
            text, file_type = _extract_scanned_pdf(filepath), "scanned_pdf"
        else:
            text, file_type = _extract_pdf(filepath), "pdf"
    else:
        raise ValueError(f"Unsupported file type '{ext}'. Accepted: .pdf, .docx, .doc")
    log.info(f"    Extracted {len(text):,} characters from {path.name}")
    return {"filename": path.name, "text": text, "type": file_type}


# ═════════════════════════════════════════════════════════════════════════════
# STAGE 2 — Map-Reduce  (synchronous — runs in executor)
# ═════════════════════════════════════════════════════════════════════════════

def chunk_text(text: str, chunk_size: int = 3000, overlap: int = 200) -> list[str]:
    words = text.split()
    step  = chunk_size - overlap
    return [
        " ".join(words[i : i + chunk_size])
        for i in range(0, len(words), step)
        if " ".join(words[i : i + chunk_size]).strip()
    ]


def _map_chunk(client, model, chunk, filename, idx) -> str:
    return _llm_call(
        client, model,
        system_message=(
                f"You are extracting key information from a section of '{filename}', "
                    "a procurement or regulatory document. "
                    "Extract and preserve ALL of the following if present: "
                    "1. Monetary values, budgets, fees (exact figures) "
                    "2. Dates, deadlines, validity periods "
                    "3. Eligibility criteria and required qualifications "
                    "4. Compliance requirements and legal references "
                    "5. Named parties, organizations, roles "
                    "6. Numbered clauses or article references "
                    "Be concise but do NOT omit specific figures or requirements."
        ),
        user_message=f"Section {idx + 1}:\n{chunk}",
        #max_tokens=500,
        max_tokens=800,
    )


def _reduce_summaries(client, model, summaries, filename) -> str:
    numbered = "\n\n".join(f"[Section {i+1}]\n{s}" for i, s in enumerate(summaries))
    return _llm_call(
        client, model,
        system_message=(
            "Produce one coherent, well-structured summary of the entire document. "
            "Preserve key facts, figures, and decisions."
        ),
        user_message=f"Summaries of all sections of '{filename}':\n\n{numbered}",
        #max_tokens=800,
        max_tokens=1200,
    )


def _compute_similarity(source_text, partial_summaries, final_summary) -> dict:
    partial_concat = "\n\n".join(partial_summaries)
    scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
    rouge  = scorer.score(source_text, final_summary)
    vec    = TfidfVectorizer(stop_words="english")
    mat    = vec.fit_transform([partial_concat, final_summary])
    tfidf  = float(cosine_similarity(mat[0:1], mat[1:2])[0][0])
    comp   = round(len(source_text) / max(len(final_summary), 1), 1)
    return {
        "rouge_l_precision": round(rouge["rougeL"].precision, 4),
        "rouge_l_recall":    round(rouge["rougeL"].recall, 4),
        "tfidf_cosine":      round(tfidf, 4),
        "compression_ratio": comp,
        "overall":           round(0.65 * tfidf + 0.35 * rouge["rougeL"].precision, 4),
    }


def _interpret_scores(scores) -> str:
    overall = scores["overall"]
    ratio   = scores["compression_ratio"]
    tag     = "aggressive" if ratio > 20 else "moderate"
    note    = f" (compression {ratio}x — {tag})"
    if overall >= 0.70: return f"Excellent — faithfully captures source content.{note}"
    if overall >= 0.55: return f"Good — covers most key points; minor gaps possible.{note}"
    if overall >= 0.40: return f"Fair — relevant but may miss some details.{note}"
    return f"Poor — diverges significantly from source.{note}"


def map_reduce_sync(client, model, extracted: dict, notify_sync) -> dict:
    """
    Fully synchronous Map-Reduce.  notify_sync(msg) is a plain callable
    that queues messages — the async tool awaits them after each stage.
    """
    filename = extracted["filename"]
    text     = extracted["text"]
    chunks   = chunk_text(text)

    notify_sync(f"Map-Reduce: {filename}  ({len(chunks)} chunks)")

    partial_summaries = []
    for i, chunk in enumerate(chunks):
        notify_sync(f"  Mapping chunk {i + 1}/{len(chunks)} ...")
        partial_summaries.append(_map_chunk(client, model, chunk, filename, i))

    notify_sync(f"  Reducing {len(partial_summaries)} chunk summaries ...")
    final_summary = _reduce_summaries(client, model, partial_summaries, filename)
    notify_sync(f"  Summary ready: {len(final_summary):,} chars")

    notify_sync("  Computing similarity scores ...")
    scores  = _compute_similarity(text, partial_summaries, final_summary)
    quality = _interpret_scores(scores)
    notify_sync(
        f"  Scores — ROUGE-L: {scores['rouge_l_precision']:.3f}  "
        f"TF-IDF: {scores['tfidf_cosine']:.3f}  Overall: {scores['overall']:.3f}"
    )
    notify_sync(f"  Quality: {quality}")

    return {
        "filename":          filename,
        "type":              extracted["type"],
        "summary":           final_summary,
        "similarity_scores": scores,
        "quality_label":     quality,
    }


# ═════════════════════════════════════════════════════════════════════════════
# STAGE 3 — Final LLM query
# ═════════════════════════════════════════════════════════════════════════════

def _build_context(reduced_docs: list[dict]) -> str:
    return "\n\n".join(
        f"=== Document: {d['filename']} (type: {d['type']}) ===\n{d['summary']}"
        for d in reduced_docs
    )


def query_llm_sync(client, model, reduced_docs, user_query) -> str:
    context = _build_context(reduced_docs)
    log.info(f"Final query context: {len(context):,} chars")
    return _llm_call(
        client, model,
        system_message=(
                        "You are a procurement law expert. Answer based strictly on the provided "
                        "document summaries. For each claim: "
                        "1. Cite the specific document name and section "
                        "2. Quote or paraphrase the relevant text "
                        "3. If the answer is not in the documents, say so explicitly — do not invent. "
                        "Structure your answer with clear headings."
        ),
        user_message=f"Summaries:\n\n{context}\n\n---\nQuestion: {user_query}",
        max_tokens=1500,
    )


# ═════════════════════════════════════════════════════════════════════════════
# MCP TOOL — async so every ctx.info() is awaited before the result is sent
# ═════════════════════════════════════════════════════════════════════════════

@mcp.tool()
async def summarize_documents(
    files: list[dict],
    query: str = "Give me a detailed summary of all documents",
    model: str = "gpt-4o-mini",
    ctx: Context = None,
) -> str:
    """
    Summarize one or more documents using a Map-Reduce + LLM pipeline.

    Args:
        files: List of file objects, each with:
                 - filename       (str) original file name; extension used for routing
                 - content_base64 (str) base64-encoded file bytes
        query: Question or instruction to answer across all documents.
        model: LLM model name (default: gpt-4o-mini).

    Returns:
        JSON string with keys:
            status       "ok" | "error"
            documents    list of per-document results
            final_answer the answer to `query` synthesized across all documents
    """
    # ── Safe ctx.info() wrapper ───────────────────────────────────────────────
    # Being async lets us AWAIT ctx.info() directly.
    # This guarantees every notification is sent and acknowledged by the MCP
    # transport BEFORE execution continues — which means:
    #   • No BrokenResourceError flood  (connection is still open at send time)
    #   • Every _notify message reaches the Java McpSseClient and is forwarded
    #     via progressListener → sendMessage(out, "⚙️ " + msg) → JSP onmessage
    async def _notify(msg: str):
        log.info(msg)
        if ctx:
            try:
                await ctx.info(msg)
            except Exception:
                # Client already disconnected — swallow silently.
                # This can only happen on the very last notification if the
                # client closes the connection the instant the result arrives.
                pass

    loop = asyncio.get_event_loop()

    try:
        await _notify("=== STAGE 0 — Initializing ===")
        client_obj, model_name = _make_client(model)

        with tempfile.TemporaryDirectory() as tmp_dir:

            # ── Decode & save uploaded files ──────────────────────────────────
            await _notify("=== STAGE 1 — TEXT EXTRACTION ===")
            saved_paths: list[str] = []
            for f in files:
                fname   = f["filename"]
                content = base64.b64decode(f["content_base64"])
                path    = os.path.join(tmp_dir, fname)
                with open(path, "wb") as fp:
                    fp.write(content)
                saved_paths.append(path)
                await _notify(f"Saved: {fname}  ({len(content):,} bytes)")

            # ── Extract text (blocking I/O → thread pool) ─────────────────────
            extracted_docs: list[dict] = []
            for path in saved_paths:
                await _notify(f"Extracting: {Path(path).name}")
                extracted = await loop.run_in_executor(
                    None, extract_text, path, tmp_dir
                )
                extracted_docs.append(extracted)

            # ── Map-Reduce (blocking CPU + network → thread pool) ─────────────
            # We use a message queue: the sync worker puts messages into it,
            # and we drain the queue with await between CPU-bound calls so that
            # ctx.info() is always awaited on the async side.
            await _notify("=== STAGE 2 — MAP-REDUCE (per document) ===")
            reduced_docs: list[dict] = []

            for doc in extracted_docs:
                # Queue for messages produced by the sync worker
                msg_queue: asyncio.Queue[str] = asyncio.Queue()

                def _sync_notify(msg: str, q=msg_queue, lp=loop):
                    """Called from the thread pool — puts msg on the queue."""
                    lp.call_soon_threadsafe(q.put_nowait, msg)

                # Run the blocking map_reduce in a thread pool
                future = loop.run_in_executor(
                    None, map_reduce_sync, client_obj, model_name, doc, _sync_notify
                )

                # Drain the queue while the executor is running
                while not future.done():
                    try:
                        msg = await asyncio.wait_for(msg_queue.get(), timeout=0.5)
                        await _notify(msg)
                    except asyncio.TimeoutError:
                        pass   # nothing in queue yet; check future.done() again

                # Flush any messages that arrived just as future completed
                while not msg_queue.empty():
                    await _notify(msg_queue.get_nowait())

                reduced_docs.append(await future)

            # ── Final LLM query ───────────────────────────────────────────────
            await _notify("=== STAGE 3 — FINAL LLM QUERY ===")
            final_answer = await loop.run_in_executor(
                None, query_llm_sync, client_obj, model_name, reduced_docs, query
            )
            await _notify("Done.")

        result = {
            "status":       "ok",
            "query":        query,
            "documents":    reduced_docs,
            "final_answer": final_answer,
        }

    except Exception as exc:
        log.exception("Pipeline error")
        result = {"status": "error", "error": str(exc)}

    return json.dumps(result, ensure_ascii=False, indent=2)


# ═════════════════════════════════════════════════════════════════════════════
# Entry point
# ═════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    import uvicorn

    port = int(os.environ.get("MCP_PORT", 7860))
    log.info(f"Starting MCP SSE server on http://0.0.0.0:{port}")
    log.info("  SSE endpoint : GET  /sse")
    log.info("  RPC endpoint : POST /messages/")

    _local_host: bytes = f"localhost:{port}".encode()

    class _ProxyHostFix:
        """Raw ASGI middleware — rewrites Host to localhost before MCP sees it."""
        def __init__(self, asgi_app):
            self._app = asgi_app

        async def __call__(self, scope, receive, send):
            if scope["type"] in ("http", "websocket"):
                scope["headers"] = [
                    (name, _local_host)
                    if name in (b"host", b"x-forwarded-host")
                    else (name, value)
                    for name, value in scope.get("headers", [])
                ]
            try:
                await self._app(scope, receive, send)
            except Exception as exc:
                exc_type = type(exc).__name__
                if "ClientDisconnect" in exc_type or "Disconnect" in exc_type:
                    log.debug(f"Client disconnected ({exc_type}) — ignored.")
                else:
                    raise

    uvicorn.run(_ProxyHostFix(mcp.sse_app()), host="0.0.0.0", port=port)