Coverage for src/local_deep_research/citation_handlers/base_citation_handler.py: 96%
171 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Base class for all citation handlers.
3"""
5from abc import ABC, abstractmethod
6from typing import Any, Callable, Dict, List, Optional, Sequence, Union
8from langchain_core.documents import Document
9from loguru import logger
11from ..utilities.type_utils import unwrap_setting
12from ..utilities.json_utils import _coerce_content_blocks, get_llm_response_text
15class _HeadingRepetitionGuard:
16 """Detects a model re-emitting its own headings instead of stopping.
18 A local model that misses its EOS token can attend back to the start of
19 what it just wrote and reproduce the same section cycle until the context
20 window hard-clips — 30+ iterations of identical subheadings, tables and
21 prose, all of which then land in the report (#6452). Nothing was watching
22 the stream, so the only bound was the context window.
24 The signal is a markdown heading line repeating verbatim. It is chosen
25 over n-gram or block similarity because it is O(1) per line, needs no
26 buffer of the text so far, and is nearly impossible to trip by accident:
27 a legitimate subsection does not emit one identical ``### Heading`` four
28 times. Prose repetition without repeated headings is NOT caught — the
29 guard is a bound on the reported failure, not a general degeneration
30 detector, and a rule loose enough to catch every loop would truncate real
31 documents.
33 Headings are compared after stripping trailing whitespace and leading
34 ``#``/space, so ``## Costs`` and ``### Costs`` count together: the loop
35 reproduces its own text, and a level change between cycles is still the
36 same cycle.
38 A repeated heading alone is not the signal, because a legitimate
39 comparative report repeats them by design: ``## Study A`` /
40 ``### Methodology`` / ``### Results``, then ``## Study B`` and the same
41 two subsections again. A four-item report reaches four
42 ``### Methodology`` with nothing degenerate happening, and truncating it
43 is exactly the silent quality loss this guard exists to prevent.
45 What separates the two is whether the document is still PRODUCING. A
46 comparative report emits a heading nobody has seen before at every item
47 (``## Study B``); a loop has stopped emitting new ones entirely and only
48 replays what it already wrote. So the abort needs both a heading at
49 ``LIMIT`` and a run of ``CONSECUTIVE_KNOWN`` headings in a row that the
50 model has already emitted, and any previously-unseen heading resets that
51 run. Both are O(1) per line and need no buffer of the text so far.
53 That also bounds a loop repeating one heading and nothing else: every
54 occurrence after the first is a known heading, so the run climbs and the
55 abort follows a cycle later than for a two-heading loop.
57 Lines inside a fenced code block are not headings. ``# comment`` is the
58 first line of half the bash and Python snippets there are, and a technical
59 report quoting four similar config blocks would otherwise trip the guard.
61 Chunk boundaries fall anywhere, so a partial last line is held back until
62 its newline arrives rather than being counted as its own heading.
63 """
65 #: Occurrences of one heading before it can end the stream. Three is
66 #: reachable by a document that legitimately revisits a title (a summary
67 #: repeating a section name); four is not, and a loop reaches it in two
68 #: cycles.
69 LIMIT = 4
71 #: Headings in a row that the model has already emitted. A comparative
72 #: report breaks this run at every item with a heading nobody has seen
73 #: before; a loop cannot, because it has stopped producing new ones.
74 CONSECUTIVE_KNOWN = 4
76 def __init__(self) -> None:
77 self._counts: Dict[str, int] = {}
78 self._pending = ""
79 self._in_fence = False
80 self._consecutive_known = 0
82 def feed(self, texts: Sequence[str]) -> Optional[str]:
83 """Consume newly streamed text; return the repeated heading, or None."""
84 for text in texts:
85 self._pending += text
86 *lines, self._pending = self._pending.split("\n")
87 for line in lines:
88 repeated = self._count(line)
89 if repeated is not None:
90 return repeated
91 return None
93 def _count(self, line: str) -> Optional[str]:
94 stripped = line.strip()
95 if stripped.startswith("```"):
96 self._in_fence = not self._in_fence
97 return None
98 if self._in_fence or not stripped.startswith("#"):
99 return None
100 heading = stripped.lstrip("#").strip()
101 if not heading: 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true
102 return None
103 seen = self._counts.get(heading, 0) + 1
104 self._counts[heading] = seen
105 if seen == 1:
106 self._consecutive_known = 0
107 return None
108 self._consecutive_known += 1
109 if seen < self.LIMIT:
110 return None
111 if self._consecutive_known < self.CONSECUTIVE_KNOWN:
112 return None
113 return heading
116class BaseCitationHandler(ABC):
117 """Abstract base class for citation handlers."""
119 def __init__(self, llm, settings_snapshot=None):
120 self.llm = llm
121 self.settings_snapshot = settings_snapshot or {}
122 self._fact_checking_logged = False
123 self.stream_callback: Optional[Callable[[str], None]] = None
125 def set_stream_callback(self, callback: Callable[[str], None]):
126 """Set a callback that receives each streamed LLM token."""
127 self.stream_callback = callback
129 def _new_repetition_guard(self) -> "_HeadingRepetitionGuard":
130 """The streaming degeneration guard, as a seam a subclass can widen."""
131 return _HeadingRepetitionGuard()
133 def _handle_chunk(self, chunk: Any, chunks: List[str]) -> None:
134 """Normalize one streamed chunk, record it and push it to the callback.
136 Shared by the sync and async streaming paths so the two differ only
137 in the LLM call line.
138 """
139 text = (
140 chunk
141 if isinstance(chunk, str)
142 else getattr(chunk, "content", str(chunk))
143 )
144 if isinstance(text, list):
145 # Keep spaces and split reasoning tags intact until chunks join.
146 text = _coerce_content_blocks(text)
147 if text:
148 chunks.append(text)
149 # Bind to a local so mypy sees the Optional narrowed — the
150 # attribute is Optional[Callable] and this helper is called
151 # from both streaming paths after their own None checks.
152 callback = self.stream_callback
153 if callback is not None: 153 ↛ exitline 153 didn't return from function '_handle_chunk' because the condition on line 153 was always true
154 try:
155 callback(text)
156 except Exception:
157 logger.debug(
158 "stream_callback failed"
159 ) # Non-critical: don't break synthesis
161 def _log_repetition_abort(self, heading: str, chunks: List[str]) -> None:
162 """One line a reader can act on: what repeated, and how much was kept."""
163 logger.warning(
164 "Stream aborted after a degenerate repetition loop: the heading "
165 "{!r} was emitted {} times. Keeping the {} chunk(s) produced so "
166 "far; the model was re-emitting its own output instead of "
167 "stopping.",
168 heading,
169 _HeadingRepetitionGuard.LIMIT,
170 len(chunks),
171 )
173 def _join_chunks(self, chunks: List[str]) -> str:
174 """Normalize the joined chunks exactly like the invoke() path does.
176 ``.stream()`` / ``.astream()`` bypass the normalization in
177 ProcessingLLMWrapper.invoke / ainvoke, so without
178 this a reasoning model's ``<think>…</think>`` would leak into the
179 persisted answer. The live token stream is a separate concern.
180 """
181 return get_llm_response_text("".join(chunks))
183 def _partial_after_stream_failure(self, chunks: List[str]) -> Optional[str]:
184 """Resolve a mid-stream failure: partial text, or ``None`` to fall back.
186 If any chunks already crossed the wire to the client, restarting via
187 a full ``invoke()`` would (a) double-bill the LLM and (b) cause the
188 frontend's accumulated streamed text to diverge from the new full
189 response — the chat bubble then shows partial chunks while the DB row
190 carries the invoke()-result. Only fall back when nothing was emitted.
191 """
192 if chunks:
193 logger.warning(
194 "Stream errored after {} chunks; returning partial "
195 "content (no invoke() fallback to avoid double-bill "
196 "and UI/DB divergence)",
197 len(chunks),
198 )
199 return self._join_chunks(chunks)
200 logger.debug("Stream failed before any chunk; falling back to invoke()")
201 return None
203 def _invoke_with_streaming(self, prompt: str) -> str:
204 """
205 Invoke the LLM, streaming tokens through the callback if set.
207 Falls back to a single ``invoke()`` call when no callback is
208 registered or when the LLM does not support ``.stream()``.
210 Stays on the synchronous LangChain API: citation handlers run inside
211 per-run research threads with no event loop, and langchain's async
212 httpx client is process-cached and loop-bound, so driving the async
213 core on a throwaway per-call loop would break or re-send every call
214 after the first (see #6293).
216 Returns:
217 The complete response text.
218 """
219 if self.stream_callback and hasattr(self.llm, "stream"):
220 chunks: List[str] = []
221 guard = self._new_repetition_guard()
222 try:
223 for chunk in self.llm.stream(prompt):
224 before = len(chunks)
225 self._handle_chunk(chunk, chunks)
226 # Cooperative abort for unwrapped/custom LLMs whose
227 # own stream() does not check the flag (wrapped LLMs
228 # already raise per chunk via ProcessingLLMWrapper).
229 # Resolved from the thread's search context so handlers
230 # holding a raw model still stop promptly on Stop.
231 # NOTE: ResearchTerminatedException inherits
232 # BaseException, so the ``except Exception`` below lets
233 # it propagate while swallowing only lookup errors.
234 try:
235 from ..llm.cancellation import (
236 get_current_research_id,
237 raise_if_terminated,
238 )
240 _rid = get_current_research_id()
241 if _rid:
242 raise_if_terminated(_rid)
243 except Exception:
244 logger.debug(
245 "Termination-flag lookup failed (non-critical)"
246 )
247 repeated = guard.feed(chunks[before:])
248 if repeated is not None:
249 self._log_repetition_abort(repeated, chunks)
250 break
251 return self._join_chunks(chunks)
252 except Exception:
253 partial = self._partial_after_stream_failure(chunks)
254 if partial is not None:
255 return partial
257 # No callback (non-chat research), stream unavailable, or the stream
258 # died before the first chunk: delegate to the same normalization
259 # _invoke_text uses, so <think> blocks are stripped and str/object
260 # responses are handled uniformly.
261 return self._invoke_text(prompt)
263 async def _invoke_with_streaming_async(self, prompt: str) -> str:
264 """Async counterpart of :meth:`_invoke_with_streaming`.
266 Additive: no production caller awaits it yet. It is intended to be
267 submitted to the process's one shared event loop via
268 ``asyncio.run_coroutine_threadsafe`` (ADR-0013,
269 docs/decisions/0013-research-runs-keep-daemon-thread.md); which
270 thread hosts that loop is not decided yet (#6467).
271 :meth:`_invoke_with_streaming` remains the canonical production
272 path on the research thread, and this must never go through
273 ``asyncio.run`` or a per-thread loop (#6293). The streaming
274 semantics are the sync path's, shared through the same helpers —
275 including returning the PARTIAL text after a mid-stream failure
276 instead of restarting with a full call.
278 Not loop-safe yet, for three reasons: ``TokenCountingCallback``
279 writes call-usage rows to the user's encrypted database from the
280 loop's default executor rather than the daemon thread (rule 5;
281 #6294 is related); the egress audit context is thread-local and
282 must be passed in and re-armed on the loop side; and the submit
283 helper must run the coroutine in an explicit minimal context so it
284 does not inherit the daemon thread's ambient search context,
285 including the password (rule 4).
286 """
287 if self.stream_callback and hasattr(self.llm, "astream"):
288 chunks: List[str] = []
289 guard = self._new_repetition_guard()
290 try:
291 async for chunk in self.llm.astream(prompt):
292 before = len(chunks)
293 self._handle_chunk(chunk, chunks)
294 repeated = guard.feed(chunks[before:])
295 if repeated is not None:
296 self._log_repetition_abort(repeated, chunks)
297 break
298 return self._join_chunks(chunks)
299 except Exception:
300 partial = self._partial_after_stream_failure(chunks)
301 if partial is not None:
302 return partial
304 return await self._invoke_text_async(prompt)
306 def _invoke_text(self, prompt: str) -> str:
307 """Invoke the LLM and return normalized text.
309 Handles both message objects (``.content``) and raw string responses,
310 and strips ``<think>`` reasoning blocks via ``get_llm_response_text``.
312 Synchronous for the same reason as :meth:`_invoke_with_streaming`.
313 """
314 return get_llm_response_text(self.llm.invoke(prompt))
316 async def _invoke_text_async(self, prompt: str) -> str:
317 """Async counterpart of :meth:`_invoke_text`.
319 Additive: no production caller awaits it yet (see
320 :meth:`_invoke_with_streaming_async`). Configured models expose
321 ``ainvoke`` through ``ProcessingLLMWrapper``. Custom injected LLMs
322 must also provide ``ainvoke`` to use this async API; it does not
323 fall back to blocking synchronous invocation.
324 """
325 return get_llm_response_text(await self.llm.ainvoke(prompt))
327 def get_setting(self, key: str, default=None):
328 """Get a setting value from the snapshot."""
329 if key in self.settings_snapshot:
330 return unwrap_setting(self.settings_snapshot[key])
331 return default
333 def is_fact_checking_enabled(self) -> bool:
334 """Check if fact-checking is enabled and log the state once."""
335 enabled = self.get_setting("general.enable_fact_checking", False)
336 if not self._fact_checking_logged:
337 handler_name = type(self).__name__
338 if enabled:
339 logger.info(
340 f"[{handler_name}] Fact-checking is ENABLED — "
341 f"extra LLM call per synthesis"
342 )
343 else:
344 logger.info(f"[{handler_name}] Fact-checking is DISABLED")
345 self._fact_checking_logged = True
346 return bool(enabled)
348 def _get_output_instruction_prefix(self) -> str:
349 """
350 Get formatted output instructions from settings if present.
352 This allows users to customize output language, tone, style, and formatting
353 for research answers and reports. Instructions are prepended to prompts
354 sent to the LLM.
356 Returns:
357 str: Formatted instruction prefix if custom instructions are set,
358 empty string otherwise.
360 Examples:
361 - "Respond in Spanish with formal academic tone"
362 - "Use simple language suitable for beginners"
363 - "Be concise with bullet points"
364 """
365 output_instructions = self.get_setting(
366 "general.output_instructions", ""
367 ).strip()
369 if output_instructions:
370 return f"User-Specified Output Style: {output_instructions}\n\n"
371 return ""
373 def _create_documents(
374 self, search_results: Union[str, List[Dict]], nr_of_links: int = 0
375 ) -> List[Document]:
376 """
377 Convert search results to LangChain documents format and add index
378 to original search results.
379 """
380 documents: List[Document] = []
381 if isinstance(search_results, str):
382 return documents
384 for i, result in enumerate(search_results):
385 if isinstance(result, dict): 385 ↛ 384line 385 didn't jump to line 384 because the condition on line 385 was always true
386 # Add index to the original search result dictionary if it doesn't exist
387 # This preserves indices that were already set (e.g., for topic organization)
388 if "index" not in result:
389 result["index"] = str(i + nr_of_links + 1)
391 content = result.get("full_content")
392 if not content:
393 content = result.get("snippet", "")
394 # Use the index from the result if it exists, otherwise calculate it
395 doc_index = int(result.get("index", i + nr_of_links + 1))
396 documents.append(
397 Document(
398 page_content=content,
399 metadata={
400 "source": result.get("link", f"source_{i + 1}"),
401 "title": result.get("title", f"Source {i + 1}"),
402 "index": doc_index,
403 },
404 )
405 )
406 return documents
408 def _format_sources(self, documents: List[Document]) -> str:
409 """Override to customize source text; keep each source's [N] first."""
410 sources = []
411 for doc in documents:
412 source_id = doc.metadata["index"]
413 sources.append(f"[{source_id}] {doc.page_content}")
414 return "\n\n".join(sources)
416 def _no_sources_response(self, question: str) -> Dict[str, Any]:
417 """
418 Explicit no-sources result returned instead of invoking the LLM.
420 Prompting the LLM to "answer with citations [1], [2]…" while the
421 sources section is empty makes it fall back on its training data
422 and fabricate references. Handlers call this to refuse synthesis
423 instead.
424 """
425 logger.warning(
426 f"[{type(self).__name__}] No sources available for synthesis of "
427 f"'{question[:100]}' — skipping LLM call to avoid fabricated "
428 f"citations"
429 )
430 content = (
431 "No sources were found for this question. The selected search "
432 "engines or document collections returned no results; this can "
433 "also happen when a search fails with an error (check the "
434 "research logs). No answer was generated because, without "
435 "sources, it would have to rely on the language model's "
436 "built-in knowledge and could contain fabricated citations."
437 )
438 return {"content": content, "documents": []}
440 @abstractmethod
441 def analyze_initial(
442 self, query: str, search_results: Union[str, List[Dict]]
443 ) -> Dict[str, Any]:
444 """Process initial analysis with citations."""
445 pass
447 @abstractmethod
448 def analyze_followup(
449 self,
450 question: str,
451 search_results: Union[str, List[Dict]],
452 previous_knowledge: str,
453 nr_of_links: int,
454 ) -> Dict[str, Any]:
455 """Process follow-up analysis with citations."""
456 pass