305 lines
17 KiB
Python
305 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(ROOT / "backend"))
|
|
|
|
from app.models.contracts import ChatMessage, MessagePart # noqa: E402
|
|
from app.services.agent_service import AgentService, TOOL_SCHEMAS, system_prompt # noqa: E402
|
|
from app.services.library import CdslLibrary # noqa: E402
|
|
from app.services.part_skills import PartSkillLibrary # noqa: E402
|
|
from app.services.storage import WorkspaceStore # noqa: E402
|
|
from app.settings import ProviderConfig, ProviderModel, Settings # noqa: E402
|
|
|
|
|
|
BACKEND = ROOT / "backend"
|
|
PART_SKILL_ROOT = BACKEND / "agent" / "skills" / "cad-engine" / "references" / "part-skills"
|
|
|
|
|
|
def workplane(z: float = 0.0) -> dict[str, list[float]]:
|
|
return {
|
|
"origin_mm": [0.0, 0.0, z],
|
|
"x_dir": [1.0, 0.0, 0.0],
|
|
"y_dir": [0.0, 1.0, 0.0],
|
|
"normal": [0.0, 0.0, 1.0],
|
|
}
|
|
|
|
|
|
def document(part_id: str, sketches: list[dict], features: list[dict]) -> dict:
|
|
return {
|
|
"schema": "cad.cdsl.llm.v1",
|
|
"schema_version": "1.1.0",
|
|
"kind": "part",
|
|
"part_id": part_id,
|
|
"meta": {"unit": "mm"},
|
|
"geometry": {"sketches": sketches},
|
|
"features": features,
|
|
}
|
|
|
|
|
|
def rectangle(center: list[float], width: float, height: float) -> dict:
|
|
half_width, half_height = width / 2, height / 2
|
|
return {"type": "polygon", "vertices": [
|
|
[center[0] - half_width, center[1] - half_height], [center[0] + half_width, center[1] - half_height],
|
|
[center[0] + half_width, center[1] + half_height], [center[0] - half_width, center[1] + half_height],
|
|
]}
|
|
|
|
|
|
def circles(items: list[dict]) -> dict:
|
|
return {"type": "analytic_contours", "contours": [
|
|
{"role": "outer", "closed": True, "segments": [{"type": "circle", "center": item["center"], "radius_mm": item["radius_mm"]}]}
|
|
for item in items
|
|
]}
|
|
|
|
|
|
def circle_grid(radius: float, count_x: int, count_y: int, spacing_x: float, spacing_y: float) -> dict:
|
|
return circles([
|
|
{"center": [(column - (count_x - 1) / 2) * spacing_x, (row - (count_y - 1) / 2) * spacing_y], "radius_mm": radius}
|
|
for row in range(count_y) for column in range(count_x)
|
|
])
|
|
|
|
|
|
def obround(length: float, width: float) -> dict:
|
|
radius, left, right = width / 2, -length / 2 + width / 2, length / 2 - width / 2
|
|
return {"type": "analytic_contours", "contours": [{"role": "outer", "closed": True, "segments": [
|
|
{"type": "line", "start": [right, radius], "end": [left, radius]},
|
|
{"type": "arc", "start": [left, radius], "end": [left, -radius], "center": [left, 0], "radius_mm": radius},
|
|
{"type": "line", "start": [left, -radius], "end": [right, -radius]},
|
|
{"type": "arc", "start": [right, -radius], "end": [right, radius], "center": [right, 0], "radius_mm": radius},
|
|
]}]}
|
|
|
|
|
|
def annulus(inner_radius: float, outer_radius: float) -> dict:
|
|
return circles([
|
|
{"center": [0, 0], "radius_mm": outer_radius},
|
|
{"center": [0, 0], "radius_mm": inner_radius},
|
|
])
|
|
|
|
|
|
def mounting_plate_fixture() -> dict:
|
|
return document("golden-mounting-plate", [
|
|
{"id": "base", "workplane": workplane(), "profile": rectangle([0, 0], 80, 60)},
|
|
{"id": "grid", "workplane": workplane(10), "profile": circle_grid(3, 2, 2, 50, 30)},
|
|
], [
|
|
{"id": "base_add", "atomic_id": "extrude_add_blind", "depends_on": [], "params": {"distance_mm": 10}, "sketch_id": "base"},
|
|
{"id": "mount_pattern", "atomic_id": "extrude_cut_blind", "depends_on": ["base_add"], "params": {"distance_mm": 10, "reverse": True}, "sketch_id": "grid"},
|
|
{"id": "counterbores", "atomic_id": "hole_counterbore", "depends_on": ["mount_pattern"], "params": {"diameter_mm": 6, "depth_mm": 10, "counterbore_diameter_mm": 12, "counterbore_depth_mm": 4, "positions": [{"mm": [-15, 0, 0]}, {"mm": [15, 0, 0]}], "host_face": {"frame": workplane(10)}}, "sketch_id": "base"},
|
|
])
|
|
|
|
|
|
def mounting_bracket_fixture() -> dict:
|
|
web_plane = {"origin_mm": [0, -20, 0], "x_dir": [1, 0, 0], "y_dir": [0, 0, -1], "normal": [0, 1, 0]}
|
|
return document("golden-mounting-bracket", [
|
|
{"id": "base", "workplane": workplane(), "profile": rectangle([0, 0], 80, 40)},
|
|
{"id": "web", "workplane": web_plane, "profile": rectangle([0, -15], 50, 30)},
|
|
{"id": "slot", "workplane": workplane(6), "profile": obround(24, 8)},
|
|
{"id": "symmetric_holes", "workplane": workplane(6), "profile": circles([{"center": [-25, 0], "radius_mm": 3}, {"center": [25, 0], "radius_mm": 3}])},
|
|
], [
|
|
{"id": "base_add", "atomic_id": "extrude_add_blind", "depends_on": [], "params": {"distance_mm": 6}, "sketch_id": "base"},
|
|
{"id": "web_add", "atomic_id": "extrude_add_blind", "depends_on": ["base_add"], "params": {"distance_mm": 6}, "sketch_id": "web"},
|
|
{"id": "slot_cut", "atomic_id": "extrude_cut_blind", "depends_on": ["web_add"], "params": {"distance_mm": 6, "reverse": True}, "sketch_id": "slot"},
|
|
{"id": "symmetric_cut", "atomic_id": "extrude_cut_blind", "depends_on": ["slot_cut"], "params": {"distance_mm": 6, "reverse": True}, "sketch_id": "symmetric_holes"},
|
|
])
|
|
|
|
|
|
def flange_fixture() -> dict:
|
|
return document("golden-flange", [
|
|
{"id": "base", "workplane": workplane(), "profile": annulus(10, 40)},
|
|
{"id": "bolt_circle", "workplane": workplane(12), "profile": circles([{"center": [25, 0], "radius_mm": 3}, {"center": [0, 25], "radius_mm": 3}, {"center": [-25, 0], "radius_mm": 3}, {"center": [0, -25], "radius_mm": 3}])},
|
|
], [
|
|
{"id": "base_add", "atomic_id": "extrude_add_blind", "depends_on": [], "params": {"distance_mm": 12}, "sketch_id": "base"},
|
|
{"id": "bolt_holes", "atomic_id": "extrude_cut_blind", "depends_on": ["base_add"], "params": {"distance_mm": 12, "reverse": True}, "sketch_id": "bolt_circle"},
|
|
])
|
|
|
|
|
|
def shaft_fixture() -> dict:
|
|
end_plane = {"origin_mm": [0, 35, 0], "x_dir": [1, 0, 0], "y_dir": [0, 0, -1], "normal": [0, 1, 0]}
|
|
return document("golden-stepped-shaft", [
|
|
{"id": "shaft_profile", "workplane": workplane(), "profile": {"type": "polygon", "vertices": [[0, 0], [10, 0], [10, 15], [7, 15], [7, 35], [0, 35]]}},
|
|
], [
|
|
{"id": "shaft_add", "atomic_id": "revolve_add", "depends_on": [], "params": {"angle_deg": 360, "axis": {"origin_mm": [0, 0, 0], "direction": [0, 1, 0]}}, "sketch_id": "shaft_profile"},
|
|
{"id": "coaxial_bore", "atomic_id": "hole_blind", "depends_on": ["shaft_add"], "params": {"diameter_mm": 4, "depth_mm": 35, "positions": [{"mm": [0, 0, 0]}], "host_face": {"frame": end_plane}}, "sketch_id": "shaft_profile"},
|
|
])
|
|
|
|
|
|
def bearing_housing_fixture() -> dict:
|
|
return document("golden-bearing-housing", [
|
|
{"id": "base", "workplane": workplane(), "profile": rectangle([0, 0], 100, 60)},
|
|
{"id": "housing", "workplane": workplane(8), "profile": {"type": "circle", "radius_mm": 28}},
|
|
{"id": "base_holes", "workplane": workplane(8), "profile": circle_grid(4, 2, 2, 80, 40)},
|
|
], [
|
|
{"id": "base_add", "atomic_id": "extrude_add_blind", "depends_on": [], "params": {"distance_mm": 8}, "sketch_id": "base"},
|
|
{"id": "housing_add", "atomic_id": "extrude_add_blind", "depends_on": ["base_add"], "params": {"distance_mm": 25}, "sketch_id": "housing"},
|
|
{"id": "bearing_seat", "atomic_id": "hole_counterbore", "depends_on": ["housing_add"], "params": {"diameter_mm": 12, "depth_mm": 25, "counterbore_diameter_mm": 30, "counterbore_depth_mm": 10, "positions": [{"mm": [0, 0, 0]}], "host_face": {"frame": workplane(33)}}, "sketch_id": "housing"},
|
|
{"id": "base_mount_holes", "atomic_id": "extrude_cut_blind", "depends_on": ["bearing_seat"], "params": {"distance_mm": 8, "reverse": True}, "sketch_id": "base_holes"},
|
|
])
|
|
|
|
|
|
def hex_nut_fixture() -> dict:
|
|
return document("golden-hex-nut", [
|
|
{"id": "hex", "workplane": workplane(), "profile": {"type": "polygon", "vertices": [[10, 0], [5, 8.660254], [-5, 8.660254], [-10, 0], [-5, -8.660254], [5, -8.660254]]}},
|
|
{"id": "thread_bore", "workplane": workplane(10), "profile": {"type": "circle", "radius_mm": 3}},
|
|
], [
|
|
{"id": "nut_add", "atomic_id": "extrude_add_blind", "depends_on": [], "params": {"distance_mm": 10}, "sketch_id": "hex"},
|
|
{"id": "m6_bore", "atomic_id": "extrude_cut_blind", "depends_on": ["nut_add"], "params": {"distance_mm": 10, "reverse": True}, "sketch_id": "thread_bore"},
|
|
])
|
|
|
|
|
|
class PartSkillLibraryTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.library = PartSkillLibrary(PART_SKILL_ROOT)
|
|
|
|
def test_english_and_chinese_triggers(self) -> None:
|
|
english = self.library.select("Create a mounting plate with counterbored hole pattern")
|
|
chinese = self.library.select("带沉孔和孔阵列的安装底板")
|
|
|
|
self.assertIn("planning/mounting-plate", english["skill_ids"])
|
|
self.assertIn("planning/mounting-plate", chinese["skill_ids"])
|
|
self.assertIn("atomic/counterbored-hole-creation", chinese["skill_ids"])
|
|
|
|
def test_exclusions_and_selection_limits(self) -> None:
|
|
selection = self.library.select("Create a flange nut with an M6 threaded hole")
|
|
|
|
self.assertNotIn("planning/flange", selection["skill_ids"])
|
|
self.assertLessEqual(len(selection["planning_ids"]), 1)
|
|
self.assertLessEqual(len(selection["support_ids"]), 3)
|
|
self.assertEqual(len(selection["skill_ids"]), len(set(selection["skill_ids"])))
|
|
|
|
def test_inheritance_and_primary_family_conflict(self) -> None:
|
|
inherited = ["planning/mounting-plate", "atomic/counterbored-hole-creation"]
|
|
addition = self.library.select("Add one M6 threaded hole", inherited)
|
|
conflict = self.library.select("Add a flange bolt circle", inherited)
|
|
replacement = self.library.select("Replace the whole part with a flange bolt circle", inherited)
|
|
|
|
self.assertIn("planning/mounting-plate", addition["skill_ids"])
|
|
self.assertIn("atomic/threaded-hole-creation", addition["skill_ids"])
|
|
self.assertEqual(conflict["conflict"]["current"], "planning/mounting-plate")
|
|
self.assertEqual(conflict["conflict"]["matched"], "planning/flange")
|
|
self.assertEqual(replacement["planning_ids"], ["planning/flange"])
|
|
self.assertIsNone(replacement["conflict"])
|
|
|
|
def test_catalog_and_bridge_files_are_complete(self) -> None:
|
|
payload = json.loads((PART_SKILL_ROOT / "catalog.json").read_text(encoding="utf-8"))
|
|
categories = {kind: 0 for kind in ("planning", "functional", "atomic")}
|
|
for skill in payload["skills"]:
|
|
categories[skill["kind"]] += 1
|
|
self.assertTrue((PART_SKILL_ROOT / skill["bridge"]).is_file())
|
|
self.assertTrue((PART_SKILL_ROOT / skill["source"]).is_file())
|
|
self.assertTrue(skill["triggers"])
|
|
self.assertNotIn("capability_translation_rules", skill)
|
|
self.assertEqual(categories, {"planning": 6, "functional": 7, "atomic": 8})
|
|
self.assertTrue((PART_SKILL_ROOT / "LICENSE").is_file())
|
|
self.assertTrue((PART_SKILL_ROOT / "PROVENANCE.md").is_file())
|
|
|
|
def test_audit_records_selection_and_assumptions_without_interpreting_geometry(self) -> None:
|
|
fixture = mounting_plate_fixture()
|
|
audit = self.library.audit(self.library.select("安装底板"), fixture, ["孔位置由用户确认"])
|
|
self.assertIn("planning/mounting-plate", audit["skill_ids"])
|
|
self.assertEqual(audit["assumptions"], ["孔位置由用户确认"])
|
|
self.assertTrue(all(item.get("version") and item.get("summary") for item in audit["skills"]))
|
|
self.assertNotIn("capability_translations", audit)
|
|
self.assertNotIn("evidence", audit)
|
|
|
|
|
|
class AgentPartSkillTests(unittest.TestCase):
|
|
def test_agent_stream_selects_and_injects_part_skill_before_first_model_call(self) -> None:
|
|
class CaptureAgent(AgentService):
|
|
def __init__(self, *args: object, **kwargs: object) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self.first_messages: list[dict[str, object]] = []
|
|
|
|
async def _complete(self, messages: list[dict[str, object]], *args: object, **kwargs: object) -> dict[str, object]:
|
|
self.first_messages = [dict(message) for message in messages]
|
|
return {"choices": [{"message": {"role": "assistant", "content": "已分析", "tool_calls": []}}]}
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
settings = self._settings(Path(directory))
|
|
library = PartSkillLibrary(PART_SKILL_ROOT)
|
|
agent = CaptureAgent(settings, WorkspaceStore(settings), CdslLibrary(settings), library)
|
|
message = ChatMessage(id="user_1", role="user", parts=[MessagePart(type="text", text="生成一个法兰螺栓圆")])
|
|
|
|
async def consume() -> None:
|
|
async for _ in agent.stream([message], None, None):
|
|
pass
|
|
|
|
asyncio.run(consume())
|
|
self.assertIn("[planning/flange]", str(agent.first_messages[0]["content"]))
|
|
self.assertIn("[functional/flange-bolt-circle]", str(agent.first_messages[0]["content"]))
|
|
|
|
def test_prompt_contains_bridge_and_direct_cdsl_tools(self) -> None:
|
|
library = PartSkillLibrary(PART_SKILL_ROOT)
|
|
selection = library.select("Create a flange with a bolt circle")
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
settings = self._settings(Path(directory))
|
|
prompt = system_prompt(settings, "Create a flange with a bolt circle", part_skill_context=library.render_context(selection))
|
|
|
|
self.assertIn("[planning/flange]", prompt)
|
|
self.assertIn("circular-pattern atomic", prompt)
|
|
self.assertEqual([tool["function"]["name"] for tool in TOOL_SCHEMAS], ["analyze_image_reference", "extract_image_sketch_candidates", "search_cdsl_library", "read_cdsl_reference", "describe_design_intent", "read_current_cdsl", "generate_cdsl_model", "patch_cdsl_model"])
|
|
|
|
def test_generation_persists_assumptions_and_selected_skills(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
settings = self._settings(Path(directory))
|
|
store = WorkspaceStore(settings)
|
|
library = PartSkillLibrary(PART_SKILL_ROOT)
|
|
agent = AgentService(settings, store, CdslLibrary(settings), library)
|
|
request = "M6 六角螺母"
|
|
selection = library.select(request)
|
|
planning_state = {"phase": "INTAKE", "design_brief": ""}
|
|
planned, _ = asyncio.run(agent._run_tool(
|
|
"describe_design_intent",
|
|
{"plan": "Create an M6 hex nut with a centered cylindrical bore representing the thread.", "assumptions": []},
|
|
"",
|
|
request,
|
|
[],
|
|
part_skill_selection=selection,
|
|
planning_state=planning_state,
|
|
))
|
|
self.assertTrue(planned["ok"])
|
|
result, generated = asyncio.run(agent._run_tool(
|
|
"generate_cdsl_model",
|
|
{"cdsl": hex_nut_fixture(), "summary": "M6 nut", "assumptions": ["M6 thread is represented as a cylindrical bore"]},
|
|
"",
|
|
request,
|
|
["reference-fixture"],
|
|
part_skill_selection=selection,
|
|
planning_state=planning_state,
|
|
))
|
|
|
|
self.assertTrue(result["ok"])
|
|
self.assertIsNotNone(generated)
|
|
task = store.read_task(str(generated["task_id"]))
|
|
revision = task["revisions"][0]
|
|
audit = json.loads(store.artifact_path(task["task_id"], revision["part_skills_path"]).read_text(encoding="utf-8"))
|
|
report = json.loads(store.artifact_path(task["task_id"], revision["report_path"]).read_text(encoding="utf-8"))
|
|
self.assertIn("M6 thread is represented as a cylindrical bore", audit["assumptions"])
|
|
self.assertIn("atomic/threaded-hole-creation", revision["part_skill_ids"])
|
|
self.assertEqual(report["generation_context"]["cdsl_reference_ids"], ["reference-fixture"])
|
|
|
|
@staticmethod
|
|
def _settings(root: Path) -> Settings:
|
|
provider = ProviderConfig("test", "Test", "https://example.invalid/v1", "test-key", (ProviderModel("test-model"),))
|
|
return Settings(
|
|
task_root=root / "tasks",
|
|
conversation_root=root / "conversations",
|
|
library_root=BACKEND / "cdsl_library",
|
|
engine_root=BACKEND / "engine" / "cdsl_engine",
|
|
llm_base_url=provider.base_url,
|
|
llm_api_key=provider.api_key,
|
|
llm_model="test-model",
|
|
llm_timeout_s=1,
|
|
default_provider_id="test",
|
|
providers=(provider,),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|