Coverage for src/local_deep_research/web_search_engines/rate_limiting/llm/wrapper.py: 92%
286 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Rate-limited wrapper for LLM calls.
3"""
5import inspect
6from typing import Optional
7from urllib.parse import urlparse
9from ....security import (
10 redact_url_for_log,
11 scrub_error,
12 url_authority_without_userinfo,
13)
14from ....security.secure_logging import logger
15from tenacity import (
16 retry,
17 stop_after_attempt,
18 retry_if_exception,
19)
20from tenacity.wait import wait_base
22from ....llm.providers.base import normalize_provider
24from ..tracker import get_tracker
25from ..exceptions import RateLimitError
26from .detection import is_llm_rate_limit_error, extract_retry_after
29class AdaptiveLLMWait(wait_base):
30 """Adaptive wait strategy for LLM rate limiting."""
32 def __init__(self, tracker, engine_type: str):
33 self.tracker = tracker
34 self.engine_type = engine_type
35 self.last_error = None
37 def __call__(self, retry_state) -> float:
38 # Store the error for potential retry-after extraction
39 if retry_state.outcome and retry_state.outcome.failed:
40 self.last_error = retry_state.outcome.exception()
42 # Get adaptive wait time from tracker
43 wait_time: float = self.tracker.get_wait_time(self.engine_type)
45 # If we have a retry-after from the error, use it
46 if self.last_error:
47 retry_after = extract_retry_after(self.last_error)
48 if retry_after > 0:
49 wait_time = max(wait_time, float(retry_after))
51 logger.info(
52 f"LLM rate limit wait for {self.engine_type}: {wait_time:.2f}s"
53 )
54 return wait_time
57def create_rate_limited_llm_wrapper(base_llm, provider: Optional[str] = None):
58 """
59 Create a rate-limited wrapper around an LLM instance.
61 Args:
62 base_llm: The base LLM instance to wrap
63 provider: Optional provider name (e.g., 'openai', 'anthropic')
65 Returns:
66 A wrapped LLM instance with rate limiting capabilities
67 """
69 class RateLimitedLLMWrapper:
70 """Wrapper that adds rate limiting to LLM calls."""
72 # invoke, ainvoke, stream, astream, batch and abatch pass their
73 # arguments, the call's RunnableConfig included, to base_llm
74 # unchanged. ProcessingLLMWrapper relies on this to see which
75 # runnable receives that config (config/llm_config.py).
76 _ldr_forwards_call_arguments = True
78 def __init__(self, llm, provider_name: Optional[str] = None):
79 self.base_llm = llm
80 self.provider = provider_name
81 self.rate_limiter = None
82 # Cancellation scope (see llm/cancellation.py). Bound via
83 # ``bind_llm_to_research`` by the outer ProcessingLLMWrapper;
84 # falls back to the thread's search context per call.
85 self._ldr_research_id = getattr(llm, "_ldr_research_id", None)
87 # Only setup rate limiting if enabled
88 if self._should_rate_limit(): 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true
89 self.rate_limiter = get_tracker()
90 logger.info(
91 f"Rate limiting enabled for LLM provider: {self._get_rate_limit_key()}"
92 )
94 def _effective_research_id(self):
95 """Resolve cancellation scope: binding > thread context."""
96 try:
97 from ....llm.cancellation import effective_research_id
99 return effective_research_id(self, self._ldr_research_id)
100 except Exception:
101 return getattr(self, "_ldr_research_id", None)
103 def _check_terminated(self, research_id=None):
104 """Raise ``ResearchTerminatedException`` if Stop was requested."""
105 from ....llm.cancellation import raise_if_terminated
107 rid = (
108 research_id
109 if research_id is not None
110 else self._effective_research_id()
111 )
112 if rid:
113 raise_if_terminated(rid)
115 def _should_rate_limit(self) -> bool:
116 """Check if rate limiting should be applied to this LLM."""
117 # Rate limiting for LLMs is currently disabled by default
118 # TODO: Pass settings_snapshot to enable proper configuration
119 #
120 # Blocker for flipping this on, now that ainvoke() is a real
121 # async path: tenacity's AsyncRetrying._run_wait awaits the wait
122 # strategy through wrap_to_async_func, which for a SYNC callable
123 # just calls it inline — so AdaptiveLLMWait.__call__ runs ON the
124 # event loop, and it reaches tracker.get_wait_time ->
125 # _ensure_estimates_loaded, which opens the user's encrypted
126 # (SQLCipher) DB. Only the sleep between attempts is non-blocking.
127 # Do not enable this without first making AdaptiveLLMWait
128 # loop-safe (async wait strategy, or offload the tracker read).
129 return False
131 def _check_if_local_model(self) -> bool:
132 """Check if the LLM is a local model that shouldn't be rate limited."""
133 # Don't rate limit local models
134 local_providers = [
135 "ollama",
136 "lmstudio",
137 "llamacpp",
138 "local",
139 "none",
140 ]
141 if normalize_provider(self.provider) in local_providers:
142 logger.debug(
143 f"Skipping rate limiting for local provider: {self.provider}"
144 )
145 return True
147 # Check if base URL indicates local model
148 if hasattr(self.base_llm, "base_url"):
149 base_url = str(self.base_llm.base_url)
150 if any(
151 local in base_url
152 # Classify a configured local service URL; no listening socket is opened.
153 for local in ["localhost", "127.0.0.1", "0.0.0.0"] # nosec B104
154 ):
155 logger.debug(
156 "Skipping rate limiting for local URL: "
157 f"{redact_url_for_log(base_url)}"
158 )
159 return True
161 return False
163 def _get_rate_limit_key(self) -> str:
164 """Build composite key: provider-url-model"""
165 provider = self.provider or "unknown"
167 # Extract URL
168 url = "unknown"
169 if hasattr(self.base_llm, "base_url"):
170 url = str(self.base_llm.base_url)
171 elif hasattr(self.base_llm, "_client") and hasattr(
172 self.base_llm._client, "base_url"
173 ):
174 url = str(self.base_llm._client.base_url)
176 # Clean URL: remove protocol and trailing slashes
177 if url != "unknown":
178 parsed = urlparse(url)
179 # Authority minus userinfo: the key is logged. A
180 # scheme-less value keeps the old path fallback unless
181 # that path could hold userinfo.
182 if parsed.netloc:
183 url = url_authority_without_userinfo(url) or "unknown"
184 elif "@" in parsed.path:
185 url = "unknown"
186 else:
187 url = parsed.path
188 url = url.rstrip("/")
190 # Extract model
191 model = "unknown"
192 if hasattr(self.base_llm, "model_name"):
193 model = str(self.base_llm.model_name)
194 elif hasattr(self.base_llm, "model"):
195 model = str(self.base_llm.model)
197 # Clean model name
198 model = model.replace("/", "-").replace(":", "-")
200 return f"{provider}-{url}-{model}"
202 def invoke(self, *args, **kwargs):
203 """Invoke the LLM with rate limiting if enabled."""
204 self._check_terminated()
205 result = self._call_rate_limited(
206 lambda: self.base_llm.invoke(*args, **kwargs)
207 )
208 self._check_terminated()
209 return result
211 async def ainvoke(self, *args, **kwargs):
212 """Async invoke with the same rate limiting and scrub as invoke().
214 Without this, ``ainvoke`` falls through ``__getattr__`` to the
215 base LLM and silently skips the adaptive wait/retry loop and
216 the RateLimitError wrapping — a behavioural hole for any caller
217 that reaches this async twin, including the fan-out path daemon
218 code submits to the process's one shared event loop (ADR-0013,
219 docs/decisions/0013-research-runs-keep-daemon-thread.md); the
220 sync ``invoke`` stays the canonical production path on the
221 research thread. The ``@retry`` decorator detects the coroutine
222 and sleeps between attempts
223 with ``asyncio.sleep`` (tenacity AsyncRetrying), so the sleep
224 itself does not block the loop. The wait STRATEGY that
225 computes that sleep still runs synchronously on the loop —
226 see the note on ``_should_rate_limit``. ``_acall_rate_limited``
227 also calls ``tracker.record_outcome`` (and from it
228 ``_update_estimate``) synchronously on the loop after each
229 successful call and each rate-limit failure; those can read
230 the settings context and the search context and persist
231 through the metrics database. Both are moot today because
232 ``self.rate_limiter`` is always ``None``.
233 See #6232/#6246/#6249/#6268/#6347 for the history.
234 """
235 self._check_terminated()
236 result = await self._acall_rate_limited(
237 lambda: self.base_llm.ainvoke(*args, **kwargs)
238 )
239 self._check_terminated()
240 return result
242 def _scrub_or_raise(self, error):
243 """Re-raise with scrubbing: provider 429 bodies can echo the
244 Authorization header, so the body must not reach logs or
245 exception chains unscrubbed. Non-rate-limit errors pass
246 through unchanged."""
247 if is_llm_rate_limit_error(error):
248 logger.warning("LLM rate limit error detected")
249 safe_msg = scrub_error(error)
250 raise RateLimitError(f"LLM rate limit: {safe_msg}") from None
251 raise error
253 def _call_rate_limited(self, call):
254 """Run one base-LLM call with retries and tracker bookkeeping
255 when rate limiting is enabled; scrub rate-limit errors either
256 way."""
257 if not self.rate_limiter:
258 try:
259 return call()
260 except Exception as e:
261 self._scrub_or_raise(e)
263 rate_limit_key = self._get_rate_limit_key()
264 tracker = self.rate_limiter
265 if tracker is None: # defensive: narrowed above in this branch 265 ↛ 266line 265 didn't jump to line 266 because the condition on line 265 was never true
266 return call()
268 @retry(
269 wait=AdaptiveLLMWait(tracker, rate_limit_key),
270 stop=stop_after_attempt(3),
271 retry=retry_if_exception(is_llm_rate_limit_error),
272 )
273 def _with_retry():
274 try:
275 return call()
276 except Exception as e:
277 self._scrub_or_raise(e)
279 try:
280 result = _with_retry()
281 tracker.record_outcome(
282 engine_type=rate_limit_key,
283 wait_time=0,
284 success=True,
285 retry_count=0,
286 )
287 return result
288 except Exception as e:
289 if is_llm_rate_limit_error(e):
290 tracker.record_outcome(
291 engine_type=rate_limit_key,
292 wait_time=0,
293 success=False,
294 retry_count=0,
295 )
296 raise
298 def stream(self, *args, **kwargs):
299 """Stream through the base LLM, scrubbing failures wherever they
300 surface.
302 `BaseChatModel.stream` is a generator function, so calling it
303 returns a generator without making the request: a 429 is raised on
304 the first `next()`, after the call has already returned. Wrapping
305 only the call therefore protects nothing, which is why this
306 iterates the source here and catches during iteration.
308 Chunks are forwarded untouched: a credential echoed in a 429 body
309 can straddle chunk boundaries, so scrubbing individual chunks is
310 unsound. Only exceptions go through the scrub.
311 """
313 def _source():
314 try:
315 return self.base_llm.stream(*args, **kwargs)
316 except Exception as e:
317 self._scrub_or_raise(e)
319 research_id = self._effective_research_id()
320 if research_id:
321 self._check_terminated(research_id)
322 chunks = iter(self._call_rate_limited(_source))
323 try:
324 while True:
325 # Only the advance is guarded, and `yield` stays outside the
326 # Only the advance is guarded, and `yield` stays outside the
327 # inner `try`: a handler that encloses the yield also catches
328 # whatever the CONSUMER throws back in, which would scrub an
329 # unrelated caller-side failure as if it were a provider 429.
330 # The outer handler only closes the upstream and
331 # re-raises the same object it caught, so a consumer
332 # `throw()` still propagates untouched to the caller.
333 # The termination check sits after next() and before yield:
334 # a Stop aborts without yielding the just-pulled chunk and
335 # without draining the rest of the stream.
336 try:
337 chunk = next(chunks)
338 except StopIteration:
339 break
340 except Exception as e:
341 self._scrub_or_raise(e)
342 if research_id:
343 self._check_terminated(research_id)
344 yield chunk
345 except BaseException as exc:
346 # Close the upstream generator on ANY exit — early consumer
347 # close (GeneratorExit lands at the yield above) or an
348 # exception — so its connection does not linger until GC
349 # (#6607).
350 self._close_upstream(chunks, exc)
351 raise
352 self._close_upstream(chunks, None)
354 def _close_upstream(self, chunks, body_exc):
355 """Close a sync upstream iterator after stream() stops.
357 The upstream's own cleanup can raise (e.g. a LangChain callback
358 handler failing in on_llm_error), and that text may carry
359 provider or credential material, so any close failure is
360 handled by ``_handle_close_error`` rather than left to
361 propagate raw. Only an explicit consumer close (GeneratorExit)
362 still surfaces the (scrubbed) close failure; a normal
363 exhaustion or an unrelated body exception both keep the close
364 failure out of the caller's way. The attribute lookup itself
365 is inside the guarded try: a proxy whose ``__getattr__`` raises
366 must not replace whatever is already propagating. Tolerant of
367 non-generator iterables (mocks, list-backed sources) that have
368 no close().
369 """
370 try:
371 closer = getattr(chunks, "close", None)
372 if closer is None:
373 return
374 closer()
375 except Exception as close_exc:
376 self._handle_close_error(close_exc, body_exc)
378 async def _aclose_upstream(self, chunks, body_exc):
379 """Async twin of _close_upstream() for astream()."""
380 # Skip an upstream that is already finished or mid-close: at
381 # loop shutdown / cyclic GC the event loop's asyncgen finalizer
382 # may already be running its aclose(), and a second one raises
383 # "aclose(): asynchronous generator is already running".
384 # Only the order where the loop's aclose starts first can be
385 # guarded here. When this aclose starts first and the
386 # upstream's cleanup suspends (e.g. an async callback handler
387 # awaiting in on_llm_error), the loop's own later aclose hits
388 # that RuntimeError inside asyncio and is logged by the loop's
389 # exception handler. That is benign: the upstream is still
390 # closed exactly once, and the log carries no provider text.
391 # It needs a consumer that abandons astream() without aclose().
392 #
393 # All the lookups below (ag_running, ag_frame, aclose) live
394 # inside this same guarded try as the close call itself: a
395 # proxy whose __getattr__ raises must not replace whatever
396 # exception is already propagating.
397 try:
398 if getattr(chunks, "ag_running", False) is True:
399 return
400 if inspect.isasyncgen(chunks) and chunks.ag_frame is None:
401 return
402 acloser = getattr(chunks, "aclose", None)
403 if acloser is None: 403 ↛ 404line 403 didn't jump to line 404 because the condition on line 403 was never true
404 return
405 result = acloser()
406 # Plain callables (e.g. a Mock source) return a
407 # non-awaitable; only real coroutines are awaited.
408 if inspect.isawaitable(result):
409 await result
410 except Exception as close_exc:
411 self._handle_close_error(close_exc, body_exc)
413 def _handle_close_error(self, close_exc, body_exc):
414 """Decide whether a close() failure should surface.
416 Only an explicit consumer close (``body_exc`` is
417 ``GeneratorExit``) keeps the current scrub-and-raise
418 behaviour. A normal exhaustion (``body_exc`` is ``None``) or
419 any other in-flight exception must not be replaced by the
420 close failure: it is logged at debug with the close
421 exception's type name only and dropped.
422 """
423 if isinstance(body_exc, GeneratorExit):
424 self._scrub_or_raise(close_exc)
425 return
426 logger.debug(
427 "Suppressed {} while closing the upstream stream",
428 type(close_exc).__name__,
429 )
431 async def astream(self, *args, **kwargs):
432 """Async generator implementing LangChain's astream contract:
433 ``async for chunk in model.astream(...)``.
435 Named `astream`, not `_astream`: `__getattr__` forwards any name
436 this class does not define straight to the raw model, so an
437 underscore-prefixed twin is never reached and the surface stays
438 unscrubbed while looking covered.
440 `BaseChatModel.astream` is an async generator function, so the
441 provider error arrives on the first `__anext__()` rather than at
442 call time — each advance is guarded for that reason.
443 Chunks are forwarded untouched: a credential echoed in a 429 body
444 can straddle chunk boundaries, so per-chunk scrubbing is unsound.
445 """
446 research_id = self._effective_research_id()
447 if research_id:
448 self._check_terminated(research_id)
449 try:
450 chunks = self.base_llm.astream(*args, **kwargs).__aiter__()
451 except Exception as e:
452 self._scrub_or_raise(e)
454 try:
455 while True:
456 # `yield` deliberately outside the inner `try`; see
457 # stream(). The outer handler re-raises the same object,
458 # so a consumer `throw()` still propagates untouched.
459 # Termination is checked after each advance so a Stop
460 # aborts without yielding the just-pulled chunk.
461 try:
462 chunk = await chunks.__anext__()
463 except StopAsyncIteration:
464 break
465 except Exception as e:
466 self._scrub_or_raise(e)
467 if research_id:
468 self._check_terminated(research_id)
469 yield chunk
470 except BaseException as exc:
471 # Async twin of stream()'s cleanup: close the upstream
472 # async generator on any exit instead of leaking the
473 # connection until GC or loop shutdown (#6607).
474 await self._aclose_upstream(chunks, exc)
475 raise
476 await self._aclose_upstream(chunks, None)
478 async def _acall_rate_limited(self, call):
479 """Async variant of _call_rate_limited."""
480 if not self.rate_limiter:
481 try:
482 return await call()
483 except Exception as e:
484 self._scrub_or_raise(e)
486 rate_limit_key = self._get_rate_limit_key()
487 tracker = self.rate_limiter
488 if tracker is None: # defensive: narrowed above in this branch 488 ↛ 489line 488 didn't jump to line 489 because the condition on line 488 was never true
489 return await call()
491 @retry(
492 wait=AdaptiveLLMWait(tracker, rate_limit_key),
493 stop=stop_after_attempt(3),
494 retry=retry_if_exception(is_llm_rate_limit_error),
495 )
496 async def _with_retry():
497 try:
498 return await call()
499 except Exception as e:
500 self._scrub_or_raise(e)
502 try:
503 result = await _with_retry()
504 tracker.record_outcome(
505 engine_type=rate_limit_key,
506 wait_time=0,
507 success=True,
508 retry_count=0,
509 )
510 return result
511 except Exception as e:
512 if is_llm_rate_limit_error(e):
513 tracker.record_outcome(
514 engine_type=rate_limit_key,
515 wait_time=0,
516 success=False,
517 retry_count=0,
518 )
519 raise
521 @staticmethod
522 def _configs_for_inputs(config, inputs):
523 """Resolve per-item configs from a shared config or a list of
524 per-input configs (LangChain Runnable.batch takes either)."""
525 if isinstance(config, list):
526 if len(config) != len(inputs):
527 raise ValueError(
528 "configs must match the number of inputs "
529 f"({len(config)} != {len(inputs)})"
530 )
531 return list(config)
532 return [config] * len(inputs)
534 def batch(
535 self, inputs, config=None, *, return_exceptions=False, **kwargs
536 ):
537 """Batch through invoke(), forwarding config (shared or
538 per-input) so caller callbacks/tags/metadata survive.
540 Only ``Exception`` is captured for ``return_exceptions=True`` —
541 ``ResearchTerminatedException`` (a ``BaseException``) always
542 propagates so a Stop is never swallowed into a result list.
543 """
544 configs = self._configs_for_inputs(config, inputs)
545 results = []
546 for msg, conf in zip(inputs, configs):
547 try:
548 results.append(self.invoke(msg, conf, **kwargs))
549 except Exception as e:
550 if return_exceptions:
551 results.append(e)
552 else:
553 raise
554 return results
556 async def abatch(
557 self, inputs, config=None, *, return_exceptions=False, **kwargs
558 ):
559 """Async counterpart of batch() (same BaseException guarantee)."""
560 configs = self._configs_for_inputs(config, inputs)
561 results = []
562 for msg, conf in zip(inputs, configs):
563 try:
564 results.append(await self.ainvoke(msg, conf, **kwargs))
565 except Exception as e:
566 if return_exceptions:
567 results.append(e)
568 else:
569 raise
570 return results
572 def bind_tools(self, tools, **kwargs):
573 """Bind tools and re-wrap so the bound runnable keeps rate
574 limiting and scrubbing (same shape as
575 ProcessingLLMWrapper.bind_tools, #4804)."""
576 bound = self.base_llm.bind_tools(tools, **kwargs)
577 wrapped = create_rate_limited_llm_wrapper(bound, self.provider)
578 try:
579 wrapped._ldr_research_id = self._effective_research_id()
580 inner = getattr(wrapped, "base_llm", None)
581 if inner is not None: 581 ↛ 596line 581 didn't jump to line 596 because the condition on line 581 was always true
582 try:
583 setattr(
584 inner,
585 "_ldr_research_id",
586 wrapped._ldr_research_id,
587 )
588 except Exception:
589 logger.debug(
590 "bind_tools inner scope propagation failed (non-critical)"
591 )
592 except Exception:
593 logger.debug(
594 "bind_tools scope propagation failed (non-critical)"
595 )
596 return wrapped
598 def with_structured_output(self, schema, **kwargs):
599 """Wrap the structured-output runnable the same way as
600 bind_tools."""
601 bound = self.base_llm.with_structured_output(schema, **kwargs)
602 wrapped = create_rate_limited_llm_wrapper(bound, self.provider)
603 try:
604 wrapped._ldr_research_id = self._effective_research_id()
605 except Exception:
606 logger.debug(
607 "with_structured_output scope propagation failed (non-critical)"
608 )
609 return wrapped
611 # Pass through any other attributes to the base LLM
612 def __getattr__(self, name):
613 return getattr(self.base_llm, name)
615 def close(self):
616 """Close underlying HTTP clients held by this LLM. Idempotent."""
617 try:
618 from ....utilities.llm_utils import _close_base_llm
620 _close_base_llm(self.base_llm)
621 except Exception:
622 logger.debug(
623 "best-effort cleanup of HTTP clients on shutdown",
624 exc_info=True,
625 )
627 def __str__(self):
628 return f"RateLimited({str(self.base_llm)})"
630 def __repr__(self):
631 return f"RateLimitedLLMWrapper({repr(self.base_llm)})"
633 return RateLimitedLLMWrapper(base_llm, provider)