phase 0
This commit is contained in:
@@ -0,0 +1,335 @@
|
||||
"""
|
||||
query_engine.py — Natural language geometric query handler.
|
||||
|
||||
Supports both single-query (--query "...") and interactive REPL (--repl).
|
||||
REPL keeps the model in memory between queries for speed.
|
||||
All output is formatted ASCII tables.
|
||||
|
||||
Supported query types (see SKILL.md for full reference):
|
||||
bounding box overall model extents
|
||||
face count total faces by type
|
||||
all parts full assembly listing
|
||||
list all holes all cylindrical through-features
|
||||
list all mounting holes cylinders dia < 15mm, axis ⊥ primary face
|
||||
holes diameter N filter by diameter
|
||||
wall thickness min distance between opposing parallel faces
|
||||
largest face largest planar face area
|
||||
help list supported queries
|
||||
exit / quit exit REPL
|
||||
"""
|
||||
import logging
|
||||
import re
|
||||
import textwrap
|
||||
from typing import Optional
|
||||
|
||||
from .loader import StepModel
|
||||
|
||||
logger = logging.getLogger("step_processor.query")
|
||||
|
||||
# Regex patterns for query routing
|
||||
_BBOX_RE = re.compile(r"\b(bounding.?box|extents|dimensions|overall.?size)\b", re.I)
|
||||
_FACECOUNT_RE = re.compile(r"\b(face.?count|how many faces|number of faces)\b", re.I)
|
||||
_PARTS_RE = re.compile(r"\ball.?parts|part.?list|assembly|components\b", re.I)
|
||||
_HOLES_RE = re.compile(r"\bholes?\b", re.I)
|
||||
_MOUNTING_RE = re.compile(r"\bmounting\b", re.I)
|
||||
_DIA_RE = re.compile(r"diameter\s+([\d.]+)\s*mm?", re.I)
|
||||
_WALL_RE = re.compile(r"\bwall.?thickness\b", re.I)
|
||||
_LARGEST_RE = re.compile(r"\blargest.?face\b", re.I)
|
||||
_HELP_RE = re.compile(r"\bhelp\b", re.I)
|
||||
_EXIT_RE = re.compile(r"\b(exit|quit|q)\b", re.I)
|
||||
|
||||
|
||||
def run_query(model: StepModel, query: str) -> str:
|
||||
"""Dispatch a query string and return formatted output."""
|
||||
q = query.strip()
|
||||
if _EXIT_RE.search(q):
|
||||
return "EXIT"
|
||||
if _HELP_RE.search(q):
|
||||
return _help_text()
|
||||
if _BBOX_RE.search(q):
|
||||
return _query_bounding_box(model)
|
||||
if _FACECOUNT_RE.search(q):
|
||||
return _query_face_count(model)
|
||||
if _PARTS_RE.search(q):
|
||||
return _query_all_parts(model)
|
||||
if _HOLES_RE.search(q):
|
||||
dia_match = _DIA_RE.search(q)
|
||||
dia_filter = float(dia_match.group(1)) if dia_match else None
|
||||
mounting_only = bool(_MOUNTING_RE.search(q))
|
||||
return _query_holes(model, mounting_only=mounting_only, dia_filter=dia_filter)
|
||||
if _WALL_RE.search(q):
|
||||
return _query_wall_thickness(model)
|
||||
if _LARGEST_RE.search(q):
|
||||
return _query_largest_face(model)
|
||||
return (f"Query not recognized: '{q}'\n"
|
||||
f"Type 'help' to see supported queries.")
|
||||
|
||||
|
||||
def repl(model: StepModel, step_path):
|
||||
"""Launch interactive REPL. Returns when user types exit/quit."""
|
||||
print(f"\nSTEP Query REPL — {step_path.name}")
|
||||
print("Type 'help' for supported queries, 'exit' to quit.\n")
|
||||
while True:
|
||||
try:
|
||||
q = input("> ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print()
|
||||
break
|
||||
if not q:
|
||||
continue
|
||||
result = run_query(model, q)
|
||||
if result == "EXIT":
|
||||
break
|
||||
print(result)
|
||||
print()
|
||||
|
||||
|
||||
# ── Query implementations ─────────────────────────────────────────────────────
|
||||
|
||||
def _query_bounding_box(model: StepModel) -> str:
|
||||
try:
|
||||
if model.backend == "build123d":
|
||||
bb = model.shape.bounding_box()
|
||||
x = round(bb.size.X, 2)
|
||||
y = round(bb.size.Y, 2)
|
||||
z = round(bb.size.Z, 2)
|
||||
else:
|
||||
bb = model.shape.bounding_box
|
||||
x = round(bb.XMax - bb.XMin, 2)
|
||||
y = round(bb.YMax - bb.YMin, 2)
|
||||
z = round(bb.ZMax - bb.ZMin, 2)
|
||||
return _table(
|
||||
f"BOUNDING BOX — {model.path.name}",
|
||||
["Axis", "Dimension"],
|
||||
[["Width (X)", f"{x} mm ({x/25.4:.3f} in)"],
|
||||
["Depth (Y)", f"{y} mm ({y/25.4:.3f} in)"],
|
||||
["Height (Z)", f"{z} mm ({z/25.4:.3f} in)"]]
|
||||
)
|
||||
except Exception as e:
|
||||
return f"Bounding box query failed: {e}"
|
||||
|
||||
|
||||
def _query_face_count(model: StepModel) -> str:
|
||||
if model.backend != "build123d":
|
||||
return f"Face count query requires build123d (loaded via {model.backend})"
|
||||
try:
|
||||
from build123d import Compound
|
||||
all_faces = model.shape.faces()
|
||||
planar = sum(1 for f in all_faces if f.geom_type() == "PLANE")
|
||||
cylindrical = sum(1 for f in all_faces if f.geom_type() == "CYLINDER")
|
||||
other = len(all_faces) - planar - cylindrical
|
||||
return _table(
|
||||
f"FACE COUNT — {model.path.name}",
|
||||
["Type", "Count"],
|
||||
[["Planar", str(planar)],
|
||||
["Cylindrical", str(cylindrical)],
|
||||
["Other", str(other)],
|
||||
["Total", str(len(all_faces))]]
|
||||
)
|
||||
except Exception as e:
|
||||
return f"Face count failed: {e}"
|
||||
|
||||
|
||||
def _query_all_parts(model: StepModel) -> str:
|
||||
if not model.parts:
|
||||
return (f"No assembly structure found in {model.path.name}.\n"
|
||||
f"File appears to be a single solid body.")
|
||||
rows = []
|
||||
for p in model.parts:
|
||||
rows.append([
|
||||
p.get("part_number", "—"),
|
||||
p.get("name", "—"),
|
||||
str(p.get("quantity", 1)),
|
||||
str(p.get("level", 0)),
|
||||
p.get("parent", ""),
|
||||
])
|
||||
return _table(
|
||||
f"ALL PARTS — {model.path.name}",
|
||||
["#", "Name", "Qty", "Level", "Parent"],
|
||||
rows
|
||||
) + f"\nTotal: {len(model.parts)} parts"
|
||||
|
||||
|
||||
def _query_holes(model: StepModel, mounting_only: bool = False,
|
||||
dia_filter: Optional[float] = None) -> str:
|
||||
if model.backend != "build123d":
|
||||
return f"Hole query requires build123d (loaded via {model.backend})"
|
||||
try:
|
||||
holes = _find_holes(model, mounting_only=mounting_only, dia_filter=dia_filter)
|
||||
if not holes:
|
||||
label = "mounting holes" if mounting_only else "holes"
|
||||
qualifier = f" ≈{dia_filter}mm" if dia_filter else ""
|
||||
return f"No {label}{qualifier} found in {model.path.name}."
|
||||
header = "MOUNTING HOLES" if mounting_only else "ALL HOLES"
|
||||
# Group by diameter bucket for summary view
|
||||
from collections import Counter
|
||||
dia_counts = Counter(round(h["dia"], 1) for h in holes)
|
||||
MAX_ROWS = 50
|
||||
display_holes = holes[:MAX_ROWS]
|
||||
rows = []
|
||||
for i, h in enumerate(display_holes, 1):
|
||||
rows.append([
|
||||
str(i),
|
||||
f"{h['dia']:.2f} mm",
|
||||
f"{h['depth']:.2f} mm" if h["depth"] else "—",
|
||||
f"({h['x']:.1f}, {h['y']:.1f}, {h['z']:.1f})",
|
||||
])
|
||||
result = _table(
|
||||
f"{header} — {model.path.name}",
|
||||
["#", "Diameter", "Depth", "Position (x,y,z)"],
|
||||
rows
|
||||
)
|
||||
result += f"\nShowing {len(display_holes)} of {len(holes)} unique hole locations"
|
||||
result += "\n\nDIAMETER SUMMARY"
|
||||
result += "\n" + "─" * 30
|
||||
for dia, count in sorted(dia_counts.items()):
|
||||
result += f"\n {dia:.1f} mm ×{count}"
|
||||
result += "\n" + "─" * 30
|
||||
return result
|
||||
except Exception as e:
|
||||
return f"Hole query failed: {e}"
|
||||
|
||||
|
||||
def _find_holes(model: StepModel, mounting_only: bool, dia_filter):
|
||||
"""Extract and deduplicate cylindrical faces from build123d model.
|
||||
|
||||
Deduplication: round axis position to 1mm grid, group by (dia_bucket, x, y, z).
|
||||
This collapses multiple cylindrical faces from the same physical hole
|
||||
(e.g. inner + outer surface of same cylinder) into one entry.
|
||||
"""
|
||||
from OCP.BRepAdaptor import BRepAdaptor_Surface
|
||||
from OCP.GeomAbs import GeomAbs_Cylinder
|
||||
|
||||
seen = {} # key → best entry
|
||||
try:
|
||||
faces = model.shape.faces()
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
for face in faces:
|
||||
try:
|
||||
adaptor = BRepAdaptor_Surface(face.wrapped)
|
||||
if adaptor.GetType() != GeomAbs_Cylinder:
|
||||
continue
|
||||
cyl = adaptor.Cylinder()
|
||||
r = cyl.Radius()
|
||||
dia = round(r * 2, 2)
|
||||
# Diameter filters
|
||||
if mounting_only and dia > 15.0:
|
||||
continue
|
||||
if dia_filter and abs(dia - dia_filter) > 0.5:
|
||||
continue
|
||||
axis_pt = cyl.Location()
|
||||
# Round to 1mm grid for deduplication
|
||||
gx = round(axis_pt.X())
|
||||
gy = round(axis_pt.Y())
|
||||
gz = round(axis_pt.Z())
|
||||
dia_bucket = round(dia, 1)
|
||||
key = (dia_bucket, gx, gy, gz)
|
||||
if key not in seen:
|
||||
bb = face.bounding_box()
|
||||
depth = round(max(bb.size.X, bb.size.Y, bb.size.Z), 2)
|
||||
seen[key] = {
|
||||
"dia": dia,
|
||||
"depth": depth,
|
||||
"x": round(axis_pt.X(), 1),
|
||||
"y": round(axis_pt.Y(), 1),
|
||||
"z": round(axis_pt.Z(), 1),
|
||||
}
|
||||
except Exception:
|
||||
continue
|
||||
return list(seen.values())
|
||||
|
||||
|
||||
def _query_wall_thickness(model: StepModel) -> str:
|
||||
if model.backend != "build123d":
|
||||
return f"Wall thickness query requires build123d (loaded via {model.backend})"
|
||||
try:
|
||||
faces = model.shape.faces()
|
||||
planar = [f for f in faces if f.geom_type() == "PLANE"]
|
||||
if len(planar) < 2:
|
||||
return "Insufficient planar faces to determine wall thickness."
|
||||
# Heuristic: find minimum non-zero distance between parallel opposing faces
|
||||
min_t = None
|
||||
for i, f1 in enumerate(planar):
|
||||
n1 = f1.normal_at()
|
||||
for f2 in planar[i+1:]:
|
||||
n2 = f2.normal_at()
|
||||
# Parallel if normals are anti-parallel
|
||||
dot = abs(n1.dot(n2))
|
||||
if dot > 0.99:
|
||||
c1 = f1.center()
|
||||
c2 = f2.center()
|
||||
dist = round(abs((c1 - c2).dot(n1)), 3)
|
||||
if dist > 0.01:
|
||||
if min_t is None or dist < min_t:
|
||||
min_t = dist
|
||||
if min_t is None:
|
||||
return "Could not determine wall thickness from available faces."
|
||||
return _table(
|
||||
f"WALL THICKNESS — {model.path.name}",
|
||||
["Measurement", "Value"],
|
||||
[["Minimum wall thickness",
|
||||
f"{min_t} mm ({min_t/25.4:.3f} in)"]]
|
||||
)
|
||||
except Exception as e:
|
||||
return f"Wall thickness query failed: {e}"
|
||||
|
||||
|
||||
def _query_largest_face(model: StepModel) -> str:
|
||||
if model.backend != "build123d":
|
||||
return f"Largest face query requires build123d (loaded via {model.backend})"
|
||||
try:
|
||||
faces = model.shape.faces()
|
||||
planar = [(f, f.area()) for f in faces if f.geom_type() == "PLANE"]
|
||||
if not planar:
|
||||
return "No planar faces found."
|
||||
largest, area = max(planar, key=lambda x: x[1])
|
||||
bb = largest.bounding_box()
|
||||
return _table(
|
||||
f"LARGEST PLANAR FACE — {model.path.name}",
|
||||
["Property", "Value"],
|
||||
[["Area", f"{round(area, 2)} mm²"],
|
||||
["Width", f"{round(bb.size.X, 2)} mm"],
|
||||
["Height", f"{round(bb.size.Z, 2)} mm"]]
|
||||
)
|
||||
except Exception as e:
|
||||
return f"Largest face query failed: {e}"
|
||||
|
||||
|
||||
# ── Formatting helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
def _table(title: str, headers: list, rows: list) -> str:
|
||||
col_widths = [len(h) for h in headers]
|
||||
for row in rows:
|
||||
for i, cell in enumerate(row):
|
||||
col_widths[i] = max(col_widths[i], len(str(cell)))
|
||||
sep = "─" * (sum(col_widths) + 3 * len(headers) - 1)
|
||||
lines = [title, sep]
|
||||
header_line = " ".join(h.ljust(col_widths[i]) for i, h in enumerate(headers))
|
||||
lines.append(header_line)
|
||||
lines.append(sep)
|
||||
for row in rows:
|
||||
lines.append(" ".join(str(cell).ljust(col_widths[i]) for i, cell in enumerate(row)))
|
||||
lines.append(sep)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _help_text() -> str:
|
||||
return textwrap.dedent("""\
|
||||
SUPPORTED QUERIES
|
||||
─────────────────────────────────────────────────────────────
|
||||
bounding box Overall model extents (W×D×H in mm)
|
||||
face count Faces by type (planar, cylindrical, other)
|
||||
all parts Full assembly listing with quantities
|
||||
list all holes All cylindrical through-features
|
||||
list all mounting holes Holes smaller than 15mm diameter
|
||||
holes diameter 4.2mm Filter holes by specific diameter
|
||||
wall thickness Minimum wall thickness estimate
|
||||
largest face Largest planar face area
|
||||
help Show this message
|
||||
exit Exit the REPL
|
||||
─────────────────────────────────────────────────────────────
|
||||
Tip: geometry queries require build123d backend.
|
||||
If the file loaded via FreeCAD fallback, only bounding box
|
||||
and parts list are available.""")
|
||||
Reference in New Issue
Block a user