Spaces:
Sleeping
Sleeping
Commit ·
2ebc6f1
1
Parent(s): b20b4ef
refactor: convert PipelineResult from dataclass to Pydantic BaseModel
Browse files- Replace @dataclass with Pydantic BaseModel
- Add model_config with arbitrary_types_allowed for ExecutionResult.cq.Workplane
- Change exported_files default to Field(default_factory=dict)
- Add TestPipelineResultModel with model_dump and default_exported_files tests
- All 9 tests pass
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- core/pipeline.py +6 -4
- tests/test_pipeline.py +21 -0
core/pipeline.py
CHANGED
|
@@ -9,10 +9,11 @@ Pipeline stages:
|
|
| 9 |
5. (Optional) Auto-retry with error feedback if execution fails
|
| 10 |
"""
|
| 11 |
|
| 12 |
-
from dataclasses import dataclass
|
| 13 |
from pathlib import Path
|
| 14 |
from typing import Optional
|
| 15 |
|
|
|
|
|
|
|
| 16 |
from core.cadquery_prompts import build_messages
|
| 17 |
from core.executor import ExecutionResult, execute_cadquery, export_all
|
| 18 |
from core.validator import validate_for_cnc, CNCValidationResult
|
|
@@ -25,13 +26,14 @@ from core.backends import (
|
|
| 25 |
# ── Pipeline ──────────────────────────────────────────────────────────────
|
| 26 |
|
| 27 |
|
| 28 |
-
|
| 29 |
-
|
|
|
|
| 30 |
prompt: str
|
| 31 |
generated_code: str
|
| 32 |
execution: ExecutionResult
|
| 33 |
validation: Optional[CNCValidationResult] = None
|
| 34 |
-
exported_files: dict[str, Path] =
|
| 35 |
retry_count: int = 0
|
| 36 |
|
| 37 |
def summary(self) -> str:
|
|
|
|
| 9 |
5. (Optional) Auto-retry with error feedback if execution fails
|
| 10 |
"""
|
| 11 |
|
|
|
|
| 12 |
from pathlib import Path
|
| 13 |
from typing import Optional
|
| 14 |
|
| 15 |
+
from pydantic import BaseModel, Field
|
| 16 |
+
|
| 17 |
from core.cadquery_prompts import build_messages
|
| 18 |
from core.executor import ExecutionResult, execute_cadquery, export_all
|
| 19 |
from core.validator import validate_for_cnc, CNCValidationResult
|
|
|
|
| 26 |
# ── Pipeline ──────────────────────────────────────────────────────────────
|
| 27 |
|
| 28 |
|
| 29 |
+
class PipelineResult(BaseModel):
|
| 30 |
+
model_config = {"arbitrary_types_allowed": True}
|
| 31 |
+
|
| 32 |
prompt: str
|
| 33 |
generated_code: str
|
| 34 |
execution: ExecutionResult
|
| 35 |
validation: Optional[CNCValidationResult] = None
|
| 36 |
+
exported_files: dict[str, Path] = Field(default_factory=dict)
|
| 37 |
retry_count: int = 0
|
| 38 |
|
| 39 |
def summary(self) -> str:
|
tests/test_pipeline.py
CHANGED
|
@@ -7,6 +7,7 @@ import pytest
|
|
| 7 |
from pathlib import Path
|
| 8 |
from core.pipeline import run_pipeline, PipelineResult
|
| 9 |
from core.backends import MockBackend
|
|
|
|
| 10 |
|
| 11 |
pytestmark = pytest.mark.requires_cadquery
|
| 12 |
|
|
@@ -76,3 +77,23 @@ class TestRunPipeline:
|
|
| 76 |
output_dir=tmp_output_dir,
|
| 77 |
)
|
| 78 |
assert result.execution.success is True
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 7 |
from pathlib import Path
|
| 8 |
from core.pipeline import run_pipeline, PipelineResult
|
| 9 |
from core.backends import MockBackend
|
| 10 |
+
from core.executor import ExecutionResult
|
| 11 |
|
| 12 |
pytestmark = pytest.mark.requires_cadquery
|
| 13 |
|
|
|
|
| 77 |
output_dir=tmp_output_dir,
|
| 78 |
)
|
| 79 |
assert result.execution.success is True
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
class TestPipelineResultModel:
|
| 83 |
+
def test_model_dump(self):
|
| 84 |
+
exec_result = ExecutionResult(success=True, volume=1000.0, bounding_box=(10, 10, 10), face_count=6, edge_count=12)
|
| 85 |
+
pr = PipelineResult(
|
| 86 |
+
prompt="test",
|
| 87 |
+
generated_code="code",
|
| 88 |
+
execution=exec_result,
|
| 89 |
+
retry_count=0,
|
| 90 |
+
)
|
| 91 |
+
d = pr.model_dump()
|
| 92 |
+
assert d["prompt"] == "test"
|
| 93 |
+
assert d["retry_count"] == 0
|
| 94 |
+
|
| 95 |
+
def test_default_exported_files(self):
|
| 96 |
+
exec_result = ExecutionResult(success=False, error="fail")
|
| 97 |
+
pr = PipelineResult(prompt="test", generated_code="code", execution=exec_result)
|
| 98 |
+
assert pr.exported_files == {}
|
| 99 |
+
assert pr.validation is None
|