Coverage for src/local_deep_research/database/models/research.py: 94%
156 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Core research models for tasks, queries, and results.
3"""
5import enum
6import json
8from loguru import logger
9from sqlalchemy import (
10 JSON,
11 Column,
12 Enum,
13 Float,
14 ForeignKey,
15 Index,
16 Integer,
17 String,
18 Text,
19 text,
20)
21from sqlalchemy.orm import relationship
22from sqlalchemy.sql import operators
23from sqlalchemy.types import TypeDecorator
25from sqlalchemy_utc import UtcDateTime, utcnow
27from ...constants import ResearchStatus
28from .base import Base
31class LenientJSON(TypeDecorator):
32 """Read malformed legacy JSON text as absent without changing writes.
34 Releases before v0.5 persisted this column as TEXT through
35 ``json.dumps``, and the stock ``JSON`` result processor raises
36 ``JSONDecodeError`` while such a row is being loaded whenever the stored
37 text is not valid JSON (the empty string from an interrupted write, a
38 truncated ``json.dumps`` output, a ``str(dict)`` repr). That raise
39 happens during the query, above every handler ``try``, so
40 ``GET /history/status/{research_id}`` answered 500 for the row instead of
41 degrading to an empty log or metadata (#6533, #7017). A malformed
42 ``research_meta`` value could also make one row take down the entire
43 history list.
45 Reads therefore return ``None`` for an unparseable value — the absent
46 value readers of both columns already handle — and the load can no
47 longer fail. Writes are untouched: the stock ``JSON`` bind path still
48 serializes Python objects, so the DDL and the stored format are
49 unchanged and no migration is needed.
50 """
52 impl = JSON
53 cache_ok = True
55 def coerce_compared_value(self, op, value):
56 # TypeDecorator otherwise JSON-encodes a LIKE pattern as a JSON
57 # string. Subscription lookups compare against the stored TEXT form,
58 # so their pattern must be bound as an ordinary SQL string. Delegate
59 # other operators to JSON to preserve its key/index coercion rules.
60 if op in ( 60 ↛ 67line 60 didn't jump to line 67 because the condition on line 60 was always true
61 operators.like_op,
62 operators.not_like_op,
63 operators.ilike_op,
64 operators.not_ilike_op,
65 ):
66 return String()
67 return self.impl.coerce_compared_value(op, value)
69 def result_processor(self, dialect, coltype):
70 def process(value):
71 # Only TEXT carriers need decoding. A backend with a native
72 # JSON type hands back an already-parsed object — list, dict,
73 # or a scalar such as 5 / true / null — and must not be fed to
74 # json.loads, which would reject the non-string scalars and
75 # turn a perfectly valid value into None. LDR ships on
76 # SQLite/SQLCipher only, so this is defensive, not load-bearing.
77 if not isinstance(value, (str, bytes, bytearray)):
78 return value
79 try:
80 return json.loads(value)
81 except (TypeError, ValueError):
82 logger.warning(
83 "Unparseable JSON value in a research_history column; "
84 "treating it as absent"
85 )
86 return None
88 return process
91class ResearchTask(Base):
92 """
93 Main research tasks that users create.
94 This is the top-level object that contains all research activities.
95 """
97 __tablename__ = "research_tasks"
99 id = Column(Integer, primary_key=True)
100 title = Column(String(500), nullable=False)
101 description = Column(Text)
102 status = Column(
103 String(50), default="pending"
104 ) # pending, in_progress, completed, failed
105 priority = Column(Integer, default=0) # Higher number = higher priority
106 tags = Column(JSON) # List of tags for categorization
107 research_metadata = Column(JSON) # Flexible field for additional data
109 # Timestamps
110 created_at = Column(UtcDateTime, default=utcnow())
111 updated_at = Column(UtcDateTime, default=utcnow(), onupdate=utcnow())
112 started_at = Column(UtcDateTime)
113 completed_at = Column(UtcDateTime)
115 # Relationships
116 searches = relationship(
117 "SearchQuery",
118 back_populates="research_task",
119 cascade="all, delete-orphan",
120 )
121 results = relationship(
122 "SearchResult",
123 back_populates="research_task",
124 cascade="all, delete-orphan",
125 )
126 reports = relationship(
127 "Report", back_populates="research_task", cascade="all, delete-orphan"
128 )
130 def __repr__(self):
131 return f"<ResearchTask(title='{self.title}', status='{self.status}')>"
134class SearchQuery(Base):
135 """
136 Individual search queries executed as part of research tasks.
137 Tracks what was searched and when.
138 """
140 __tablename__ = "search_queries"
142 id = Column(Integer, primary_key=True)
143 research_task_id = Column(
144 Integer, ForeignKey("research_tasks.id", ondelete="CASCADE")
145 )
146 query = Column(Text, nullable=False)
147 search_engine = Column(String(50)) # google, bing, duckduckgo, etc.
148 search_type = Column(String(50)) # web, academic, news, etc.
149 parameters = Column(JSON) # Additional search parameters
151 # Status tracking
152 status = Column(
153 String(50), default="pending"
154 ) # pending, executing, completed, failed
155 error_message = Column(Text)
156 retry_count = Column(Integer, default=0)
158 # Timestamps
159 created_at = Column(UtcDateTime, default=utcnow())
160 executed_at = Column(UtcDateTime)
161 completed_at = Column(UtcDateTime)
163 # Relationships
164 research_task = relationship("ResearchTask", back_populates="searches")
165 results = relationship(
166 "SearchResult",
167 back_populates="search_query",
168 cascade="all, delete-orphan",
169 )
171 # Indexes for performance
172 __table_args__ = (
173 Index("idx_research_task_status", "research_task_id", "status"),
174 Index("idx_search_engine", "search_engine", "status"),
175 )
177 def __repr__(self):
178 return f"<SearchQuery(query='{self.query[:50]}...', status='{self.status}')>"
181class SearchResult(Base):
182 """
183 Individual search results from queries.
184 Stores both the initial result and any fetched content.
185 """
187 __tablename__ = "search_results"
189 id = Column(Integer, primary_key=True)
190 research_task_id = Column(
191 Integer, ForeignKey("research_tasks.id", ondelete="CASCADE")
192 )
193 search_query_id = Column(
194 Integer, ForeignKey("search_queries.id", ondelete="CASCADE")
195 )
197 # Basic result information
198 title = Column(String(500))
199 url = Column(Text, index=True) # Indexed for deduplication
200 snippet = Column(Text)
202 # Extended content
203 content = Column(Text) # Full content if fetched
204 content_type = Column(String(50)) # html, pdf, text, etc.
205 content_hash = Column(String(64)) # For deduplication
207 # Metadata
208 relevance_score = Column(Float) # Calculated relevance
209 position = Column(Integer) # Position in search results
210 domain = Column(String(255), index=True)
211 language = Column(String(10))
212 published_date = Column(UtcDateTime)
213 author = Column(String(255))
215 # Status tracking
216 fetch_status = Column(String(50)) # pending, fetched, failed, skipped
217 fetch_error = Column(Text)
219 # Timestamps
220 created_at = Column(UtcDateTime, default=utcnow())
221 fetched_at = Column(UtcDateTime)
223 # Relationships
224 research_task = relationship("ResearchTask", back_populates="results")
225 search_query = relationship("SearchQuery", back_populates="results")
227 # Indexes for performance
228 __table_args__ = (
229 Index("idx_task_relevance", "research_task_id", "relevance_score"),
230 Index("idx_content_hash", "content_hash"),
231 Index("idx_domain_task", "domain", "research_task_id"),
232 )
234 def __repr__(self):
235 return f"<SearchResult(title='{self.title[:50] if self.title else 'No title'}...', score={self.relevance_score})>"
238class ResearchMode(enum.Enum):
239 """Research modes available."""
241 QUICK = "quick"
242 DETAILED = "detailed"
245class ResearchResource(Base):
246 """Resources associated with research projects."""
248 __tablename__ = "research_resources"
250 id = Column(Integer, primary_key=True, autoincrement=True)
251 # Named Index() in __table_args__ below (NOT index=True on this
252 # column) so both the create_all path (fresh installs, test fixtures)
253 # and the migration path (0006:286-289) produce an identically-
254 # named index: ix_research_resources_research_id. Using index=True
255 # would give create_all an auto-named index that diverges from the
256 # migration's named one; that asymmetry previously meant fresh
257 # installs had no research_id index and ran full-table scans on
258 # 20+ call sites.
259 research_id = Column(
260 String(36),
261 ForeignKey("research_history.id", ondelete="CASCADE"),
262 nullable=False,
263 )
264 title = Column(Text)
265 url = Column(Text)
266 content_preview = Column(Text)
267 source_type = Column(Text)
268 resource_metadata = Column("metadata", JSON)
269 created_at = Column(String, nullable=False)
270 # `use_alter=True` breaks the model-level circular FK with Document
271 # (Document.resource_id → research_resources.id, ResearchResource.document_id
272 # → documents.id). Without it, every cold start emits an SAWarning from
273 # `Base.metadata.sorted_tables` because SQLAlchemy can't topologically
274 # order the two tables. With `use_alter=True`, this FK is emitted as a
275 # post-CREATE ALTER TABLE so both tables exist by the time the
276 # constraint is added. Migration 0005 already chose not to enforce
277 # this FK at the DB level on existing DBs (SQLite batch-alter
278 # limitation), so this only changes how `create_all()` emits the FK
279 # for fresh installs going forward.
280 document_id = Column(
281 String(36),
282 ForeignKey(
283 "documents.id",
284 ondelete="SET NULL",
285 use_alter=True,
286 name="fk_research_resources_document_id",
287 ),
288 nullable=True,
289 index=True,
290 )
292 # Relationships
293 research = relationship("ResearchHistory", back_populates="resources")
294 document = relationship("Document", foreign_keys=[document_id])
295 paper_appearance = relationship(
296 "PaperAppearance",
297 back_populates="resource",
298 uselist=False,
299 cascade="all, delete-orphan",
300 )
302 __table_args__ = (
303 Index("ix_research_resources_research_id", "research_id"),
304 )
306 def __repr__(self):
307 return f"<ResearchResource(title='{self.title}', url='{self.url}')>"
310class ResearchHistory(Base):
311 """
312 Research history table.
313 Tracks research sessions and their progress.
314 """
316 __tablename__ = "research_history"
318 # UUID as primary key
319 id = Column(String(36), primary_key=True)
320 # The search query.
321 query = Column(Text, nullable=False)
322 # The mode of research (e.g., 'quick_summary', 'detailed_report').
323 mode = Column(Text, nullable=False)
324 # Current status of the research.
325 status = Column(Text, nullable=False)
326 # The timestamp when the research started.
327 created_at = Column(Text, nullable=False)
328 # The timestamp when the research was completed.
329 completed_at = Column(Text)
330 # Duration of the research in seconds.
331 duration_seconds = Column(Integer)
332 # Path to the generated report.
333 report_path = Column(Text)
334 # Report content stored in database
335 report_content = Column(Text)
336 # Additional metadata about the research.
337 research_meta = Column(LenientJSON)
338 # Latest progress log message.
339 progress_log = Column(LenientJSON)
340 # Current progress of the research (as a percentage).
341 progress = Column(Integer)
342 # Title of the research report.
343 title = Column(Text)
345 # Optional link to chat session that triggered this research.
346 # Named Index() in __table_args__ below (NOT index=True on this
347 # column) so both the create_all path (fresh installs, test
348 # fixtures) and the migration path (in migration 0010) produce an
349 # identically-named index: ix_research_history_chat_session_id.
350 # The matching ResearchResource pattern above documents the same
351 # divergence risk in detail.
352 chat_session_id = Column(
353 String(36),
354 ForeignKey("chat_sessions.id", ondelete="SET NULL"),
355 nullable=True,
356 )
358 # Atomic counter for ChatService.add_progress_step's per-research
359 # sequence_number allocation (mirrors ChatSession.message_count).
360 step_count = Column(Integer, nullable=False, default=0, server_default="0")
362 # Relationships
363 resources = relationship(
364 "ResearchResource",
365 back_populates="research",
366 cascade="all, delete-orphan",
367 )
368 chat_session = relationship(
369 "ChatSession",
370 foreign_keys=[chat_session_id],
371 back_populates="researches",
372 )
373 chat_messages = relationship(
374 "ChatMessage",
375 back_populates="research",
376 passive_deletes=True,
377 )
378 progress_steps = relationship(
379 "ChatProgressStep",
380 back_populates="research",
381 passive_deletes=True,
382 )
384 __table_args__ = (
385 Index("ix_research_history_chat_session_id", "chat_session_id"),
386 # Partial unique index closing the SELECT-then-INSERT race in
387 # chat/routes.py::send_message — only one in-progress research
388 # per chat session, NULL chat_session_id rows unconstrained.
389 # Migration 0010 mirrors this index exactly.
390 Index(
391 "ux_research_history_chat_session_in_progress",
392 "chat_session_id",
393 unique=True,
394 sqlite_where=text(
395 "status = 'in_progress' AND chat_session_id IS NOT NULL"
396 ),
397 postgresql_where=text(
398 "status = 'in_progress' AND chat_session_id IS NOT NULL"
399 ),
400 ),
401 )
403 def __repr__(self):
404 return f"<ResearchHistory(query='{self.query[:50]}...', status={self.status})>"
407class Research(Base):
408 """
409 Modern research tracking with better type safety.
410 """
412 __tablename__ = "research"
414 id = Column(Integer, primary_key=True)
415 query = Column(String, nullable=False)
416 status = Column(
417 Enum(ResearchStatus), default=ResearchStatus.PENDING, nullable=False
418 )
419 mode = Column(
420 Enum(ResearchMode), default=ResearchMode.QUICK, nullable=False
421 )
422 created_at = Column(UtcDateTime, server_default=utcnow(), nullable=False)
423 updated_at = Column(
424 UtcDateTime, server_default=utcnow(), onupdate=utcnow(), nullable=False
425 )
426 progress = Column(Float, default=0.0, nullable=False)
427 start_time = Column(UtcDateTime, nullable=True)
428 end_time = Column(UtcDateTime, nullable=True)
429 error_message = Column(Text, nullable=True)
431 def __repr__(self):
432 return f"<Research(query='{self.query[:50]}...', status={self.status.value})>"
435class ResearchStrategy(Base):
436 """
437 Track which search strategy was used for each research.
438 """
440 __tablename__ = "research_strategies"
442 id = Column(Integer, primary_key=True)
443 # FK targets research_history.id (the live UUID-keyed table) rather than
444 # the dormant `research` table. save_research_strategy passes the
445 # research_history UUID, so the prior Integer FK to research.id raised
446 # FOREIGN KEY constraint failed on every commit once v1.6.0 turned on
447 # PRAGMA foreign_keys.
448 research_id = Column(
449 String(36),
450 ForeignKey("research_history.id", ondelete="CASCADE"),
451 nullable=False,
452 unique=True,
453 index=True,
454 )
455 strategy_name = Column(String(100), nullable=False, index=True)
456 created_at = Column(UtcDateTime, server_default=utcnow(), nullable=False)
458 def __repr__(self):
459 return f"<ResearchStrategy(research_id={self.research_id}, strategy={self.strategy_name})>"