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")