Coverage for src/local_deep_research/advanced_search_system/parallel_search.py: 100%

1"""Parallel search execution helper. 

2 

3Provides :func:`run_parallel_searches`, a small helper that runs a sequence 

4of search queries concurrently using :class:`~concurrent.futures.ThreadPoolExecutor` 

5while preserving the research/Flask contexts expected by worker threads. 

6 

7The helper is intentionally generic: 

8 

9* ``search_fn`` is a callable accepting the query string and returning 

10 whatever per-query payload the caller needs (a list of results, a dict 

11 with metadata, etc.). Callers wrap it with context-preserving 

12 decorators (e.g. :func:`preserve_research_context`) before passing it 

13 in; this helper stays focused on the concurrency concern. 

14* ``context_factory`` is an optional zero-arg callable (e.g. 

15 :func:`~local_deep_research.utilities.threading_utils.thread_context`) 

16 that returns a **fresh** Flask :class:`~flask.ctx.AppContext` for one 

17 worker thread. It is invoked once per query *on the calling thread* and 

18 the resulting context is pushed into the matching worker via 

19 :func:`~local_deep_research.utilities.threading_utils.thread_with_app_context`. 

20 A single :class:`~flask.ctx.AppContext` instance must never be shared 

21 across concurrent threads — Flask tracks the context-var token stack on 

22 the instance, so concurrent push/pop from multiple threads raises 

23 ``ValueError: <Token> was created in a different Context``. ``None`` 

24 disables context propagation, preserving the behavior of strategies that 

25 never propagated a context (e.g. focused-iteration, progressive). 

26 

27Returns a list of ``(query, payload)`` tuples in completion order. 

28Callers that need a question-keyed dict build it from this list; callers 

29that only need a flat list of results flatten it directly. 

30""" 

31 

32from __future__ import annotations 

33 

34import concurrent.futures 

35from typing import Callable, List, Optional, Tuple, TypeVar 

36 

37from flask.ctx import AppContext 

38from loguru import logger 

39 

40from ..utilities.threading_utils import thread_with_app_context 

41 

42T = TypeVar("T") 

43 

44 

45def run_parallel_searches( 

46 queries: List[str], 

47 search_fn: Callable[[str], T], 

48 max_workers: Optional[int] = None, 

49 context_factory: Optional[Callable[[], Optional[AppContext]]] = None, 

50) -> List[Tuple[str, T]]: 

51 """Run ``search_fn`` for each query in parallel. 

52 

53 Args: 

54 queries: Queries to search. If empty, returns an empty list 

55 immediately (and logs a warning). 

56 search_fn: Callable invoked as ``search_fn(query)`` inside a worker 

57 thread. Callers are responsible for wrapping it with any 

58 context-preserving decorators (e.g. 

59 :func:`preserve_research_context`) before passing it in, and 

60 for their own error handling (the callable should never raise 

61 — return an empty payload on failure instead, matching the 

62 pre-existing contract of the strategies this was extracted 

63 from). 

64 max_workers: Size of the thread pool. Defaults to ``len(queries)`` 

65 when ``None``, matching the historical behavior of the 

66 source-based, focused-iteration, and progressive strategies. 

67 context_factory: Optional zero-arg callable returning a **fresh** 

68 Flask app context per call (e.g. 

69 :func:`~local_deep_research.utilities.threading_utils.thread_context`). 

70 Invoked once per query on the calling thread; each worker gets 

71 its own context. Passing a single shared context instead would 

72 crash under concurrency, so a factory (not a context) is 

73 required. ``None`` disables context propagation, preserving the 

74 behavior of strategies that never propagated a context. 

75 

76 Returns: 

77 A list of ``(query, payload)`` tuples in completion order. Each 

78 ``payload`` is whatever ``search_fn`` returned for that query. 

79 """ 

80 if not queries: 

81 logger.warning("No queries provided for parallel search") 

82 return [] 

83 

84 if max_workers is None: 

85 max_workers = len(queries) 

86 

87 @thread_with_app_context 

88 def _worker(query: str) -> Tuple[str, T]: 

89 return (query, search_fn(query)) 

90 

91 with concurrent.futures.ThreadPoolExecutor( 

92 max_workers=max_workers 

93 ) as executor: 

94 # context_factory() is called here on the calling thread, once per 

95 # query, so every worker receives its own fresh AppContext. Sharing 

96 # one instance across workers is what raises "Token was created in a 

97 # different Context". 

98 futures = [ 

99 executor.submit( 

100 _worker, 

101 context_factory() if context_factory is not None else None, 

102 q, 

103 ) 

104 for q in queries 

105 ] 

106 return [ 

107 future.result() 

108 for future in concurrent.futures.as_completed(futures) 

109 ]