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