Coverage for src/local_deep_research/database/models/chat.py: 95%
60 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Chat models for conversations and messages.
3Domain-driven organization following news.py, research.py patterns.
5Chat schema:
6- ChatMessage.content is NOT NULL — every row stores its rendered text inline.
7- No CHECK constraint (content NOT NULL is sufficient).
8- Step rows live in ChatProgressStep, NOT in ChatMessage.
9- ChatSession.status is a typed Enum (ChatSessionStatus).
10"""
12import enum
13from datetime import datetime
15from sqlalchemy import (
16 JSON,
17 Column,
18 Enum,
19 ForeignKey,
20 Index,
21 Integer,
22 String,
23 Text,
24 UniqueConstraint,
25)
26from sqlalchemy.orm import Mapped, mapped_column, relationship
27from sqlalchemy_utc import UtcDateTime, utcnow
29from .base import Base
32class ChatRole(str, enum.Enum):
33 """Who authored a chat message."""
35 USER = "user"
36 ASSISTANT = "assistant"
39class ChatMessageType(str, enum.Enum):
40 """Role a message plays in the conversation flow.
42 Step rows (transient research progress) live in ChatProgressStep, not
43 in ChatMessage. The chat_message_type Enum no longer includes "step".
44 """
46 QUERY = "query"
47 FOLLOWUP = "followup"
48 RESPONSE = "response"
51class ChatSessionStatus(str, enum.Enum):
52 """Lifecycle state of a chat session."""
54 ACTIVE = "active"
55 ARCHIVED = "archived"
56 DELETED = "deleted"
59class ChatSession(Base):
60 """
61 Chat conversation session.
62 A session contains multiple messages and can span multiple research operations.
63 """
65 __tablename__ = "chat_sessions"
67 # UUID as primary key (consistent with ResearchHistory)
68 id: Mapped[str] = mapped_column(String(36), primary_key=True)
70 # Session metadata
71 title = Column(String(500))
72 # No create_constraint=True: the migration intentionally omits the
73 # CHECK constraint (matches the project pattern; enforcement is
74 # at the ORM/service layer via ChatSessionStatus(value)). Adding it
75 # here would create schema drift — fresh installs would have the
76 # CHECK, migrated DBs would not.
77 status: Mapped[ChatSessionStatus] = mapped_column(
78 Enum(
79 ChatSessionStatus,
80 values_callable=lambda obj: [e.value for e in obj],
81 name="chat_session_status",
82 ),
83 nullable=False,
84 default=ChatSessionStatus.ACTIVE,
85 server_default="active",
86 )
88 # Accumulated context across messages — used to build the LLM prompt for
89 # follow-up turns. Updated under row-level lock by ChatService.
90 accumulated_context = Column(JSON)
92 created_at = Column(UtcDateTime, nullable=False, default=utcnow())
94 # message_count counts surviving chat messages. Its atomic increment
95 # acquires the write lock, while sequence_number is allocated above both
96 # persisted and pending messages. Deleting earlier messages therefore
97 # reduces the count without reusing a surviving sequence number.
98 # Step rows do NOT increment this counter (they live in
99 # chat_progress_steps with their own per-research counter on
100 # research_history.step_count).
101 message_count = Column(
102 Integer, nullable=False, default=0, server_default="0"
103 )
105 # Relationships
106 messages = relationship(
107 "ChatMessage",
108 back_populates="session",
109 cascade="all, delete-orphan",
110 order_by="ChatMessage.sequence_number",
111 )
112 researches = relationship(
113 "ResearchHistory",
114 back_populates="chat_session",
115 # DB-level FK (research_history.chat_session_id ONDELETE SET NULL)
116 # handles the cleanup; no need for the ORM to emit a redundant
117 # UPDATE before the parent DELETE.
118 passive_deletes=True,
119 )
120 progress_steps = relationship(
121 "ChatProgressStep",
122 back_populates="session",
123 cascade="all, delete-orphan",
124 order_by="ChatProgressStep.created_at",
125 )
127 # Indexes for performance
128 __table_args__ = (
129 Index("idx_chat_session_status", "status"),
130 Index("idx_chat_session_created", "created_at"),
131 # Composite for the sidebar list_sessions hot query
132 # (WHERE status=? ORDER BY created_at DESC).
133 Index("idx_chat_session_status_created", "status", "created_at"),
134 )
136 def __repr__(self):
137 return (
138 f"<ChatSession(id='{self.id}', title='{self.title}', "
139 f"messages={self.message_count})>"
140 )
143class ChatMessage(Base):
144 """
145 Individual durable message within a chat session.
146 Content is always inline; research_id is a back-reference,
147 not a content source.
148 """
150 __tablename__ = "chat_messages"
152 # UUID as primary key
153 id: Mapped[str] = mapped_column(String(36), primary_key=True)
155 # Foreign keys
156 # No `index=True` here — see __table_args__ note below: single-column
157 # indexes are named explicitly so create_all and the 0010 migration
158 # produce identically-named indexes.
159 session_id: Mapped[str] = mapped_column(
160 String(36),
161 ForeignKey("chat_sessions.id", ondelete="CASCADE"),
162 nullable=False,
163 )
165 # Back-reference to research, if this turn triggered a research run.
166 # ON DELETE SET NULL: deleting the underlying research preserves the
167 # chat row (content is stored inline; the link just becomes stale).
168 # This is the snapshot semantic the schema delivers.
169 research_id: Mapped[str | None] = mapped_column(
170 String(36),
171 ForeignKey("research_history.id", ondelete="SET NULL"),
172 nullable=True,
173 )
175 # Message content
176 # See ChatSession.status note: no create_constraint=True to match the
177 # migration's deliberately-no-CHECK shape.
178 role = Column(
179 Enum(
180 ChatRole,
181 values_callable=lambda obj: [e.value for e in obj],
182 name="chat_role",
183 ),
184 nullable=False,
185 )
186 message_type = Column(
187 Enum(
188 ChatMessageType,
189 values_callable=lambda obj: [e.value for e in obj],
190 name="chat_message_type",
191 ),
192 nullable=False,
193 )
194 content = Column(Text, nullable=False)
196 # Sequence for ordering within a session
197 sequence_number = Column(Integer, nullable=False)
199 created_at: Mapped[datetime] = mapped_column(
200 UtcDateTime, nullable=False, default=utcnow()
201 )
203 # Relationships
204 session = relationship("ChatSession", back_populates="messages")
205 research = relationship("ResearchHistory", back_populates="chat_messages")
207 # Named Index in __table_args__ (NOT index=True on the column) so both
208 # the create_all path (fresh installs, test fixtures) and the migration
209 # path (0010 CHAT_INDEXES list) produce identically-named indexes —
210 # mirroring the ResearchResource pattern in models/research.py.
211 __table_args__ = (
212 UniqueConstraint(
213 "session_id",
214 "sequence_number",
215 name="uq_chat_message_session_seq",
216 ),
217 # Named indexes mirroring 0010 CHAT_INDEXES so create_all and the
218 # migration produce identically-named indexes (avoids drift between
219 # fresh-install and migrated DBs).
220 Index("ix_chat_messages_session_id", "session_id"),
221 Index("ix_chat_messages_research_id", "research_id"),
222 Index(
223 "ix_chat_messages_session_created",
224 "session_id",
225 "created_at",
226 ),
227 )
229 def __repr__(self):
230 return (
231 f"<ChatMessage(session='{self.session_id}', "
232 f"seq={self.sequence_number}, role='{self.role}')>"
233 )
236class ChatProgressStep(Base):
237 """
238 Transient research-progress step. Streamed live during a running research,
239 persisted so a session reload during in-flight research can rehydrate the
240 progress accordion.
242 Lives in its own table so step rows do not pollute conversation
243 sequencing or the chat_messages.message_count counter.
244 """
246 __tablename__ = "chat_progress_steps"
248 id = Column(String(36), primary_key=True)
250 # Step belongs to a research run. CASCADE: deleting the research drops
251 # all its steps (they're transient progress; nothing to preserve).
252 # No `index=True` — see __table_args__ below.
253 research_id = Column(
254 String(36),
255 ForeignKey("research_history.id", ondelete="CASCADE"),
256 nullable=False,
257 )
259 # Denormalized session FK for fast session-scoped queries (the resume
260 # path on chat reload reads steps by session). CASCADE: deleting the
261 # session drops all its steps.
262 # No `index=True` — see __table_args__ below.
263 session_id = Column(
264 String(36),
265 ForeignKey("chat_sessions.id", ondelete="CASCADE"),
266 nullable=False,
267 )
269 # Phase from research_service._STEP_PHASES (e.g. "search", "observation").
270 # Nullable so legacy rows migrated from chat_messages can be carried
271 # forward without phase data.
272 phase = Column(String(64), nullable=True)
274 content = Column(Text, nullable=False)
276 # Per-research sequence — allocated atomically by ChatService.add_progress_step
277 # via UPDATE...RETURNING on research_history.step_count.
278 sequence_number = Column(Integer, nullable=False)
280 created_at = Column(UtcDateTime, nullable=False, default=utcnow())
282 # Relationships
283 research = relationship("ResearchHistory", back_populates="progress_steps")
284 session = relationship("ChatSession", back_populates="progress_steps")
286 __table_args__ = (
287 UniqueConstraint(
288 "research_id",
289 "sequence_number",
290 name="uq_chat_progress_step_research_seq",
291 ),
292 # Named indexes mirroring 0010 CHAT_INDEXES so create_all and the
293 # migration produce identically-named indexes.
294 Index("ix_chat_progress_steps_research_id", "research_id"),
295 Index("ix_chat_progress_steps_session_id", "session_id"),
296 Index(
297 "ix_chat_progress_steps_session_created",
298 "session_id",
299 "created_at",
300 ),
301 )
303 def __repr__(self):
304 return (
305 f"<ChatProgressStep(research='{self.research_id}', "
306 f"seq={self.sequence_number}, phase='{self.phase}')>"
307 )