Coverage for backend/ahuora-builder-types/src/ahuora_builder_types/payloads/ml_request_schema.py: 88%

1from typing import Any, Literal 

2 

3from pydantic import BaseModel, Field, model_validator 

4 

5 

6class MLTrainRequestPayload(BaseModel): 

7 """Payload describing an object-storage-backed ML training request.""" 

8 

9 csv_bucket: str 

10 csv_key: str 

11 csv_delimiter: str | None = None 

12 input_labels: list[str] 

13 output_labels: list[str] 

14 task_id: int 

15 

16 

17class MLTrainingRegressionMetricPayload(BaseModel): 

18 """Regression metrics produced for one trained output column.""" 

19 

20 mean_squared_error: float 

21 r2_score: float 

22 

23 

24class MLTrainingChartPayload(BaseModel): 

25 """Chart payload produced for one trained output column.""" 

26 

27 min: float 

28 max: float 

29 qq_plot_data: str 

30 output_label: str 

31 

32 

33class MLTrainingResultPayload(BaseModel): 

34 """Structured ML training output stored on the task and ML model records.""" 

35 

36 surrogate_model: dict[str, Any] 

37 charts: list[MLTrainingChartPayload] 

38 metrics: list[MLTrainingRegressionMetricPayload] 

39 test_results_bucket: str 

40 test_results_key: str 

41 timing: dict[str, Any] = Field(default_factory=dict) 

42 

43 

44class MLTrainingCompletionPayload(BaseModel): 

45 """Result returned by the ML worker after training succeeds or fails.""" 

46 

47 json_response: MLTrainingResultPayload | None = None 

48 error: str | None = None 

49 log: str | None = None 

50 traceback: str | None = None 

51 task_id: int 

52 status: Literal["success", "error"] 

53 

54 @model_validator(mode="after") 

55 def validate_status_payload(self) -> "MLTrainingCompletionPayload": 

56 """Ensure success and error payloads carry the expected companion fields.""" 

57 if self.status == "success": 

58 if self.json_response is None: 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true

59 raise ValueError("Successful ML training payloads must include json_response.") 

60 if self.error is not None: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true

61 raise ValueError("Successful ML training payloads cannot include error text.") 

62 elif self.error is None: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 raise ValueError("Failed ML training payloads must include error text.") 

64 return self 

65