Coverage for src/local_deep_research/database/models/research.py: 95%

1""" 

2Core research models for tasks, queries, and results. 

3""" 

4 

5import enum 

6 

7from sqlalchemy import ( 

8 JSON, 

9 Column, 

10 Enum, 

11 Float, 

12 ForeignKey, 

13 Index, 

14 Integer, 

15 String, 

16 Text, 

17 text, 

18) 

19from sqlalchemy.orm import relationship 

20 

21from sqlalchemy_utc import UtcDateTime, utcnow 

22 

23from ...constants import ResearchStatus 

24from .base import Base 

25 

26 

27class ResearchTask(Base): 

28 """ 

29 Main research tasks that users create. 

30 This is the top-level object that contains all research activities. 

31 """ 

32 

33 __tablename__ = "research_tasks" 

34 

35 id = Column(Integer, primary_key=True) 

36 title = Column(String(500), nullable=False) 

37 description = Column(Text) 

38 status = Column( 

39 String(50), default="pending" 

40 ) # pending, in_progress, completed, failed 

41 priority = Column(Integer, default=0) # Higher number = higher priority 

42 tags = Column(JSON) # List of tags for categorization 

43 research_metadata = Column(JSON) # Flexible field for additional data 

44 

45 # Timestamps 

46 created_at = Column(UtcDateTime, default=utcnow()) 

47 updated_at = Column(UtcDateTime, default=utcnow(), onupdate=utcnow()) 

48 started_at = Column(UtcDateTime) 

49 completed_at = Column(UtcDateTime) 

50 

51 # Relationships 

52 searches = relationship( 

53 "SearchQuery", 

54 back_populates="research_task", 

55 cascade="all, delete-orphan", 

56 ) 

57 results = relationship( 

58 "SearchResult", 

59 back_populates="research_task", 

60 cascade="all, delete-orphan", 

61 ) 

62 reports = relationship( 

63 "Report", back_populates="research_task", cascade="all, delete-orphan" 

64 ) 

65 

66 def __repr__(self): 

67 return f"<ResearchTask(title='{self.title}', status='{self.status}')>" 

68 

69 

70class SearchQuery(Base): 

71 """ 

72 Individual search queries executed as part of research tasks. 

73 Tracks what was searched and when. 

74 """ 

75 

76 __tablename__ = "search_queries" 

77 

78 id = Column(Integer, primary_key=True) 

79 research_task_id = Column( 

80 Integer, ForeignKey("research_tasks.id", ondelete="CASCADE") 

81 ) 

82 query = Column(Text, nullable=False) 

83 search_engine = Column(String(50)) # google, bing, duckduckgo, etc. 

84 search_type = Column(String(50)) # web, academic, news, etc. 

85 parameters = Column(JSON) # Additional search parameters 

86 

87 # Status tracking 

88 status = Column( 

89 String(50), default="pending" 

90 ) # pending, executing, completed, failed 

91 error_message = Column(Text) 

92 retry_count = Column(Integer, default=0) 

93 

94 # Timestamps 

95 created_at = Column(UtcDateTime, default=utcnow()) 

96 executed_at = Column(UtcDateTime) 

97 completed_at = Column(UtcDateTime) 

98 

99 # Relationships 

100 research_task = relationship("ResearchTask", back_populates="searches") 

101 results = relationship( 

102 "SearchResult", 

103 back_populates="search_query", 

104 cascade="all, delete-orphan", 

105 ) 

106 

107 # Indexes for performance 

108 __table_args__ = ( 

109 Index("idx_research_task_status", "research_task_id", "status"), 

110 Index("idx_search_engine", "search_engine", "status"), 

111 ) 

112 

113 def __repr__(self): 

114 return f"<SearchQuery(query='{self.query[:50]}...', status='{self.status}')>" 

115 

116 

117class SearchResult(Base): 

118 """ 

119 Individual search results from queries. 

120 Stores both the initial result and any fetched content. 

121 """ 

122 

123 __tablename__ = "search_results" 

124 

125 id = Column(Integer, primary_key=True) 

126 research_task_id = Column( 

127 Integer, ForeignKey("research_tasks.id", ondelete="CASCADE") 

128 ) 

129 search_query_id = Column( 

130 Integer, ForeignKey("search_queries.id", ondelete="CASCADE") 

131 ) 

132 

133 # Basic result information 

134 title = Column(String(500)) 

135 url = Column(Text, index=True) # Indexed for deduplication 

136 snippet = Column(Text) 

137 

138 # Extended content 

139 content = Column(Text) # Full content if fetched 

140 content_type = Column(String(50)) # html, pdf, text, etc. 

141 content_hash = Column(String(64)) # For deduplication 

142 

143 # Metadata 

144 relevance_score = Column(Float) # Calculated relevance 

145 position = Column(Integer) # Position in search results 

146 domain = Column(String(255), index=True) 

147 language = Column(String(10)) 

148 published_date = Column(UtcDateTime) 

149 author = Column(String(255)) 

150 

151 # Status tracking 

152 fetch_status = Column(String(50)) # pending, fetched, failed, skipped 

153 fetch_error = Column(Text) 

154 

155 # Timestamps 

156 created_at = Column(UtcDateTime, default=utcnow()) 

157 fetched_at = Column(UtcDateTime) 

158 

159 # Relationships 

160 research_task = relationship("ResearchTask", back_populates="results") 

161 search_query = relationship("SearchQuery", back_populates="results") 

162 

163 # Indexes for performance 

164 __table_args__ = ( 

165 Index("idx_task_relevance", "research_task_id", "relevance_score"), 

166 Index("idx_content_hash", "content_hash"), 

167 Index("idx_domain_task", "domain", "research_task_id"), 

168 ) 

169 

170 def __repr__(self): 

171 return f"<SearchResult(title='{self.title[:50] if self.title else 'No title'}...', score={self.relevance_score})>" 

172 

173 

174class ResearchMode(enum.Enum): 

175 """Research modes available.""" 

176 

177 QUICK = "quick" 

178 DETAILED = "detailed" 

179 

180 

181class ResearchResource(Base): 

182 """Resources associated with research projects.""" 

183 

184 __tablename__ = "research_resources" 

185 

186 id = Column(Integer, primary_key=True, autoincrement=True) 

187 # Named Index() in __table_args__ below (NOT index=True on this 

188 # column) so both the create_all path (fresh installs, test fixtures) 

189 # and the migration path (0006:286-289) produce an identically- 

190 # named index: ix_research_resources_research_id. Using index=True 

191 # would give create_all an auto-named index that diverges from the 

192 # migration's named one; that asymmetry previously meant fresh 

193 # installs had no research_id index and ran full-table scans on 

194 # 20+ call sites. 

195 research_id = Column( 

196 String(36), 

197 ForeignKey("research_history.id", ondelete="CASCADE"), 

198 nullable=False, 

199 ) 

200 title = Column(Text) 

201 url = Column(Text) 

202 content_preview = Column(Text) 

203 source_type = Column(Text) 

204 resource_metadata = Column("metadata", JSON) 

205 created_at = Column(String, nullable=False) 

206 # `use_alter=True` breaks the model-level circular FK with Document 

207 # (Document.resource_id → research_resources.id, ResearchResource.document_id 

208 # → documents.id). Without it, every cold start emits an SAWarning from 

209 # `Base.metadata.sorted_tables` because SQLAlchemy can't topologically 

210 # order the two tables. With `use_alter=True`, this FK is emitted as a 

211 # post-CREATE ALTER TABLE so both tables exist by the time the 

212 # constraint is added. Migration 0005 already chose not to enforce 

213 # this FK at the DB level on existing DBs (SQLite batch-alter 

214 # limitation), so this only changes how `create_all()` emits the FK 

215 # for fresh installs going forward. 

216 document_id = Column( 

217 String(36), 

218 ForeignKey( 

219 "documents.id", 

220 ondelete="SET NULL", 

221 use_alter=True, 

222 name="fk_research_resources_document_id", 

223 ), 

224 nullable=True, 

225 index=True, 

226 ) 

227 

228 # Relationships 

229 research = relationship("ResearchHistory", back_populates="resources") 

230 document = relationship("Document", foreign_keys=[document_id]) 

231 paper_appearance = relationship( 

232 "PaperAppearance", 

233 back_populates="resource", 

234 uselist=False, 

235 cascade="all, delete-orphan", 

236 ) 

237 

238 __table_args__ = ( 

239 Index("ix_research_resources_research_id", "research_id"), 

240 ) 

241 

242 def __repr__(self): 

243 return f"<ResearchResource(title='{self.title}', url='{self.url}')>" 

244 

245 

246class ResearchHistory(Base): 

247 """ 

248 Research history table. 

249 Tracks research sessions and their progress. 

250 """ 

251 

252 __tablename__ = "research_history" 

253 

254 # UUID as primary key 

255 id = Column(String(36), primary_key=True) 

256 # The search query. 

257 query = Column(Text, nullable=False) 

258 # The mode of research (e.g., 'quick_summary', 'detailed_report'). 

259 mode = Column(Text, nullable=False) 

260 # Current status of the research. 

261 status = Column(Text, nullable=False) 

262 # The timestamp when the research started. 

263 created_at = Column(Text, nullable=False) 

264 # The timestamp when the research was completed. 

265 completed_at = Column(Text) 

266 # Duration of the research in seconds. 

267 duration_seconds = Column(Integer) 

268 # Path to the generated report. 

269 report_path = Column(Text) 

270 # Report content stored in database 

271 report_content = Column(Text) 

272 # Additional metadata about the research. 

273 research_meta = Column(JSON) 

274 # Latest progress log message. 

275 progress_log = Column(JSON) 

276 # Current progress of the research (as a percentage). 

277 progress = Column(Integer) 

278 # Title of the research report. 

279 title = Column(Text) 

280 

281 # Optional link to chat session that triggered this research. 

282 # Named Index() in __table_args__ below (NOT index=True on this 

283 # column) so both the create_all path (fresh installs, test 

284 # fixtures) and the migration path (in migration 0010) produce an 

285 # identically-named index: ix_research_history_chat_session_id. 

286 # The matching ResearchResource pattern above documents the same 

287 # divergence risk in detail. 

288 chat_session_id = Column( 

289 String(36), 

290 ForeignKey("chat_sessions.id", ondelete="SET NULL"), 

291 nullable=True, 

292 ) 

293 

294 # Atomic counter for ChatService.add_progress_step's per-research 

295 # sequence_number allocation (mirrors ChatSession.message_count). 

296 step_count = Column(Integer, nullable=False, default=0, server_default="0") 

297 

298 # Relationships 

299 resources = relationship( 

300 "ResearchResource", 

301 back_populates="research", 

302 cascade="all, delete-orphan", 

303 ) 

304 chat_session = relationship( 

305 "ChatSession", 

306 foreign_keys=[chat_session_id], 

307 back_populates="researches", 

308 ) 

309 chat_messages = relationship( 

310 "ChatMessage", 

311 back_populates="research", 

312 passive_deletes=True, 

313 ) 

314 progress_steps = relationship( 

315 "ChatProgressStep", 

316 back_populates="research", 

317 passive_deletes=True, 

318 ) 

319 

320 __table_args__ = ( 

321 Index("ix_research_history_chat_session_id", "chat_session_id"), 

322 # Partial unique index closing the SELECT-then-INSERT race in 

323 # chat/routes.py::send_message — only one in-progress research 

324 # per chat session, NULL chat_session_id rows unconstrained. 

325 # Migration 0010 mirrors this index exactly. 

326 Index( 

327 "ux_research_history_chat_session_in_progress", 

328 "chat_session_id", 

329 unique=True, 

330 sqlite_where=text( 

331 "status = 'in_progress' AND chat_session_id IS NOT NULL" 

332 ), 

333 postgresql_where=text( 

334 "status = 'in_progress' AND chat_session_id IS NOT NULL" 

335 ), 

336 ), 

337 ) 

338 

339 def __repr__(self): 

340 return f"<ResearchHistory(query='{self.query[:50]}...', status={self.status})>" 

341 

342 

343class Research(Base): 

344 """ 

345 Modern research tracking with better type safety. 

346 """ 

347 

348 __tablename__ = "research" 

349 

350 id = Column(Integer, primary_key=True) 

351 query = Column(String, nullable=False) 

352 status = Column( 

353 Enum(ResearchStatus), default=ResearchStatus.PENDING, nullable=False 

354 ) 

355 mode = Column( 

356 Enum(ResearchMode), default=ResearchMode.QUICK, nullable=False 

357 ) 

358 created_at = Column(UtcDateTime, server_default=utcnow(), nullable=False) 

359 updated_at = Column( 

360 UtcDateTime, server_default=utcnow(), onupdate=utcnow(), nullable=False 

361 ) 

362 progress = Column(Float, default=0.0, nullable=False) 

363 start_time = Column(UtcDateTime, nullable=True) 

364 end_time = Column(UtcDateTime, nullable=True) 

365 error_message = Column(Text, nullable=True) 

366 

367 def __repr__(self): 

368 return f"<Research(query='{self.query[:50]}...', status={self.status.value})>" 

369 

370 

371class ResearchStrategy(Base): 

372 """ 

373 Track which search strategy was used for each research. 

374 """ 

375 

376 __tablename__ = "research_strategies" 

377 

378 id = Column(Integer, primary_key=True) 

379 # FK targets research_history.id (the live UUID-keyed table) rather than 

380 # the dormant `research` table. save_research_strategy passes the 

381 # research_history UUID, so the prior Integer FK to research.id raised 

382 # FOREIGN KEY constraint failed on every commit once v1.6.0 turned on 

383 # PRAGMA foreign_keys. 

384 research_id = Column( 

385 String(36), 

386 ForeignKey("research_history.id", ondelete="CASCADE"), 

387 nullable=False, 

388 unique=True, 

389 index=True, 

390 ) 

391 strategy_name = Column(String(100), nullable=False, index=True) 

392 created_at = Column(UtcDateTime, server_default=utcnow(), nullable=False) 

393 

394 def __repr__(self): 

395 return f"<ResearchStrategy(research_id={self.research_id}, strategy={self.strategy_name})>"