Upload drawio_engine.py with huggingface_hub

drawio_engine.py

@@ -0,0 +1,179 @@
+#!/usr/bin/env python3
+"""Agentic draw.io build engine.
+
+extract_plan(path) : full .drawio  -> compact build-plan (style table + nodes +
+                     containers + edges, real geometry & SAP styles, base64
+                     image payloads stripped).
+build_drawio(plan) : build-plan -> complete, valid .drawio XML (the executor).
+
+The model is trained to emit the build-plan; the engine renders the full file.
+Round-tripping every plan guarantees the target is reconstructible & valid.
+"""
+import base64, zlib, urllib.parse, html, re, json
+import xml.etree.ElementTree as ET
+from xml.sax.saxutils import quoteattr
+
+IMAGE_DATA_RE = re.compile(r"image=data:[^;,]+;base64,[^;]+;?")
+IMG_RE = re.compile(r"image=(data:[^;]+)")  # SAP icons: data:image/png,<b64>
+WS_RE = re.compile(r"\s+")
+
+def signature(style):
+    """Group styles by icon identity (image payload) else by full string."""
+    if not style:
+        return "_plain_"
+    m = IMG_RE.search(style)
+    return ("img:" + m.group(1)) if m else ("sty:" + style.strip())
+
+def _inflate(text):
+    text = (text or "").strip()
+    if text.startswith("<"):
+        return text
+    try:
+        return urllib.parse.unquote(zlib.decompress(base64.b64decode(text), -15).decode("utf-8"))
+    except Exception:
+        return None
+
+def _label(val):
+    if not val:
+        return ""
+    return WS_RE.sub(" ", re.sub(r"<[^>]+>", " ", html.unescape(val))).strip()
+
+def _clean_style(style):
+    """Strip heavy base64 image payloads but keep SAP shapes/colors/fonts."""
+    if not style:
+        return "rounded=0;whiteSpace=wrap;html=1;"
+    s = IMAGE_DATA_RE.sub("", style)
+    s = re.sub(r";{2,}", ";", s).strip(";")
+    return (s + ";") if s else "rounded=0;whiteSpace=wrap;html=1;"
+
+def _models(path):
+    data = open(path, "r", encoding="utf-8", errors="ignore").read()
+    try:
+        tree = ET.fromstring(data)
+    except ET.ParseError:
+        return
+    for di, diagram in enumerate(tree.iter("diagram")):
+        gm = diagram.find("mxGraphModel")
+        if gm is None:
+            inf = _inflate(diagram.text)
+            if not inf:
+                continue
+            try:
+                gm = ET.fromstring(inf)
+            except ET.ParseError:
+                continue
+        yield diagram.get("name", f"Page-{di+1}"), gm
+
+def extract_plan(path, sig2key, max_pages=1):
+    """Return list of build-plans (one per page). Styles -> library keys."""
+    plans = []
+    for name, gm in _models(path):
+        root = gm.find("root")
+        if root is None:
+            continue
+        def sid(style):
+            return sig2key.get(signature(style), "box")
+        nodes, edges, idmap = [], [], {}
+        order = []
+        for cell in root.iter("mxCell"):
+            cid = cell.get("id")
+            if cid in (None, "0", "1"):
+                continue
+            order.append(cell)
+        # assign compact ids
+        for i, cell in enumerate(order):
+            idmap[cell.get("id")] = i + 2
+        for cell in order:
+            cid = idmap[cell.get("id")]
+            if cell.get("vertex") == "1":
+                geo = cell.find("mxGeometry")
+                if geo is None:
+                    continue
+                try:
+                    x = int(float(geo.get("x", 0) or 0)); y = int(float(geo.get("y", 0) or 0))
+                    w = int(float(geo.get("width", 80) or 80)); h = int(float(geo.get("height", 60) or 60))
+                except ValueError:
+                    x = y = 0; w, h = 80, 60
+                parent = idmap.get(cell.get("parent"), 1)
+                n = {"i": cid, "l": _label(cell.get("value")), "k": sid(cell.get("style", "")),
+                     "x": x, "y": y, "w": w, "h": h}
+                if parent != 1:
+                    n["p"] = parent
+                nodes.append(n)
+            elif cell.get("edge") == "1":
+                s = idmap.get(cell.get("source")); t = idmap.get(cell.get("target"))
+                if s is None or t is None:
+                    continue
+                e = {"a": s, "b": t, "k": sid(cell.get("style", ""))}
+                lab = _label(cell.get("value"))
+                if lab:
+                    e["l"] = lab
+                edges.append(e)
+        if len(nodes) < 2:
+            continue
+        plans.append({"page": name, "nodes": nodes, "edges": edges})
+    return plans[:max_pages] if max_pages else plans
+
+def build_drawio(plan, lib):
+    """Executor: build-plan + style library -> full valid .drawio XML string."""
+    cells = ['<mxCell id="0" />', '<mxCell id="1" parent="0" />']
+    def st(k):
+        return lib.get(k, "rounded=0;whiteSpace=wrap;html=1;")
+    for n in plan.get("nodes", []):
+        parent = n.get("p", 1)
+        cells.append(
+            f'<mxCell id="{n["i"]}" value={quoteattr(str(n.get("l","")))} '
+            f'style={quoteattr(st(n.get("s",0)))} vertex="1" parent="{parent}">'
+            f'<mxGeometry x="{n.get("x",0)}" y="{n.get("y",0)}" '
+            f'width="{n.get("w",80)}" height="{n.get("h",60)}" as="geometry" /></mxCell>')
+    eid = 100000
+    for e in plan.get("edges", []):
+        lab = e.get("l", "")
+        vattr = f' value={quoteattr(str(lab))}' if lab else ''
+        cells.append(
+            f'<mxCell id="{eid}"{vattr} style={quoteattr(st(e.get("s",0)))} '
+            f'edge="1" parent="1" source="{e["a"]}" target="{e["b"]}">'
+            f'<mxGeometry relative="1" as="geometry" /></mxCell>')
+        eid += 1
+    name = plan.get("page", "Architecture")
+    return (f'<mxfile host="archgen"><diagram name={quoteattr(name)}>'
+            f'<mxGraphModel dx="1200" dy="800" grid="0" gridSize="10" guides="1" '
+            f'tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" '
+            f'pageWidth="1600" pageHeight="1100" math="0" shadow="0"><root>'
+            + "".join(cells) +
+            '</root></mxGraphModel></diagram></mxfile>')
+
+def plan_to_text(plan):
+    """Compact one-line-per-record serialization the model emits."""
+    return json.dumps(plan, ensure_ascii=False, separators=(",", ":"))
+
+def validate(xml):
+    """True if well-formed and structurally a real diagram."""
+    try:
+        t = ET.fromstring(xml)
+    except ET.ParseError:
+        return False, "parse-error"
+    cells = list(t.iter("mxCell"))
+    verts = [c for c in cells if c.get("vertex") == "1"]
+    return (len(verts) >= 2), f"{len(verts)} nodes, {len([c for c in cells if c.get('edge')=='1'])} edges"
+
+if __name__ == "__main__":
+    import glob, os
+    lib = json.load(open("dataset/style_library.json"))
+    sig2key = json.load(open("dataset/sig2key.json"))
+    toks = []
+    for f in sorted(glob.glob("data_refarch/**/*.drawio", recursive=True)):
+        if "/archive/" in f:
+            continue
+        for plan in extract_plan(f, sig2key, max_pages=2):
+            txt = plan_to_text(plan)
+            xml = build_drawio(plan, lib)
+            ok, info = validate(xml)
+            t = len(txt) // 3
+            toks.append(t)
+            print(f"{t:6d}tok plan | nodes={len(plan['nodes']):3d} edges={len(plan['edges']):3d} "
+                  f"| {'OK ' if ok else 'BAD'} {info:18s} | {os.path.basename(f)[:42]}")
+    if toks:
+        print(f"\npages={len(toks)} max~tok={max(toks)} "
+              f"median~={sorted(toks)[len(toks)//2]} mean~={sum(toks)//len(toks)} "
+              f"| >7000tok: {sum(1 for t in toks if t>7000)}")