refactor: update CrewOrchestrator to return ChatTurnResponse

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

agents/crew_orchestrator.py

@@ -15,7 +15,7 @@ from agents.base import BaseOrchestrator
 from core.utils import derive_part_name
 from agents.design_state import DesignState, DesignPlan, extract_decisions, compute_score
 from agents.gap_analyzer import analyze_gaps, generate_question_cards
-from agents.agent_flow import AgentResponse
+from agents.agent_flow import AgentResponse, ChatTurnResponse, PreviewData
 from config.settings import settings
 
 logger = logging.getLogger(__name__)
@@ -109,44 +109,39 @@ class CrewOrchestrator(BaseOrchestrator):
         history: list[dict],
         mentions: list[str] | None = None,
         max_history: int = 30,
-        design_state: dict | None = None,
+        design_state: DesignState | None = None,
         plan_context: bool = False,
-    ) -> dict:
+    ) -> ChatTurnResponse:
+        # Backward compat: accept dict during migration
+        state = design_state if isinstance(design_state, DesignState) else DesignState(**(design_state or {}))
+
         # Phase: manual plan trigger (before crew/fallback dispatch)
-        state = DesignState(**(design_state or {}))
         if state.phase == "exploring" and _is_plan_trigger(message):
             score = compute_score(state)
             plan = DesignPlan.from_state(state, confidence_score=score)
             state.phase = "planning"
             state.plan = plan
-            return {
-                "responses": [],
-                "preview": None,
-                "design_state": state.model_dump(),
-                "question_cards": [],
-            }
+            return ChatTurnResponse(design_state=state)
 
         if not self._crew_available:
-            return self._fallback(message, history, mentions, max_history, design_state, plan_context)
+            return self._fallback(message, history, mentions, max_history, state, plan_context)
 
         try:
-            return self._run_crew(message, history, mentions, max_history, design_state, plan_context)
+            return self._run_crew(message, history, mentions, max_history, state, plan_context)
         except Exception as exc:
             logger.warning("CrewAI run failed (%s), falling back", exc, exc_info=True)
             try:
-                return self._fallback(message, history, mentions, max_history, design_state, plan_context)
+                return self._fallback(message, history, mentions, max_history, state, plan_context)
             except Exception as fallback_exc:
                 logger.error("Fallback also failed: %s", fallback_exc, exc_info=True)
-                return {
-                    "responses": [AgentResponse.from_agent(
+                return ChatTurnResponse(
+                    responses=[AgentResponse.from_agent(
                         "design",
                         f"Backend error: {exc}. Fallback also failed: {fallback_exc}. "
                         f"Please check that your API key is set correctly.",
-                    ).model_dump()],
-                    "preview": None,
-                    "design_state": design_state or {},
-                    "question_cards": [],
-                }
+                    )],
+                    design_state=state,
+                )
 
     def _run_crew(
         self,
@@ -154,13 +149,13 @@ class CrewOrchestrator(BaseOrchestrator):
         history: list[dict],
         mentions: list[str] | None,
         max_history: int,
-        design_state_dict: dict | None,
+        design_state: DesignState | None,
         plan_context: bool = False,
-    ) -> dict:
+    ) -> ChatTurnResponse:
         from agents.agent_flow import AgentFlowState, AgentDispatchFlow, collect_responses
         from agents.tools import set_design_state
 
-        state = DesignState(**(design_state_dict or {}))
+        state = design_state if isinstance(design_state, DesignState) else DesignState(**(design_state or {}))
 
         # Phase: if in planning and user sends a non-plan message, reset to exploring.
         # When plan_context=True the message is a plan-field query (Ask Agent),
@@ -204,8 +199,6 @@ class CrewOrchestrator(BaseOrchestrator):
         cad_code = flow.state.cad_code
         cam_plan = flow.state.cam_plan
 
-        responses = [r.model_dump() for r in agent_responses]
-
         # ── Post-processing ─────────────────────────────────────────────
         preview = None
 
@@ -236,18 +229,18 @@ class CrewOrchestrator(BaseOrchestrator):
                     pass
 
                 validation = validate_for_cnc(shape, part_name=part_name)
-                preview = {
-                    "success": True,
-                    "part_name": part_name,
-                    "stl_url": f"/api/models/{part_name}.stl",
-                    "step_url": f"/api/models/{part_name}.step",
-                    "threemf_url": f"/api/models/{part_name}.3mf",
-                    "execution": execution_data,
-                    "validation": validation.model_dump(),
-                }
+                preview = PreviewData(
+                    success=True,
+                    part_name=part_name,
+                    stl_url=f"/api/models/{part_name}.stl",
+                    step_url=f"/api/models/{part_name}.step",
+                    threemf_url=f"/api/models/{part_name}.3mf",
+                    execution=execution_data,
+                    validation=validation.model_dump(),
+                )
 
         # G-code generation
-        if preview and preview.get("success") and cam_plan:
+        if preview and preview.success and cam_plan:
             from core.cam import generate_gcode
             from agents.tools import get_last_shape
 
@@ -259,12 +252,11 @@ class CrewOrchestrator(BaseOrchestrator):
                     tool_config=cam_plan.to_tool_config(),
                     post_processor=cam_plan.post_processor,
                 )
-                preview["cam"] = cam_result.model_dump()
+                preview.cam = cam_result.model_dump()
                 if cam_result.success and cam_result.gcode:
-                    part_name = preview["part_name"]
-                    gcode_path = self.output_dir / f"{part_name}.gcode"
+                    gcode_path = self.output_dir / f"{preview.part_name}.gcode"
                     gcode_path.write_text(cam_result.gcode)
-                    preview["gcode_url"] = f"/api/models/{part_name}.gcode"
+                    preview.gcode_url = f"/api/models/{preview.part_name}.gcode"
 
         # Update design state
         updated_state = extract_decisions(agent_responses, state, message)
@@ -273,8 +265,7 @@ class CrewOrchestrator(BaseOrchestrator):
         gap_result = analyze_gaps(agent_responses)
         question_cards = []
         if gap_result.has_gaps:
-            cards = generate_question_cards(gap_result, updated_state, user_message=message)
-            question_cards = [c.model_dump() for c in cards]
+            question_cards = generate_question_cards(gap_result, updated_state, user_message=message)
 
         # Auto-trigger plan if score crosses threshold
         if updated_state.phase == "exploring":
@@ -292,12 +283,12 @@ class CrewOrchestrator(BaseOrchestrator):
                     updated_state.plan = None
                     break
 
-        return {
-            "responses": responses,
-            "preview": preview,
-            "design_state": updated_state.model_dump(),
-            "question_cards": question_cards,
-        }
+        return ChatTurnResponse(
+            responses=agent_responses,
+            preview=preview,
+            design_state=updated_state,
+            question_cards=question_cards,
+        )
 
     def _fallback(
         self,
@@ -305,30 +296,21 @@ class CrewOrchestrator(BaseOrchestrator):
         history: list[dict],
         mentions: list[str] | None,
         max_history: int,
-        design_state: dict | None,
+        design_state: DesignState | None,
         plan_context: bool = False,
-    ) -> dict:
+    ) -> ChatTurnResponse:
         """Fall back to MockChatBackend."""
         from agents.tools import set_design_state
         from agents.orchestrator import MockChatBackend
 
-        # Keep design state consistent even in fallback path — pre-extract
-        # from user message so gap analysis reflects current turn data.
-        state = DesignState(**(design_state or {}))
+        state = design_state if isinstance(design_state, DesignState) else DesignState(**(design_state or {}))
         state = state.update_from_messages([], user_message=message)
         set_design_state(state)
 
         mock = MockChatBackend(output_dir=self.output_dir)
-        turn = mock.chat_turn(message, history, mentions, design_state=state, plan_context=plan_context)
-
-        # Gap analysis on the typed responses
-        gap_result = analyze_gaps(turn.responses)
-        question_cards = []
-        if gap_result.has_gaps:
-            cards = generate_question_cards(gap_result, state, user_message=message)
-            question_cards = [c.model_dump() for c in cards]
-
-        # Convert to dict for CrewOrchestrator's current dict-based return
-        result: dict = turn.model_dump()
-        result["question_cards"] = question_cards
+        result = mock.chat_turn(message, history, mentions, design_state=state, plan_context=plan_context)
+        if not result.question_cards:
+            gap_result = analyze_gaps(result.responses)
+            if gap_result.has_gaps:
+                result.question_cards = generate_question_cards(gap_result, state, user_message=message)
         return result

agents/gap_analyzer.py

@@ -6,11 +6,15 @@ and generates structured question cards for the UI to present to the user.
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from pydantic import BaseModel, Field
 
 from config.settings import settings
 from agents.definitions import AGENTS
-from agents.agent_flow import AgentResponse
+
+if TYPE_CHECKING:
+    from agents.agent_flow import AgentResponse
 
 
 # ── Models ────────────────────────────────────────────────────────────────────

tests/test_crew_orchestrator.py

@@ -1,7 +1,8 @@
 """Tests for CrewOrchestrator — CrewAI-based multi-agent orchestrator."""
 
 from agents.crew_orchestrator import CrewOrchestrator, _get_crewai_model, _is_plan_trigger
-from agents.agent_flow import WIKI_DIR
+from agents.agent_flow import ChatTurnResponse, PreviewData, WIKI_DIR
+from agents.design_state import DesignState, DesignPlan
 
 
 class TestGetCrewaiModel:
@@ -42,16 +43,14 @@ class TestCrewOrchestratorFallback:
         orch = CrewOrchestrator(backend_name="gemini", output_dir=tmp_output_dir)
         orch._crew_available = False
         result = orch.chat_turn("test", history=[])
-        assert "responses" in result
-        assert "preview" in result
-        assert "design_state" in result
+        assert isinstance(result, ChatTurnResponse)
 
     def test_response_format(self, tmp_output_dir):
         orch = CrewOrchestrator(backend_name="gemini", output_dir=tmp_output_dir)
         orch._crew_available = False
         result = orch.chat_turn("I need a bracket", history=[])
-        assert isinstance(result["responses"], list)
-        assert isinstance(result["design_state"], dict)
+        assert isinstance(result.responses, list)
+        assert isinstance(result.design_state, DesignState)
 
 
 class TestGetOrchestrator:
@@ -79,32 +78,24 @@ class TestGetOrchestrator:
 class TestGapAnalysis:
     def test_not_ready_produces_question_cards(self):
         orch = CrewOrchestrator(backend_name="mock")
-        result = orch.chat_turn(
-            message="generate a bracket",
-            history=[],
-            design_state={},
-        )
-        assert "question_cards" in result
+        result = orch.chat_turn(message="generate a bracket", history=[])
+        assert isinstance(result.question_cards, list)
 
     def test_no_question_cards_when_no_gaps(self):
         orch = CrewOrchestrator(backend_name="mock")
         result = orch.chat_turn(
-            message="I need a bracket",
-            history=[],
-            design_state={"material": "aluminum", "dimensions": {"width": 60}},
+            message="I need a bracket", history=[],
+            design_state=DesignState(material="aluminum", dimensions={"width": 60}),
         )
-        assert "question_cards" in result
-        assert isinstance(result["question_cards"], list)
+        assert isinstance(result.question_cards, list)
 
     def test_plan_trigger_includes_question_cards_key(self):
         orch = CrewOrchestrator(backend_name="mock")
         result = orch.chat_turn(
-            message="show plan",
-            history=[],
-            design_state={"material": "aluminum"},
+            message="show plan", history=[],
+            design_state=DesignState(material="aluminum"),
         )
-        assert "question_cards" in result
-        assert result["question_cards"] == []
+        assert result.question_cards == []
 
 
 class TestPlanningPhase:
@@ -113,62 +104,36 @@ class TestPlanningPhase:
     def test_manual_plan_trigger(self):
         """User typing a trigger keyword returns plan without running crew."""
         orch = CrewOrchestrator(backend_name="mock")
-        state_dict = {
-            "part_name": "bracket",
-            "material": "aluminum 6061",
-            "dimensions": {"width": 60, "height": 40, "depth": 20},
-            "axis_recommendation": "3-axis",
-        }
-        result = orch.chat_turn(
-            message="show plan",
-            history=[],
-            design_state=state_dict,
+        state = DesignState(
+            part_name="bracket", material="aluminum 6061",
+            dimensions={"width": 60, "height": 40, "depth": 20},
+            axis_recommendation="3-axis",
         )
-        assert result["design_state"]["phase"] == "planning"
-        assert result["design_state"]["plan"] is not None
-        assert result["design_state"]["plan"]["material"] == "aluminum 6061"
+        result = orch.chat_turn(message="show plan", history=[], design_state=state)
+        assert result.design_state.phase == "planning"
+        assert result.design_state.plan is not None
+        assert result.design_state.plan.material == "aluminum 6061"
 
     def test_approved_phase_keeps_approved(self):
         """When phase is approved, orchestrator keeps it for agent run."""
         orch = CrewOrchestrator(backend_name="mock")
-        plan_dict = {
-            "part_name": "bracket", "description": "test", "material": "aluminum",
-            "dimensions": {"width": 60}, "features": [], "constraints": [],
-            "axis_recommendation": "3-axis", "machining_notes": [],
-            "confidence_score": 9.0,
-        }
-        state_dict = {
-            "phase": "approved",
-            "plan": plan_dict,
-            "material": "aluminum",
-            "dimensions": {"width": 60},
-        }
-        result = orch.chat_turn(
-            message="Generate the approved design",
-            history=[],
-            design_state=state_dict,
+        plan = DesignPlan(
+            part_name="bracket", description="test", material="aluminum",
+            dimensions={"width": 60}, features=[], constraints=[],
+            axis_recommendation="3-axis", machining_notes=[], confidence_score=9.0,
         )
-        assert "responses" in result
+        state = DesignState(phase="approved", plan=plan, material="aluminum", dimensions={"width": 60})
+        result = orch.chat_turn(message="Generate the approved design", history=[], design_state=state)
+        assert isinstance(result.responses, list)
 
     def test_planning_phase_resets_on_message(self):
         """If phase is planning and user sends regular message, reset to exploring."""
         orch = CrewOrchestrator(backend_name="mock")
-        plan_dict = {
-            "part_name": "bracket", "description": "", "material": "steel",
-            "dimensions": {}, "features": [], "constraints": [],
-            "axis_recommendation": "", "machining_notes": [],
-            "confidence_score": 5.0,
-        }
-        state_dict = {
-            "phase": "planning",
-            "plan": plan_dict,
-            "material": "steel",
-        }
-        result = orch.chat_turn(
-            message="actually change the material",
-            history=[],
-            design_state=state_dict,
+        plan = DesignPlan(
+            part_name="bracket", description="", material="steel",
+            dimensions={}, features=[], constraints=[],
+            axis_recommendation="", machining_notes=[], confidence_score=5.0,
         )
-        # Should reset to exploring since user sent a regular message while in planning
-        ds = result["design_state"]
-        assert ds["phase"] in ("exploring", "planning")  # may auto-trigger again
+        state = DesignState(phase="planning", plan=plan, material="steel")
+        result = orch.chat_turn(message="actually change the material", history=[], design_state=state)
+        assert result.design_state.phase in ("exploring", "planning")