Coverage for src/local_deep_research/search_system.py: 91%
126 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1# src/local_deep_research/search_system.py
2from typing import Any, Callable, Dict
4from langchain_core.language_models import BaseChatModel
5from loguru import logger
7from .advanced_search_system.findings.repository import FindingsRepository
8from .advanced_search_system.questions.standard_question import (
9 StandardQuestionGenerator,
10)
11from .advanced_search_system.strategies.followup.enhanced_contextual_followup import (
12 EnhancedContextualFollowUpStrategy,
13)
15from .utilities.type_utils import unwrap_setting
16from .citation_handler import CitationHandler
17from .web_search_engines.search_engine_base import BaseSearchEngine
18from .constants import DEFAULT_SEARCH_TOOL
21def _ensure_snapshot_username(settings_snapshot, username):
22 """Ensure the run's settings snapshot carries the username under the
23 ``_username`` key that snapshot-driven consumers look it up by.
25 The LangGraph agent re-instantiates the search engine *per tool call*
26 from this snapshot (``_make_web_search_tool`` -> ``create_search_engine``),
27 and registering a user's document collections needs the username
28 (``get_user_db_session``), which collection registration reads only from
29 ``settings_snapshot["_username"]``. The snapshot is built without it (the
30 user lives under ``system.user``), so without this a collection/library
31 primary fails inside the agent with "Unknown search engine 'collection_…'".
33 Injected here in ``AdvancedSearchSystem.__init__`` — the narrowest common
34 consumer of the strategy-running paths — so the web run and the
35 programmatic API are repaired in one place. (No-ops when ``username`` is
36 absent, e.g. benchmarks/news, which don't use per-user collections anyway.)
37 Mirrors the convention ``get_search()`` already applies to its own local
38 snapshot copy.
40 Returns the snapshot unchanged when there is no username, it is not a dict,
41 or ``_username`` is already set (never overwrites an explicit value);
42 otherwise a shallow copy with ``_username`` added (no mutation).
43 """
44 if (
45 username
46 and isinstance(settings_snapshot, dict)
47 and not settings_snapshot.get("_username")
48 ):
49 return {**settings_snapshot, "_username": username}
50 return settings_snapshot
53class AdvancedSearchSystem:
54 """
55 Advanced search system that coordinates different search strategies.
56 """
58 def __init__(
59 self,
60 llm: BaseChatModel,
61 search: BaseSearchEngine,
62 strategy_name: str = "source-based", # Default to comprehensive research strategy
63 include_text_content: bool = True,
64 use_cross_engine_filter: bool = True,
65 max_iterations: int | None = None,
66 questions_per_iteration: int | None = None,
67 use_atomic_facts: bool = False,
68 username: str | None = None,
69 settings_snapshot: dict | None = None,
70 research_id: str | None = None,
71 research_context: dict | None = None,
72 programmatic_mode: bool = False,
73 search_original_query: bool = True,
74 ):
75 """Initialize the advanced search system.
77 Args:
78 llm: LLM to use for the search strategy.
79 search: Search engine to use for queries.
80 strategy_name: The name of the search strategy to use. Options:
81 - "standard": Basic iterative search strategy
82 - "iterdrag": Iterative Dense Retrieval Augmented Generation
83 - "source-based": Focuses on finding and extracting from sources
84 - "parallel": Runs multiple search queries in parallel
85 - "rapid": Quick single-pass search
86 - "recursive": Recursive decomposition of complex queries
87 - "iterative": Loop-based reasoning with persistent knowledge
88 - "adaptive": Adaptive step-by-step reasoning
89 - "smart": Automatically chooses best strategy based on query
90 - "browsecomp": Optimized for BrowseComp-style puzzle queries
91 - "evidence": Enhanced evidence-based verification with improved candidate discovery
92 - "constrained": Progressive constraint-based search that narrows candidates step by step
93 - "parallel-constrained": Parallel constraint-based search with combined constraint execution
94 - "early-stop-constrained": Parallel constraint search with immediate evaluation and early stopping at 99% confidence
95 - "dual-confidence": Dual confidence scoring with positive/negative/uncertainty
96 - "dual-confidence-with-rejection": Dual confidence with early rejection of poor candidates
97 - "concurrent-dual-confidence": Concurrent search & evaluation with progressive constraint relaxation
98 - "modular": Modular architecture using constraint checking and candidate exploration modules
99 - "browsecomp-entity": Entity-focused search for BrowseComp questions with knowledge graph building
100 - "iterative-refinement": Iteratively refines results using LLM evaluation and follow-up queries
101 include_text_content: If False, only includes metadata and links in search results
102 use_cross_engine_filter: Whether to filter results across search
103 engines.
104 max_iterations: The maximum number of search iterations to
105 perform. Will be read from the settings if not specified.
106 questions_per_iteration: The number of questions to include in
107 each iteration. Will be read from the settings if not specified.
108 use_atomic_facts: Whether to use atomic fact decomposition for
109 complex queries when using the source-based strategy.
110 programmatic_mode: If True, disables database operations and metrics tracking.
111 This is useful for running searches without database dependencies.
112 search_original_query: Whether to include the original query in the first iteration
113 of search. Set to False for news searches to avoid sending long subscription
114 prompts to search engines.
116 """
117 # Store research context for strategies
118 self.research_id = research_id
119 self.research_context = research_context
120 self.username = username
122 # Store required components
123 self.model = llm
124 self.search = search
126 # Store settings snapshot. Inject the username under "_username" so
127 # snapshot-driven consumers (notably the LangGraph agent's per-call
128 # search-engine creation, which must register the user's document
129 # collections) can find it. See _ensure_snapshot_username.
130 self.settings_snapshot = _ensure_snapshot_username(
131 settings_snapshot or {}, username
132 )
134 # Store programmatic mode
135 self.programmatic_mode = programmatic_mode
137 # Store search original query setting
138 self.search_original_query = search_original_query
140 # Log if running in programmatic mode
141 if self.programmatic_mode:
142 logger.warning(
143 "Running in programmatic mode - database operations and metrics tracking disabled. "
144 "Rate limiting, search metrics, and persistence features will not be available."
145 )
147 # Get iterations setting
148 self.max_iterations = max_iterations
149 if self.max_iterations is None:
150 # Use settings from snapshot
151 if "search.iterations" in self.settings_snapshot:
152 self.max_iterations = unwrap_setting(
153 self.settings_snapshot["search.iterations"]
154 )
155 else:
156 self.max_iterations = 1 # Default
158 self.questions_per_iteration = questions_per_iteration
159 if self.questions_per_iteration is None:
160 # Use settings from snapshot
161 if "search.questions_per_iteration" in self.settings_snapshot:
162 self.questions_per_iteration = unwrap_setting(
163 self.settings_snapshot["search.questions_per_iteration"]
164 )
165 else:
166 self.questions_per_iteration = 3 # Default
168 # Log the strategy name that's being used
169 logger.info(
170 f"Initializing AdvancedSearchSystem with strategy_name='{strategy_name}'"
171 )
173 # Initialize components
174 self.citation_handler = CitationHandler(
175 self.model, settings_snapshot=self.settings_snapshot
176 )
177 self.question_generator = StandardQuestionGenerator(self.model)
178 self.findings_repository = FindingsRepository(self.model)
179 # For backward compatibility
180 self.questions_by_iteration: dict[Any, Any] = {}
181 self.progress_callback = lambda _1, _2, _3: None
182 self.all_links_of_system: list[dict[Any, Any]] = []
184 # Initialize strategy using factory
185 from .search_system_factory import create_strategy
187 # Special handling for follow-up strategy which needs different logic
188 if strategy_name.lower() in [
189 "enhanced-contextual-followup",
190 "enhanced_contextual_followup",
191 "contextual-followup",
192 "contextual_followup",
193 ]:
194 logger.info("Creating EnhancedContextualFollowUpStrategy instance")
195 # Get delegate strategy from research context
196 # This should be the user's preferred strategy from settings
197 delegate_strategy_name = (
198 self.research_context.get("delegate_strategy", "source-based")
199 if self.research_context
200 else "source-based"
201 )
203 delegate = create_strategy(
204 strategy_name=delegate_strategy_name,
205 model=self.model,
206 search=self.search,
207 all_links_of_system=[],
208 settings_snapshot=self.settings_snapshot,
209 knowledge_accumulation_mode=True,
210 search_original_query=self.search_original_query,
211 )
213 # Create the contextual follow-up strategy with the delegate
214 self.strategy = EnhancedContextualFollowUpStrategy(
215 model=self.model,
216 search=self.search,
217 delegate_strategy=delegate,
218 all_links_of_system=self.all_links_of_system,
219 settings_snapshot=self.settings_snapshot,
220 research_context=self.research_context,
221 )
222 else:
223 # Use factory for all other strategies
224 logger.info(f"Creating {strategy_name} strategy using factory")
225 self.strategy = create_strategy(
226 strategy_name=strategy_name,
227 model=self.model,
228 search=self.search,
229 all_links_of_system=self.all_links_of_system,
230 settings_snapshot=self.settings_snapshot,
231 # Pass strategy-specific parameters
232 include_text_content=include_text_content,
233 use_cross_engine_filter=use_cross_engine_filter,
234 use_atomic_facts=use_atomic_facts,
235 max_iterations=self.max_iterations,
236 questions_per_iteration=self.questions_per_iteration,
237 # Special parameters for iterative strategy
238 search_iterations_per_round=self.max_iterations or 1,
239 questions_per_search=self.questions_per_iteration,
240 # Special parameters for adaptive strategy
241 max_steps=self.max_iterations,
242 source_questions_per_iteration=self.questions_per_iteration,
243 # Special parameters for evidence and constrained strategies
244 max_search_iterations=self.max_iterations,
245 # Special parameters for focused iteration
246 use_browsecomp_optimization=True,
247 # Pass search original query parameter
248 search_original_query=self.search_original_query,
249 # Forwarded so strategies that create engines per tool call
250 # (e.g. langgraph-agent) can match the system's mode.
251 programmatic_mode=self.programmatic_mode,
252 )
254 # Log the actual strategy class
255 logger.info(f"Created strategy of type: {type(self.strategy).__name__}")
257 # Configure the strategy with our attributes
258 if ( 258 ↛ exitline 258 didn't return from function '__init__' because the condition on line 258 was always true
259 hasattr(self, "progress_callback")
260 and self.progress_callback is not None
261 ):
262 self.strategy.set_progress_callback(self.progress_callback)
264 def close(self):
265 """Close resources held by the search system.
267 Cascades close to the strategy, which may hold persistent
268 ThreadPoolExecutor instances (e.g. ConstraintParallelStrategy holds
269 search_executor and evaluation_executor,
270 ConcurrentDualConfidenceStrategy holds evaluation_executor).
272 NOTE: Does NOT close self.search (the search engine) — the caller
273 (run_research_process) manages search engine lifecycle separately
274 because the search engine may be shared or reused.
276 Most strategies (including the default source-based) use
277 context-managed ThreadPoolExecutors that clean up automatically.
278 The close() call here is a safety net for the two constraint-based
279 strategies that hold persistent executors in __init__. Those
280 strategies also shut down their executors in
281 find_relevant_information()'s finally block, so this is a second
282 line of defense for the edge case where the method is never called.
283 """
284 from .utilities.resource_utils import safe_close
286 if hasattr(self, "strategy"):
287 safe_close(self.strategy, "search strategy")
289 def _progress_callback(
290 self, message: str, progress: int, metadata: dict
291 ) -> None:
292 """Handle progress updates from the strategy."""
293 logger.info(f"Progress: {progress}% - {message}")
294 if hasattr(self, "progress_callback"):
295 self.progress_callback(message, progress, metadata)
297 def set_progress_callback(
298 self, callback: Callable[[str, int, dict], None]
299 ) -> None:
300 """Set a callback function to receive progress updates."""
301 self.progress_callback = callback
302 if hasattr(self, "strategy"): 302 ↛ exitline 302 didn't return from function 'set_progress_callback' because the condition on line 302 was always true
303 self.strategy.set_progress_callback(callback)
305 def analyze_topic(
306 self,
307 query: str,
308 is_user_search: bool = True,
309 is_news_search: bool = False,
310 user_id: str = "anonymous",
311 search_id: str | None = None,
312 **kwargs,
313 ) -> Dict:
314 """Analyze a topic using the current strategy.
316 Args:
317 query: The research query to analyze
318 is_user_search: Whether this is a user-initiated search
319 is_news_search: Whether this is a news search
320 user_id: The user ID for tracking
321 search_id: The search ID (auto-generated if not provided)
322 **kwargs: Additional arguments
323 """
325 # Generate search ID if not provided
326 if search_id is None:
327 import uuid
329 search_id = str(uuid.uuid4())
331 # Defense-in-depth: arm the PEP-578 egress audit-hook backstop for this
332 # run if no caller already armed it. The web worker arms it in
333 # research_service before calling us; CLI, the news scheduler and the
334 # programmatic API construct AdvancedSearchSystem directly and would
335 # otherwise run the full pipeline with the secondary net inactive.
336 _armed_egress = self._arm_egress_backstop()
337 try:
338 # Perform the search
339 return self._perform_search(
340 query, search_id, is_user_search, is_news_search, user_id
341 )
342 finally:
343 if _armed_egress:
344 # Only clear what WE armed, so a reused thread doesn't leak the
345 # context to unrelated subsequent work (the web worker owns its
346 # own teardown via @thread_cleanup).
347 from .security.egress.audit_hook import clear_active_context
349 clear_active_context()
351 def _arm_egress_backstop(self) -> bool:
352 """Arm the egress audit-hook context for this run when no caller has.
354 Returns True only when THIS call armed the context (so ``analyze_topic``
355 clears it on exit). Returns False when it was already armed (web worker)
356 or could not be built. Never raises — a backstop failure must not break
357 a research run; the explicit PEPs remain the primary enforcement.
358 """
359 try:
360 from .security.egress.audit_hook import (
361 get_active_context,
362 set_active_context,
363 )
364 from .security.egress.policy import (
365 PolicyDeniedError,
366 context_from_snapshot,
367 )
368 except Exception:
369 return False
371 if get_active_context() is not None:
372 return False # already armed (e.g. by the web research worker)
373 if not self.settings_snapshot:
374 return False
376 try:
377 primary = unwrap_setting(
378 self.settings_snapshot.get("search.tool", DEFAULT_SEARCH_TOOL)
379 )
380 ctx = context_from_snapshot(
381 self.settings_snapshot,
382 primary or DEFAULT_SEARCH_TOOL,
383 username=self.username,
384 )
385 except (PolicyDeniedError, ValueError):
386 # Corrupted scope: the explicit PEPs and the run-start
387 # precheck handle these; don't arm a partial backstop.
388 logger.bind(policy_audit=True).debug(
389 "egress backstop not armed: policy unevaluable", exc_info=True
390 )
391 return False
392 except Exception:
393 return False
395 set_active_context(ctx)
396 return True
398 def _perform_search(
399 self,
400 query: str,
401 search_id: str,
402 is_user_search: bool,
403 is_news_search: bool,
404 user_id: str,
405 ) -> Dict:
406 """Perform the actual search."""
407 # Send progress message with LLM info
408 # Get settings from snapshot if available
409 llm_provider = "unknown"
410 llm_model = "unknown"
411 search_tool = "unknown"
413 if self.settings_snapshot:
414 # Extract values from settings snapshot
415 provider_setting = self.settings_snapshot.get("llm.provider", {})
416 llm_provider = (
417 provider_setting.get("value", "unknown")
418 if isinstance(provider_setting, dict)
419 else provider_setting
420 )
422 model_setting = self.settings_snapshot.get("llm.model", {})
423 llm_model = (
424 model_setting.get("value", "unknown")
425 if isinstance(model_setting, dict)
426 else model_setting
427 )
429 tool_setting = self.settings_snapshot.get("search.tool", {})
430 search_tool = (
431 tool_setting.get("value", DEFAULT_SEARCH_TOOL)
432 if isinstance(tool_setting, dict)
433 else tool_setting
434 )
436 self.progress_callback(
437 f"Using {llm_provider} model: {llm_model}",
438 1, # Low percentage to show this as an early step
439 {
440 "phase": "setup",
441 "llm_info": {
442 "name": llm_model,
443 "provider": llm_provider,
444 },
445 },
446 )
447 # Send progress message with search strategy info
448 self.progress_callback(
449 f"Using search tool: {search_tool}",
450 1.5, # Between setup and processing steps
451 {
452 "phase": "setup",
453 "search_info": {
454 "tool": search_tool,
455 },
456 },
457 )
459 # Use the strategy to analyze the topic
460 result = self.strategy.analyze_topic(query)
462 # Update our attributes for backward compatibility
464 self.questions_by_iteration = (
465 self.strategy.questions_by_iteration.copy()
466 )
467 # Send progress message with search info
469 # Only extend if they're different objects in memory to avoid duplication
470 # This check prevents doubling the list when they reference the same object
471 # Fix for issue #301: "too many links in detailed report mode"
472 if id(self.all_links_of_system) != id(
473 self.strategy.all_links_of_system
474 ):
475 self.all_links_of_system.extend(self.strategy.all_links_of_system)
477 # Include the search system instance for access to citations
478 result["search_system"] = self
479 result["all_links_of_system"] = self.all_links_of_system
481 # Ensure query is included in the result
482 if "query" not in result:
483 result["query"] = query
484 result["questions_by_iteration"] = self.questions_by_iteration
486 # Call news callback
487 try:
488 from .news.core.search_integration import NewsSearchCallback
490 callback = NewsSearchCallback()
491 context = {
492 "is_user_search": is_user_search,
493 "is_news_search": is_news_search,
494 "user_id": user_id,
495 "search_id": search_id,
496 }
497 callback(query, result, context)
498 except Exception:
499 logger.exception("Error in news callback")
501 return result