traceforge-pcb 1.0.0__py3-none-any.whl

traceforge/bench/checks.py

@@ -0,0 +1,410 @@
+"""Machine-checkable requirements, evaluated on a :class:`Netlist` and a :class:`Board`.
+
+A task names the parts it cares about as *roles* (``mcu``, ``usbc``, ``ldo``)
+with selectors, then lists checks over those roles. Selectors and checks only
+look at what KiCad writes (values, library ids, footprints, pin names, nets), so
+they are blind to which tool produced the project.
+"""
+
+from __future__ import annotations
+
+import re
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Any
+
+from traceforge.bench.design import Board, Component, Netlist, clean_net, edge_distance, is_ground
+
+_PREFIX = {
+    "p": 1e-12,
+    "n": 1e-9,
+    "u": 1e-6,
+    "µ": 1e-6,
+    "μ": 1e-6,
+    "m": 1e-3,
+    "k": 1e3,
+    "K": 1e3,
+    "M": 1e6,
+    "G": 1e9,
+    "R": 1.0,
+    "r": 1.0,
+}
+_VALUE_RE = re.compile(r"(\d+(?:[.,]\d+)?)\s*([pnuµμmkKMGRr]?)(\d*)")
+PASSIVE_KINDS = "R|C|L|D|LED|F|FB|FUSE|RN|VD|TVS|Z"
+
+
+def parse_value(text: str | float | int | None) -> float | None:
+    """'5.1k' / '5k1' / '4R7' / '100nF' / '0.1 uF 16V' / 5100 -> number (SI)."""
+    if text is None:
+        return None
+    if isinstance(text, (int, float)):
+        return float(text)
+    match = _VALUE_RE.search(str(text).replace("Ω", "").replace("ohm", ""))
+    if not match:
+        return None
+    whole, prefix, frac = match.groups()
+    number = float(whole.replace(",", "."))
+    if frac and prefix:
+        number = float(f"{whole}.{frac}")
+    return number * _PREFIX.get(prefix, 1.0)
+
+
+def value_matches(text: str, want: Any, tolerance: float = 0.03) -> bool:
+    got, target = parse_value(text), parse_value(want)
+    if got is None or target is None:
+        return False
+    return abs(got - target) <= tolerance * abs(target) + 1e-18
+
+
+@dataclass
+class Selector:
+    """Which components play a role. All given fields must match."""
+
+    match: str = ""  # regex searched in value / lib id / MPN / description
+    kind: str = ""  # regex fully matching the reference prefix (R, C, U, J, ...)
+    footprint: str = ""  # regex searched in the footprint id
+    value: str | float | None = None  # numeric value (5.1k, 100n) within tolerance
+    tolerance: float = 0.03
+    exclude: str = ""  # regex that must NOT be found in the part's text (e.g. LED)
+
+    @classmethod
+    def build(cls, raw: Any) -> Selector:
+        if isinstance(raw, str):
+            return cls(match=raw)
+        return cls(**{k: v for k, v in dict(raw).items() if k in cls.__dataclass_fields__})
+
+    def accepts(self, comp: Component) -> bool:
+        if self.kind and not re.fullmatch(self.kind, comp.kind, re.I):
+            return False
+        if self.match and not re.search(self.match, comp.text(), re.I):
+            return False
+        if self.exclude and re.search(self.exclude, comp.text(), re.I):
+            return False
+        if self.footprint and not re.search(self.footprint, comp.footprint, re.I):
+            return False
+        return not (
+            self.value is not None and not value_matches(comp.value, self.value, self.tolerance)
+        )
+
+
+@dataclass
+class Result:
+    id: str
+    label: str
+    score: float  # 0..1
+    detail: str = ""
+    weight: float = 1.0
+
+    @property
+    def passed(self) -> bool:
+        return self.score >= 0.999
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "label": self.label,
+            "passed": self.passed,
+            "score": round(self.score, 4),
+            "weight": self.weight,
+            "detail": self.detail,
+        }
+
+
+@dataclass
+class Context:
+    netlist: Netlist
+    board: Board | None
+    roles: dict[str, Selector] = field(default_factory=dict)
+    base: Netlist | None = None  # edit tasks: the project before the edit
+
+    def selector(self, role_or_selector: Any) -> Selector:
+        if isinstance(role_or_selector, Selector):
+            return role_or_selector
+        if isinstance(role_or_selector, str) and role_or_selector in self.roles:
+            return self.roles[role_or_selector]
+        return Selector.build(role_or_selector)
+
+    def candidates(self, role_or_selector: Any) -> list[Component]:
+        selector = self.selector(role_or_selector)
+        comps = dict(self.netlist.components)
+        if self.board:  # PCB-only items (mounting holes, logos) count too
+            for ref, comp in self.board.netlist().components.items():
+                comps.setdefault(ref, comp)
+        return [
+            c for ref, c in sorted(comps.items()) if not ref.startswith("#") and selector.accepts(c)
+        ]
+
+    def endpoint_nets(self, raw: Any) -> set[str]:
+        """``{role, pin}`` -> nets on that pin; ``{net: regex}`` or ``GND`` -> named nets."""
+        if isinstance(raw, list):
+            return set().union(*(self.endpoint_nets(r) for r in raw)) if raw else set()
+        if isinstance(raw, str):
+            raw = {"net": raw}
+        nets = set(self.netlist.nets)
+        if "net" in raw:
+            pattern = str(raw["net"])
+            if pattern.upper() == "GND":
+                return {n for n in nets if is_ground(n)}
+            return {n for n in nets if re.fullmatch(pattern, clean_net(n).lstrip("/"), re.I)}
+        pin_re = str(raw.get("pin", ".*"))
+        out = set()
+        selector = (
+            raw.get("role") or raw.get("part") or {k: v for k, v in raw.items() if k != "pin"}
+        )
+        for comp in self.candidates(selector):
+            for pin, net in self.netlist.pins_of(comp.ref):
+                if pin_matches(pin.pin, pin.function, pin_re):
+                    out.add(net)
+        return out
+
+
+def pin_matches(number: str, function: str, pattern: str) -> bool:
+    if re.fullmatch(pattern, number, re.I):
+        return True
+    names = [function, *re.split(r"[/,\s]+", function)] if function else []
+    return any(n and re.fullmatch(pattern, n, re.I) for n in names)
+
+
+def _describe(raw: Any) -> str:
+    if isinstance(raw, list):
+        return " or ".join(_describe(r) for r in raw)
+    if isinstance(raw, str):
+        return raw
+    if "net" in raw:
+        return f"net {raw['net']}"
+    who = raw.get("role") or raw.get("part") or raw.get("match") or raw.get("kind") or "part"
+    return f"{who}.{raw['pin']}" if "pin" in raw else str(who)
+
+
+# ---------------------------------------------------------------------- checks
+
+
+def check_part(ctx: Context, spec: dict) -> tuple[float, str]:
+    found = ctx.candidates(spec.get("role") or spec.get("part"))
+    need = int(spec.get("min", 1))
+    refs = ", ".join(c.ref for c in found[:8]) or "none"
+    if "max" in spec and len(found) > int(spec["max"]):
+        return 0.0, f"{len(found)} found (max {spec['max']}): {refs}"
+    return min(1.0, len(found) / need) if need else 1.0, f"{len(found)}/{need}: {refs}"
+
+
+def check_connected(ctx: Context, spec: dict) -> tuple[float, str]:
+    """All endpoints share at least one net."""
+    sets = [ctx.endpoint_nets(e) for e in spec["pins"]]
+    if not all(sets):
+        missing = [_describe(e) for e, s in zip(spec["pins"], sets, strict=True) if not s]
+        return 0.0, "not found: " + ", ".join(missing)
+    common = set.intersection(*sets)
+    if common:
+        return 1.0, "on " + ", ".join(sorted(common))
+    return 0.0, " / ".join(",".join(sorted(s)) for s in sets)
+
+
+def check_separate(ctx: Context, spec: dict) -> tuple[float, str]:
+    """Endpoints exist and are NOT on a common net (e.g. something sits in series)."""
+    sets = [ctx.endpoint_nets(e) for e in spec["pins"]]
+    if not all(sets):
+        return 0.0, "endpoint not found"
+    common = set.intersection(*sets)
+    return (0.0, "shorted on " + ", ".join(sorted(common))) if common else (1.0, "separate nets")
+
+
+def _two_sided(ctx: Context, sel: Selector, a: set[str], b: set[str]) -> list[str]:
+    hits = []
+    for comp in ctx.candidates(sel):
+        nets = {net for _, net in ctx.netlist.pins_of(comp.ref)}
+        if any(x != y for x in nets & a for y in nets & b):
+            hits.append(comp.ref)
+    return hits
+
+
+def check_through(ctx: Context, spec: dict) -> tuple[float, str]:
+    """A component matching ``via`` bridges ``from`` and ``to`` (e.g. CC1 -> 5.1k -> GND)."""
+    a, b = ctx.endpoint_nets(spec["from"]), ctx.endpoint_nets(spec["to"])
+    if not a or not b:
+        return 0.0, f"endpoint not found ({_describe(spec['from'])} / {_describe(spec['to'])})"
+    via = ctx.selector(spec["via"])
+    hits = _two_sided(ctx, via, a, b)
+    need = int(spec.get("min", 1))
+    return min(1.0, len(hits) / need), (", ".join(hits) if hits else "no bridging part")
+
+
+def check_path(ctx: Context, spec: dict) -> tuple[float, str]:
+    """A chain of at most ``max_parts`` components links ``from`` to ``to``.
+
+    ``allow`` restricts which components may be traversed (default: passives
+    and diodes); every selector in ``through`` must appear on the chain.
+    """
+    starts, goals = ctx.endpoint_nets(spec["from"]), ctx.endpoint_nets(spec["to"])
+    if not starts or not goals:
+        return 0.0, "endpoint not found"
+    allow = ctx.selector(spec.get("allow", {"kind": PASSIVE_KINDS}))
+    through = [ctx.selector(t) for t in spec.get("through", [])]
+    max_parts = int(spec.get("max_parts", 3))
+    if not through and starts & goals:
+        return 1.0, "same net: " + ", ".join(sorted(starts & goals))
+    comps = ctx.candidates(allow)
+    comp_nets = {c.ref: {n for _, n in ctx.netlist.pins_of(c.ref)} for c in comps}
+    marks = {c.ref: frozenset(i for i, t in enumerate(through) if t.accepts(c)) for c in comps}
+    full = frozenset(range(len(through)))
+    queue = deque((n, frozenset(), (), frozenset()) for n in starts)
+    while queue:
+        net, done, chain, used = queue.popleft()
+        if len(chain) >= max_parts:
+            continue
+        for ref, nets in comp_nets.items():
+            if ref in used or net not in nets:
+                continue
+            got = done | marks[ref]
+            for nxt in nets - {net}:
+                if nxt in goals and got == full:
+                    return 1.0, " -> ".join([*chain, ref])
+                queue.append((nxt, got, (*chain, ref), used | {ref}))
+    return 0.0, "no path"
+
+
+def check_preserve(ctx: Context, spec: dict) -> tuple[float, str]:
+    """Edit tasks: the original connectivity survives (node sets, net names may change)."""
+    if ctx.base is None:
+        return 0.0, "no base project"
+    ignore_refs = set(spec.get("ignore_refs", []))
+    ignore_nets = re.compile(spec["ignore_nets"], re.I) if spec.get("ignore_nets") else None
+    base = {
+        n: s
+        for n, s in ctx.base.node_sets(ignore_refs).items()
+        if not (ignore_nets and ignore_nets.fullmatch(n))
+    }
+    now = set(ctx.netlist.node_sets(ignore_refs).values())
+    kept = [n for n, s in base.items() if s in now]
+    lost = sorted(set(base) - set(kept))
+    score = len(kept) / len(base) if base else 1.0
+    return score, f"{len(kept)}/{len(base)} nets kept" + (
+        f"; changed: {', '.join(lost[:8])}" if lost else ""
+    )
+
+
+def check_matched_length(ctx: Context, spec: dict) -> tuple[float, str]:
+    """A routed differential pair: both nets have copper and their length skew is small."""
+    if ctx.board is None:
+        return 0.0, "no board"
+    nets = []
+    for end in spec["nets"]:
+        found = sorted(ctx.endpoint_nets(end))
+        if len(found) != 1:
+            return 0.0, f"{_describe(end)} -> {found or 'no net'}"
+        nets.append(found[0])
+    lengths = [ctx.board.track_length.get(n, 0.0) for n in nets]
+    if min(lengths) <= 0:
+        return 0.0, f"unrouted: {dict(zip(nets, lengths, strict=True))}"
+    skew = max(lengths) - min(lengths)
+    limit = float(spec.get("max_skew_mm", 1.0))
+    detail = f"{nets[0]} {lengths[0]:.2f} mm / {nets[1]} {lengths[1]:.2f} mm, skew {skew:.2f} mm"
+    return (1.0 if skew <= limit else 0.5), detail
+
+
+def check_count(ctx: Context, spec: dict) -> tuple[float, str]:
+    """Component count of a kind, e.g. at least 4 mounting holes."""
+    return check_part(ctx, spec)
+
+
+CHECKS = {
+    "part": check_part,
+    "count": check_count,
+    "connected": check_connected,
+    "separate": check_separate,
+    "through": check_through,
+    "path": check_path,
+    "preserve": check_preserve,
+    "matched_length": check_matched_length,
+}
+
+
+def evaluate(ctx: Context, checks: list[dict]) -> list[Result]:
+    results = []
+    for i, spec in enumerate(checks):
+        kind = spec.get("type", "part")
+        label = spec.get("label") or f"{kind} {i + 1}"
+        try:
+            score, detail = CHECKS[kind](ctx, spec)
+        except Exception as exc:  # a broken project must not crash the scorer
+            score, detail = 0.0, f"error: {type(exc).__name__}: {exc}"
+        results.append(
+            Result(
+                spec.get("id", f"r{i + 1}"), label, score, detail, float(spec.get("weight", 1.0))
+            )
+        )
+    return results
+
+
+# --------------------------------------------------------------- board checks
+
+
+def board_checks(board: Board | None, spec: dict) -> list[Result]:
+    out: list[Result] = []
+    if not spec:
+        return out
+    if board is None:
+        return [Result("board", "board exists", 0.0, "no .kicad_pcb")]
+    size = board.size
+    if "max_width" in spec or "max_height" in spec:
+        limit = sorted([float(spec.get("max_width", 1e9)), float(spec.get("max_height", 1e9))])
+        if size is None:
+            out.append(Result("size", "board size", 0.0, "no Edge.Cuts outline"))
+        else:
+            ok = all(s <= lim + 0.1 for s, lim in zip(sorted(size), limit, strict=True))
+            out.append(
+                Result(
+                    "size",
+                    f"board within {spec.get('max_width')} x {spec.get('max_height')} mm",
+                    float(ok),
+                    f"{size[0]} x {size[1]} mm",
+                )
+            )
+    if "layers" in spec:
+        got = len(board.copper_layers)
+        out.append(
+            Result(
+                "layers",
+                f"{spec['layers']} copper layers",
+                float(got == int(spec["layers"])),
+                f"{got} layers",
+            )
+        )
+    return out
+
+
+def edge_checks(ctx: Context, roles: list[Any], max_inset_mm: float = 3.0) -> list[Result]:
+    out = []
+    for role in roles:
+        label = f"{role if isinstance(role, str) else _describe(role)} on board edge"
+        if ctx.board is None:
+            out.append(Result("edge", label, 0.0, "no board"))
+            continue
+        refs = {c.ref for c in ctx.candidates(role)}
+        fps = [fp for fp in ctx.board.footprints if fp.ref in refs]
+        if not fps:
+            out.append(Result("edge", label, 0.0, "connector not on the board"))
+            continue
+        dist = {fp.ref: edge_distance(ctx.board, fp) for fp in fps}
+        on_edge = [r for r, d in dist.items() if d is not None and d <= max_inset_mm]
+        out.append(
+            Result(
+                "edge",
+                label,
+                len(on_edge) / len(fps),
+                ", ".join(f"{r}: {d} mm" for r, d in dist.items()),
+            )
+        )
+    return out
+
+
+def refdes_check(board: Board | None) -> Result:
+    if board is None:
+        return Result("refdes", "reference designators visible", 0.0, "no board")
+    fps = [fp for fp in board.footprints if fp.ref and not fp.ref.startswith("#") and fp.pads]
+    hidden = [fp.ref for fp in fps if not fp.ref_visible]
+    score = 1 - len(hidden) / len(fps) if fps else 0.0
+    detail = f"{len(fps) - len(hidden)}/{len(fps)} visible on silkscreen"
+    if hidden:
+        detail += "; hidden: " + ", ".join(sorted(hidden)[:10])
+    return Result("refdes", "reference designators visible", score, detail)

traceforge/bench/cli.py

@@ -0,0 +1,121 @@
+"""``traceforge bench``: list, prepare, run and score benchmark tasks."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from traceforge.bench.tasks import default_root, load_tasks
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="traceforge bench", description="Open benchmark for PCB design agents"
+    )
+    p.add_argument("--root", type=Path, help="benchmark folder (default: benchmarks/)")
+    sub = p.add_subparsers(dest="action", required=True)
+
+    ls = sub.add_parser("list", help="list tasks")
+    ls.add_argument("select", nargs="*", help="task ids, id prefixes, tiers or kinds")
+
+    prep = sub.add_parser(
+        "prepare", help="(re)generate planted review/edit projects from their specs"
+    )
+    prep.add_argument("select", nargs="*")
+
+    run = sub.add_parser("run", help="run tasks against an agent and score them")
+    run.add_argument("select", nargs="*")
+    run.add_argument("--agent", default="traceforge", help="traceforge | peer | command")
+    run.add_argument(
+        "--cmd", help="for --agent command: CLI with {workspace} and {prompt} placeholders"
+    )
+    run.add_argument("--model", help="model for the traceforge adapter")
+    run.add_argument("--exe", help="executable of the peer agent CLI")
+    run.add_argument(
+        "--agent-arg", action="append", default=[], help="extra argument for the agent CLI"
+    )
+    run.add_argument(
+        "--out", type=Path, help="results folder (default: bench-results/<agent>-<time>)"
+    )
+    run.add_argument("--timeout", type=float, default=3600, help="seconds per task")
+    run.add_argument(
+        "--prices", type=Path, help="price table YAML (default: benchmarks/prices.yaml)"
+    )
+    run.add_argument("--offline", action="store_true", help="skip LCSC verification")
+
+    sc = sub.add_parser("score", help="score an existing project folder for one task")
+    sc.add_argument("task")
+    sc.add_argument("workspace", type=Path)
+    sc.add_argument("--answer", type=Path, help="file with the agent's final answer (review tasks)")
+    sc.add_argument("--out", type=Path, help="write the result JSON here")
+    sc.add_argument("--offline", action="store_true")
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    root = args.root or default_root()
+    tasks = load_tasks(root, getattr(args, "select", None) or None)
+    if args.action == "list":
+        for t in tasks:
+            print(f"{t.id:<34} {t.tier:<7} {t.kind:<7} {t.title}")
+        print(f"{len(tasks)} tasks")
+        return 0
+    if args.action == "prepare":
+        from traceforge.bench.review import build_planted
+
+        for t in tasks:
+            if t.planted and t.setup_project:
+                result = build_planted(root, t.planted, t.setup_project)
+                print(f"{t.id}: {'ok' if result['ok'] else 'issues'} -> {t.setup_project}")
+        return 0
+    if args.action == "score":
+        from traceforge.bench.score import score_project
+
+        match = [t for t in tasks if t.id == args.task] or [
+            t for t in tasks if t.id.startswith(args.task)
+        ]
+        if len(match) != 1:
+            print(f"task {args.task!r} matches {len(match)} tasks", file=sys.stderr)
+            return 2
+        answer = args.answer.read_text(encoding="utf-8") if args.answer else ""
+        if args.out:
+            work = args.out.parent / f"{args.out.stem}.files"
+        else:
+            import tempfile
+
+            work = Path(tempfile.mkdtemp(prefix=f"bench-{match[0].id}-"))
+        result = score_project(
+            match[0], args.workspace.resolve(), answer=answer, work_dir=work, offline=args.offline
+        )
+        text = json.dumps(result, indent=1, default=str)
+        if args.out:
+            args.out.write_text(text, encoding="utf-8")
+        print(text if not args.out else f"{match[0].id}: {result['score']} -> {args.out}")
+        return 0
+    from datetime import datetime
+
+    from traceforge.bench.adapters import make_adapter
+    from traceforge.bench.runner import run_suite
+
+    adapter = make_adapter(
+        args.agent,
+        command=args.cmd,
+        model=args.model,
+        executable=args.exe,
+        extra_args=args.agent_arg,
+    )
+    out = args.out or Path("bench-results") / f"{adapter.name}-{datetime.now():%Y%m%d-%H%M%S}"
+    suite = run_suite(
+        tasks,
+        adapter,
+        out.resolve(),
+        timeout=args.timeout,
+        prices_path=args.prices or root / "prices.yaml",
+        offline=args.offline,
+        progress=lambda msg: print(msg, file=sys.stderr),
+    )
+    print((out / "results.md").read_text(encoding="utf-8"))
+    return 0 if suite["results"] else 1