610 lines
29 KiB
Python
610 lines
29 KiB
Python
from __future__ import annotations
|
||
|
||
import ast
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import pathlib
|
||
import re
|
||
import resource
|
||
import sqlite3
|
||
import subprocess
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
import uuid
|
||
from datetime import datetime, timezone
|
||
from typing import Callable
|
||
|
||
|
||
ProgressCallback = Callable[[str, str, int, int], None]
|
||
|
||
|
||
class AI3DModelError(RuntimeError):
|
||
pass
|
||
|
||
|
||
class AI3DModelEngine:
|
||
"""STEP-first text/image-to-CAD runner with an untrusted-code sandbox."""
|
||
|
||
TERMINAL_STAGES = {"completed", "failed"}
|
||
ALLOWED_IMPORT_ROOTS = {"build123d", "math"}
|
||
BANNED_CALLS = {
|
||
"compile", "eval", "exec", "getattr", "globals", "input", "locals",
|
||
"open", "setattr", "vars", "__import__",
|
||
}
|
||
|
||
def __init__(self) -> None:
|
||
self.state_dir = pathlib.Path(
|
||
os.environ.get("DALONGXIA_AI_3D_STATE_DIR", "/var/lib/dalongxia-chat/ai-3d-model")
|
||
)
|
||
self.runtime_dir = pathlib.Path(
|
||
os.environ.get("DALONGXIA_AI_3D_RUNTIME_DIR", "/opt/dalongxia-cad-runtime")
|
||
)
|
||
self.skill_dir = pathlib.Path(
|
||
os.environ.get("DALONGXIA_AI_3D_SKILL_DIR", "/opt/dalongxia-chat-api/ai-3d-model/cad")
|
||
)
|
||
self.model = os.environ.get("DALONGXIA_AI_3D_MODEL", "qwen3-coder-plus")
|
||
self.vision_model = os.environ.get("DALONGXIA_AI_3D_VISION_MODEL", "qwen-vl-plus")
|
||
self.api_base = os.environ.get(
|
||
"DALONGXIA_AI_3D_API_BASE",
|
||
"https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||
).rstrip("/")
|
||
self.timeout = max(60, int(os.environ.get("DALONGXIA_AI_3D_TIMEOUT", "300")))
|
||
self.state_dir.mkdir(parents=True, exist_ok=True)
|
||
(self.state_dir / "jobs").mkdir(parents=True, exist_ok=True)
|
||
self.db_path = self.state_dir / "jobs.sqlite3"
|
||
self._init_db()
|
||
|
||
def _connect(self) -> sqlite3.Connection:
|
||
conn = sqlite3.connect(self.db_path, timeout=20)
|
||
conn.row_factory = sqlite3.Row
|
||
return conn
|
||
|
||
def _init_db(self) -> None:
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS jobs (
|
||
job_id TEXT PRIMARY KEY,
|
||
context_key TEXT NOT NULL,
|
||
prompt TEXT NOT NULL,
|
||
stage TEXT NOT NULL,
|
||
display_text TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
result_json TEXT NOT NULL DEFAULT '{}',
|
||
error TEXT NOT NULL DEFAULT ''
|
||
)
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_ai3d_context_updated ON jobs(context_key, updated_at DESC)"
|
||
)
|
||
|
||
@staticmethod
|
||
def _now() -> str:
|
||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||
|
||
@staticmethod
|
||
def context_key(conversation_id: str, user_id: str, session_key: str) -> str:
|
||
raw = "|".join([str(user_id or ""), str(conversation_id or ""), str(session_key or "")])
|
||
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||
|
||
def _insert_job(self, context_key: str, prompt: str) -> dict:
|
||
job_id = f"cad_{datetime.now().strftime('%Y%m%d%H%M%S')}_{uuid.uuid4().hex[:8]}"
|
||
now = self._now()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"INSERT INTO jobs(job_id,context_key,prompt,stage,display_text,created_at,updated_at) VALUES(?,?,?,?,?,?,?)",
|
||
(job_id, context_key, prompt, "queued", "任务已进入CAD生成队列", now, now),
|
||
)
|
||
return self.get_job(job_id) or {}
|
||
|
||
def _update_job(self, job_id: str, stage: str, display_text: str, *, result=None, error="") -> dict:
|
||
now = self._now()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"UPDATE jobs SET stage=?,display_text=?,updated_at=?,result_json=?,error=? WHERE job_id=?",
|
||
(stage, display_text, now, json.dumps(result or {}, ensure_ascii=False), str(error or "")[:4000], job_id),
|
||
)
|
||
return self.get_job(job_id) or {}
|
||
|
||
def get_job(self, job_id: str) -> dict | None:
|
||
with self._connect() as conn:
|
||
row = conn.execute("SELECT * FROM jobs WHERE job_id=?", (job_id,)).fetchone()
|
||
if not row:
|
||
return None
|
||
value = dict(row)
|
||
try:
|
||
value["result"] = json.loads(value.pop("result_json") or "{}")
|
||
except json.JSONDecodeError:
|
||
value["result"] = {}
|
||
return value
|
||
|
||
def latest_job(self, context_key: str) -> dict | None:
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
"SELECT job_id FROM jobs WHERE context_key=? ORDER BY updated_at DESC LIMIT 1",
|
||
(context_key,),
|
||
).fetchone()
|
||
return self.get_job(str(row["job_id"])) if row else None
|
||
|
||
@staticmethod
|
||
def _is_entry_only(message: str) -> bool:
|
||
compact = re.sub(r"[\s，。！？!?、:：]+", "", str(message or "").lower())
|
||
compact = re.sub(r"【技能路由[:：]ai-3d-model】", "", compact)
|
||
return compact in {"", "ai画3d模型", "ai画3d", "生成3d模型", "文字生成cad"}
|
||
|
||
@staticmethod
|
||
def _is_progress_request(message: str) -> bool:
|
||
text = str(message or "").strip()
|
||
return bool(re.fullmatch(r"(?:继续|进度|查看进度|现在怎么样|完成了吗|结果呢)[？?。\s]*", text))
|
||
|
||
def _api_key(self) -> str:
|
||
direct = os.environ.get("DASHSCOPE_API_KEY", "").strip()
|
||
if direct:
|
||
return direct
|
||
key_path = os.environ.get("DASHSCOPE_API_KEY_FILE", "").strip()
|
||
if key_path:
|
||
try:
|
||
return pathlib.Path(key_path).read_text(encoding="utf-8").strip()
|
||
except OSError as exc:
|
||
raise AI3DModelError("CAD模型服务密钥文件不可读") from exc
|
||
raise AI3DModelError("CAD模型服务尚未配置密钥")
|
||
|
||
@staticmethod
|
||
def _has_image_intent(prompt: str) -> bool:
|
||
text = re.sub(r"\s+", "", str(prompt or "")).lower()
|
||
return any(
|
||
marker in text
|
||
for marker in (
|
||
"根据图片", "根据上传", "按图", "参考图", "示意图", "工程图",
|
||
"图中标注", "照着图", "图片创建", "图片生成", "image-to-cad",
|
||
)
|
||
)
|
||
|
||
@staticmethod
|
||
def _vision_content(image_inputs: list[dict], prompt: str) -> list[dict]:
|
||
content: list[dict] = []
|
||
for item in image_inputs[:4]:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
kind = str(item.get("kind") or "")
|
||
value = str(item.get("value") or "").strip()
|
||
if kind == "image_url" and re.match(r"^https?://", value, re.I):
|
||
content.append({"type": "image_url", "image_url": {"url": value}})
|
||
elif kind == "image_base64" and value.startswith("data:image/"):
|
||
content.append({"type": "image_url", "image_url": {"url": value}})
|
||
content.append(
|
||
{
|
||
"type": "text",
|
||
"text": (
|
||
"请作为机械CAD读图工程师，只提取图像中有证据的几何与尺寸。"
|
||
"识别外形、孔/槽/凸台、对称关系、截面、标注数值与单位。"
|
||
"不要凭空添加图中没有的结构；看不清的内容必须放入uncertainties。"
|
||
"返回一个JSON对象，字段为summary、geometry、dimensions、uncertainties、modeling_instruction。\n"
|
||
f"用户要求：{prompt.strip()}"
|
||
),
|
||
}
|
||
)
|
||
return content
|
||
|
||
def _visual_spec(self, prompt: str, image_inputs: list[dict]) -> dict:
|
||
content = self._vision_content(image_inputs, prompt)
|
||
if not any(item.get("type") == "image_url" for item in content):
|
||
raise AI3DModelError("图片内容不完整，请上传可读取的PNG/JPG/WebP图片")
|
||
body = {
|
||
"model": self.vision_model,
|
||
"messages": [{"role": "user", "content": content}],
|
||
"temperature": 0.0,
|
||
}
|
||
request = urllib.request.Request(
|
||
f"{self.api_base}/chat/completions",
|
||
data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
|
||
headers={"Authorization": f"Bearer {self._api_key()}", "Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
||
payload = json.loads(response.read().decode("utf-8"))
|
||
content_value = payload["choices"][0]["message"]["content"]
|
||
if isinstance(content_value, list):
|
||
content_value = "".join(
|
||
str(item.get("text") or "") for item in content_value if isinstance(item, dict)
|
||
)
|
||
raw = str(content_value or "").strip()
|
||
raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw, flags=re.S)
|
||
match = re.search(r"\{.*\}", raw, flags=re.S)
|
||
spec = json.loads(match.group(0) if match else raw)
|
||
except (urllib.error.URLError, TimeoutError, KeyError, IndexError, TypeError, json.JSONDecodeError) as exc:
|
||
raise AI3DModelError(f"CAD图片理解失败：{exc}") from exc
|
||
if not isinstance(spec, dict) or not str(spec.get("modeling_instruction") or "").strip():
|
||
raise AI3DModelError("CAD图片理解未返回可建模的结构和尺寸")
|
||
return spec
|
||
|
||
def _model_source(
|
||
self,
|
||
prompt: str,
|
||
repair_error: str = "",
|
||
previous_source: str = "",
|
||
visual_spec: dict | None = None,
|
||
) -> tuple[str, dict]:
|
||
system = """你是参数化机械CAD工程师。把中文需求转换成安全的 build123d Python。
|
||
只返回一个JSON对象：{"title":"简短中文名","assumptions":["..."],"source":"..."}。
|
||
source必须：1) 只允许 import math 或 from build123d import ...；2) 定义无参数函数 gen_step()；
|
||
3) 返回一个闭合的正体积 Part、Solid、Compound 或带标签的装配；4) 单位毫米；5) 使用明确命名参数；
|
||
6) 不读写文件、不访问网络、环境变量或系统，不使用动态执行；7) 不调用 show_object；8) 简单需求也要完整建模。
|
||
优先使用 Box、Cylinder、Sphere、Cone、BuildPart、BuildSketch、Locations、extrude、fillet、chamfer、Hole。
|
||
不要返回Markdown代码围栏，不要臆造工程认证或承载能力。"""
|
||
user = f"用户CAD需求：\n{prompt.strip()}"
|
||
if visual_spec:
|
||
user += (
|
||
"\n\n以下是视觉模型从用户图片提取的读图结果。"
|
||
"建模必须以其中可见标注为准，不得用常见尺寸替换：\n"
|
||
+ json.dumps(visual_spec, ensure_ascii=False)
|
||
)
|
||
if repair_error:
|
||
user += f"\n\n上一次源码构建失败，请仅修复导致失败的部分。错误：\n{repair_error[-3000:]}"
|
||
user += f"\n\n上一次源码：\n{previous_source[-10000:]}"
|
||
body = {
|
||
"model": self.model,
|
||
"messages": [{"role": "system", "content": system}, {"role": "user", "content": user}],
|
||
"temperature": 0.1,
|
||
"response_format": {"type": "json_object"},
|
||
}
|
||
request = urllib.request.Request(
|
||
f"{self.api_base}/chat/completions",
|
||
data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
|
||
headers={"Authorization": f"Bearer {self._api_key()}", "Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
||
payload = json.loads(response.read().decode("utf-8"))
|
||
except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
|
||
raise AI3DModelError(f"CAD模型服务调用失败：{exc}") from exc
|
||
try:
|
||
content = payload["choices"][0]["message"]["content"]
|
||
if isinstance(content, list):
|
||
content = "".join(str(x.get("text") or "") for x in content if isinstance(x, dict))
|
||
text = str(content).strip()
|
||
if text.startswith("```"):
|
||
text = re.sub(r"^```(?:json)?\s*|\s*```$", "", text, flags=re.S)
|
||
result = json.loads(text)
|
||
source = str(result.get("source") or "").strip()
|
||
except (KeyError, IndexError, TypeError, json.JSONDecodeError) as exc:
|
||
raise AI3DModelError("CAD模型服务未返回有效源码") from exc
|
||
self._validate_source(source)
|
||
return source, result
|
||
|
||
def _validate_source(self, source: str) -> None:
|
||
if not source or len(source) > 50_000:
|
||
raise AI3DModelError("生成的CAD源码为空或过长")
|
||
try:
|
||
tree = ast.parse(source)
|
||
except SyntaxError as exc:
|
||
raise AI3DModelError(f"生成源码语法错误：{exc.msg}") from exc
|
||
has_gen_step = False
|
||
for node in ast.walk(tree):
|
||
if isinstance(node, (ast.Import, ast.ImportFrom)):
|
||
names = [alias.name.split(".")[0] for alias in node.names] if isinstance(node, ast.Import) else [(node.module or "").split(".")[0]]
|
||
if any(name not in self.ALLOWED_IMPORT_ROOTS for name in names):
|
||
raise AI3DModelError("生成源码包含不允许的模块")
|
||
elif isinstance(node, ast.FunctionDef) and node.name == "gen_step" and not node.args.args:
|
||
has_gen_step = True
|
||
elif isinstance(node, ast.Call):
|
||
name = node.func.id if isinstance(node.func, ast.Name) else ""
|
||
if name in self.BANNED_CALLS:
|
||
raise AI3DModelError("生成源码包含不安全调用")
|
||
elif isinstance(node, ast.Attribute) and str(node.attr).startswith("__"):
|
||
raise AI3DModelError("生成源码包含不安全属性访问")
|
||
if not has_gen_step:
|
||
raise AI3DModelError("生成源码缺少无参数 gen_step()")
|
||
|
||
@staticmethod
|
||
def _limits() -> None:
|
||
resource.setrlimit(resource.RLIMIT_CPU, (180, 180))
|
||
resource.setrlimit(resource.RLIMIT_AS, (5 * 1024**3, 5 * 1024**3))
|
||
resource.setrlimit(resource.RLIMIT_FSIZE, (512 * 1024**2, 512 * 1024**2))
|
||
resource.setrlimit(resource.RLIMIT_NOFILE, (256, 256))
|
||
# RLIMIT_NPROC is per Linux uid, and the shared Chat API uid already owns
|
||
# browser workers. Keep a finite ceiling without blocking namespace setup.
|
||
resource.setrlimit(resource.RLIMIT_NPROC, (2048, 2048))
|
||
|
||
def _sandbox_command(self, job_dir: pathlib.Path, argv: list[str]) -> list[str]:
|
||
command = [
|
||
"bwrap", "--unshare-net", "--unshare-pid", "--unshare-ipc", "--unshare-uts",
|
||
"--die-with-parent", "--new-session",
|
||
"--tmpfs", "/", "--ro-bind", "/proc", "/proc", "--dev", "/dev", "--tmpfs", "/dev/shm",
|
||
"--dir", "/opt", "--dir", "/opt/dalongxia-cad-runtime",
|
||
"--dir", "/opt/dalongxia-chat-api", "--dir", "/opt/dalongxia-chat-api/ai-3d-model",
|
||
"--dir", "/work", "--tmpfs", "/tmp",
|
||
"--ro-bind", "/usr", "/usr", "--ro-bind", "/lib", "/lib",
|
||
"--ro-bind", "/etc", "/etc", "--ro-bind", "/sys", "/sys",
|
||
]
|
||
if pathlib.Path("/lib64").exists():
|
||
command += ["--ro-bind", "/lib64", "/lib64"]
|
||
command += [
|
||
"--ro-bind", str(self.runtime_dir), str(self.runtime_dir),
|
||
"--ro-bind", str(self.skill_dir), str(self.skill_dir),
|
||
"--bind", str(job_dir), "/work", "--chdir", "/work",
|
||
"--setenv", "HOME", "/tmp", "--setenv", "TMPDIR", "/tmp",
|
||
"--setenv", "PATH", f"{self.runtime_dir}/venv/bin:/usr/bin:/bin",
|
||
"--clearenv",
|
||
]
|
||
# Environment is rebuilt after clearenv by bwrap in command order.
|
||
command += [
|
||
"--setenv", "HOME", "/tmp", "--setenv", "TMPDIR", "/tmp",
|
||
"--setenv", "PATH", f"{self.runtime_dir}/venv/bin:/usr/bin:/bin",
|
||
"--setenv", "PLAYWRIGHT_BROWSERS_PATH", f"{self.runtime_dir}/ms-playwright",
|
||
]
|
||
return command + argv
|
||
|
||
def _run(self, job_dir: pathlib.Path, argv: list[str], *, stdin: str = "", timeout: int = 240) -> subprocess.CompletedProcess:
|
||
proc = subprocess.run(
|
||
self._sandbox_command(job_dir, argv), input=stdin, text=True,
|
||
stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout,
|
||
preexec_fn=self._limits, env={"PATH": "/usr/bin:/bin"},
|
||
)
|
||
if proc.returncode != 0:
|
||
detail = (proc.stderr or proc.stdout or "CAD工具执行失败")[-5000:]
|
||
raise AI3DModelError(detail)
|
||
return proc
|
||
|
||
def _run_trusted_snapshot(
|
||
self,
|
||
job_dir: pathlib.Path,
|
||
argv: list[str],
|
||
*,
|
||
stdin: str,
|
||
timeout: int,
|
||
) -> subprocess.CompletedProcess:
|
||
"""Render an already-validated STEP outside bwrap.
|
||
|
||
Playwright cannot spawn Chromium from this host's unprivileged user
|
||
namespace. The renderer receives only the inert STEP artifact, never the
|
||
generated Python source, and gets a minimal environment.
|
||
"""
|
||
env = {
|
||
"HOME": "/tmp",
|
||
"TMPDIR": "/tmp",
|
||
"PATH": f"{self.runtime_dir}/venv/bin:/usr/bin:/bin",
|
||
"PLAYWRIGHT_BROWSERS_PATH": str(self.runtime_dir / "ms-playwright"),
|
||
}
|
||
proc = subprocess.run(
|
||
argv, input=stdin, text=True, cwd=job_dir,
|
||
stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout,
|
||
env=env,
|
||
)
|
||
if proc.returncode != 0:
|
||
detail = (proc.stderr or proc.stdout or "CAD预览执行失败")[-5000:]
|
||
raise AI3DModelError(detail)
|
||
return proc
|
||
|
||
@staticmethod
|
||
def _requested_exports(prompt: str) -> list[str]:
|
||
upper = str(prompt or "").upper()
|
||
return [fmt for fmt in ("STL", "3MF", "GLB") if fmt in upper]
|
||
|
||
def _build_and_validate(
|
||
self,
|
||
job_dir: pathlib.Path,
|
||
source_name: str,
|
||
requested_exports: list[str] | None = None,
|
||
progress: ProgressCallback | None = None,
|
||
) -> dict:
|
||
python = str(self.runtime_dir / "venv/bin/python")
|
||
gen = str(self.skill_dir / "scripts/gen")
|
||
inspect_tool = str(self.skill_dir / "scripts/inspect")
|
||
snapshot = str(self.skill_dir / "scripts/snapshot")
|
||
self._run(job_dir, [python, gen, source_name, "--write", "--json"], timeout=self.timeout)
|
||
step_name = source_name[:-3]
|
||
step_path = job_dir / step_name
|
||
if not step_path.is_file() or step_path.stat().st_size < 100:
|
||
raise AI3DModelError("STEP文件没有生成")
|
||
if progress:
|
||
progress("validating", "STEP几何已生成，正在复核闭合性、体积和尺寸边界", 4, 5)
|
||
validation = self._run(job_dir, [python, inspect_tool, "validate", step_name], timeout=self.timeout)
|
||
# Per-solid validity does not detect two individually valid parts
|
||
# occupying the same volume. Run the geometry kernel's pairwise
|
||
# boolean-intersection audit before producing a preview or claiming a
|
||
# usable assembly. The inspect command exits non-zero on a clash, so
|
||
# _run turns it into a repairable CAD failure instead of a false pass.
|
||
interference = self._run(
|
||
job_dir,
|
||
# The bundled inspect CLI exposes structured output through
|
||
# ``--format json`` (and already defaults to JSON on newer
|
||
# versions). ``--json`` is not accepted by the runtime currently
|
||
# installed on the server.
|
||
[python, inspect_tool, "interfere", step_name, "--format", "json"],
|
||
timeout=self.timeout,
|
||
)
|
||
refs = self._run(
|
||
job_dir,
|
||
[python, inspect_tool, "refs", step_name, "--facts", "--planes", "--positioning"],
|
||
timeout=self.timeout,
|
||
)
|
||
preview_base = str(job_dir / "preview.png")
|
||
snapshot_job = {
|
||
"input": str(step_path), "mode": "view",
|
||
"outputs": [{"path": preview_base, "camera": "iso"}],
|
||
"render": {"viewLabels": True, "padding": 0.12, "sizeProfile": "diagnostic"},
|
||
}
|
||
if progress:
|
||
progress("rendering", "几何校验通过，正在生成等轴预览图", 4, 5)
|
||
self._run_trusted_snapshot(
|
||
job_dir,
|
||
[python, snapshot, "--job", "-", "--json"],
|
||
stdin=json.dumps(snapshot_job),
|
||
timeout=self.timeout,
|
||
)
|
||
previews = sorted(job_dir.glob("preview*.png"), key=lambda p: p.stat().st_mtime, reverse=True)
|
||
if not previews:
|
||
raise AI3DModelError("模型预览图没有生成")
|
||
result = {
|
||
"step_path": str(step_path), "source_path": str(job_dir / source_name),
|
||
"preview_path": str(previews[0]), "validation": validation.stdout.strip()[-5000:],
|
||
"interference": interference.stdout.strip()[-12000:],
|
||
"facts": refs.stdout.strip()[-12000:],
|
||
}
|
||
formats = [str(item).upper() for item in (requested_exports or [])]
|
||
if formats:
|
||
export_tool = str(self.skill_dir / "scripts/export")
|
||
flags = [f"--{fmt.lower()}" for fmt in formats]
|
||
self._run(job_dir, [python, export_tool, step_name, *flags, "--json"], timeout=self.timeout)
|
||
for fmt in formats:
|
||
export_path = step_path.with_suffix(f".{fmt.lower()}")
|
||
if not export_path.is_file() or export_path.stat().st_size < 100:
|
||
raise AI3DModelError(f"{fmt}文件没有生成")
|
||
result[f"{fmt.lower()}_path"] = str(export_path)
|
||
return result
|
||
|
||
def run(
|
||
self,
|
||
message: str,
|
||
*,
|
||
conversation_id: str,
|
||
user_id: str,
|
||
session_key: str,
|
||
progress: ProgressCallback | None = None,
|
||
image_inputs: list[dict] | None = None,
|
||
) -> dict:
|
||
image_inputs = [item for item in (image_inputs or []) if isinstance(item, dict)][:4]
|
||
context_key = self.context_key(conversation_id, user_id, session_key)
|
||
if self._is_entry_only(message):
|
||
return {
|
||
"status": "waiting",
|
||
"answer": "### ai画3d模型\n\n请直接在对话框描述要画的模型和关键尺寸。默认单位为毫米，主文件为 STEP。\n\n例如：`画一个80×50×4毫米的L形安装支架，两边各有2个直径5毫米安装孔，圆角2毫米，同时导出STL。`",
|
||
"files": [],
|
||
}
|
||
|
||
if self._is_progress_request(message):
|
||
latest = self.latest_job(context_key)
|
||
if not latest:
|
||
return {"status": "waiting", "answer": "### ai画3d模型\n\n当前会话还没有CAD任务，请直接描述模型和尺寸。", "files": []}
|
||
result = latest.get("result") or {}
|
||
if latest["stage"] == "completed":
|
||
return {"status": "completed", "answer": self._completed_answer(latest), "files": result.get("files") or []}
|
||
if latest["stage"] == "failed":
|
||
return {"status": "failed", "answer": f"### ai画3d模型未完成\n\n{latest['display_text']}\n\n失败原因：{latest['error']}", "files": []}
|
||
return {"status": "waiting", "answer": f"### ai画3d模型进度\n\n{latest['display_text']}\n\n任务号：`{latest['job_id']}`", "files": []}
|
||
|
||
if self._has_image_intent(message) and not image_inputs:
|
||
return {
|
||
"status": "waiting",
|
||
"answer": (
|
||
"### ai画3d模型\n\n我收到了按图建模的要求，但当前请求没有可读取的图片内容。"
|
||
"请直接上传PNG、JPG或WebP图片，不会在没有读图证据时自行猜测尺寸。"
|
||
),
|
||
"files": [],
|
||
}
|
||
|
||
job = self._insert_job(context_key, message)
|
||
job_id = str(job["job_id"])
|
||
job_dir = self.state_dir / "jobs" / job_id
|
||
job_dir.mkdir(parents=True, exist_ok=False)
|
||
source_name = "model.step.py"
|
||
assumptions: list[str] = []
|
||
|
||
def emit(stage: str, text: str, current: int) -> None:
|
||
self._update_job(job_id, stage, text)
|
||
if progress:
|
||
progress(stage, text, current, 5)
|
||
|
||
def build_progress(stage: str, text: str, current: int, total: int) -> None:
|
||
self._update_job(job_id, stage, text)
|
||
if progress:
|
||
progress(stage, text, current, total)
|
||
|
||
try:
|
||
visual_spec: dict | None = None
|
||
if image_inputs:
|
||
emit("vision", "正在识别图片中的轮廓、孔位、尺寸标注和建模约束", 1)
|
||
visual_spec = self._visual_spec(message, image_inputs)
|
||
emit("specifying", "图片读图完成，正在把可见标注整理为参数化建模要求", 2)
|
||
else:
|
||
emit("specifying", "正在把文字需求整理为尺寸、特征和建模约束", 1)
|
||
source = ""
|
||
last_error = ""
|
||
meta: dict = {}
|
||
for attempt in range(1, 3):
|
||
emit("coding", "正在生成参数化CAD源码" if attempt == 1 else "首次构建未通过，正在自动修复源码", 2)
|
||
source, meta = self._model_source(
|
||
message,
|
||
last_error,
|
||
source,
|
||
visual_spec=visual_spec,
|
||
)
|
||
(job_dir / source_name).write_text(source + "\n", encoding="utf-8")
|
||
assumptions = [str(x) for x in (meta.get("assumptions") or []) if str(x).strip()][:10]
|
||
try:
|
||
emit("building", "正在无网络隔离沙箱中构建STEP几何", 3)
|
||
built = self._build_and_validate(
|
||
job_dir,
|
||
source_name,
|
||
requested_exports=self._requested_exports(message),
|
||
progress=build_progress,
|
||
)
|
||
break
|
||
except Exception as exc:
|
||
last_error = str(exc)
|
||
if attempt >= 2:
|
||
raise
|
||
result = {
|
||
**built, "assumptions": assumptions, "title": str(meta.get("title") or "3D模型"),
|
||
"files": [],
|
||
}
|
||
emit("delivering", "模型、源码和预览已生成，正在上传交付文件", 4)
|
||
self._update_job(job_id, "delivering", "模型已完成并通过几何校验，正在上传交付文件", result=result)
|
||
completed = self.get_job(job_id) or {}
|
||
return {
|
||
"status": "completed", "answer": self._completed_answer(completed),
|
||
"files": [], "local_files": built, "job_id": job_id,
|
||
}
|
||
except Exception as exc:
|
||
friendly = re.sub(r"\s+", " ", str(exc)).strip()[-1000:] or "CAD生成失败"
|
||
self._update_job(job_id, "failed", "模型生成未完成", error=friendly)
|
||
return {
|
||
"status": "failed",
|
||
"answer": f"### ai画3d模型未完成\n\n任务没有被错误标记为成功。系统已自动生成并修复两轮，但几何校验仍未通过。\n\n失败原因：{friendly}",
|
||
"files": [], "job_id": job_id,
|
||
}
|
||
|
||
def attach_files(self, job_id: str, files: list[dict]) -> dict | None:
|
||
job = self.get_job(job_id)
|
||
if not job:
|
||
return None
|
||
result = dict(job.get("result") or {})
|
||
result["files"] = [dict(item) for item in files if isinstance(item, dict)]
|
||
return self._update_job(
|
||
job_id,
|
||
"completed",
|
||
"模型已完成并通过几何校验，交付文件已上传",
|
||
result=result,
|
||
error=str(job.get("error") or ""),
|
||
)
|
||
|
||
def mark_delivery_failed(self, job_id: str, error: str) -> dict | None:
|
||
job = self.get_job(job_id)
|
||
if not job:
|
||
return None
|
||
return self._update_job(
|
||
job_id,
|
||
"failed",
|
||
"模型已生成，但交付文件上传失败",
|
||
result=dict(job.get("result") or {}),
|
||
error=str(error or "文件上传失败")[:1000],
|
||
)
|
||
|
||
@staticmethod
|
||
def _completed_answer(job: dict) -> str:
|
||
result = job.get("result") or {}
|
||
assumptions = result.get("assumptions") or []
|
||
assumptions_text = "\n".join(f"- {x}" for x in assumptions) or "- 未使用额外假设"
|
||
return (
|
||
f"### ai画3d模型完成\n\n模型已生成，并完成 STEP 几何校验与预览复核。\n\n"
|
||
f"- 任务号：`{job.get('job_id')}`\n- 模型：{result.get('title') or '3D模型'}\n"
|
||
f"- 主交付：STEP\n- 可编辑源码：build123d Python\n\n使用的建模假设：\n{assumptions_text}"
|
||
)
|