297 lines
15 KiB
Python
297 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
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 # 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 mounting_plate_cdsl() -> dict:
|
|
holes = [
|
|
{"role": "outer", "closed": True, "segments": [{"type": "circle", "center": center, "radius_mm": 3}]}
|
|
for center in [[-40, -20], [40, -20], [-40, 20], [40, 20]]
|
|
]
|
|
return {
|
|
"schema": "cad.cdsl.llm.v1",
|
|
"schema_version": "1.0",
|
|
"kind": "part",
|
|
"part_id": "mounting-plate",
|
|
"geometry": {"sketches": [
|
|
{"id": "base_sketch", "workplane": _workplane(), "profile": {"type": "polygon", "vertices": [[-50, -30], [50, -30], [50, 30], [-50, 30]]}},
|
|
{"id": "holes_sketch", "workplane": _workplane(10), "profile": {"type": "analytic_contours", "contours": holes}},
|
|
]},
|
|
"features": [
|
|
{"id": "base_add", "atomic_id": "extrude_add_blind", "depends_on": [], "sketch_id": "base_sketch", "params": {"distance_mm": 10}},
|
|
{"id": "hole_cut", "atomic_id": "extrude_cut_blind", "depends_on": ["base_add"], "sketch_id": "holes_sketch", "params": {"distance_mm": 10, "reverse": True}},
|
|
],
|
|
}
|
|
|
|
|
|
class DirectCdslFlowTests(unittest.TestCase):
|
|
def settings(self, 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,),
|
|
)
|
|
|
|
def test_design_brief_is_required_before_library_or_cdsl(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
settings = self.settings(Path(directory))
|
|
store = WorkspaceStore(settings)
|
|
agent = AgentService(settings, store, CdslLibrary(settings), PartSkillLibrary(PART_SKILL_ROOT))
|
|
state = {"phase": "INTAKE", "design_brief": ""}
|
|
searched, _ = asyncio.run(agent._run_tool("search_cdsl_library", {"query": "mounting plate"}, "", "mounting plate", [], planning_state=state))
|
|
generated, _ = asyncio.run(agent._run_tool("generate_cdsl_model", {"cdsl": {}, "summary": "x", "assumptions": []}, "", "mounting plate", [], planning_state=state))
|
|
|
|
self.assertEqual(searched["code"], "DESIGN_BRIEF_REQUIRED")
|
|
self.assertEqual(generated["code"], "DESIGN_BRIEF_REQUIRED")
|
|
self.assertEqual(list(settings.task_root.glob("cad_*")), [])
|
|
|
|
def test_generation_normalizes_legacy_llm_cdsl_before_building(self) -> None:
|
|
legacy_cdsl = {
|
|
"schema": "cad.cdsl.llm.v1",
|
|
"part_id": "legacy-flange-base",
|
|
"geometry": {"sketches": [{
|
|
"id": "base_sketch",
|
|
"plane": "XY",
|
|
"offset_mm": 12,
|
|
"profile": {"type": "circle", "radius_mm": 20},
|
|
}]},
|
|
"features": [{
|
|
"id": "base_add",
|
|
"atomic_id": "extrude_add_blind",
|
|
"sketch": "base_sketch",
|
|
"params": {"distance_mm": 8},
|
|
}],
|
|
}
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
settings = self.settings(Path(directory))
|
|
store = WorkspaceStore(settings)
|
|
agent = AgentService(settings, store, CdslLibrary(settings), PartSkillLibrary(PART_SKILL_ROOT))
|
|
state = {"phase": "PLANNED", "design_brief": "Create a flange base."}
|
|
captured_build: dict[str, object] = {}
|
|
|
|
def fake_build_revision(**kwargs: object) -> dict[str, object]:
|
|
captured_build.update(kwargs)
|
|
return {"task_id": "cad_aaaaaaaaaaaa", "revision_id": "rev_001"}
|
|
|
|
with patch("app.services.agent_service.build_revision", side_effect=fake_build_revision):
|
|
result, _ = asyncio.run(agent._run_tool(
|
|
"generate_cdsl_model",
|
|
{"cdsl": legacy_cdsl, "summary": "flange base", "assumptions": []},
|
|
"", "Create a flange base", [], planning_state=state,
|
|
))
|
|
|
|
built_cdsl = captured_build["cdsl"]
|
|
self.assertTrue(result["ok"])
|
|
self.assertEqual(built_cdsl["features"][0]["sketch_id"], "base_sketch")
|
|
self.assertEqual(built_cdsl["features"][0]["depends_on"], [])
|
|
self.assertNotIn("sketch", built_cdsl["features"][0])
|
|
self.assertIn("workplane", built_cdsl["geometry"]["sketches"][0])
|
|
self.assertEqual(len(result["normalization_repairs"]), 3)
|
|
|
|
def test_successful_generation_ends_the_agent_tool_loop(self) -> None:
|
|
class CaptureAgent(AgentService):
|
|
def __init__(self, *args: object, **kwargs: object) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self.responses = [
|
|
{
|
|
"choices": [{"message": {
|
|
"role": "assistant", "content": "", "tool_calls": [{
|
|
"id": "brief", "type": "function", "function": {
|
|
"name": "describe_design_intent",
|
|
"arguments": json.dumps({"plan": "Create a cylindrical part.", "assumptions": []}),
|
|
},
|
|
}],
|
|
}}],
|
|
},
|
|
{
|
|
"choices": [{"message": {
|
|
"role": "assistant", "content": "", "tool_calls": [{
|
|
"id": "generate", "type": "function", "function": {
|
|
"name": "generate_cdsl_model",
|
|
"arguments": json.dumps({
|
|
"cdsl": mounting_plate_cdsl(),
|
|
"summary": "mounting plate",
|
|
"assumptions": [],
|
|
}),
|
|
},
|
|
}],
|
|
}}],
|
|
},
|
|
]
|
|
|
|
async def _complete(self, *args: object, **kwargs: object) -> dict[str, object]:
|
|
return self.responses.pop(0)
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
settings = self.settings(Path(directory))
|
|
store = WorkspaceStore(settings)
|
|
agent = CaptureAgent(settings, store, CdslLibrary(settings), PartSkillLibrary(PART_SKILL_ROOT))
|
|
|
|
def fake_build_revision(**kwargs: object) -> dict[str, object]:
|
|
return {
|
|
"task_id": "cad_aaaaaaaaaaaa", "revision_id": "rev_001",
|
|
"cdsl_path": "revisions/rev_001/model.cdsl.json",
|
|
"step_path": "revisions/rev_001/model.step",
|
|
"glb_path": "revisions/rev_001/model.glb",
|
|
"report_path": "revisions/rev_001/rebuild-report.json",
|
|
"summary": str(kwargs["summary"]), "reference_ids": [], "engine": "cdsl_only",
|
|
}
|
|
|
|
with patch("app.services.agent_service.build_revision", side_effect=fake_build_revision):
|
|
async def consume() -> None:
|
|
message = ChatMessage(id="user_1", role="user", parts=[MessagePart(type="text", text="Create a mounting plate")])
|
|
async for _ in agent.stream([message], None, None):
|
|
pass
|
|
|
|
asyncio.run(consume())
|
|
|
|
saved = store.read_conversation(next(item.name for item in settings.conversation_root.iterdir()))
|
|
parts = saved["messages"][-1]["parts"]
|
|
self.assertEqual(agent.responses, [])
|
|
self.assertTrue(any(part["type"] == "data-cad-result" for part in parts))
|
|
self.assertFalse(any(part["type"] == "data-cad-error" for part in parts))
|
|
|
|
def test_text_brief_is_returned_to_model_and_cdsl_is_the_only_contract(self) -> None:
|
|
class CaptureAgent(AgentService):
|
|
def __init__(self, *args: object, **kwargs: object) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self.seen_messages: list[list[dict[str, object]]] = []
|
|
self.responses = [
|
|
{
|
|
"choices": [{"message": {
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": [{
|
|
"id": "brief",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "describe_design_intent",
|
|
"arguments": json.dumps({
|
|
"plan": "Create a rectangular mounting plate, then cut four mounting holes and a center slot.",
|
|
"assumptions": ["Use millimetres."],
|
|
}),
|
|
},
|
|
}],
|
|
}}],
|
|
},
|
|
{
|
|
"choices": [{"message": {
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": [{
|
|
"id": "generate",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "generate_cdsl_model",
|
|
"arguments": json.dumps({
|
|
"cdsl": mounting_plate_cdsl(),
|
|
"summary": "mounting plate",
|
|
"assumptions": ["Use millimetres."],
|
|
}),
|
|
},
|
|
}],
|
|
}}],
|
|
},
|
|
{"choices": [{"message": {"role": "assistant", "content": "已生成。", "tool_calls": []}}]},
|
|
]
|
|
|
|
async def _complete(self, messages: list[dict[str, object]], *args: object, **kwargs: object) -> dict[str, object]:
|
|
self.seen_messages.append([dict(message) for message in messages])
|
|
return self.responses.pop(0)
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
settings = self.settings(Path(directory))
|
|
store = WorkspaceStore(settings)
|
|
skills = PartSkillLibrary(PART_SKILL_ROOT)
|
|
agent = CaptureAgent(settings, store, CdslLibrary(settings), skills)
|
|
captured_build: dict[str, object] = {}
|
|
|
|
def fake_build_revision(**kwargs: object) -> dict[str, object]:
|
|
captured_build.update(kwargs)
|
|
return {
|
|
"task_id": "cad_aaaaaaaaaaaa",
|
|
"revision_id": "rev_001",
|
|
"cdsl_path": "revisions/rev_001/model.cdsl.json",
|
|
"step_path": "revisions/rev_001/model.step",
|
|
"glb_path": "revisions/rev_001/model.glb",
|
|
"report_path": "revisions/rev_001/rebuild-report.json",
|
|
"summary": str(kwargs["summary"]),
|
|
"reference_ids": list(kwargs["reference_ids"]),
|
|
"engine": "cdsl_only",
|
|
}
|
|
|
|
with patch("app.services.agent_service.build_revision", side_effect=fake_build_revision):
|
|
async def consume() -> None:
|
|
message = ChatMessage(id="user_1", role="user", parts=[MessagePart(type="text", text="Create a mounting plate")])
|
|
async for _ in agent.stream([message], None, None):
|
|
pass
|
|
|
|
asyncio.run(consume())
|
|
|
|
brief_result = json.loads(str(agent.seen_messages[1][-1]["content"]))
|
|
self.assertEqual(brief_result["plan"], "Create a rectangular mounting plate, then cut four mounting holes and a center slot.")
|
|
self.assertNotIn("structures", brief_result)
|
|
self.assertNotIn("design_intent", captured_build)
|
|
self.assertEqual(captured_build["parent_revision_id"], "")
|
|
|
|
def test_revision_parent_is_taken_from_the_current_successful_cdsl_revision(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
settings = self.settings(Path(directory))
|
|
store = WorkspaceStore(settings)
|
|
agent = AgentService(settings, store, CdslLibrary(settings), PartSkillLibrary(PART_SKILL_ROOT))
|
|
state = {"phase": "INTAKE", "design_brief": "", "current_model_read": True}
|
|
asyncio.run(agent._run_tool(
|
|
"describe_design_intent",
|
|
{"plan": "Increase the plate thickness and preserve the existing hole layout.", "assumptions": []},
|
|
"cad_aaaaaaaaaaaa", "Revise the plate", [], planning_state=state,
|
|
))
|
|
store.read_task = lambda _task_id: {"current_revision": "rev_007"} # type: ignore[method-assign]
|
|
captured_build: dict[str, object] = {}
|
|
|
|
def fake_build_revision(**kwargs: object) -> dict[str, object]:
|
|
captured_build.update(kwargs)
|
|
return {"task_id": "cad_aaaaaaaaaaaa", "revision_id": "rev_008"}
|
|
|
|
with patch("app.services.agent_service.build_revision", side_effect=fake_build_revision):
|
|
result, _ = asyncio.run(agent._run_tool(
|
|
"generate_cdsl_model",
|
|
{"cdsl": mounting_plate_cdsl(), "summary": "revised plate", "assumptions": []},
|
|
"cad_aaaaaaaaaaaa", "Revise the plate", [], planning_state=state,
|
|
))
|
|
|
|
self.assertTrue(result["ok"])
|
|
self.assertEqual(captured_build["parent_revision_id"], "rev_007")
|