Coverage for src/local_deep_research/citation_handlers/base_citation_handler.py: 96%
81 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""
2Base class for all citation handlers.
3"""
5from abc import ABC, abstractmethod
6from typing import Any, Callable, Dict, List, Optional, Union
8from langchain_core.documents import Document
9from loguru import logger
11from ..utilities.type_utils import unwrap_setting
12from ..utilities.json_utils import get_llm_response_text
15class BaseCitationHandler(ABC):
16 """Abstract base class for citation handlers."""
18 def __init__(self, llm, settings_snapshot=None):
19 self.llm = llm
20 self.settings_snapshot = settings_snapshot or {}
21 self._fact_checking_logged = False
22 self.stream_callback: Optional[Callable[[str], None]] = None
24 def set_stream_callback(self, callback: Callable[[str], None]):
25 """Set a callback that receives each streamed LLM token."""
26 self.stream_callback = callback
28 def _invoke_with_streaming(self, prompt: str) -> str:
29 """
30 Invoke the LLM, streaming tokens through the callback if set.
32 Falls back to a single ``invoke()`` call when no callback is
33 registered or when the LLM does not support ``.stream()``.
35 Returns:
36 The complete response text.
37 """
38 if self.stream_callback and hasattr(self.llm, "stream"):
39 chunks = []
40 try:
41 for chunk in self.llm.stream(prompt):
42 text = (
43 chunk
44 if isinstance(chunk, str)
45 else getattr(chunk, "content", str(chunk))
46 )
47 if text:
48 chunks.append(text)
49 try:
50 self.stream_callback(text)
51 except Exception:
52 logger.debug(
53 "stream_callback failed"
54 ) # Non-critical: don't break synthesis
55 # Normalize the joined chunks exactly like the invoke() path
56 # below: .stream() bypasses ProcessingLLMWrapper.invoke (which
57 # is the only place <think> blocks are stripped), so without
58 # this a reasoning model's <think>…</think> would leak into the
59 # persisted answer. The live token stream is a separate concern.
60 return get_llm_response_text("".join(chunks))
61 except Exception:
62 # If any chunks already crossed the wire to the client,
63 # restarting via .invoke() would (a) double-bill the LLM
64 # and (b) cause the frontend's accumulated streamed text
65 # to diverge from the new full response — the chat bubble
66 # then shows partial chunks while the DB row carries the
67 # invoke()-result. Only fall back when nothing was emitted.
68 if chunks:
69 logger.warning(
70 "Stream errored after {} chunks; returning partial "
71 "content (no invoke() fallback to avoid double-bill "
72 "and UI/DB divergence)",
73 len(chunks),
74 )
75 return get_llm_response_text("".join(chunks))
76 logger.debug(
77 "Stream failed before any chunk; falling back to invoke()"
78 )
80 # No callback (non-chat research) or stream unavailable: delegate to
81 # the same normalization _invoke_text uses, so <think> blocks
82 # are stripped and str/object responses are handled uniformly.
83 return get_llm_response_text(self.llm.invoke(prompt))
85 def _invoke_text(self, prompt: str) -> str:
86 """Invoke the LLM and return normalized text.
88 Handles both message objects (``.content``) and raw string responses,
89 and strips ``<think>`` reasoning blocks via ``get_llm_response_text``.
90 """
91 return get_llm_response_text(self.llm.invoke(prompt))
93 def get_setting(self, key: str, default=None):
94 """Get a setting value from the snapshot."""
95 if key in self.settings_snapshot:
96 return unwrap_setting(self.settings_snapshot[key])
97 return default
99 def is_fact_checking_enabled(self) -> bool:
100 """Check if fact-checking is enabled and log the state once."""
101 enabled = self.get_setting("general.enable_fact_checking", False)
102 if not self._fact_checking_logged:
103 handler_name = type(self).__name__
104 if enabled:
105 logger.info(
106 f"[{handler_name}] Fact-checking is ENABLED — "
107 f"extra LLM call per synthesis"
108 )
109 else:
110 logger.info(f"[{handler_name}] Fact-checking is DISABLED")
111 self._fact_checking_logged = True
112 return bool(enabled)
114 def _get_output_instruction_prefix(self) -> str:
115 """
116 Get formatted output instructions from settings if present.
118 This allows users to customize output language, tone, style, and formatting
119 for research answers and reports. Instructions are prepended to prompts
120 sent to the LLM.
122 Returns:
123 str: Formatted instruction prefix if custom instructions are set,
124 empty string otherwise.
126 Examples:
127 - "Respond in Spanish with formal academic tone"
128 - "Use simple language suitable for beginners"
129 - "Be concise with bullet points"
130 """
131 output_instructions = self.get_setting(
132 "general.output_instructions", ""
133 ).strip()
135 if output_instructions:
136 return f"User-Specified Output Style: {output_instructions}\n\n"
137 return ""
139 def _create_documents(
140 self, search_results: Union[str, List[Dict]], nr_of_links: int = 0
141 ) -> List[Document]:
142 """
143 Convert search results to LangChain documents format and add index
144 to original search results.
145 """
146 documents: List[Document] = []
147 if isinstance(search_results, str):
148 return documents
150 for i, result in enumerate(search_results):
151 if isinstance(result, dict): 151 ↛ 150line 151 didn't jump to line 150 because the condition on line 151 was always true
152 # Add index to the original search result dictionary if it doesn't exist
153 # This preserves indices that were already set (e.g., for topic organization)
154 if "index" not in result:
155 result["index"] = str(i + nr_of_links + 1)
157 content = result.get("full_content", result.get("snippet", ""))
158 # Use the index from the result if it exists, otherwise calculate it
159 doc_index = int(result.get("index", i + nr_of_links + 1))
160 documents.append(
161 Document(
162 page_content=content,
163 metadata={
164 "source": result.get("link", f"source_{i + 1}"),
165 "title": result.get("title", f"Source {i + 1}"),
166 "index": doc_index,
167 },
168 )
169 )
170 return documents
172 def _format_sources(self, documents: List[Document]) -> str:
173 """Format sources with numbers for citation."""
174 sources = []
175 for doc in documents:
176 source_id = doc.metadata["index"]
177 sources.append(f"[{source_id}] {doc.page_content}")
178 return "\n\n".join(sources)
180 def _no_sources_response(self, question: str) -> Dict[str, Any]:
181 """
182 Explicit no-sources result returned instead of invoking the LLM.
184 Prompting the LLM to "answer with citations [1], [2]…" while the
185 sources section is empty makes it fall back on its training data
186 and fabricate references. Handlers call this to refuse synthesis
187 instead.
188 """
189 logger.warning(
190 f"[{type(self).__name__}] No sources available for synthesis of "
191 f"'{question[:100]}' — skipping LLM call to avoid fabricated "
192 f"citations"
193 )
194 content = (
195 "No sources were found for this question. The selected search "
196 "engines or document collections returned no results; this can "
197 "also happen when a search fails with an error (check the "
198 "research logs). No answer was generated because, without "
199 "sources, it would have to rely on the language model's "
200 "built-in knowledge and could contain fabricated citations."
201 )
202 return {"content": content, "documents": []}
204 @abstractmethod
205 def analyze_initial(
206 self, query: str, search_results: Union[str, List[Dict]]
207 ) -> Dict[str, Any]:
208 """Process initial analysis with citations."""
209 pass
211 @abstractmethod
212 def analyze_followup(
213 self,
214 question: str,
215 search_results: Union[str, List[Dict]],
216 previous_knowledge: str,
217 nr_of_links: int,
218 ) -> Dict[str, Any]:
219 """Process follow-up analysis with citations."""
220 pass