pdf-quiz-generator/backend/app/services/illustrate.py

"""A diagram for a section, written as SVG.

Some things are a picture and are taught as prose because prose is what fits
in a database column: a timeline, a branching decision, a table of what
distinguishes four look-alike conditions, the sequence a hormone axis runs in.
A reader remembers the shape of those long after the sentences have gone.

So: the model is shown one section and asked, first, whether a diagram is
genuinely the content here — and only then to draw one. "No" is the expected
answer for most sections and is not a failure. A picture of a paragraph is
worse than the paragraph.

The SVG is checked before it is stored. Markup written by a model and rendered
in somebody's browser is the same shape of risk as markup written by anybody
else, and `uploads` already serves SVG under a sandbox CSP — this is the belt
to that brace.
"""
import logging
import re
import uuid
import xml.etree.ElementTree as ET

logger = logging.getLogger(__name__)

SVG_NS = "http://www.w3.org/2000/svg"
#: Nine kilobytes is a rich hand-drawn diagram; the milestone timeline is that
#: size. Much past this and the model is emitting a traced bitmap or a wall of
#: repeated paths, neither of which anybody wants in a section.
MAX_BYTES = 60_000
#: What to draw with, when nothing is configured for it. Measured: the article
#: model of the day, ds-deepseek-v4.1-flash, returns an empty completion for a
#: long SVG prompt — the same model draws a circle happily, so it is the size
#: of the ask rather than the ask itself. These two draw clean SVG.
DRAWING_MODELS = ("openrouter-claude-sonnet-4", "openrouter-gpt-4.1")
#: The section, trimmed. A diagram is of one mechanism, not of everything a
#: section says, and a shorter prompt is the difference between a drawing and
#: an empty reply.
BODY_CHARS = 2_500
#: Anything that executes, loads, or reaches outside the document.
FORBIDDEN_TAGS = {"script", "foreignobject", "iframe", "image", "use",
                  "animate", "set", "handler"}
FORBIDDEN_ATTR = re.compile(r"^on", re.I)
#: Anything that fetches from outside the file. `url(#arrowhead)` is not that
#: — it is how every marker, gradient and clip path in SVG refers to a <defs>
#: entry a few lines above, and refusing it threw away otherwise good drawings
#: for pointing at their own arrowheads. Only a url() that leaves the document
#: is refused.
REMOTE = re.compile(r"(?:https?:)?//|url\s*\(\s*(?!#)|javascript:|data:", re.I)

PROMPT = """You are illustrating one section of a pediatric reference article.

Article: {title}
Section: {section}

{body}

Find the one mechanism in this section that is a picture — a sequence, a
timeline, a branching decision, an anatomical relationship, or a comparison of
things that are confused with each other — and draw only that. Not the section:
one mechanism from it. Most sections are prose and hold no such thing; saying
so is the right answer and costs nothing, and a picture that restates a
paragraph is worse than the paragraph.

If no diagram belongs here, reply with exactly one line:

USEFUL: no

If one does, reply in exactly this shape and nothing else:

USEFUL: yes
TITLE: a short figure title
ALT: one full sentence describing what is drawn
<svg ...>...</svg>

Rules for the SVG:
- One root <svg> with a viewBox, no width or height, role="img", and an
  aria-label repeating the ALT sentence.
- Plain shapes and text only: rect, circle, ellipse, line, polyline, polygon,
  path, text, g, defs, marker, style. No script, no foreignObject, no <image>,
  no <use>, no external fonts, no links, no animation.
- Legible: system-ui text at 12px or larger, dark text on light fill, and
  nothing that relies on colour alone to carry meaning.
- Correct before it is pretty. Every label a fact from the section above.
- Nothing overlaps: a label on a connector goes above or below the line with
  clear space around it, never across a box, and two boxes never touch. Leave
  at least 20 units between anything and anything else.
- Do not wrap the reply in markdown fences and do not write JSON."""


def read_reply(raw: str) -> dict | None:
    """The four things a reply carries, out of plain text.

    Asking for JSON was the first attempt and it failed on seven sections out
    of eight: an SVG inside a JSON string needs every quote and newline
    escaped, and a model that draws a good diagram will still get that wrong.
    Nothing here needs escaping, so nothing here can be escaped incorrectly.
    """
    text = (raw or "").strip()
    if text.startswith("```"):
        text = text.split("\n", 1)[1] if "\n" in text else text[3:]
        if text.rstrip().endswith("```"):
            text = text.rstrip()[:-3]
    lines = text.strip().splitlines()
    head = {}
    for line in lines[:4]:
        for key in ("USEFUL", "TITLE", "ALT"):
            if line.upper().startswith(f"{key}:"):
                head[key] = line.split(":", 1)[1].strip()
    if head.get("USEFUL", "").lower().startswith("n"):
        return {"useful": False}
    at = text.find("<svg")
    end = text.rfind("</svg>")
    if at == -1 or end == -1:
        return None
    return {"useful": True, "title": head.get("TITLE", ""),
            "alt": head.get("ALT", ""), "svg": text[at:end + 6]}


def check(svg: str) -> str | None:
    """The SVG if it is safe and sane, otherwise None with a reason logged."""
    text = (svg or "").strip()
    if not text.startswith("<"):
        return None
    if len(text.encode()) > MAX_BYTES:
        logger.info("Illustrate: refused an SVG of %d bytes", len(text.encode()))
        return None
    try:
        root = ET.fromstring(text)
    except ET.ParseError as broken:
        logger.info("Illustrate: refused unparseable SVG (%s)", broken)
        return None
    if root.tag not in ("svg", f"{{{SVG_NS}}}svg"):
        return None
    if not root.get("viewBox"):
        return None
    # The namespace, added rather than demanded.
    #
    # An SVG in an <img> is a standalone document, and a browser will not draw
    # one without xmlns — it is not optional there, only inside HTML. Models
    # supply it about half the time, and a drawing that is otherwise correct
    # should not be thrown away over an attribute that can simply be written
    # in. This was the whole of "some figures render and some show the alt
    # text": nothing to do with the guard, the thumbnailer, or the file size.
    if root.tag == "svg" and 'xmlns=' not in text[:text.index(">") + 1]:
        text = text.replace("<svg", f'<svg xmlns="{SVG_NS}"', 1)
    for node in root.iter():
        tag = node.tag.split("}")[-1].lower()
        if tag in FORBIDDEN_TAGS:
            logger.info("Illustrate: refused an SVG containing <%s>", tag)
            return None
        for name, value in node.attrib.items():
            plain = name.split("}")[-1]
            if FORBIDDEN_ATTR.match(plain):
                logger.info("Illustrate: refused an SVG with an %s handler", plain)
                return None
            if REMOTE.search(value or "") and plain not in ("d", "points"):
                logger.info("Illustrate: refused an SVG reaching outside itself (%s)", plain)
                return None
    return text


def figure_line(alt: str, path: str) -> str:
    """The markdown an article section carries a figure as."""
    clean = " ".join((alt or "Figure").split()).replace("]", ")").replace("[", "(")
    return f"\n\n![{clean}](/uploads/{path})"


def already_illustrated(content: str) -> bool:
    return "![" in (content or "")


def key_for() -> str:
    return f"media/{uuid.uuid4().hex}.svg"