simit 0.1.0__py3-none-any.whl

simit/skills/routing.py

@@ -0,0 +1,89 @@
+"""Visualization routing: pick the skill that should realize a description.
+
+The router prompt is assembled from the active skills: each contributes its
+category definition (``route_hint``) and example requests. With exactly the
+built-in skills active, the assembled prompt is identical to the original
+SIMIT router prompt (``prompts.CLASSIFY_PROMPT``).
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Sequence
+
+from . import prompts as P
+
+_NUMBER_WORDS = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten",
+                 "eleven", "twelve", "thirteen", "fourteen", "fifteen", "sixteen", "seventeen", "eighteen",
+                 "nineteen", "twenty", "twenty-one", "twenty-two", "twenty-three", "twenty-four",
+                 "twenty-five", "twenty-six", "twenty-seven", "twenty-eight", "twenty-nine", "thirty"]
+
+
+def _parse_builtin_prompt(text: str):
+    """Split CLASSIFY_PROMPT into (definitions{LABEL: paragraph}, label order,
+    few-shot [(request, LABEL)])."""
+    head, rest = text.split("\n\n", 1)
+    defs_part, tail = rest.split("\n\nAnswer with exactly one word:", 1)
+    definitions, order = {}, []
+    for para in defs_part.split("\n\n"):
+        label, body = para.split(":", 1)
+        definitions[label] = body.strip()
+        order.append(label)
+    shots = re.findall(r'Request: "(.*?)"\nAnswer: (\w+)\n', tail, re.DOTALL)
+    shots = [(r, a) for r, a in shots if r != "{request}"]
+    return head, definitions, order, shots
+
+
+_HEAD, BUILTIN_DEFINITIONS, BUILTIN_ORDER, BUILTIN_SHOTS = _parse_builtin_prompt(P.CLASSIFY_PROMPT)
+
+
+def _number_word(n: int) -> str:
+    return _NUMBER_WORDS[n] if n < len(_NUMBER_WORDS) else str(n)
+
+
+def build_router_prompt(skills: Sequence) -> str:
+    """Router prompt (with a ``{request}`` placeholder) for ``skills``."""
+    labels = [s.label for s in skills]
+    head = _HEAD.replace("seventeen", _number_word(len(labels)))
+    defs = "\n\n".join(f"{s.label}: {_wrap_hint(s)}" for s in skills)
+    if len(labels) == 1:
+        answer = f"Answer with exactly one word: {labels[0]}."
+    else:
+        answer = "Answer with exactly one word: " + ", ".join(labels[:-1]) + f", or {labels[-1]}."
+    active = set(labels)
+    shots = [(req, lab) for req, lab in BUILTIN_SHOTS if lab in active]
+    for s in skills:
+        if s.label not in BUILTIN_DEFINITIONS:
+            shots += [(req, s.label) for req in s.get_route_examples()]
+    body = "".join(f'Request: "{req}"\nAnswer: {lab}\n\n' for req, lab in shots)
+    prompt = f"{head}\n\n{defs}\n\n{answer}\n\n{body}" + 'Request: "{request}"\nAnswer:'
+    return prompt
+
+
+def _wrap_hint(skill) -> str:
+    hint = BUILTIN_DEFINITIONS.get(skill.label) if skill.route_hint == "" else skill.route_hint
+    return (hint or skill.name.replace("_", " ")).strip()
+
+
+def parse_route(answer: str, skills: Sequence, default: str) -> str:
+    """Map the router's free-text answer to a skill name (longest label first,
+    so HTML_COMPOSITE beats HTML)."""
+    answer = answer.strip().upper()
+    for skill in sorted(skills, key=lambda s: len(s.label), reverse=True):
+        if skill.label in answer:
+            return skill.name
+    return default
+
+
+def pre_route(request: str, skills: Sequence) -> str | None:
+    """Deterministic routing rules that bypass the model (``Skill.route_pattern``)."""
+    ordered = sorted(skills, key=lambda s: -getattr(s, "route_priority", 0))
+    for skill in ordered:
+        pattern = getattr(skill, "route_pattern", None)
+        if pattern is None or not pattern.search(request):
+            continue
+        exclude = getattr(skill, "route_exclude_pattern", None)
+        if exclude is not None and exclude.search(request):
+            continue
+        return skill.name
+    return None

simit/skills/sandbox.py

@@ -0,0 +1,188 @@
+"""Isolated execution of model-written matplotlib code (the FIGURE skill)."""
+
+from __future__ import annotations
+
+import os
+import re
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+from PIL import Image
+
+from ..utils import CJK_RE
+from .helpers import GEOMETRY_HELPERS_SRC, TEXT_OVERLAP_FIX_SRC
+
+CODE_BLOCK_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.DOTALL)
+
+# The model is only ever asked for matplotlib/numpy code; anything touching
+# the OS is refused rather than executed.
+DENYLIST = (
+    "import os", "import sys", "import subprocess", "import socket",
+    "import shutil", "__import__", "eval(", "exec(", "open(", "os.", "sys.",
+)
+
+MAX_DIM = 2000
+MAX_ASPECT_RATIO = 9.0
+TIMEOUT_S = 20
+
+
+def extract_code(text: str) -> str | None:
+    m = CODE_BLOCK_RE.search(text)
+    return m.group(1).strip() if m else None
+
+
+def is_safe(code: str) -> bool:
+    return not any(pattern in code for pattern in DENYLIST)
+
+
+def find_non_ascii_issue(code: str) -> str | None:
+    if CJK_RE.search(code):
+        return (
+            "The code mixes Chinese characters into Python identifiers/calls "
+            "(e.g. `ax.plot共同点列表,...`), which is invalid syntax. Rewrite the "
+            "entire script using only English words and standard ASCII "
+            "characters for every variable name, function call, and argument."
+        )
+    return None
+
+
+def find_known_api_issues(code: str) -> str | None:
+    if "boxplot(" in code and "tick_labels" not in code and re.search(r"\blabels\s*=", code):
+        return (
+            "ax.boxplot's label argument is called tick_labels in this matplotlib "
+            "version, not labels. Replace labels= with tick_labels= in the boxplot() call."
+        )
+    if re.search(r"\bArc\(", code) and re.search(r"\b(angle1|angle2|start_angle|end_angle)\s*=", code):
+        return (
+            "matplotlib's Arc patch takes theta1 and theta2 for its sweep angles -- "
+            "angle1/angle2/start_angle/end_angle are not real Arc parameters. Either fix the "
+            "Arc(...) call to use theta1=/theta2=, or better, use the pre-defined "
+            "angle_arc(ax, vertex, theta1, theta2, radius=..., label=...) helper instead of "
+            "constructing Arc directly."
+        )
+    return None
+
+
+_ZYGOTE_SRC = r"""
+import os, sys, signal, struct, pickle, time, traceback
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot, numpy  # noqa: F401  (preloaded: forked jobs start warm)
+rd, wr = os.fdopen(0, "rb"), os.fdopen(1, "wb")
+def send(obj):
+    data = pickle.dumps(obj); wr.write(struct.pack("<I", len(data)) + data); wr.flush()
+while True:
+    head = rd.read(4)
+    if len(head) < 4:
+        break
+    script, workdir, timeout = pickle.loads(rd.read(struct.unpack("<I", head)[0]))
+    err_path = os.path.join(workdir, "stderr.txt")
+    pid = os.fork()
+    if pid == 0:
+        try:
+            os.chdir(workdir)
+            fd = os.open(err_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC)
+            os.dup2(fd, 2)
+            code = compile(script, "script.py", "exec")
+            exec(code, {"__name__": "__main__"})
+            os._exit(0)
+        except SystemExit as e:
+            os._exit(e.code if isinstance(e.code, int) else 1)
+        except BaseException:
+            traceback.print_exc()
+            sys.stderr.flush()
+            os._exit(1)
+    deadline = time.time() + timeout
+    status = None
+    while time.time() < deadline:
+        done, st = os.waitpid(pid, os.WNOHANG)
+        if done:
+            status = st
+            break
+        time.sleep(0.01)
+    if status is None:
+        os.kill(pid, signal.SIGKILL)
+        os.waitpid(pid, 0)
+        send((False, f"Execution timed out after {timeout}s"))
+        continue
+    err = ""
+    if os.path.exists(err_path):
+        with open(err_path) as f:
+            err = f.read()[-2000:]
+    send((os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0, err))
+"""
+
+
+class _Zygote:
+    """A warm interpreter (matplotlib/numpy imported) that forks one child per
+    script, so each FIGURE render skips the ~1.5s interpreter+import start."""
+
+    def __init__(self):
+        self.proc = None
+
+    def _start(self):
+        env = {"MPLBACKEND": "Agg", "PATH": "/usr/bin:/bin"}
+        if "HOME" in os.environ:
+            env["HOME"] = os.environ["HOME"]
+        self.proc = subprocess.Popen([sys.executable, "-c", _ZYGOTE_SRC], stdin=subprocess.PIPE,
+                                     stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, env=env)
+
+    def run(self, script: str, workdir: str, timeout: int) -> tuple[bool, str]:
+        import pickle
+        import struct
+        if self.proc is None or self.proc.poll() is not None:
+            self._start()
+        data = pickle.dumps((script, workdir, timeout))
+        self.proc.stdin.write(struct.pack("<I", len(data)) + data)
+        self.proc.stdin.flush()
+        head = self.proc.stdout.read(4)
+        if len(head) < 4:
+            self.proc = None
+            raise RuntimeError("figure sandbox died")
+        return pickle.loads(self.proc.stdout.read(struct.unpack("<I", head)[0]))
+
+
+_ZYGOTE = _Zygote() if hasattr(os, "fork") and sys.platform.startswith("linux") else None
+
+
+def run_script(code: str) -> tuple[Image.Image | None, str]:
+    """Run ``code`` isolated from this process; it must save ``output.png``.
+    Returns (image, "") on success or (None, error_text)."""
+    with tempfile.TemporaryDirectory(prefix="simit_fig_") as tmp:
+        workdir = Path(tmp)
+        script = GEOMETRY_HELPERS_SRC + "\n" + TEXT_OVERLAP_FIX_SRC + "\n" + code
+        ok, err = None, ""
+        if _ZYGOTE is not None:
+            try:
+                ok, err = _ZYGOTE.run(script, tmp, TIMEOUT_S)
+            except Exception:
+                ok = None
+        if ok is None:  # no fork available (or the zygote died): fresh interpreter
+            (workdir / "script.py").write_text(script)
+            env = {"MPLBACKEND": "Agg", "PATH": "/usr/bin:/bin", "MPLCONFIGDIR": tmp}
+            if "HOME" in os.environ:
+                env["HOME"] = os.environ["HOME"]
+            try:
+                proc = subprocess.run([sys.executable, "script.py"], cwd=workdir, capture_output=True,
+                                      text=True, timeout=TIMEOUT_S, env=env)
+                ok, err = proc.returncode == 0, proc.stderr[-2000:]
+            except subprocess.TimeoutExpired:
+                ok, err = False, f"Execution timed out after {TIMEOUT_S}s"
+        if not ok:
+            return None, err or "script failed"
+        png = workdir / "output.png"
+        if not png.exists():
+            return None, "Script exited cleanly but did not produce output.png"
+        image = Image.open(png)
+        image.load()
+        w, h = image.size
+        if w > MAX_DIM or h > MAX_DIM or max(w, h) / max(1, min(w, h)) > MAX_ASPECT_RATIO:
+            return None, (
+                f"Rendered image has pathological dimensions {w}x{h}px. This usually "
+                "means a shape was added with ax.add_patch(...) without also calling "
+                "ax.set_xlim(...)/ax.set_ylim(...) explicitly, or bbox_inches='tight' picked up "
+                "an artist placed far outside the intended view. Fix the axis limits."
+            )
+        return image, ""

simit/skills/web.py

@@ -0,0 +1,253 @@
+"""Browser-backed rendering (HTML, Mermaid) through Playwright's Chromium.
+
+One headless browser is kept alive per process, so a render costs one page
+load instead of a browser launch. The Chromium binary is downloaded on first
+use (``python -m playwright install chromium``) -- no Node.js, Mermaid CLI,
+or system Chromium is needed: Mermaid runs from the ``mermaid.min.js`` bundled
+with this package. Network access from rendered pages is blocked.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import subprocess
+import sys
+import threading
+from pathlib import Path
+
+from .base import SpecError
+
+_ASSETS = Path(__file__).parent / "assets"
+_MERMAID_JS = _ASSETS / "mermaid.min.js"
+
+_SCRIPT_RE = re.compile(r"<script[\s>]", re.IGNORECASE)
+_HANDLER_RE = re.compile(r"\bon\w+\s*=", re.IGNORECASE)
+
+
+def is_safe_markup(code: str) -> bool:
+    """Model-written HTML/SVG must not run scripts."""
+    return not (_SCRIPT_RE.search(code) or _HANDLER_RE.search(code))
+
+
+_LAUNCH_ARGS = ["--no-sandbox", "--disable-gpu", "--disable-dev-shm-usage"]
+_SYSTEM_BROWSERS = ("chromium", "chromium-browser", "google-chrome", "google-chrome-stable", "chrome")
+
+
+def _system_browser() -> str | None:
+    if os.environ.get("SIMIT_CHROMIUM"):
+        return os.environ["SIMIT_CHROMIUM"]
+    import shutil
+
+    for name in _SYSTEM_BROWSERS:
+        path = shutil.which(name)
+        if path:
+            return path
+    return None
+
+
+def _install_chromium():
+    """Download Playwright's headless Chromium once, serialized across processes."""
+    import fcntl
+    import tempfile
+
+    lock_path = Path(tempfile.gettempdir()) / "simit_playwright_install.lock"
+    with open(lock_path, "w") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        print("[simit] downloading headless Chromium for HTML/Mermaid rendering (one-time)...",
+              file=sys.stderr, flush=True)
+        r = subprocess.run([sys.executable, "-m", "playwright", "install", "chromium-headless-shell"],
+                           check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=900)
+        if r.returncode != 0:  # Playwright's (Node) downloader fails on some networks
+            _download_headless_shell()
+
+
+def _download_headless_shell():
+    """Fetch the build Playwright expects with plain Python and unpack it where
+    Playwright looks for it."""
+    import shutil
+    import tempfile
+    import urllib.request
+    import zipfile
+
+    out = subprocess.run([sys.executable, "-m", "playwright", "install", "--dry-run", "chromium-headless-shell"],
+                         capture_output=True, text=True, timeout=120).stdout
+    block = out.split("\n\n")[0]
+    loc = re.search(r"Install location:\s*(.+)", block)
+    urls = re.findall(r"Download (?:url|fallback \d+):\s*(\S+)", block)
+    if not loc or not urls:
+        return
+    dest = Path(loc.group(1).strip())
+    with tempfile.TemporaryDirectory() as tmp:
+        archive = Path(tmp) / "browser.zip"
+        for url in urls:
+            try:
+                with urllib.request.urlopen(url, timeout=60) as resp, open(archive, "wb") as f:
+                    shutil.copyfileobj(resp, f)
+                break
+            except Exception:
+                continue
+        else:
+            return
+        staging = dest.with_name(dest.name + ".partial")
+        shutil.rmtree(staging, ignore_errors=True)
+        with zipfile.ZipFile(archive) as z:
+            for info in z.infolist():
+                path = z.extract(info, staging)
+                mode = (info.external_attr >> 16) & 0o777
+                if mode:
+                    os.chmod(path, mode)
+        shutil.rmtree(dest, ignore_errors=True)
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        os.replace(staging, dest)
+        (dest / "INSTALLATION_COMPLETE").touch()
+
+
+class _Browser:
+    def __init__(self):
+        self._pw = None
+        self._browser = None
+        self._owner = None
+
+    def _launch(self):
+        """Playwright's own Chromium, else a system Chrome/Chromium, else
+        download Playwright's Chromium and retry."""
+        from playwright.sync_api import sync_playwright
+
+        if self._pw is not None:  # a previous launch failed or another thread owned it
+            try:
+                self._pw.stop()
+            except Exception:
+                pass
+        self._pw = sync_playwright().start()
+        errors = []
+        for attempt in ("bundled", "system", "install"):
+            try:
+                if attempt == "system":
+                    exe = _system_browser()
+                    if exe is None:
+                        continue
+                    self._browser = self._pw.chromium.launch(executable_path=exe, args=_LAUNCH_ARGS)
+                else:
+                    if attempt == "install":
+                        _install_chromium()
+                    self._browser = self._pw.chromium.launch(args=_LAUNCH_ARGS)
+                return
+            except Exception as e:
+                errors.append(f"{attempt}: {str(e).splitlines()[0]}")
+        raise SpecError("no usable Chromium for HTML/Mermaid rendering (" + "; ".join(errors) +
+                        "). Run `python -m playwright install chromium` or set SIMIT_CHROMIUM.")
+
+    def page(self, width: int, height: int):
+        # Playwright's sync API is bound to the thread that started it.
+        if self._browser is None or self._owner != (os.getpid(), threading.get_ident()):
+            self._launch()
+            self._owner = (os.getpid(), threading.get_ident())
+        page = self._browser.new_page(viewport={"width": width, "height": height})
+        page.route("**/*", lambda route: route.abort()
+                   if route.request.url.startswith(("http://", "https://")) else route.continue_())
+        return page
+
+
+_BROWSER = _Browser()
+
+
+def render_html(html: str, width: int = 900, height: int = 700) -> bytes:
+    """Screenshot of the ``width`` x ``height`` viewport (PNG bytes)."""
+    if not is_safe_markup(html):
+        raise SpecError("HTML contains script elements or event handlers")
+    page = _BROWSER.page(width, height)
+    try:
+        page.set_content(html, wait_until="load", timeout=30_000)
+        return page.screenshot(timeout=30_000)
+    finally:
+        page.close()
+
+
+_MERMAID_PAGE = """<!DOCTYPE html><html><head><meta charset="utf-8">
+<style>body{margin:0;background:white;} #out{width:100%;background:white;}</style>
+</head><body><div id="out"></div></body></html>"""
+
+_MERMAID_RENDER_JS = """async (code) => {
+  try {
+    mermaid.initialize({startOnLoad: false, securityLevel: 'strict', theme: 'default'});
+    const {svg} = await mermaid.render('simit_graph', code);
+    const out = document.getElementById('out');
+    out.innerHTML = svg;
+    return {ok: true};
+  } catch (e) {
+    return {ok: false, error: String((e && e.message) || e)};
+  }
+}"""
+
+
+def render_mermaid(code: str, width: int = 900) -> bytes:
+    """Render Mermaid code to PNG (same page width as ``mmdc -w 900``)."""
+    page = _BROWSER.page(width, 600)
+    try:
+        page.set_content(_MERMAID_PAGE, wait_until="load")
+        page.add_script_tag(path=str(_MERMAID_JS))
+        result = page.evaluate(_MERMAID_RENDER_JS, code)
+        if not result.get("ok"):
+            msg = result.get("error", "Mermaid could not render this diagram")
+            raise SpecError(msg.strip().splitlines()[0][:400] if msg.strip() else "Mermaid syntax error")
+        # Like mmdc: screenshot the SVG's bounding box (mermaid caps its width
+        # at the natural size, and the page width caps it at ``width``).
+        element = page.query_selector("#out svg")
+        if element is None:
+            raise SpecError("Mermaid produced no diagram")
+        return element.screenshot(timeout=30_000)
+    finally:
+        page.close()
+
+
+def render_svg(svg_code: str, output_size: int = 600) -> bytes:
+    """Rasterize SVG onto a white ``output_size`` square, scaled to fit and
+    centred (as cairosvg's output_width/output_height did), via resvg in
+    vl-convert -- no system cairo needed."""
+    import io
+
+    import vl_convert
+    from PIL import Image
+
+    if not is_safe_markup(svg_code):
+        raise SpecError("SVG contains script elements or event handlers")
+    m = re.search(r"<svg\b[^>]*>", svg_code, re.IGNORECASE | re.DOTALL)
+    if not m:
+        raise SpecError("no <svg> element found")
+    w, h = _svg_size(m.group(0))
+    scale = output_size / max(w, h) if w and h else 1.0
+    try:
+        png = vl_convert.svg_to_png(svg_code, scale=scale)
+    except Exception as e:  # malformed XML etc.
+        raise SpecError(f"SVG could not be rendered: {str(e).splitlines()[0][:300]}") from None
+    drawing = Image.open(io.BytesIO(png)).convert("RGBA")
+    drawing.thumbnail((output_size, output_size))
+    canvas = Image.new("RGB", (output_size, output_size), "white")
+    canvas.paste(drawing, ((output_size - drawing.width) // 2, (output_size - drawing.height) // 2), drawing)
+    buf = io.BytesIO()
+    canvas.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def _svg_size(tag: str) -> tuple[float, float]:
+    def num(attr):
+        m = re.search(rf'\b{attr}\s*=\s*["\']\s*([\d.]+)', tag)
+        return float(m.group(1)) if m else 0.0
+
+    w, h = num("width"), num("height")
+    if not (w and h):
+        vb = re.search(r'viewBox\s*=\s*["\']\s*[-\d.]+[\s,]+[-\d.]+[\s,]+([\d.]+)[\s,]+([\d.]+)', tag)
+        if vb:
+            w, h = float(vb.group(1)), float(vb.group(2))
+    return w, h
+
+
+def render_vegalite(spec: dict) -> bytes:
+    import vl_convert
+
+    try:
+        return vl_convert.vegalite_to_png(vl_spec=json.dumps(spec), scale=2)
+    except Exception as e:
+        raise SpecError(str(e).strip().splitlines()[0][:400] or "invalid Vega-Lite spec") from None