simit 0.1.0__py3-none-any.whl

simit/skills/builtin.py

@@ -0,0 +1,614 @@
+"""The built-in skill library (17 visualization categories).
+
+Each skill mirrors one category of the SIMIT pipeline: a narrow
+specification prompt, a parser, and a deterministic renderer. ``natural``
+uses the model's own (or a tool) image generator instead.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import io
+import json
+import math
+import re
+import tempfile
+from pathlib import Path
+from typing import Any, Optional
+
+from PIL import Image, ImageDraw, ImageFont
+
+from ..utils import CJK_RE, png_bytes
+from . import prompts as P
+from .base import Realization, Skill, SpecError, extract_fenced_block
+from .routing import BUILTIN_DEFINITIONS
+
+
+def _render_file(fn, *args, suffix: str = ".png") -> Image.Image:
+    """Run a renderer that writes to a path and load the result."""
+    with tempfile.TemporaryDirectory(prefix="simit_render_") as tmp:
+        out = str(Path(tmp) / f"out{suffix}")
+        try:
+            fn(*args, out)
+        except SpecError:
+            raise
+        except Exception as e:  # deterministic renderers raise ValueError & co on bad specs
+            raise SpecError(str(e) or type(e).__name__) from None
+        image = Image.open(out)
+        image.load()
+        return image
+
+
+def _png(data: bytes) -> Image.Image:
+    image = Image.open(io.BytesIO(data))
+    image.load()
+    return image
+
+
+class _Builtin(Skill):
+    """Built-in skills take their router definition from the original prompt."""
+
+    repair_template: str = ""
+
+    def __init__(self):
+        if not self.route_hint:
+            self.route_hint = BUILTIN_DEFINITIONS[self.label]
+
+    def build_repair_prompt(self, request, previous, reason):
+        return self.repair_template.format(request=request, reason=reason)
+
+
+# --------------------------------------------------------------------- natural
+class NaturalSkill(_Builtin):
+    """Photos/illustrations of real-world content: native image generation."""
+
+    name = "natural"
+    verify_rubric = "strict"
+    # Body-part / botanical descriptions are routed here before asking the
+    # model, which often mistakes them for GEOMETRY.
+    route_pattern = re.compile(
+        r'\b(?:foot|feet|hand|palm|finger|thumb|toe|knee|ankle|elbow|wrist|shoulder|hip'
+        r'|eye|nose|mouth|ear|neck|chest|spine|rib|skull|bone|muscle|skin'
+        r'|brain|lung|liver|kidney|stomach|intestine|artery|vein|nerve'
+        r'|leaf|leaves|petal|sepal|flower|blossom|bloom'
+        r'|moss|algae|fern|mushroom|fungi|pollen|spore'
+        r'|fur|feather|claw|hoof|paw|snout|beak|gill|whisker)\b',
+        re.IGNORECASE,
+    )
+    route_exclude_pattern = re.compile(
+        r'\b(?:diagram|schematic|anatomical|cross.?section|anatomy|flowchart|food\s+chain)\b', re.IGNORECASE)
+
+    async def realize(self, ctx, request, retries=2):
+        image = await ctx.generate_image(request)
+        return None if image is None else Realization(image, self.name, spec_text=request)
+
+
+# ----------------------------------------------------------- DSL-based skills
+class DiagramSkill(_Builtin):
+    name = "diagram"
+    prompt = P.DIAGRAM_DSL_PROMPT
+    repair_template = P.DSL_REPAIR_PROMPT
+    retry_temperature = 0.0
+
+    def parse(self, text):
+        from .renderers import parse_dsl
+        nodes, edges = parse_dsl(text)
+        if not nodes:
+            raise SpecError("no NODE: lines found")
+        return nodes, edges
+
+    def render(self, spec):
+        from .renderers import render_flowchart
+        return _render_file(render_flowchart, *spec)
+
+
+class GraphSkill(_Builtin):
+    name = "graph"
+    prompt = P.GRAPH_DSL_PROMPT
+    repair_template = P.GRAPH_REPAIR_PROMPT
+    retry_temperature = 0.0
+
+    def parse(self, text):
+        from .renderers import parse_graph_dsl
+        nodes, edges = parse_graph_dsl(text)
+        if not nodes:
+            raise SpecError("no NODE: lines found")
+        return nodes, edges
+
+    def render(self, spec):
+        from .renderers import render_graph
+        return _render_file(render_graph, *spec)
+
+
+class MoleculeSkill(_Builtin):
+    name = "molecule"
+    prompt = P.MOLECULE_PROMPT
+    max_new_tokens = 128
+    retry_temperature = 0.0
+
+    def shortcut(self, request):
+        # The model's chemistry recall is unreliable even for common molecules.
+        from .renderers import lookup_common_molecule
+        return lookup_common_molecule(request)
+
+    def parse(self, text):
+        from rdkit import Chem, RDLogger
+        from .renderers import parse_molecule
+        smiles, title = parse_molecule(text)
+        if not smiles:
+            raise SpecError("missing")
+        RDLogger.DisableLog("rdApp.*")
+        if Chem.MolFromSmiles(smiles) is None:
+            raise SpecError("RDKit could not parse it")
+        return smiles, title
+
+    def build_repair_prompt(self, request, previous, reason):
+        from .renderers import parse_molecule
+        smiles = parse_molecule(previous)[0] or "(none)"
+        return P.MOLECULE_REPAIR_PROMPT.format(request=request, smiles=smiles, reason=reason)
+
+    def render(self, spec):
+        from rdkit import Chem
+        from rdkit.Chem import Draw
+        smiles, title = spec
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            raise SpecError("RDKit could not parse it")
+        return Draw.MolToImage(mol, size=(450, 450), legend=title)
+
+
+class CircuitSkill(_Builtin):
+    name = "circuit"
+    prompt = P.CIRCUIT_DSL_PROMPT
+    repair_template = P.CIRCUIT_REPAIR_PROMPT
+    max_new_tokens = 256
+    retry_temperature = 0.0
+
+    def parse(self, text):
+        from .renderers import parse_circuit_dsl
+        slots = parse_circuit_dsl(text)
+        if not slots:
+            raise SpecError("no COMPONENT: lines found")
+        return slots
+
+    def render(self, spec):
+        from .renderers import render_circuit
+        return _render_file(render_circuit, spec)
+
+
+class VectorSkill(_Builtin):
+    name = "vector"
+    prompt = P.VECTOR_DSL_PROMPT
+    repair_template = P.VECTOR_REPAIR_PROMPT
+    retry_temperature = 0.0
+
+    def parse(self, text):
+        from .renderers import parse_vector_dsl
+        origin_box, vectors = parse_vector_dsl(text)
+        if not vectors:
+            raise SpecError("no VECTOR: lines found")
+        return origin_box, vectors
+
+    def render(self, spec):
+        from .renderers import render_vectors
+        return _render_file(render_vectors, *spec)
+
+
+class VennSkill(_Builtin):
+    name = "venn"
+    prompt = P.VENN_DSL_PROMPT
+    repair_template = P.VENN_REPAIR_PROMPT
+    retry_temperature = 0.0
+
+    def parse(self, text):
+        from .renderers import parse_venn_dsl
+        fields = parse_venn_dsl(text)
+        if not fields["SET_A"] or not fields["SET_B"]:
+            raise SpecError("missing SET_A or SET_B")
+        return fields
+
+    def render(self, spec):
+        from .renderers import render_venn
+        return _render_file(render_venn, spec)
+
+
+class TableSkill(_Builtin):
+    name = "table"
+    prompt = P.TABLE_DSL_PROMPT
+    repair_template = P.TABLE_REPAIR_PROMPT
+    retry_temperature = 0.0
+
+    def parse(self, text):
+        from .renderers import parse_table_dsl
+        title, columns, rows = parse_table_dsl(text)
+        if not columns or not rows:
+            raise SpecError("missing COLUMNS or ROW lines")
+        return title, columns, rows
+
+    def render(self, spec):
+        from .renderers import render_table
+        return _render_file(render_table, *spec)
+
+
+class SceneSkill(_Builtin):
+    name = "scene"
+    prompt = P.SCENE_DSL_PROMPT
+    repair_template = P.SCENE_REPAIR_PROMPT
+    retry_temperature = 0.0
+
+    def parse(self, text):
+        from .renderers import parse_scene_dsl
+        title, items = parse_scene_dsl(text)
+        if not items:
+            raise SpecError("no SHAPE: lines found")
+        return title, items
+
+    def render(self, spec):
+        from .renderers import render_scene
+        return _render_file(render_scene, *spec)
+
+
+class PuzzleSkill(_Builtin):
+    name = "puzzle"
+    prompt = P.PUZZLE_DSL_PROMPT
+    repair_template = P.PUZZLE_REPAIR_PROMPT
+    retry_temperature = 0.0
+
+    def parse(self, text):
+        from .renderers import parse_puzzle_dsl
+        spec = parse_puzzle_dsl(text)
+        if not spec["layout"] or not spec["panels"]:
+            raise SpecError("missing LAYOUT or PANEL lines")
+        return spec
+
+    def render(self, spec):
+        from .renderers import render_puzzle
+        return _render_file(render_puzzle, spec)
+
+
+class GeometrySkill(_Builtin):
+    """Points and references only; every length/angle/area is computed by the
+    renderer from coordinates."""
+
+    name = "geometry"
+    prompt = P.GEOMETRY_DSL_PROMPT
+    max_new_tokens = 768
+
+    def build_repair_prompt(self, request, previous, reason):
+        return P.GEOMETRY_REPAIR_PROMPT.format(request=request, previous=previous, reason=reason)
+
+    async def realize(self, ctx, request, retries=2):
+        from .renderers import parse_geometry_dsl
+        prompt = self.build_prompt(request)
+        for attempt in range(retries + 1):
+            temperature = self.first_temperature if attempt == 0 else self.retry_temperature
+            cleaned = self.clean(await ctx.generate(prompt, max_new_tokens=self.max_new_tokens,
+                                                    temperature=temperature))
+            # Corruption often hits only a few lines: drop them and keep the
+            # rest, unless so much is lost that a retry is the better bet.
+            raw = cleaned.splitlines()
+            usable = [ln for ln in raw if not CJK_RE.search(ln)]
+            if len(raw) - len(usable) > len(raw) / 3 and attempt < retries:
+                prompt = self.build_repair_prompt(request, cleaned, "your answer contained non-English characters")
+                continue
+            cleaned = "\n".join(usable)
+            commands = parse_geometry_dsl(cleaned)
+            if not commands:
+                prompt = self.build_repair_prompt(request, cleaned, "no structure lines found")
+                continue
+            try:
+                image = await ctx.render(self, commands)
+            except SpecError as e:
+                prompt = self.build_repair_prompt(request, cleaned, str(e))
+                continue
+            return Realization(image, self.name, spec_text=cleaned, attempts=attempt + 1)
+        return None
+
+    def render(self, spec):
+        from .renderers import render_geometry
+        return _render_file(render_geometry, spec)
+
+
+# ------------------------------------------------------------- code-based skills
+class FigureSkill(_Builtin):
+    """Charts and other numeric figures: the model writes matplotlib code."""
+
+    name = "figure"
+    prompt = P.MASTER_PROMPT
+    max_new_tokens = 1536
+
+    async def realize(self, ctx, request, retries=2):
+        from .sandbox import extract_code, find_known_api_issues, find_non_ascii_issue, is_safe
+        prompt = self.build_prompt(request)
+        for attempt in range(retries + 1):
+            temperature = self.first_temperature if attempt == 0 else self.retry_temperature
+            code = extract_code(self.clean(await ctx.generate(prompt, max_new_tokens=self.max_new_tokens,
+                                                              temperature=temperature)))
+            if code is None:
+                prompt = self.build_prompt(request)
+                continue
+            if not is_safe(code):
+                ctx.log("[figure] generated code failed the safety check; not running it")
+                return None
+            issue = find_non_ascii_issue(code) or find_known_api_issues(code)
+            if issue:
+                prompt = P.REPAIR_PROMPT.format(request=request, code=code, error=issue)
+                continue
+            try:
+                image = await ctx.render(self, code)
+            except SpecError as e:
+                prompt = P.REPAIR_PROMPT.format(request=request, code=code, error=str(e))
+                continue
+            return Realization(image, self.name, spec_text=code, attempts=attempt + 1)
+        return None
+
+    def render(self, spec):
+        from .sandbox import run_script
+        image, error = run_script(spec)
+        if image is None:
+            raise SpecError(error or "unknown error")
+        return image
+
+
+class SvgSkill(_Builtin):
+    name = "svg"
+    prompt = P.SVG_PROMPT
+    repair_template = P.SVG_REPAIR_PROMPT
+    max_new_tokens = 1024
+
+    def parse(self, text):
+        code = extract_fenced_block(text)
+        if code is None:
+            raise SpecError("no SVG code block found in response")
+        if "<svg" not in code:
+            raise SpecError("no <svg> element found in the code block")
+        return code[code.index("<svg"):]
+
+    def render(self, spec):
+        from .web import render_svg
+        return _png(render_svg(spec))
+
+
+class MermaidSkill(_Builtin):
+    name = "mermaid"
+    prompt = P.MERMAID_PROMPT
+    repair_template = P.MERMAID_REPAIR_PROMPT
+    max_new_tokens = 768
+
+    def parse(self, text):
+        code = extract_fenced_block(text)
+        if code is None:
+            raise SpecError("no fenced code block found")
+        return code
+
+    def render(self, spec):
+        from .web import render_mermaid
+        return _png(render_mermaid(spec))
+
+
+_VEGALITE_SCHEME_FIXES = {"yelloworange": "yelloworangered", "redyellow": "redyellowblue", "bluered": "redblue"}
+_MONTHS = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
+
+
+def _clean_spec_strings(obj):
+    if isinstance(obj, str):
+        return CJK_RE.sub("", obj).strip().rstrip("（(")
+    if isinstance(obj, dict):
+        return {k: _clean_spec_strings(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_clean_spec_strings(v) for v in obj]
+    return obj
+
+
+def _fix_vegalite_spec(spec: dict) -> dict:
+    """Patch common model mistakes: invalid scheme names, 'area' used for a
+    heatmap, months sorted alphabetically."""
+    text = json.dumps(spec)
+    for wrong, right in _VEGALITE_SCHEME_FIXES.items():
+        text = text.replace(f'"scheme": "{wrong}"', f'"scheme": "{right}"')
+    spec = json.loads(text)
+    try:
+        enc = spec.get("encoding", {})
+        if isinstance(enc, dict):
+            typ = lambda k: enc.get(k, {}).get("type", "") if isinstance(enc.get(k), dict) else ""  # noqa: E731
+            if (spec.get("mark") == "area" and typ("x") in ("ordinal", "nominal")
+                    and typ("y") in ("ordinal", "nominal") and typ("color") == "quantitative"):
+                spec["mark"] = "rect"
+        data = spec.get("data", {})
+        values = data.get("values", []) if isinstance(data, dict) else []
+        month_fields = {f for row in values if isinstance(row, dict) for f, v in row.items() if v in _MONTHS}
+        if month_fields:
+            def add_sort(e):
+                for axis in ("x", "y"):
+                    ax = e.get(axis, {})
+                    if isinstance(ax, dict) and ax.get("field") in month_fields and "sort" not in ax:
+                        ax["sort"] = _MONTHS
+            if isinstance(spec.get("encoding"), dict):
+                add_sort(spec["encoding"])
+            for layer in spec.get("layer", []):
+                if isinstance(layer, dict) and isinstance(layer.get("encoding"), dict):
+                    add_sort(layer["encoding"])
+    except Exception:
+        pass
+    return spec
+
+
+class VegaLiteSkill(_Builtin):
+    name = "vegalite"
+    prompt = P.VEGALITE_PROMPT
+    repair_template = P.VEGALITE_REPAIR_PROMPT
+    max_new_tokens = 1536
+
+    def parse(self, text):
+        block = extract_fenced_block(text)
+        if block is None:
+            raise SpecError("no JSON code block found")
+        try:
+            spec = json.loads(block)
+        except json.JSONDecodeError as e:
+            raise SpecError(f"invalid JSON: {e}") from None
+        if not isinstance(spec, dict):
+            raise SpecError("the top-level JSON must be an object, not an array")
+        return _fix_vegalite_spec(_clean_spec_strings(spec))
+
+    def render(self, spec):
+        from .web import render_vegalite
+        return _png(render_vegalite(spec))
+
+
+# A CJK run inside CSS often carries a glued-on English "translation" word
+# (e.g. "#3498db 获取acquire"); drop the run together with that word.
+_CJK_WITH_ORPHAN_RE = re.compile(r"[一-鿿]+[A-Za-z]*")
+
+
+def _clean_html(html: str) -> str:
+    """Strip CJK contamination and stray fence markers from model HTML."""
+    lines = html.splitlines()
+    while lines and lines[0].strip().startswith("```"):
+        lines = lines[1:]
+    html = "\n".join(lines)
+    html = re.sub(r'style="([^"]*)"', lambda m: f'style="{_CJK_WITH_ORPHAN_RE.sub("", m.group(1))}"', html)
+    return CJK_RE.sub("", html)
+
+
+class HtmlSkill(_Builtin):
+    name = "html"
+    prompt = P.HTML_PROMPT
+    repair_template = P.HTML_REPAIR_PROMPT
+    max_new_tokens = 1536
+
+    def parse(self, text):
+        html = extract_fenced_block(text)
+        if html is None:
+            if "<html" in text.lower() or "<body" in text.lower():
+                html = text
+            else:
+                raise SpecError("no HTML code block found in response")
+        html = _clean_html(html)
+        if "<!DOCTYPE" not in html and "<html" not in html.lower():
+            html = f"<!DOCTYPE html><html><body style='margin:20px;font-family:sans-serif'>{html}</body></html>"
+        return html
+
+    def render(self, spec):
+        from .web import render_html
+        return _png(render_html(spec))
+
+
+# ------------------------------------------------------------------ composite
+_PANEL_SIZE = 280
+
+
+def _gray_panel(label: str, size: int = _PANEL_SIZE) -> bytes:
+    img = Image.new("RGB", (size, size), color="#e8e8e8")
+    draw = ImageDraw.Draw(img)
+    draw.rectangle([2, 2, size - 3, size - 3], outline="#aaaaaa", width=2)
+    try:
+        font = ImageFont.load_default(size=32)
+    except TypeError:  # Pillow < 10.1
+        font = ImageFont.load_default()
+    draw.text((size // 2, size // 2), label or "?", fill="#666666", anchor="mm", font=font)
+    return png_bytes(img)
+
+
+class CompositeSkill(_Builtin):
+    """N labeled sub-images in a grid; each panel is realized recursively by
+    whichever skill fits it, then composed as an HTML page."""
+
+    name = "html_composite"
+    prompt = P.HTML_COMPOSITE_PROMPT
+    repair_template = P.HTML_COMPOSITE_REPAIR_PROMPT
+    max_new_tokens = 768
+    route_priority = 1  # checked before natural's pattern
+    # Multi-panel labeled image grids are routed here before asking the model.
+    route_pattern = re.compile(
+        r'(?:'
+        r'\b(?:image|photo|picture)s?\s+(?:labeled\s+|showing\s+)?[A-Da-d1-5]'
+        r'|\bpanels?\s+(?:labeled|showing)\s+(?:[A-Da-d1-5]|different)'
+        r'|\bpanel\s+[A-Da-d1-5]\s+(?:shows?|is\b|depicts?|contains?)'
+        r'|\b(?:three|four|five|[2-5])\s+(?:image|photo|picture)s?\b'
+        r'|labeled\s+[A-Da-d]\s+through\s+[A-Da-d]'
+        r'|labeled\s+[A-Da-d][-–][A-Da-d]'
+        r'|\b(?:left|right|top|bottom)\s+panel\b.{0,60}labeled\s+[A-Da-d]'
+        r')',
+        re.IGNORECASE,
+    )
+
+    def parse(self, text):
+        title, panels = "", []
+        for line in text.splitlines():
+            line = line.strip()
+            if line.upper().startswith("TITLE:"):
+                title = line[6:].strip()
+            elif line.upper().startswith("PANEL:"):
+                content = line[6:].strip()
+                if "|" in content:
+                    desc, label = content.rsplit("|", 1)
+                    panels.append((desc.strip(), label.strip()))
+                else:
+                    panels.append((content, ""))
+        if not panels:
+            raise SpecError("no PANEL: lines found")
+        return title, panels
+
+    async def realize(self, ctx, request, retries=2):
+        prompt = self.build_prompt(request)
+        layout = None
+        for attempt in range(retries + 1):
+            temperature = self.first_temperature if attempt == 0 else self.retry_temperature
+            text = self.clean(await ctx.generate(prompt, max_new_tokens=self.max_new_tokens,
+                                                 temperature=temperature))
+            try:
+                layout = self.parse(text)
+                break
+            except SpecError as e:
+                prompt = self.build_repair_prompt(request, text, str(e))
+        if layout is None:
+            return None
+        title, panels = layout
+        subs = await asyncio.gather(*[
+            ctx.realize(desc, retries=1, allow_composite=False, size=_PANEL_SIZE) for desc, _ in panels
+        ])
+        rendered = []
+        for (desc, label), sub in zip(panels, subs):
+            rendered.append((png_bytes(sub.image) if sub is not None else _gray_panel(label), label))
+        image = await ctx.render(self, {"title": title, "panels": rendered})
+        return Realization(image, self.name, spec_text=text, attempts=attempt + 1,
+                           extra={"panels": [(d, l, s.skill if s else None) for (d, l), s in zip(panels, subs)]})
+
+    def render(self, spec):
+        from .web import render_html
+        title, panels = spec["title"], spec["panels"]
+        n = len(panels)
+        n_cols = min(4, n)
+        n_rows = math.ceil(n / n_cols)
+        width = n_cols * (_PANEL_SIZE + 30) + 60
+        height = 80 + (50 if title else 0) + n_rows * (_PANEL_SIZE + 60) + 40
+        divs = ""
+        for data, label in panels:
+            b64 = base64.b64encode(data).decode()
+            label_html = f'<div class="label">{label}</div>' if label else ""
+            divs += f'\n<div class="panel"><img src="data:image/png;base64,{b64}" alt="{label}">{label_html}</div>'
+        title_html = f'<h2 style="text-align:center;margin:0 0 16px;font-size:20px">{title}</h2>' if title else ""
+        html = f"""<!DOCTYPE html><html><head><meta charset="UTF-8"><style>
+  body {{ background:#fff; margin:20px; font-family:Arial,sans-serif; }}
+  .grid {{ display:flex; flex-wrap:wrap; gap:20px; justify-content:center; align-items:flex-start; }}
+  .panel {{ text-align:center; }}
+  .panel img {{ width:{_PANEL_SIZE}px; height:{_PANEL_SIZE}px; object-fit:contain;
+               border:1px solid #ccc; background:#fafafa; display:block; }}
+  .label {{ font-size:22px; font-weight:bold; margin-top:8px; color:#111; }}
+</style></head><body>{title_html}<div class="grid">{divs}
+</div></body></html>"""
+        return _png(render_html(html, width=width, height=height))
+
+
+BUILTIN_SKILLS = [
+    NaturalSkill, DiagramSkill, GraphSkill, MoleculeSkill, CircuitSkill, VectorSkill, VennSkill,
+    TableSkill, SceneSkill, PuzzleSkill, GeometrySkill, FigureSkill, MermaidSkill, SvgSkill,
+    VegaLiteSkill, HtmlSkill, CompositeSkill,
+]
+
+
+def default_skills() -> list[Skill]:
+    """Fresh instances of the 17 built-in skills, in router order."""
+    return [cls() for cls in BUILTIN_SKILLS]