Spaces:
Sleeping
Sleeping
File size: 8,072 Bytes
edd81ef a9d0aec eb89dca edd81ef a9d0aec edd81ef a9d0aec edd81ef a9d0aec edd81ef a9d0aec edd81ef a9d0aec edd81ef a9d0aec edd81ef eb89dca e43a92b eb89dca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 | """Tests for agents/design_state.py — state tracking and decision extraction."""
from agents.agent_flow import AgentResponse
from agents.design_state import DesignPlan, DesignState, compute_score, extract_decisions
class TestDesignState:
def test_empty_render(self):
state = DesignState()
assert state.render() == ""
def test_render_with_fields(self):
state = DesignState(
part_name="bracket",
material="aluminum 6061",
dimensions={"width": 60.0, "height": 40.0},
)
rendered = state.render()
assert "bracket" in rendered
assert "aluminum 6061" in rendered
assert "width=60.0mm" in rendered
def test_render_features(self):
state = DesignState(features=["4x M6 holes", "fillet"])
rendered = state.render()
assert "4x M6 holes" in rendered
def test_render_decisions_capped_at_5(self):
state = DesignState(decisions=[f"decision {i}" for i in range(10)])
rendered = state.render()
assert "decision 9" in rendered
assert "decision 4" not in rendered
class TestExtractDecisions:
def test_extracts_material(self):
responses = [
AgentResponse.from_agent("engineering", "I recommend aluminum 6061 for this application.")
]
state = extract_decisions(responses, DesignState())
assert "aluminum" in state.material.lower()
def test_extracts_dimensions_from_user(self):
responses = []
state = extract_decisions(responses, DesignState(), user_message="Make it 60mm wide and 40mm high")
assert state.dimensions.get("width") == 60.0
assert state.dimensions.get("height") == 40.0
def test_extracts_fastener_features(self):
responses = [
AgentResponse.from_agent("engineering", "I'll add 4x M6 clearance holes for mounting.")
]
state = extract_decisions(responses, DesignState())
assert any("M6" in f for f in state.features)
def test_extracts_axis_recommendation(self):
responses = [
AgentResponse.from_agent("cnc", "This part needs 5-axis machining due to the undercut.")
]
state = extract_decisions(responses, DesignState())
assert "5-axis" in state.axis_recommendation
def test_extracts_part_name(self):
responses = []
state = extract_decisions(responses, DesignState(), user_message="I need a servo bracket with M4 holes")
assert "servo bracket" in state.part_name.lower()
def test_preserves_existing_state(self):
existing = DesignState(material="steel", dimensions={"width": 50.0})
responses = [
AgentResponse.from_agent("engineering", "Height should be 30mm.")
]
updated = extract_decisions(responses, existing, user_message="add height")
assert updated.material == "steel"
assert updated.dimensions.get("width") == 50.0
def test_extracts_decisions_from_agreement(self):
responses = [
AgentResponse.from_agent("design", "I'd recommend an L-bracket form factor for this.")
]
state = extract_decisions(responses, DesignState())
assert len(state.decisions) > 0
def test_no_duplicate_features(self):
existing = DesignState(features=["4x M6 holes"])
responses = [
AgentResponse.from_agent("engineering", "The 4x M6 holes are properly specified.")
]
updated = extract_decisions(responses, existing)
m6_count = sum(1 for f in updated.features if "M6" in f)
assert m6_count == 1
class TestDesignPlan:
def test_create_from_state(self):
state = DesignState(
part_name="bracket",
description="mounting bracket",
material="aluminum 6061",
dimensions={"width": 60.0, "height": 40.0, "depth": 20.0},
features=["4x M6 holes"],
constraints=["min wall 3mm"],
axis_recommendation="3-axis",
decisions=["use aluminum"],
)
plan = DesignPlan.from_state(state, confidence_score=9.0)
assert plan.part_name == "bracket"
assert plan.material == "aluminum 6061"
assert plan.dimensions == {"width": 60.0, "height": 40.0, "depth": 20.0}
assert plan.confidence_score == 9.0
assert plan.machining_notes == []
def test_plan_render(self):
plan = DesignPlan(
part_name="bracket",
description="test",
material="aluminum 6061",
dimensions={"width": 60.0},
features=["4x M6 holes"],
constraints=[],
axis_recommendation="3-axis",
machining_notes=["No undercuts"],
confidence_score=9.0,
)
rendered = plan.render_approved()
assert "APPROVED DESIGN PLAN" in rendered
assert "aluminum 6061" in rendered
assert "No undercuts" in rendered
def test_plan_notes_default_empty(self):
plan = DesignPlan(part_name="test")
assert plan.notes == ""
def test_plan_notes_in_render(self):
plan = DesignPlan(
part_name="bracket",
material="aluminum",
notes="Check if 304 is overkill",
)
rendered = plan.render_approved()
assert "User Notes" in rendered
assert "Check if 304 is overkill" in rendered
def test_plan_notes_empty_not_in_render(self):
plan = DesignPlan(part_name="bracket", material="aluminum", notes="")
rendered = plan.render_approved()
assert "User Notes" not in rendered
def test_plan_from_state_no_notes(self):
state = DesignState(part_name="bracket", material="steel")
plan = DesignPlan.from_state(state, confidence_score=5.0)
assert plan.notes == ""
class TestComputeScore:
def test_empty_state_scores_zero(self):
assert compute_score(DesignState()) == 0.0
def test_material_scores_3(self):
state = DesignState(material="aluminum")
assert compute_score(state) == 3.0
def test_full_state_above_threshold(self):
state = DesignState(
part_name="bracket",
description="test bracket",
material="aluminum 6061",
dimensions={"width": 60.0, "height": 40.0, "depth": 20.0},
features=["4x M6 holes"],
constraints=["min wall 3mm"],
axis_recommendation="3-axis",
)
score = compute_score(state)
assert score >= 8.0
def test_dimension_cap_at_4(self):
state = DesignState(dimensions={
"width": 60, "height": 40, "depth": 20,
"length": 100, "diameter": 10, "radius": 5,
})
score = compute_score(state)
assert score == 4.0
def test_feature_cap_at_4(self):
state = DesignState(features=["a", "b", "c", "d", "e", "f"])
score = compute_score(state)
assert score == 4.0
def test_constraint_cap_at_2(self):
state = DesignState(constraints=["a", "b", "c", "d"])
score = compute_score(state)
assert score == 2.0
class TestDesignStatePhase:
def test_default_phase_exploring(self):
state = DesignState()
assert state.phase == "exploring"
assert state.plan is None
def test_phase_serialization(self):
plan = DesignPlan(
part_name="b", description="", material="steel",
dimensions={}, features=[], constraints=[],
axis_recommendation="", machining_notes=[],
confidence_score=5.0,
)
state = DesignState(phase="planning", plan=plan)
d = state.model_dump()
assert d["phase"] == "planning"
assert d["plan"]["material"] == "steel"
def test_roundtrip_from_dict(self):
state = DesignState(phase="approved", material="brass")
d = state.model_dump()
restored = DesignState(**d)
assert restored.phase == "approved"
assert restored.material == "brass"
|