Coverage for src/local_deep_research/database/models/benchmark.py: 100%

1"""Database models for benchmark system.""" 

2 

3import enum 

4 

5from sqlalchemy import ( 

6 JSON, 

7 Boolean, 

8 Column, 

9 Enum, 

10 Float, 

11 ForeignKey, 

12 Index, 

13 Integer, 

14 String, 

15 Text, 

16 UniqueConstraint, 

17) 

18from sqlalchemy.orm import relationship 

19from sqlalchemy_utc import UtcDateTime, utcnow 

20 

21 

22# Use the same base as the main app 

23from . import Base 

24 

25 

26class BenchmarkStatus(enum.Enum): 

27 """Status of a benchmark run.""" 

28 

29 PENDING = "pending" 

30 IN_PROGRESS = "in_progress" 

31 COMPLETED = "completed" 

32 FAILED = "failed" 

33 CANCELLED = "cancelled" 

34 PAUSED = "paused" 

35 

36 

37class DatasetType(enum.Enum): 

38 """Supported dataset types.""" 

39 

40 SIMPLEQA = "simpleqa" 

41 BROWSECOMP = "browsecomp" 

42 XBENCH_DEEPSEARCH = "xbench_deepsearch" 

43 CUSTOM = "custom" 

44 

45 

46class BenchmarkRun(Base): 

47 """Main benchmark run metadata.""" 

48 

49 __tablename__ = "benchmark_runs" 

50 

51 id = Column(Integer, primary_key=True) 

52 

53 # Run identification 

54 run_name = Column(String(255), nullable=True) # User-friendly name 

55 config_hash = Column( 

56 String(16), nullable=False, index=True 

57 ) # For compatibility matching 

58 query_hash_list = Column( 

59 JSON, nullable=False 

60 ) # List of query hashes to avoid duplication 

61 

62 # Configuration 

63 search_config = Column( 

64 JSON, nullable=False 

65 ) # Complete search configuration 

66 evaluation_config = Column(JSON, nullable=False) # Evaluation settings 

67 datasets_config = Column( 

68 JSON, nullable=False 

69 ) # Dataset selection and quantities 

70 

71 # Status and timing 

72 status = Column( 

73 Enum(BenchmarkStatus), default=BenchmarkStatus.PENDING, nullable=False 

74 ) 

75 created_at = Column(UtcDateTime, default=utcnow(), nullable=False) 

76 updated_at = Column( 

77 UtcDateTime, default=utcnow(), onupdate=utcnow(), nullable=False 

78 ) 

79 start_time = Column(UtcDateTime, nullable=True) 

80 end_time = Column(UtcDateTime, nullable=True) 

81 

82 # Provenance — captured at start_benchmark time so the YAML download 

83 # reflects the version/settings that ran the benchmark, not the version 

84 # at download time. NULL on rows created before migration 0014. 

85 ldr_version = Column(String(32), nullable=True) 

86 settings_snapshot = Column(JSON, nullable=True) 

87 

88 # Progress tracking 

89 total_examples = Column(Integer, default=0, nullable=False) 

90 completed_examples = Column(Integer, default=0, nullable=False) 

91 failed_examples = Column(Integer, default=0, nullable=False) 

92 

93 # Results summary 

94 overall_accuracy = Column(Float, nullable=True) 

95 processing_rate = Column(Float, nullable=True) # Examples per minute 

96 

97 # Error handling 

98 error_message = Column(Text, nullable=True) 

99 

100 # Relationships 

101 results = relationship( 

102 "BenchmarkResult", 

103 back_populates="benchmark_run", 

104 cascade="all, delete-orphan", 

105 lazy="dynamic", 

106 ) 

107 progress_updates = relationship( 

108 "BenchmarkProgress", 

109 back_populates="benchmark_run", 

110 cascade="all, delete-orphan", 

111 lazy="dynamic", 

112 ) 

113 

114 # Indexes for performance and extend existing 

115 __table_args__ = ( 

116 Index("idx_benchmark_runs_config_hash", "config_hash"), 

117 Index("idx_benchmark_runs_status_created", "status", "created_at"), 

118 {"extend_existing": True}, 

119 ) 

120 

121 

122class BenchmarkResult(Base): 

123 """Individual benchmark result for a single question.""" 

124 

125 __tablename__ = "benchmark_results" 

126 

127 id = Column(Integer, primary_key=True) 

128 

129 # Foreign key 

130 benchmark_run_id = Column( 

131 Integer, 

132 ForeignKey("benchmark_runs.id", ondelete="CASCADE"), 

133 nullable=False, 

134 index=True, 

135 ) 

136 

137 # Question identification 

138 example_id = Column(String(255), nullable=False) # Original dataset ID 

139 query_hash = Column( 

140 String(32), nullable=False, index=True 

141 ) # For deduplication 

142 dataset_type = Column(Enum(DatasetType), nullable=False) 

143 research_id = Column( 

144 String(36), nullable=True, index=True 

145 ) # UUID string or converted integer 

146 

147 # Question and answer 

148 question = Column(Text, nullable=False) 

149 correct_answer = Column(Text, nullable=False) 

150 

151 # Research results 

152 response = Column(Text, nullable=True) 

153 extracted_answer = Column(Text, nullable=True) 

154 confidence = Column(String(10), nullable=True) 

155 processing_time = Column(Float, nullable=True) 

156 sources = Column(JSON, nullable=True) 

157 

158 # Evaluation results 

159 is_correct = Column(Boolean, nullable=True) 

160 graded_confidence = Column(String(10), nullable=True) 

161 grader_response = Column(Text, nullable=True) 

162 

163 # Timestamps 

164 created_at = Column(UtcDateTime, default=utcnow(), nullable=False) 

165 completed_at = Column(UtcDateTime, nullable=True) 

166 

167 # Error handling 

168 research_error = Column(Text, nullable=True) 

169 evaluation_error = Column(Text, nullable=True) 

170 

171 # Additional metadata 

172 task_index = Column(Integer, nullable=True) # Order in processing 

173 result_metadata = Column(JSON, nullable=True) # Additional data 

174 

175 # Relationships 

176 benchmark_run = relationship("BenchmarkRun", back_populates="results") 

177 

178 # Indexes for performance 

179 __table_args__ = ( 

180 Index( 

181 "idx_benchmark_results_run_dataset", 

182 "benchmark_run_id", 

183 "dataset_type", 

184 ), 

185 Index("idx_benchmark_results_query_hash", "query_hash"), 

186 Index("idx_benchmark_results_completed", "completed_at"), 

187 UniqueConstraint( 

188 "benchmark_run_id", "query_hash", name="uix_run_query" 

189 ), 

190 {"extend_existing": True}, 

191 ) 

192 

193 

194class BenchmarkConfig(Base): 

195 """Saved benchmark configurations for reuse.""" 

196 

197 __tablename__ = "benchmark_configs" 

198 

199 id = Column(Integer, primary_key=True) 

200 

201 # Configuration details 

202 name = Column(String(255), nullable=False) 

203 description = Column(Text, nullable=True) 

204 config_hash = Column(String(16), nullable=False, index=True) 

205 

206 # Configuration data 

207 search_config = Column(JSON, nullable=False) 

208 evaluation_config = Column(JSON, nullable=False) 

209 datasets_config = Column(JSON, nullable=False) 

210 

211 # Metadata 

212 created_at = Column(UtcDateTime, default=utcnow(), nullable=False) 

213 updated_at = Column( 

214 UtcDateTime, default=utcnow(), onupdate=utcnow(), nullable=False 

215 ) 

216 is_default = Column(Boolean, default=False, nullable=False) 

217 is_public = Column(Boolean, default=True, nullable=False) 

218 

219 # Usage tracking 

220 usage_count = Column(Integer, default=0, nullable=False) 

221 last_used = Column(UtcDateTime, nullable=True) 

222 

223 # Performance data (if available) 

224 best_accuracy = Column(Float, nullable=True) 

225 avg_processing_rate = Column(Float, nullable=True) 

226 

227 # Indexes 

228 __table_args__ = ( 

229 Index("idx_benchmark_configs_name", "name"), 

230 Index("idx_benchmark_configs_hash", "config_hash"), 

231 Index("idx_benchmark_configs_default", "is_default"), 

232 {"extend_existing": True}, 

233 ) 

234 

235 

236class BenchmarkProgress(Base): 

237 """Real-time progress tracking for benchmark runs.""" 

238 

239 __tablename__ = "benchmark_progress" 

240 

241 id = Column(Integer, primary_key=True) 

242 

243 # Foreign key 

244 benchmark_run_id = Column( 

245 Integer, 

246 ForeignKey("benchmark_runs.id", ondelete="CASCADE"), 

247 nullable=False, 

248 index=True, 

249 ) 

250 

251 # Progress data 

252 timestamp = Column(UtcDateTime, default=utcnow(), nullable=False) 

253 completed_examples = Column(Integer, nullable=False) 

254 total_examples = Column(Integer, nullable=False) 

255 

256 # Accuracy tracking 

257 overall_accuracy = Column(Float, nullable=True) 

258 dataset_accuracies = Column(JSON, nullable=True) # Per-dataset accuracy 

259 

260 # Performance metrics 

261 processing_rate = Column(Float, nullable=True) # Examples per minute 

262 estimated_completion = Column(UtcDateTime, nullable=True) 

263 

264 # Current status 

265 current_dataset = Column(Enum(DatasetType), nullable=True) 

266 current_example_id = Column(String(255), nullable=True) 

267 

268 # Additional metrics 

269 memory_usage = Column(Float, nullable=True) # MB 

270 cpu_usage = Column(Float, nullable=True) # Percentage 

271 

272 # Relationships 

273 benchmark_run = relationship( 

274 "BenchmarkRun", back_populates="progress_updates" 

275 ) 

276 

277 # Indexes for real-time queries 

278 __table_args__ = ( 

279 Index( 

280 "idx_benchmark_progress_run_time", "benchmark_run_id", "timestamp" 

281 ), 

282 {"extend_existing": True}, 

283 )