Coverage for src/local_deep_research/chat/context.py: 94%
110 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2ChatContextManager - Custom context building for multi-turn conversations.
4This is DIFFERENT from FollowUpResearchService (single parent-child context).
5Multi-turn chat requires:
6- Rolling window of recent messages
7- Accumulated findings across conversation
8- Source deduplication across turns
9- Context summarization for long conversations
10"""
12from typing import Dict, Any, List, Optional
14from loguru import logger
16from ..config.thread_settings import get_setting_from_snapshot
19class ChatContextManager:
20 """
21 Build context from multi-turn conversation history.
23 Handles: rolling window, summarization, context accumulation.
24 Different from follow-up: accumulates from MULTIPLE previous turns.
25 """
27 MAX_FINDINGS_TO_INCLUDE = 5 # Recent findings to include
29 # Limits for the query-focused conversation summary that becomes the
30 # follow-up prompt's "previous findings" block.
31 CONTEXT_SUMMARY_MAX_SENTENCES = 8
32 CONTEXT_SUMMARY_MAX_CHARS = 2000
33 # Transcript char budget kept below BaseSummarizer.INPUT_TRUNCATE_CHARS
34 # (8000) so the summarizer's own truncation never has to drop the most
35 # recent turns — we trim oldest-first ourselves below.
36 CONTEXT_INPUT_CHAR_BUDGET = 7500
38 # Default for the chat.followup_context_mode setting (summary | raw |
39 # full | none) — what prior context a follow-up turn receives.
40 DEFAULT_FOLLOWUP_CONTEXT_MODE = "summary"
42 def __init__(
43 self,
44 session_id: str,
45 messages: List[Dict[str, Any]],
46 accumulated_context: Optional[Dict[str, Any]] = None,
47 settings_snapshot: Optional[Dict[str, Any]] = None,
48 ):
49 """
50 Initialize context manager.
52 Args:
53 session_id: Chat session ID
54 messages: List of message dictionaries with role, content, etc.
55 accumulated_context: Previously accumulated context from session
56 settings_snapshot: Optional settings to override class-level defaults
57 """
58 self.session_id = session_id
59 # chat_messages no longer contains step rows (they live in
60 # chat_progress_steps), but get_session_messages MERGES both for
61 # client rendering. Filter out steps + non-dict entries here so
62 # accumulated context only reflects durable conversation turns.
63 self.messages = [
64 msg
65 for msg in (messages or [])
66 if isinstance(msg, dict) and msg.get("message_type") != "step"
67 ]
68 # The persisted JSON value is passed in unvalidated. Corruption or a
69 # legacy writer can leave a truthy non-dict (7, "abc", [1]) that a
70 # bare `or {}` would let through, so treat any non-dict as empty.
71 if isinstance(accumulated_context, dict):
72 self.accumulated_context = accumulated_context
73 else:
74 if accumulated_context:
75 logger.warning(
76 "Ignoring non-dict accumulated_context (type={})",
77 type(accumulated_context).__name__,
78 )
79 self.accumulated_context = {}
80 # Used by build_research_context to construct the LLM that produces
81 # the query-focused conversation summary.
82 self.settings_snapshot = settings_snapshot
84 # Override class defaults from settings if provided.
85 # Note: settings_snapshot is the 4th keyword arg; passing it positionally
86 # would bind it to the unused `username` param, silently using defaults.
87 if settings_snapshot:
88 self.MAX_FINDINGS_TO_INCLUDE = get_setting_from_snapshot(
89 "chat.max_findings_to_include",
90 self.MAX_FINDINGS_TO_INCLUDE,
91 settings_snapshot=settings_snapshot,
92 )
94 def build_research_context(self, current_query: str = "") -> Dict[str, Any]:
95 """
96 Build context for the next research query.
98 Args:
99 current_query: The user's new message. On a follow-up turn it is
100 used to focus a summary of the whole prior conversation, which
101 becomes the follow-up prompt's "previous findings". On the
102 first turn there is no prior work to summarize.
104 Returns dict with:
105 - session_id: Current session
106 - original_query: The session's first user message
107 - accumulated_findings / past_findings: Prior work for the follow-up
108 (query-focused summary on follow-ups; empty on the first turn)
109 - key_entities: Important entities mentioned
110 - topics: Topics discussed
111 - is_multi_turn: Whether this is a follow-up
112 """
113 # The follow-up strategy reads "original_query" to anchor the prompt on
114 # the topic that started the conversation; without it the contextual
115 # follow-up loses the original question. Use the session's first user
116 # message as that anchor.
117 original_query = next(
118 (
119 m.get("content", "")
120 for m in self.messages
121 if isinstance(m, dict) and m.get("role") == "user"
122 ),
123 "",
124 )
126 is_multi_turn = any(
127 isinstance(m, dict) and m.get("role") == "assistant"
128 for m in self.messages
129 )
131 findings = self._select_prior_findings(current_query, is_multi_turn)
133 return {
134 "session_id": self.session_id,
135 "original_query": original_query,
136 "accumulated_findings": findings,
137 "past_findings": findings, # Research engine expects this key
138 "key_entities": self._get_key_entities(),
139 "topics": self._get_topics(),
140 "is_multi_turn": is_multi_turn,
141 "turn_count": len(self.messages),
142 }
144 def _select_prior_findings(
145 self, current_query: str, is_multi_turn: bool
146 ) -> str:
147 """Pick the follow-up's "previous findings" per chat.followup_context_mode.
149 Modes:
150 - ``summary`` (default): query-focused LLM summary of the conversation
151 - ``raw``: recent research findings, truncated
152 - ``full``: the entire conversation transcript
153 - ``none``: no prior findings
155 Only follow-up turns carry prior work; the first turn returns "".
156 """
157 if not is_multi_turn:
158 return ""
160 mode = self.DEFAULT_FOLLOWUP_CONTEXT_MODE
161 if self.settings_snapshot:
162 mode = get_setting_from_snapshot(
163 "chat.followup_context_mode",
164 mode,
165 settings_snapshot=self.settings_snapshot,
166 )
168 if mode == "none":
169 findings = ""
170 elif mode == "raw":
171 findings = self._extract_findings_from_history()
172 elif mode == "full":
173 findings = self._build_conversation_text()
174 elif current_query:
175 # "summary" with a question to focus on.
176 findings = self._summarize_prior_work(current_query)
177 else:
178 # "summary" with no question (e.g. a no-arg build_research_context
179 # call): fall back to raw recent findings.
180 findings = self._extract_findings_from_history()
182 # Observability: the summary path is otherwise silent (no token-counter
183 # entry, since get_llm runs without a research_id), so a follow-up's
184 # prior-context build looked like an unexplained pause. One line per
185 # follow-up turn records which mode ran and how much context it built.
186 logger.info(
187 "Chat follow-up prior context: mode={}, {} chars",
188 mode,
189 len(findings),
190 )
191 return findings
193 def _build_conversation_text(self) -> str:
194 """Render the prior conversation (both roles) as a plain transcript.
196 Trims oldest-first to ``CONTEXT_INPUT_CHAR_BUDGET`` so the most recent
197 turns survive and the summarizer's input cap never has to truncate.
198 """
199 lines: List[str] = []
200 used = 0
201 for msg in reversed(self.messages):
202 if not isinstance(msg, dict): 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true
203 continue
204 content = (msg.get("content") or "").strip()
205 if not content: 205 ↛ 206line 205 didn't jump to line 206 because the condition on line 205 was never true
206 continue
207 role = (msg.get("role") or "unknown").capitalize()
208 line = f"{role}: {content}"
209 remaining = self.CONTEXT_INPUT_CHAR_BUDGET - used
210 if remaining <= 0:
211 break
212 if len(line) > remaining:
213 if lines:
214 # Budget already spent on more recent turns — stop rather
215 # than partially including an older one.
216 break
217 # The most recent turn alone exceeds the budget: keep its head
218 # so the transcript still fits the summarizer's input cap.
219 line = line[:remaining]
220 lines.append(line)
221 used += len(line)
222 lines.reverse()
223 return "\n\n".join(lines)
225 def _summarize_prior_work(self, current_query: str) -> str:
226 """Summarize the prior conversation, focused on ``current_query``.
228 Returns an empty string when there is no prior conversation, the LLM
229 cannot be constructed (e.g. a misconfigured provider), or the LLM call
230 itself fails. The summary is additive context, so a failure here must
231 not crash the follow-up request — the research dispatch that follows
232 surfaces a genuinely-broken LLM through its own error handling.
233 """
234 transcript = self._build_conversation_text()
235 if not transcript: 235 ↛ 236line 235 didn't jump to line 236 because the condition on line 235 was never true
236 return ""
238 from ..config.llm_config import get_llm
239 from ..advanced_search_system.summarization import FocusedSummarizer
241 try:
242 llm = get_llm(settings_snapshot=self.settings_snapshot)
243 except Exception:
244 logger.opt(exception=True).debug(
245 "Could not build LLM for chat context summary; skipping"
246 )
247 return ""
249 return FocusedSummarizer(
250 llm,
251 focus_query=current_query,
252 max_sentences=self.CONTEXT_SUMMARY_MAX_SENTENCES,
253 max_chars=self.CONTEXT_SUMMARY_MAX_CHARS,
254 ).summarize(transcript)
256 def _extract_findings_from_history(self) -> str:
257 """
258 Extract key findings from assistant messages with research.
260 Returns combined findings text, limited in length.
261 """
262 findings = []
264 for msg in self.messages:
265 if msg.get("role") == "assistant" and msg.get("research_id"):
266 content = msg.get("content") or ""
267 # Summarize long responses - take first part
268 if len(content) > 500:
269 # Try to find a natural break point
270 break_point = content.find("\n\n", 300)
271 if break_point == -1 or break_point > 600: 271 ↛ 273line 271 didn't jump to line 273 because the condition on line 271 was always true
272 break_point = 500
273 content = content[:break_point] + "..."
274 findings.append(content)
276 # Keep only recent findings
277 recent_findings = findings[-self.MAX_FINDINGS_TO_INCLUDE :]
278 return "\n\n---\n\n".join(recent_findings)
280 def _get_key_entities(self) -> List[str]:
281 """Get key entities from accumulated context."""
282 entities = self.accumulated_context.get("key_entities", [])
283 return entities[:20] if isinstance(entities, list) else []
285 def _get_topics(self) -> List[str]:
286 """Get topics from accumulated context."""
287 topics = self.accumulated_context.get("topics", [])
288 return topics[:10] if isinstance(topics, list) else []
290 def extract_context_updates(self, new_content: str) -> Dict[str, Any]:
291 """
292 Extract context updates from a new research response.
294 Args:
295 new_content: New assistant response content
297 Returns:
298 Dict with entity, topic, and summary updates.
299 """
300 return {
301 "new_entities": [], # Could be enhanced with NLP entity extraction
302 "new_topics": [], # Could be enhanced with NLP topic modeling
303 "summary_addition": self._create_summary(new_content),
304 }
306 def _create_summary(self, content: str) -> str:
307 """
308 Create a brief summary of content for context accumulation.
310 Returns first meaningful paragraph or truncated content.
311 """
312 if not content:
313 return ""
315 # Try to get first paragraph
316 paragraphs = content.split("\n\n")
317 for para in paragraphs:
318 para = para.strip()
319 # Skip headers and very short paragraphs
320 if para and len(para) > 50 and not para.startswith("#"):
321 if len(para) > 300:
322 return para[:300] + "..."
323 return para
325 # Fallback: just truncate
326 if len(content) > 300: 326 ↛ 327line 326 didn't jump to line 327 because the condition on line 326 was never true
327 return content[:300] + "..."
328 return content