CallMeDaniel Claude Opus 4.6 (1M context) commited on
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>

Files changed (2) hide show
  1. core/pipeline.py +6 -4
  2. 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
- @dataclass
29
- class PipelineResult:
 
30
  prompt: str
31
  generated_code: str
32
  execution: ExecutionResult
33
  validation: Optional[CNCValidationResult] = None
34
- exported_files: dict[str, Path] = None
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