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 model_type: str 

16 

17 

18class MLTrainingRegressionMetricPayload(BaseModel): 

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

20 

21 mean_squared_error: float 

22 r2_score: float 

23 

24 

25class MLTrainingChartPayload(BaseModel): 

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

27 

28 min: float 

29 max: float 

30 qq_plot_data: str 

31 output_label: str 

32 

33 

34class MLTrainingResultPayload(BaseModel): 

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

36 

37 surrogate_model: dict[str, Any] 

38 charts: list[MLTrainingChartPayload] 

39 metrics: list[MLTrainingRegressionMetricPayload] 

40 test_results_bucket: str 

41 test_results_key: str 

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

43 

44 

45class MLTrainingCompletionPayload(BaseModel): 

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

47 

48 json_response: MLTrainingResultPayload | None = None 

49 error: str | None = None 

50 log: str | None = None 

51 traceback: str | None = None 

52 task_id: int 

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

54 

55 @model_validator(mode="after") 

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

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

58 if self.status == "success": 

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

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

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

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

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

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

65 return self 

66