Coverage for src/local_deep_research/advanced_search_system/tools/search_availability.py: 97%
113 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""LangChain adapter for the standalone search engine availability tracker."""
3import asyncio
4import time
6from langchain.agents.middleware import AgentMiddleware
7from langchain_core.messages import SystemMessage
9from ...exceptions import ResearchTerminatedException
10from ...security.secure_logging import logger
11from ...utilities.resource_utils import safe_close
12from ...web_search_engines.search_engine_base import (
13 BaseSearchEngine,
14 metrics_engine_name,
15)
16from ...web_search_engines.engine_availability import (
17 SearchEngineAvailability,
18 SearchFailure,
19 available_tool_names,
20 engine_failure,
21 engine_for_tool,
22 failure_from_exception,
23 http_status,
24 is_query_error,
25 search_engines_for_tools,
26)
29class NoResearchToolsAvailable(RuntimeError):
30 """Availability filtering removed every research tool from a request.
32 Raised instead of calling the model with an empty tool list: several
33 providers reject a request that carries tool-call history but defines
34 no tools. This covers only an empty list: if a fetch tool survives
35 while every search engine is disabled, the model keeps it, since it can
36 still read URLs from the query or from earlier results. It can then also
37 answer without searching; the prompt tells it to say when evidence is
38 insufficient. The message lists only safe disable explanations.
40 A cooldown is not a reason to stop: when every search engine is only
41 cooling down, the middleware waits for the soonest one instead (see
42 ``wait_for_cooldown``). Only disables that last the rest of the run, or
43 a wait that hits its bound, end up here.
44 """
47def wait_for_cooldown(
48 availability: SearchEngineAvailability,
49 engines: set[str],
50 *,
51 sleep=None,
52 check=None,
53) -> None:
54 """Block until one of ``engines`` leaves its cooldown, within the run budget.
56 Returns at once when an engine is already available or when every
57 remaining disable lasts the rest of the run. Once the soonest retry lies
58 beyond the run's remaining cooldown-wait budget, it disables the cooling
59 ``engines`` for the rest of the run (``end_cooldowns``) and returns
60 instead of waiting. A concurrent search can start another cooldown
61 during this wait; it must not leave the model with no search tool while
62 the run can still wait for a retry. Parallel waiters spend the shared
63 budget once. ``check`` (by default the tracker's ``check_cancelled``)
64 runs before each sleep of at most one second, so cancellation stays
65 responsive; a cancel already recorded on the tracker
66 (``note_cancelled``) ends the wait as well.
67 """
68 sleep = sleep or time.sleep
69 check = check or availability.check_cancelled
70 with availability.waiting():
71 while engines and engines <= availability.disabled().keys():
72 if availability.cancelled():
73 raise ResearchTerminatedException(
74 "Research was terminated by user during a cooldown wait"
75 )
76 if check is not None:
77 try:
78 check()
79 except ResearchTerminatedException:
80 # The run's termination flag is single-use; record the
81 # cancel so parallel waiters and subagents stop too.
82 availability.note_cancelled()
83 raise
84 remaining = availability.seconds_until_retry(engines)
85 if remaining is None:
86 return
87 if remaining > availability.wait_budget_remaining():
88 availability.end_cooldowns(engines)
89 return
90 sleep(min(1.0, max(remaining, 0.01)))
93def note_metrics_name(
94 availability: SearchEngineAvailability, engine_name: str, engine
95) -> None:
96 """Tell ``availability`` which name metrics record for ``engine``."""
97 # Full-content wrappers delegate the search, and its metrics, inward.
98 inner = getattr(engine, "web_search", None)
99 availability.note_metrics_name(
100 engine_name,
101 metrics_engine_name(
102 inner if isinstance(inner, BaseSearchEngine) else engine
103 ),
104 )
107def _record_skipped_search(
108 engine_name: str,
109 query: str,
110 availability: SearchEngineAvailability,
111 explanation: str,
112) -> None:
113 """Keep outages visible in search metrics although nothing was sent."""
114 try:
115 from ...metrics.search_tracker import (
116 SEARCH_STATUS_SKIPPED,
117 SearchTracker,
118 )
120 SearchTracker.record_search(
121 engine_name=availability.metrics_name(engine_name),
122 query=query,
123 success=False,
124 error_message=explanation,
125 status=SEARCH_STATUS_SKIPPED,
126 )
127 except Exception as exc:
128 logger.warning(
129 f"Could not record skipped {engine_name} search "
130 f"({type(exc).__name__})"
131 )
134def run_search(
135 engine_name: str,
136 query: str,
137 factory,
138 availability: SearchEngineAvailability,
139 *,
140 programmatic_mode: bool = False,
141):
142 """Return results or a disable explanation, with cleanup on every path.
144 A search skipped because the engine is disabled is recorded in search
145 metrics (unless ``programmatic_mode``, like the engine's own searches).
146 """
147 disabled = availability.state(engine_name)
148 if disabled is not None:
149 explanation = disabled.describe()
150 if not programmatic_mode:
151 _record_skipped_search(
152 engine_name, query, availability, explanation
153 )
154 return explanation
155 engine = None
156 try:
157 engine = factory()
158 if engine is not None: 158 ↛ 160line 158 didn't jump to line 160 because the condition on line 158 was always true
159 note_metrics_name(availability, engine_name, engine)
160 if engine is None: 160 ↛ 161line 160 didn't jump to line 161 because the condition on line 160 was never true
161 failure = SearchFailure(
162 "Could not initialize search engine; check its configuration"
163 )
164 else:
165 failure = engine_failure(engine)
166 if failure is None:
167 results = engine.run(query)
168 # Usable results win over a failed optional enrichment,
169 # fallback attempt or full-content request.
170 if isinstance(results, list) and results:
171 availability.record_success(engine_name)
172 return results
173 failure = engine_failure(engine)
174 if failure is None:
175 availability.record_success(engine_name)
176 return []
177 return availability.disable(engine_name, failure).describe()
178 except Exception as exc:
179 # Type name only: provider exceptions can embed request URLs
180 # carrying API keys, and loguru would render them.
181 logger.warning(f"{engine_name} search failed ({type(exc).__name__})")
182 if is_query_error(exc):
183 # This query was rejected; the engine still works for others.
184 return (
185 f"{engine_name} rejected this query "
186 f"(HTTP {http_status(exc)}); rephrase it and try again"
187 )
188 failure = failure_from_exception(exc) or SearchFailure(
189 "Search error; a later attempt may succeed"
190 )
191 return availability.disable(engine_name, failure).describe()
192 finally:
193 safe_close(engine, f"{engine_name} search engine", allow_none=True)
194 try:
195 from ...database.thread_local_session import (
196 cleanup_current_thread,
197 thread_session_manager,
198 )
200 if not thread_session_manager.in_scope():
201 cleanup_current_thread()
202 except Exception:
203 logger.exception("Could not clean up search tool thread")
206class SearchAvailabilityMiddleware(AgentMiddleware):
207 """Remove unavailable tools from every model request, not just the prompt.
209 Cooldowns are re-checked on every model call, so a tool comes back as
210 soon as its cooldown ends. When every search engine behind the request
211 is disabled and at least one of them only for a cooldown, the call
212 first waits (bounded) for the soonest cooldown via ``wait``, which
213 receives the set of engine names. The default is ``wait_for_cooldown``,
214 which checks the tracker's ``check_cancelled`` every second and charges
215 the run's cooldown-wait budget. Once the cancel is recorded on the
216 tracker (``note_cancelled``), every further model call raises
217 ``ResearchTerminatedException`` instead.
218 """
220 def __init__(self, availability, primary_engine, wait=None):
221 self.availability = availability
222 self.primary_engine = primary_engine
223 self.wait = wait or (
224 lambda engines: wait_for_cooldown(availability, engines)
225 )
227 def _filter_request(self, request):
228 # The run's termination flag is gone once the run has cleaned up, so
229 # a subagent back from a slow model call would otherwise keep going.
230 if self.availability.cancelled():
231 raise ResearchTerminatedException("Research was terminated by user")
232 names = [t.name for t in request.tools]
233 relevant = {
234 engine_for_tool(name, self.primary_engine) for name in names
235 } - {None}
236 # Archive fetching tracks provider health but cannot replace URL
237 # discovery. Preserve search recovery even when it remains available.
238 disabled = self.availability.disabled()
239 waiting_engines = search_engines_for_tools(names, self.primary_engine)
240 if self.availability.seconds_until_retry(waiting_engines) is None:
241 # Discovery may be permanently unavailable while a known URL's
242 # archive can recover. Keep its bounded cooldown wait as well.
243 waiting_engines = relevant
244 if (
245 waiting_engines
246 and waiting_engines <= disabled.keys()
247 and self.availability.seconds_until_retry(waiting_engines)
248 is not None
249 ):
250 # Nothing can use these providers now, but something will again: a
251 # transient failure must not end the run or leave the model
252 # answering from memory. Wait instead of filtering.
253 self.wait(waiting_engines)
254 disabled = self.availability.disabled()
255 available = available_tool_names(names, self.primary_engine, disabled)
256 tools = [t for t in request.tools if t.name in available]
257 explanations = [
258 state.describe()
259 for name, state in disabled.items()
260 if name in relevant
261 ]
262 if not explanations:
263 return request
264 if not tools:
265 raise NoResearchToolsAvailable(
266 "No research tools remain available:\n"
267 + "\n".join(explanations)
268 )
269 guidance = (
270 "\nSearch engine availability for this step:\n"
271 + "\n".join(explanations)
272 + "\nUse only the tools currently supplied. If the primary search "
273 "is disabled, use another available source. Do not call disabled "
274 "tools or wait for cooldowns; continue with available tools and "
275 "sources already gathered. If evidence is insufficient, say so."
276 )
277 if request.system_message is None:
278 message = SystemMessage(content=guidance)
279 else:
280 # Preserve structured content and any provider-specific metadata.
281 content = request.system_message.content
282 content = (
283 content + guidance
284 if isinstance(content, str)
285 else [*content, {"type": "text", "text": guidance}]
286 )
287 message = request.system_message.model_copy(
288 update={"content": content}
289 )
290 return request.override(tools=tools, system_message=message)
292 def wrap_model_call(self, request, handler):
293 return handler(self._filter_request(request))
295 async def awrap_model_call(self, request, handler):
296 # Filtering may wait out a cooldown; keep that off the event loop.
297 filtered = await asyncio.to_thread(self._filter_request, request)
298 return await handler(filtered)