Coverage for llm_dataset_engine/core/models.py: 92%

1"""

2 Core data models for execution results and metadata.

3

4 These models represent the outputs and state information from pipeline

5 execution, following clean code principles with type safety.

6"""

7

8 from dataclasses import dataclass, field

9 from datetime import datetime

10 from decimal import Decimal

11 from typing import Any, Dict, List, Optional

12 from uuid import UUID, uuid4

13

14 import pandas as pd

15

16

17@dataclass

18 class LLMResponse:

19 """Response from a single LLM invocation."""

20

21 text: str

22 tokens_in: int

23 tokens_out: int

24 model: str

25 cost: Decimal

26 latency_ms: float

27 metadata: Dict[str, Any] = field(default_factory=dict)

28

29

30@dataclass

31 class CostEstimate:

32 """Cost estimation for pipeline execution."""

33

34 total_cost: Decimal

35 total_tokens: int

36 input_tokens: int

37 output_tokens: int

38 rows: int

39 breakdown_by_stage: Dict[str, Decimal] = field(default_factory=dict)

40 confidence: str = "estimate" # estimate, sample-based, actual

41

42

43@dataclass

44 class ProcessingStats:

45 """Statistics from pipeline execution."""

46

47 total_rows: int

48 processed_rows: int

49 failed_rows: int

50 skipped_rows: int

51 rows_per_second: float

52 total_duration_seconds: float

53 stage_durations: Dict[str, float] = field(default_factory=dict)

54

55

56@dataclass

57 class ErrorInfo:

58 """Information about an error during processing."""

59

60 row_index: int

61 stage_name: str

62 error_type: str

63 error_message: str

64 timestamp: datetime

65 context: Dict[str, Any] = field(default_factory=dict)

66

67

68@dataclass

69 class ExecutionResult:

70 """Complete result from pipeline execution."""

71

72 data: pd.DataFrame

73 metrics: ProcessingStats

74 costs: CostEstimate

75 errors: List[ErrorInfo] = field(default_factory=list)

76 execution_id: UUID = field(default_factory=uuid4)

77 start_time: datetime = field(default_factory=datetime.now)

78 end_time: Optional[datetime] = None

79 success: bool = True

80 metadata: Dict[str, Any] = field(default_factory=dict)

81

82 @property

83 def duration(self) -> float:

84 """Get execution duration in seconds."""

85 if self.end_time is None:

86 return 0.0

87 return (self.end_time - self.start_time).total_seconds()

88

89 @property

90 def error_rate(self) -> float:

91 """Get error rate as percentage."""

92 if self.metrics.total_rows == 0:

93 return 0.0

94 return (

95 self.metrics.failed_rows / self.metrics.total_rows

96 ) * 100

97

98

99@dataclass

100 class ValidationResult:

101 """Result from validation checks."""

102

103 is_valid: bool

104 errors: List[str] = field(default_factory=list)

105 warnings: List[str] = field(default_factory=list)

106

107 def add_error(self, error: str) -> None:

108 """Add an error message."""

109 self.errors.append(error)

110 self.is_valid = False

111

112 def add_warning(self, warning: str) -> None:

113 """Add a warning message."""

114 self.warnings.append(warning)

115

116

117@dataclass

118 class WriteConfirmation:

119 """Confirmation of successful data write."""

120

121 path: str

122 rows_written: int

123 success: bool

124 timestamp: datetime = field(default_factory=datetime.now)

125 metadata: Dict[str, Any] = field(default_factory=dict)

126

127

128@dataclass

129 class CheckpointInfo:

130 """Information about a checkpoint."""

131

132 session_id: UUID

133 checkpoint_path: str

134 row_index: int

135 stage_index: int

136 timestamp: datetime

137 size_bytes: int

138

139

140@dataclass

141 class RowMetadata:

142 """Metadata for a single row during processing."""

143

144 row_index: int

145 row_id: Optional[Any] = None

146 batch_id: Optional[int] = None

147 attempt: int = 1

148 custom: Dict[str, Any] = field(default_factory=dict)

149

150

151@dataclass

152 class PromptBatch:

153 """Batch of prompts for processing."""

154

155 prompts: List[str]

156 metadata: List[RowMetadata]

157 batch_id: int

158

159

160@dataclass

161 class ResponseBatch:

162 """Batch of responses from LLM."""

163

164 responses: List[str]

165 metadata: List[RowMetadata]

166 tokens_used: int

167 cost: Decimal

168 batch_id: int

169 latencies_ms: List[float] = field(default_factory=list)

170