sunxianghui
93773f3887
Includes FastAPI backend, vendored step2urdf frontend, and A7 handtuned arm JSON for URDF generation.
120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import shutil
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from app.config import settings
|
|
from app.models import ParseResult, RobotDraft
|
|
from app.services.a7_gold import enforce_a7_gold_kinematics
|
|
from app.services.profiles import apply_profile
|
|
from app.services.step_text import parse_step_products
|
|
from app.services.urdf_builder import write_urdf_package
|
|
from app.services.zhipu import propose_joints_with_zhipu
|
|
from app.services.mesh_export import export_meshes_stub, geometry_backend_status
|
|
|
|
|
|
def _job_dir(job_id: str) -> Path:
|
|
return settings.jobs_dir / job_id
|
|
|
|
|
|
def create_job_from_upload(src: Path, filename: str) -> ParseResult:
|
|
job_id = uuid.uuid4().hex[:12]
|
|
jdir = _job_dir(job_id)
|
|
jdir.mkdir(parents=True, exist_ok=True)
|
|
dest = jdir / filename
|
|
shutil.copy2(src, dest)
|
|
|
|
parsed = parse_step_products(dest)
|
|
result = ParseResult(
|
|
job_id=job_id,
|
|
filename=filename,
|
|
root_name=parsed["root_name"],
|
|
parts=parsed["parts"],
|
|
product_names=parsed["product_names"],
|
|
stats={**parsed["stats"], "geometry": geometry_backend_status()},
|
|
)
|
|
(jdir / "parse.json").write_text(result.model_dump_json(indent=2), encoding="utf-8")
|
|
(jdir / "step_path.txt").write_text(str(dest), encoding="utf-8")
|
|
return result
|
|
|
|
|
|
def load_parse(job_id: str) -> ParseResult:
|
|
path = _job_dir(job_id) / "parse.json"
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"job not found: {job_id}")
|
|
return ParseResult.model_validate_json(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def step_path_for_job(job_id: str) -> Path:
|
|
txt = _job_dir(job_id) / "step_path.txt"
|
|
return Path(txt.read_text(encoding="utf-8").strip())
|
|
|
|
|
|
async def propose_draft(
|
|
job_id: str | None,
|
|
*,
|
|
profile: str,
|
|
robot_name: str,
|
|
extra_hint: str = "",
|
|
part_names: list[str] | None = None,
|
|
solids_geom: list[dict[str, Any]] | None = None,
|
|
) -> RobotDraft:
|
|
"""A7 → ARM7 gold URDF (no LLM). generic → Zhipu full propose."""
|
|
stats: dict[str, Any] = {}
|
|
names: list[str] = []
|
|
|
|
if part_names:
|
|
names = [n.strip() for n in part_names if n and str(n).strip()]
|
|
stats = {"source": "client", "part_count": len(names)}
|
|
elif job_id:
|
|
parsed = load_parse(job_id)
|
|
names = list(parsed.product_names)
|
|
stats = dict(parsed.stats)
|
|
else:
|
|
raise ValueError("Provide part_names (preferred) or job_id")
|
|
|
|
if not names:
|
|
raise ValueError("No part names available for proposal")
|
|
|
|
seed = apply_profile(profile, names, robot_name=robot_name)
|
|
|
|
# A7: always return gold kinematics (skip LLM — avoids JSON truncation).
|
|
if profile == "a7":
|
|
draft = enforce_a7_gold_kinematics(seed)
|
|
draft.notes = list(draft.notes) + [
|
|
"已使用 ARM7_urdf.urdf 金标关节（未调用智谱）。前端会按模型最长轴摆放。"
|
|
]
|
|
else:
|
|
draft = await propose_joints_with_zhipu(
|
|
seed,
|
|
extra_hint=extra_hint,
|
|
stats=stats,
|
|
solids_geom=solids_geom,
|
|
)
|
|
|
|
if job_id:
|
|
out = _job_dir(job_id) / "draft.json"
|
|
out.write_text(draft.model_dump_json(indent=2), encoding="utf-8")
|
|
return draft
|
|
|
|
|
|
def export_package(job_id: str, draft: RobotDraft, include_meshes: bool = True) -> dict[str, Any]:
|
|
jdir = _job_dir(job_id)
|
|
out_dir = jdir / "export" / draft.name
|
|
if out_dir.exists():
|
|
shutil.rmtree(out_dir)
|
|
paths = write_urdf_package(draft, out_dir)
|
|
mesh_info: dict[str, Any] = {}
|
|
if include_meshes:
|
|
mesh_info = export_meshes_stub(
|
|
step_path_for_job(job_id),
|
|
Path(paths["meshes_dir"]),
|
|
[l.name for l in draft.links],
|
|
)
|
|
meta = {"files": paths, "meshes": mesh_info, "job_id": job_id}
|
|
(out_dir / "export_meta.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")
|
|
return meta
|