ui_helpers.py

from __future__ import annotations

import base64
import html
import json
import re
from typing import Any

from document_processing import MAX_DOCUMENT_BYTES, MAX_PAGES

EMPTY_TEXT = """
<div class="empty-result">
  <strong>Ready to parse</strong>
  <span>Run Cohere Parse to extract this document's content.</span>
</div>
"""

EMPTY_TABLES = """
<div class="empty-result">
  <strong>No rendered tables yet</strong>
  <span>Tables found by Cohere Parse will appear here.</span>
</div>
"""


def format_bytes(value: int) -> str:
    if value <= 0:
        return "0 MB"
    megabytes = value / (1024 * 1024)
    precision = 0 if megabytes >= 10 else 1
    return f"{megabytes:.{precision}f} MB"


def format_duration(milliseconds: int | float) -> str:
    if milliseconds < 1000:
        return f"{round(milliseconds)} MS"
    return f"{milliseconds / 1000:.1f} S"


def header_html() -> str:
    return """
<header class="product-header">
  <a class="brand-lockup" href="/" aria-label="Parse home">
    <img src="/gradio_api/file=assets/parse-logo.svg" alt="Parse" />
  </a>
  <a class="docs-link" href="https://docs.cohere.com/docs/parse" target="_blank"
     rel="noopener noreferrer" title="Parse documentation">
    <span>Parse documentation</span>
    <img src="/gradio_api/file=assets/up_right_arrow.svg" alt="" aria-hidden="true" />
  </a>
</header>
"""


def upload_intro_html() -> str:
    return f"""
<section class="upload-intro">
  <img class="document-stack-icon"
       src="/gradio_api/file=assets/FileStack.svg" alt="" aria-hidden="true" />
  <h1>Parse a document</h1>
  <p class="upload-subtitle">
    Extract text, tables, and document structure from PDFs and images.
  </p>
  <p class="upload-limits">
    PDF, PNG, JPEG, or WebP · {MAX_PAGES} pages ·
    {format_bytes(MAX_DOCUMENT_BYTES)} maximum
  </p>
</section>
"""


def privacy_html() -> str:
    return """
<div class="privacy-note">
  Uploaded documents are converted to page images and sent to Cohere for
  processing. Do not upload documents containing information you are not
  authorized to share.
  <a href="https://cohere.com/privacy" target="_blank" rel="noopener noreferrer">
    Cohere Privacy Policy
  </a>
</div>
"""


def progress_html(message: str, *, upload: bool = False) -> str:
    safe_message = html.escape(message)
    classes = "upload-progress" if upload else "run-status is-running"
    return f"""
<div class="{classes}" role="status" aria-live="polite">
  <img class="status-spinner" src="/gradio_api/file=assets/Loader.svg"
       alt="" aria-hidden="true" />
  <strong>{safe_message}</strong>
</div>
"""


def output_progress_html(page_number: int, page_count: int) -> str:
    return f"""
<div class="output-progress" role="status" aria-live="polite">
  <img class="status-spinner" src="/gradio_api/file=assets/Loader.svg"
       alt="" aria-hidden="true" />
  <strong>Parsing document</strong>
  <span>Processing page {page_number} / {page_count}</span>
</div>
"""


def file_strip_html(state: dict) -> str:
    completed = [result for result in state["results"] if result]
    elapsed = sum(result.get("elapsed_ms", 0) for result in completed)
    characters = sum(result.get("output_chars", 0) for result in completed)
    icon = "filetype_doc.svg" if state["kind"] == "pdf" else "filetype_image.svg"
    safe_name = html.escape(state["name"])

    if completed:
        metrics = (
            f'<span class="metric-badge">{format_duration(elapsed)}</span>'
            f'<span class="metric-badge">{len(completed)}/{len(state["pages"])} PARSED</span>'
            f'<span class="metric-badge">{characters:,} CHARS</span>'
        )
    else:
        metrics = (
            f"<span>{len(state['pages'])} page"
            f"{'s' if len(state['pages']) != 1 else ''} ready</span>"
            f"<span>{format_bytes(state['size'])}</span>"
        )

    return f"""
<div class="file-strip">
  <div class="file-name">
    <img class="file-type" src="/gradio_api/file=assets/{icon}" alt="" />
    <span title="{safe_name}">{safe_name}</span>
  </div>
  <div class="file-metrics">{metrics}</div>
</div>
"""


def page_limit_html(state: dict) -> str:
    if state["total_pages"] <= MAX_PAGES:
        return ""
    skipped = state["total_pages"] - len(state["pages"])
    suffix = "page" if skipped == 1 else "pages"
    return f"""
<div class="page-limit-alert" role="alert">
  <span class="warning-icon">!</span>
  <div>
    <strong>Only the first {len(state["pages"])} pages will be processed</strong>
    <span>
      This document has {state["total_pages"]} pages. {skipped} {suffix} will
      not be included in this run.
    </span>
  </div>
</div>
"""


def settings_html(state: dict) -> str:
    return """
<div class="settings-summary">
  <div class="settings-heading">
    <p class="eyebrow">Configuration</p>
    <h2>Parse settings</h2>
  </div>
</div>
"""


def page_context_html(state: dict) -> str:
    page_count = len(state["pages"])
    current = state["current_page"] + 1
    if page_count <= 1:
        return ""

    remaining = [number for number in range(current + 1, page_count + 1)]
    if not remaining:
        return f'<span class="page-context-muted">of {page_count}</span>'
    if len(remaining) <= 3:
        labels = [str(number) for number in remaining]
    else:
        labels = [str(remaining[0]), "…", str(page_count)]
    return "".join(f"<span>{html.escape(label)}</span>" for label in labels)


def zoom_label_html(zoom: int) -> str:
    return f'<span class="zoom-value">{max(1, min(100, zoom))}%</span>'


def status_html(
    state: dict,
    *,
    running: bool = False,
    running_page: int | None = None,
    error: str = "",
) -> str:
    if error:
        return (
            '<div class="run-status is-error">'
            '<img src="/gradio_api/file=assets/loader_octagon_x.svg" alt="" />'
            '<div><strong>Parse failed</strong>'
            f"<span>{html.escape(error)}</span></div></div>"
        )
    completed = sum(result is not None for result in state["results"])
    if running:
        page_number = running_page or state["current_page"] + 1
        return progress_html(
            f"Processing page {page_number} / {len(state['pages'])}"
        )
    if completed == len(state["pages"]) and completed:
        return (
            '<div class="run-status is-complete">'
            '<img src="/gradio_api/file=assets/loader_checkmark.svg" alt="" />'
            f"<strong>{completed} page{'s' if completed != 1 else ''} parsed</strong>"
            "</div>"
        )
    return ""


def _box_html(box: dict, index: int) -> str:
    x1, y1, x2, y2 = box["bbox"]
    left = x1 * 100
    top = y1 * 100
    width = (x2 - x1) * 100
    height = (y2 - y1) * 100
    label = html.escape(box.get("label", "element"))
    color_class = f"box-color-{index % 4}"
    return f"""
<span class="parse-box {color_class}"
      style="left:{left:.3f}%;top:{top:.3f}%;width:{width:.3f}%;height:{height:.3f}%">
  <span>{label}</span>
</span>
"""


def zoom_width(zoom: int) -> float:
    bounded = max(1, min(100, zoom))
    if bounded <= 50:
        return 50 + bounded
    return 100 + ((bounded - 50) * 1.2)


def viewer_html(state: dict, *, show_boxes: bool, zoom: int) -> str:
    page = state["pages"][state["current_page"]]
    result = state["results"][state["current_page"]]
    encoded = base64.b64encode(page["png"]).decode("ascii")
    boxes = result.get("boxes", []) if result and show_boxes else []
    overlays = "".join(_box_html(box, index) for index, box in enumerate(boxes))
    return f"""
<div class="viewer-stage">
  <div class="viewer-document" style="width:{zoom_width(zoom):.1f}%">
    <img src="data:image/png;base64,{encoded}"
         width="{page["width"]}" height="{page["height"]}"
         alt="Rendered document page {page["number"]}" />
    {overlays}
  </div>
</div>
"""


def output_values(state: dict, *, render_tables: bool) -> tuple[str, str, str]:
    result = state["results"][state["current_page"]]
    if not result:
        return EMPTY_TEXT, EMPTY_TABLES, ""

    text_output = result.get("text_html") or (
        '<div class="empty-result"><strong>No text returned</strong></div>'
    )
    if render_tables and result.get("tables_html"):
        tables = f'<div class="table-output">{result["tables_html"]}</div>'
    elif not render_tables:
        tables = """
<div class="empty-result">
  <strong>Rendered tables are hidden</strong>
  <span>Enable the table setting to inspect extracted tables.</span>
</div>
"""
    else:
        tables = EMPTY_TABLES
    return text_output, tables, result.get("raw_output", "")


def download_html(state: dict) -> str:
    if not state or not all(state["results"]):
        return ""

    payload = {
        "document": {
            "name": state["name"],
            "total_pages": state["total_pages"],
            "processed_pages": len(state["pages"]),
        },
        "pages": [
            {
                "page": page["number"],
                "text": result.get("text_output", result["raw_output"]),
                "response": result.get("response"),
                "bounding_boxes": result["boxes"],
                "elapsed_ms": result["elapsed_ms"],
                "model": result["model"],
                "usage": result["usage"],
                "request_id": result["request_id"],
            }
            for page, result in zip(state["pages"], state["results"])
        ],
    }
    encoded = base64.b64encode(
        json.dumps(payload, indent=2, ensure_ascii=True).encode("utf-8")
    ).decode("ascii")
    stem = re.sub(r"[^a-zA-Z0-9._-]+", "-", state["name"]).strip("-")
    filename = f"{stem or 'cohere-parse'}-results.json"
    return f"""
<a class="download-link"
   href="data:application/json;base64,{encoded}"
   download="{html.escape(filename, quote=True)}">
  <span>Download JSON</span>
  <img src="/gradio_api/file=assets/download_white.svg" alt="" aria-hidden="true" />
</a>
"""


def configuration_required_html(api_key_name: str | None = None) -> str:
    if api_key_name:
        detail = (
            f"Add the {html.escape(api_key_name)} secret in the Space settings "
            "before parsing."
        )
    else:
        detail = "Set COHERE_ENV to either staging or production."
    return f"""
<div class="configuration-required" role="alert">
  <strong>Configuration required</strong>
  <span>{detail}</span>
</div>
"""