sunxianghui
93773f3887
Includes FastAPI backend, vendored step2urdf frontend, and A7 handtuned arm JSON for URDF generation.
249 lines
8.5 KiB
Python
249 lines
8.5 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import re
|
||
from typing import Any
|
||
|
||
import httpx
|
||
|
||
from app.config import settings
|
||
from app.models import JointDef, LinkDef, RobotDraft
|
||
|
||
|
||
SYSTEM_PROMPT = """你是机器人 URDF 运动学专家。根据零件列表与包围盒，直接设计 link/joint 树。
|
||
|
||
约定：
|
||
- 单位：origin_xyz 用毫米；origin_rpy 用弧度。
|
||
- 第一个关节相对世界系，origin 靠近对应零件 center，不要写成 [0,0,0]。
|
||
- part_names 只能用输入里的名字（如 Solid_0），禁止发明新名字。
|
||
- 每个非 base_link 恰好一个 parent joint。
|
||
- axis 为单位向量；不要把所有轴都写成 [0,0,1]。
|
||
- rationale 最多 20 字。
|
||
- 只输出一个合法 JSON 对象，不要 Markdown、不要注释、不要尾逗号。
|
||
|
||
JSON schema:
|
||
{"robot_name":"string","links":[{"name":"string","part_names":["string"]}],"joints":[{"name":"string","joint_type":"revolute|prismatic|fixed","parent":"string","child":"string","origin_xyz":[0,0,0],"origin_rpy":[0,0,0],"axis":[0,0,1],"lower":-3.14,"upper":3.14,"effort":100,"velocity":1,"rationale":"string"}],"notes":["string"]}
|
||
"""
|
||
|
||
|
||
def _strip_fences(text: str) -> str:
|
||
text = text.strip()
|
||
if text.startswith("```"):
|
||
text = re.sub(r"^```(?:json)?\s*", "", text)
|
||
text = re.sub(r"\s*```$", "", text)
|
||
return text.strip()
|
||
|
||
|
||
def _repair_json_text(text: str) -> str:
|
||
"""Best-effort cleanup for common LLM JSON issues / truncation."""
|
||
text = _strip_fences(text)
|
||
# Chinese quotes → ASCII
|
||
text = text.replace("“", '"').replace("”", '"').replace("‘", "'").replace("’", "'")
|
||
# Remove // comments
|
||
text = re.sub(r"(?m)^\s*//.*?$", "", text)
|
||
# Trailing commas before } or ]
|
||
text = re.sub(r",\s*([}\]])", r"\1", text)
|
||
|
||
# If truncated, close open braces/brackets
|
||
start = text.find("{")
|
||
if start < 0:
|
||
return text
|
||
text = text[start:]
|
||
# Drop incomplete trailing string if odd number of unescaped quotes in last line
|
||
in_str = False
|
||
escape = False
|
||
stack: list[str] = []
|
||
last_ok = 0
|
||
for i, ch in enumerate(text):
|
||
if in_str:
|
||
if escape:
|
||
escape = False
|
||
elif ch == "\\":
|
||
escape = True
|
||
elif ch == '"':
|
||
in_str = False
|
||
continue
|
||
if ch == '"':
|
||
in_str = True
|
||
continue
|
||
if ch in "{[":
|
||
stack.append("}" if ch == "{" else "]")
|
||
last_ok = i
|
||
elif ch in "}]":
|
||
if stack and stack[-1] == ch:
|
||
stack.pop()
|
||
last_ok = i
|
||
if in_str:
|
||
# close string then truncate junk after last complete structure
|
||
text = text + '"'
|
||
if stack:
|
||
# trim to last comma / incomplete key if needed
|
||
text = text.rstrip()
|
||
if text.endswith(","):
|
||
text = text[:-1]
|
||
text = text + "".join(reversed(stack))
|
||
return text
|
||
|
||
|
||
def _extract_json(text: str) -> dict[str, Any]:
|
||
candidates = [_strip_fences(text), _repair_json_text(text)]
|
||
m = re.search(r"\{[\s\S]*\}", text)
|
||
if m:
|
||
candidates.append(m.group(0))
|
||
candidates.append(_repair_json_text(m.group(0)))
|
||
|
||
errors: list[str] = []
|
||
for cand in candidates:
|
||
try:
|
||
data = json.loads(cand)
|
||
if isinstance(data, dict):
|
||
return data
|
||
except json.JSONDecodeError as e:
|
||
errors.append(str(e))
|
||
raise json.JSONDecodeError(
|
||
errors[-1] if errors else "Unable to parse LLM JSON",
|
||
text[:200],
|
||
0,
|
||
)
|
||
|
||
|
||
def _compact_solids(solids_geom: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
|
||
"""Shrink AABB payload so the model reply is less likely to truncate."""
|
||
out: list[dict[str, Any]] = []
|
||
for s in solids_geom or []:
|
||
center = s.get("center")
|
||
if not center:
|
||
mn, mx = s.get("min"), s.get("max")
|
||
if mn and mx and len(mn) == 3 and len(mx) == 3:
|
||
center = [(mn[i] + mx[i]) / 2 for i in range(3)]
|
||
if not center:
|
||
continue
|
||
out.append(
|
||
{
|
||
"name": s.get("name") or s.get("id"),
|
||
"center": [round(float(x), 1) for x in center[:3]],
|
||
}
|
||
)
|
||
return out
|
||
|
||
|
||
def _draft_from_llm(data: dict[str, Any], fallback: RobotDraft) -> RobotDraft:
|
||
links = [
|
||
LinkDef(name=l["name"], part_names=list(l.get("part_names") or []))
|
||
for l in data.get("links") or []
|
||
]
|
||
joints = []
|
||
for j in data.get("joints") or []:
|
||
jtype = j.get("joint_type") or "revolute"
|
||
if jtype not in {"revolute", "prismatic", "fixed", "continuous"}:
|
||
jtype = "revolute"
|
||
joints.append(
|
||
JointDef(
|
||
name=j.get("name") or "joint",
|
||
joint_type=jtype, # type: ignore[arg-type]
|
||
parent=j["parent"],
|
||
child=j["child"],
|
||
origin_xyz=[float(x) for x in (j.get("origin_xyz") or [0, 0, 0])],
|
||
origin_rpy=[float(x) for x in (j.get("origin_rpy") or [0, 0, 0])],
|
||
axis=[float(x) for x in (j.get("axis") or [0, 0, 1])],
|
||
lower=float(j.get("lower", -3.14)),
|
||
upper=float(j.get("upper", 3.14)),
|
||
effort=float(j.get("effort", 100)),
|
||
velocity=float(j.get("velocity", 1)),
|
||
rationale=str(j.get("rationale") or "")[:80],
|
||
)
|
||
)
|
||
if not links:
|
||
return fallback.model_copy(
|
||
update={
|
||
"llm_raw": data,
|
||
"notes": fallback.notes + ["LLM returned no links; kept seed."],
|
||
}
|
||
)
|
||
return RobotDraft(
|
||
name=str(data.get("robot_name") or fallback.name),
|
||
profile=fallback.profile,
|
||
links=links,
|
||
joints=joints,
|
||
notes=list(data.get("notes") or []),
|
||
raw_parts=fallback.raw_parts,
|
||
axis_candidates=fallback.axis_candidates,
|
||
llm_raw=data,
|
||
)
|
||
|
||
|
||
async def _chat(messages: list[dict[str, str]], *, temperature: float = 0.2) -> str:
|
||
headers = {
|
||
"Authorization": f"Bearer {settings.zhipu_api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
body = {
|
||
"model": settings.zhipu_model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": 8192,
|
||
}
|
||
url = f"{settings.zhipu_base_url.rstrip('/')}/chat/completions"
|
||
async with httpx.AsyncClient(timeout=180.0) as client:
|
||
resp = await client.post(url, headers=headers, json=body)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
return data["choices"][0]["message"]["content"]
|
||
|
||
|
||
async def propose_joints_with_zhipu(
|
||
seed: RobotDraft,
|
||
*,
|
||
extra_hint: str = "",
|
||
stats: dict[str, Any] | None = None,
|
||
solids_geom: list[dict[str, Any]] | None = None,
|
||
) -> RobotDraft:
|
||
if not settings.zhipu_api_key:
|
||
seed.notes = list(seed.notes) + [
|
||
"ZHIPU_API_KEY 未配置：无法生成关节，请配置后重试。"
|
||
]
|
||
return seed
|
||
|
||
compact = _compact_solids(solids_geom)
|
||
payload_user = {
|
||
"profile": seed.profile,
|
||
"robot_name": seed.name,
|
||
"part_names": seed.raw_parts,
|
||
"solids_center_mm": compact,
|
||
"extra_hint": extra_hint,
|
||
"instruction": (
|
||
"直接生成完整 links+joints JSON。"
|
||
"origin_xyz 用毫米；第一个关节靠近零件 center。"
|
||
),
|
||
}
|
||
|
||
messages = [
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{"role": "user", "content": json.dumps(payload_user, ensure_ascii=False)},
|
||
]
|
||
|
||
content = await _chat(messages, temperature=0.2)
|
||
try:
|
||
parsed = _extract_json(content)
|
||
except json.JSONDecodeError:
|
||
# One repair pass: ask model to fix truncated / invalid JSON
|
||
repair_messages = [
|
||
{"role": "system", "content": "只输出修复后的合法 JSON 对象，不要解释。"},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"下面内容不是合法 JSON（可能被截断或有尾逗号）。"
|
||
"请修复并只输出完整 JSON：\n\n" + content[:12000]
|
||
),
|
||
},
|
||
]
|
||
repaired = await _chat(repair_messages, temperature=0.0)
|
||
try:
|
||
parsed = _extract_json(repaired)
|
||
except json.JSONDecodeError as e:
|
||
raise ValueError(
|
||
f"智谱返回的 JSON 无法解析（常见于输出被截断）。请重试或减少零件数量。原始错误: {e.msg}"
|
||
) from e
|
||
|
||
return _draft_from_llm(parsed, seed)
|