Coverage for src/local_deep_research/web_search_engines/rate_limiting/tracker.py: 93%
408 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Adaptive rate limit tracker that learns optimal retry wait times for each search engine.
3"""
5import random
6import threading
7import time
8from collections import deque
9from contextlib import AbstractContextManager, contextmanager, nullcontext
10from typing import (
11 TYPE_CHECKING,
12 Any,
13 Dict,
14 Iterator,
15 List,
16 Optional,
17 Tuple,
18)
19from cachetools import LRUCache
21if TYPE_CHECKING:
22 from sqlalchemy.orm import Session
23from ...security.secure_logging import logger
25from ...config.thread_settings import (
26 NoSettingsContextError,
27 get_setting_from_snapshot,
28 get_settings_context,
29)
30from ...settings.env_registry import is_ci_environment
31from ...utilities.thread_context import get_search_context
32from ...utilities.metric_write_health import record_metric_write_drop
34# Lazy imports to avoid database initialization in programmatic mode
35_db_imports = None
38def _get_db_imports():
39 """Lazy load database imports only when needed."""
40 global _db_imports
41 if _db_imports is None:
42 try:
43 from ...database.models import RateLimitAttempt, RateLimitEstimate
44 from ...database.thread_local_session import thread_session_manager
46 _db_imports = {
47 "RateLimitAttempt": RateLimitAttempt,
48 "RateLimitEstimate": RateLimitEstimate,
49 "thread_session_manager": thread_session_manager,
50 }
51 except (ImportError, RuntimeError):
52 # Database not available - programmatic mode
53 _db_imports = {}
54 return _db_imports
57class AdaptiveRateLimitTracker:
58 """
59 Tracks and learns optimal retry wait times for each search engine.
60 Persists learned patterns to the main application database using SQLAlchemy.
61 """
63 def __init__(self, settings_snapshot=None, programmatic_mode=False):
64 self.settings_snapshot = settings_snapshot or {}
65 self.programmatic_mode = programmatic_mode
67 # Helper function to get settings with defaults
68 def get_setting_or_default(key, default, type_fn=None):
69 try:
70 value = get_setting_from_snapshot(
71 key,
72 settings_snapshot=self.settings_snapshot,
73 )
74 # Handle None values - return default instead of calling type_fn(None)
75 if value is None:
76 return default
77 return type_fn(value) if type_fn else value
78 except NoSettingsContextError:
79 return default
81 # Get settings with explicit defaults
82 self.memory_window = get_setting_or_default(
83 "rate_limiting.memory_window", 100, int
84 )
85 self.exploration_rate = get_setting_or_default(
86 "rate_limiting.exploration_rate", 0.1, float
87 )
88 self.learning_rate = get_setting_or_default(
89 "rate_limiting.learning_rate", 0.3, float
90 )
91 self.decay_per_day = get_setting_or_default(
92 "rate_limiting.decay_per_day", 0.95, float
93 )
95 # In programmatic mode, default to disabled
96 self.enabled = get_setting_or_default(
97 "rate_limiting.enabled",
98 not self.programmatic_mode, # Default based on mode
99 bool,
100 )
102 profile = get_setting_or_default("rate_limiting.profile", "balanced")
104 if self.programmatic_mode and self.enabled:
105 logger.info(
106 "Rate limiting enabled in programmatic mode - using memory-only tracking without persistence"
107 )
109 # Apply rate limiting profile
110 self._apply_profile(profile)
112 # Bounded in-memory caches for fast access (prevent memory leaks)
113 # Lock required: LRUCache is not thread-safe (mutates on read via LRU reordering)
114 self._cache_lock = threading.Lock()
115 self.recent_attempts: LRUCache = LRUCache(maxsize=100)
116 self.current_estimates: LRUCache = LRUCache(maxsize=100)
118 # Initialize the _estimates_loaded flag
119 self._estimates_loaded = False
121 # Load estimates from database
122 self._load_estimates()
124 logger.info(
125 f"AdaptiveRateLimitTracker initialized: enabled={self.enabled}, profile={profile}"
126 )
128 def _apply_profile(self, profile: str) -> None:
129 """Apply rate limiting profile settings."""
130 if profile == "conservative":
131 # More conservative: lower exploration, slower learning
132 self.exploration_rate = min(
133 self.exploration_rate * 0.5, 0.05
134 ) # 5% max exploration
135 self.learning_rate = min(
136 self.learning_rate * 0.7, 0.2
137 ) # Slower learning
138 logger.info("Applied conservative rate limiting profile")
139 elif profile == "aggressive":
140 # More aggressive: higher exploration, faster learning
141 self.exploration_rate = min(
142 self.exploration_rate * 1.5, 0.2
143 ) # Up to 20% exploration
144 self.learning_rate = min(
145 self.learning_rate * 1.3, 0.5
146 ) # Faster learning
147 logger.info("Applied aggressive rate limiting profile")
148 else: # balanced
149 # Use settings as-is
150 logger.info("Applied balanced rate limiting profile")
152 def _load_estimates(self) -> None:
153 """Load estimates from database into memory."""
154 # Skip database operations in programmatic mode
155 if self.programmatic_mode:
156 logger.debug(
157 "Skipping rate limit estimate loading in programmatic mode"
158 )
159 self._estimates_loaded = (
160 True # Mark as loaded to skip future attempts
161 )
162 return
164 # During initialization, we don't have user context yet
165 # Mark that we need to load estimates when user context becomes available
166 self._estimates_loaded = False
167 logger.debug(
168 "Rate limit estimates will be loaded on-demand when user context is available"
169 )
171 def _resolve_session(
172 self, username: str, password: str
173 ) -> AbstractContextManager["Session"]:
174 """Unified session resolution helper for database operations.
176 Reuses the thread-local session when an enclosing scope is active
177 (avoiding a second QueuePool connection and SQLite busy-timeout deadlock);
178 otherwise falls back to a dedicated session that commits and closes on exit.
179 Raises ``RuntimeError`` when the active scope holds no session for
180 ``username``, for example because it belongs to another user.
181 """
182 db_imports = _get_db_imports() or {}
183 thread_session_manager = db_imports.get("thread_session_manager")
184 if thread_session_manager is None:
185 try:
186 from ...database.thread_local_session import (
187 thread_session_manager,
188 )
189 except (ImportError, RuntimeError):
190 thread_session_manager = None
192 if (
193 thread_session_manager is not None
194 and thread_session_manager.in_scope()
195 ):
196 session = thread_session_manager.get_scoped_session(username)
197 if session is None:
198 raise RuntimeError(
199 "Rate telemetry requires a matching active owner session"
200 )
201 # Borrow without revalidation, replacement, commit, or rollback.
202 return nullcontext(session)
204 from ...database.thread_metrics import metrics_writer
206 metrics_writer.set_user_password(username, password)
207 return metrics_writer.get_session(username)
209 @contextmanager
210 def _isolated_begin_nested(self, session: "Session") -> Iterator[None]:
211 """Open a SAVEPOINT without flushing pre-existing pending ORM state.
213 ``Session.begin_nested()`` flushes every pending change before it
214 issues SAVEPOINT. If another connection holds SQLite's writer lock,
215 that flush fails outside the savepoint and SQLAlchemy invalidates
216 the caller's transaction. Detach already-pending instances first so
217 tracker writes cannot flush them until after the isolation boundary
218 exists; restore them after the savepoint settles. The savepoint
219 joins the caller's open write transaction, or commits at once on
220 its own when the scope has not written (``telemetry_savepoint``).
222 Pending rate-limit records skip this operation, because loading or
223 deleting their persisted counterparts would prevent safe restoration.
224 Other pending caller state retains the existing detach/restore path.
226 Both expunge and restore handle errors: cascade-expelled
227 children (``cascade="all, delete-orphan"`` such as
228 ``ResearchHistory.resources``) are skipped instead of raising
229 ``InvalidRequestError``, and any failure during expunge or restore
230 is contained so a half-failed tracker write cannot leave caller
231 objects detached from the session.
232 """
233 from ...database.models import RateLimitAttempt, RateLimitEstimate
234 from ...database.session_context import telemetry_savepoint
236 # Detaching a pending metrics row lets the tracker load another
237 # instance with the same identity, making restoration fail. A bulk
238 # reset/delete can likewise invalidate a caller's staged metrics edit.
239 # Keep those rows attached and skip this best-effort operation.
240 pending = (*session.new, *session.dirty, *session.deleted)
241 if any(
242 isinstance(row, (RateLimitAttempt, RateLimitEstimate))
243 for row in pending
244 ):
245 raise RuntimeError(
246 "Rate telemetry cannot isolate pending rate-limit records"
247 )
249 pending_new = list(getattr(session, "new", ()) or ())
250 pending_deleted = list(getattr(session, "deleted", ()) or ())
251 skip_ids = {id(obj) for obj in pending_new}
252 skip_ids.update(id(obj) for obj in pending_deleted)
253 pending_dirty = [
254 obj
255 for obj in list(getattr(session, "dirty", ()) or ())
256 if id(obj) not in skip_ids
257 ]
259 # Dedupe by identity and skip objects already detached (e.g. by an
260 # earlier cascade expunge of a parent). ``expunge`` cascades along
261 # ``cascade="all, delete-orphan"`` relationships, so a parent +
262 # child both pending must not be expunged twice.
263 ordered: list[Any] = []
264 seen: set[int] = set()
265 for obj in (*pending_new, *pending_dirty, *pending_deleted):
266 if id(obj) in seen: 266 ↛ 267line 266 didn't jump to line 267 because the condition on line 266 was never true
267 continue
268 seen.add(id(obj))
269 ordered.append(obj)
271 for obj in ordered:
272 if obj not in session:
273 # Cascade already detached this object when its parent was
274 # expunged; it remains in pending_new/dirty/deleted by
275 # Python identity, but not by session membership.
276 continue
277 try:
278 session.expunge(obj)
279 except Exception:
280 # ``expunge`` can raise for already-detached instances or
281 # for objects the session has lost track of. Swallow and
282 # continue so any later objects can still be detached and
283 # later we still attempt full restoration below.
284 logger.debug(
285 "Tracker isolation: expunge failed during pre-flush detach",
286 exc_info=True,
287 )
289 try:
290 with telemetry_savepoint(session):
291 yield
292 finally:
293 self._restore_pending_state(
294 session, pending_new, pending_dirty, pending_deleted
295 )
297 def _restore_pending_state(
298 self,
299 session: "Session",
300 pending_new: list[Any],
301 pending_dirty: list[Any],
302 pending_deleted: list[Any],
303 ) -> None:
304 """Best-effort restoration of expunged caller objects.
306 Adds new/dirty objects back to the session and re-stages deleted
307 objects. ``add(parent)`` cascades to children for relationships
308 declared with ``cascade="all, delete-orphan"``, so we skip any
309 child already re-attached by cascade to avoid double-add errors.
310 All restoration steps are individually exception-safe: a failure
311 here must never propagate out and detach caller state.
312 """
313 for obj in (*pending_new, *pending_dirty):
314 try:
315 if obj not in session:
316 session.add(obj)
317 except Exception:
318 logger.debug(
319 "Tracker isolation: failed to re-attach caller object",
320 exc_info=True,
321 )
322 for obj in pending_deleted: 322 ↛ 323line 322 didn't jump to line 323 because the loop on line 322 never started
323 try:
324 if obj not in session:
325 session.add(obj)
326 except Exception:
327 logger.debug(
328 "Tracker isolation: failed to re-add deleted caller object",
329 exc_info=True,
330 )
331 continue
332 try:
333 session.delete(obj)
334 except Exception:
335 logger.debug(
336 "Tracker isolation: failed to re-stage deleted caller object",
337 exc_info=True,
338 )
340 def _ensure_estimates_loaded(self) -> None:
341 """Load estimates from user's encrypted database if not already loaded."""
342 # Early return if already loaded or should skip
343 if self._estimates_loaded or self.programmatic_mode:
344 if not self._estimates_loaded:
345 self._estimates_loaded = True
346 return
348 # Get database imports
349 db_imports = _get_db_imports()
350 RateLimitEstimate = (
351 db_imports.get("RateLimitEstimate") if db_imports else None
352 )
354 if not db_imports or not RateLimitEstimate:
355 # Database not available
356 self._estimates_loaded = True
357 return
359 # Try to get research context from search tracker
361 context = get_search_context()
362 if not context:
363 # No context available (e.g., in tests or programmatic access)
364 self._estimates_loaded = True
365 return
367 username = context.get("username")
368 password = context.get("user_password")
370 if username and password:
371 try:
372 session: "Session"
373 with self._resolve_session(username, password) as session:
374 # The production sessionmaker defaults to
375 # ``autoflush=True``, so a bare ``query()`` could flush
376 # the caller's pending ORM writes before the SELECT
377 # runs. Under a competing SQLite writer that flush
378 # fails and SQLAlchemy invalidates the outer
379 # transaction; ``no_autoflush`` keeps the read from
380 # touching unrelated caller state.
381 with session.no_autoflush:
382 estimates: list[Any] = session.query(
383 RateLimitEstimate
384 ).all()
386 for estimate in estimates:
387 # Apply decay for old estimates
388 age_hours = (
389 time.time() - estimate.last_updated
390 ) / 3600
391 decay = self.decay_per_day ** (age_hours / 24)
393 with self._cache_lock:
394 self.current_estimates[estimate.engine_type] = {
395 "base": estimate.base_wait_seconds,
396 "min": estimate.min_wait_seconds,
397 "max": estimate.max_wait_seconds,
398 "confidence": decay,
399 }
401 logger.debug(
402 f"Loaded estimate for {estimate.engine_type}: base={estimate.base_wait_seconds:.2f}s, confidence={decay:.2f}"
403 )
405 self._estimates_loaded = True
406 logger.info(
407 f"Loaded {len(estimates)} rate limit estimates from encrypted database"
408 )
410 except Exception:
411 logger.warning("Could not load rate limit estimates")
412 # Mark as loaded anyway to avoid repeated attempts
413 self._estimates_loaded = True
415 def get_wait_time(self, engine_type: str) -> float:
416 """
417 Get adaptive wait time for a search engine.
418 Includes exploration to discover better rates.
420 Args:
421 engine_type: Name of the search engine
423 Returns:
424 Wait time in seconds
425 """
426 # If rate limiting is disabled, return minimal wait time
427 if not self.enabled:
428 return 0.1
430 # Check if we have a user context - if not, handle appropriately
431 context = get_search_context()
432 if not context and not self.programmatic_mode:
433 # No context and not in programmatic mode - this is unexpected
434 logger.warning(
435 f"No user context available for rate limiting on {engine_type} "
436 "but programmatic_mode=False. Disabling rate limiting. "
437 "This may indicate a configuration issue."
438 )
439 return 0.0
441 # In programmatic mode, we continue with memory-only rate limiting even without context
443 # Ensure estimates are loaded from database
444 self._ensure_estimates_loaded()
446 with self._cache_lock:
447 if engine_type not in self.current_estimates:
448 # First time seeing this engine - start optimistic and learn from real responses
449 # Use engine-specific optimistic defaults only for what we know for sure
450 optimistic_defaults = {
451 "SearXNGSearchEngine": 0.1, # Self-hosted default engine
452 }
454 wait_time = optimistic_defaults.get(
455 engine_type, 0.1
456 ) # Default optimistic for others
457 logger.info(
458 f"No rate limit data for {engine_type}, starting optimistic with {wait_time}s"
459 )
460 return wait_time
462 estimate = self.current_estimates[engine_type]
463 base_wait = estimate["base"]
465 # Security: random used for non-security rate-limit jitter, not tokens or secrets
466 # Exploration vs exploitation
467 # nosemgrep: semgrep.rules.weak-random-generation, reason: Non-security scheduling or rate-limit jitter
468 if random.random() < self.exploration_rate:
469 # Explore: try a faster rate to see if API limits have relaxed
470 wait_time = base_wait * random.uniform(0.5, 0.9)
471 logger.debug(
472 f"Exploring faster rate for {engine_type}: {wait_time:.2f}s"
473 )
474 else:
475 # Exploit: use learned estimate with jitter
476 wait_time = base_wait * random.uniform(0.9, 1.1)
478 # Enforce bounds
479 wait_time = max(
480 float(estimate["min"]), min(wait_time, float(estimate["max"]))
481 )
482 return float(wait_time)
484 def apply_rate_limit(self, engine_type: str) -> float:
485 """
486 Apply rate limiting for the given engine type.
487 This is a convenience method that combines checking if rate limiting
488 is enabled, getting the wait time, and sleeping if necessary.
490 Args:
491 engine_type: The type of search engine
493 Returns:
494 The wait time that was applied (0 if rate limiting is disabled)
495 """
496 if not self.enabled:
497 return 0.0
499 wait_time = self.get_wait_time(engine_type)
500 if wait_time > 0:
501 logger.debug(
502 f"{engine_type} waiting {wait_time:.2f}s before request"
503 )
504 time.sleep(wait_time)
505 return wait_time
507 def record_outcome(
508 self,
509 engine_type: str,
510 wait_time: float,
511 success: bool,
512 retry_count: int,
513 error_type: Optional[str] = None,
514 search_result_count: Optional[int] = None,
515 ) -> None:
516 """
517 Record the outcome of a retry attempt.
519 Args:
520 engine_type: Name of the search engine
521 wait_time: How long we waited before this attempt
522 success: Whether the attempt succeeded
523 retry_count: Which retry attempt this was (1, 2, 3, etc.)
524 error_type: Type of error if failed
525 search_result_count: Number of search results returned (for quality monitoring)
526 """
527 # If rate limiting is disabled, don't record outcomes
528 if not self.enabled:
529 logger.info(
530 f"Rate limiting disabled - not recording outcome for {engine_type}"
531 )
532 return
534 logger.debug(
535 f"Recording rate limit outcome for {engine_type}: success={success}, wait_time={wait_time}s"
536 )
537 timestamp = time.time()
539 # NOTE: Database writes for rate limiting are disabled to prevent
540 # database locking issues under heavy parallel search load.
541 # Rate limiting still works via in-memory tracking below.
542 # Historical rate limit data is not persisted to DB.
544 # Update in-memory tracking
545 with self._cache_lock:
546 if engine_type not in self.recent_attempts:
547 # Get current memory window setting from thread context
548 settings_context = get_settings_context()
549 if settings_context:
550 current_memory_window = int(
551 settings_context.get_setting(
552 "rate_limiting.memory_window", self.memory_window
553 )
554 )
555 else:
556 current_memory_window = self.memory_window
558 self.recent_attempts[engine_type] = deque(
559 maxlen=current_memory_window
560 )
562 self.recent_attempts[engine_type].append(
563 {
564 "wait_time": wait_time,
565 "success": success,
566 "timestamp": timestamp,
567 "retry_count": retry_count,
568 "search_result_count": search_result_count,
569 }
570 )
572 # Update estimates
573 self._update_estimate(engine_type)
575 def _update_estimate(self, engine_type: str) -> None:
576 """Update wait time estimate based on recent attempts."""
577 with self._cache_lock:
578 if (
579 engine_type not in self.recent_attempts
580 or len(self.recent_attempts[engine_type]) < 3
581 ):
582 attempt_count = len(self.recent_attempts.get(engine_type) or [])
583 logger.info(
584 f"Not updating estimate for {engine_type} - only {attempt_count} attempts (need 3)"
585 )
586 return
588 attempts = list(self.recent_attempts[engine_type])
589 old_estimate = self.current_estimates.get(engine_type)
591 # Calculate success rate and optimal wait time
592 successful_waits = [a["wait_time"] for a in attempts if a["success"]]
593 failed_waits = [a["wait_time"] for a in attempts if not a["success"]]
595 if not successful_waits:
596 # All attempts failed - increase wait time with a cap
597 new_base = max(failed_waits) * 1.5 if failed_waits else 10.0
598 # Cap the base wait time to prevent runaway growth
599 new_base = min(new_base, 10.0) # Max 10 seconds base when all fail
600 else:
601 # Use 50th percentile (median) of successful waits for more stability
602 # This provides a balanced approach between speed and reliability
603 successful_waits.sort()
604 # Lower-median index: `(n - 1) // 2` is the middle element for odd n
605 # and the lower of the two middles for even n. The prior
606 # `int(n * 0.50) - 1` (left from a 25th-percentile version) sat below
607 # the median for odd n: at n == 3, the minimum qualifying sample, it
608 # picked index 0, the fastest wait, not the median.
609 median_index = (len(successful_waits) - 1) // 2
610 new_base = successful_waits[median_index]
612 # Update estimate with learning rate (exponential moving average)
613 if old_estimate is not None:
614 old_base = old_estimate["base"]
615 # Get current learning rate from settings context
616 settings_context = get_settings_context()
617 if settings_context:
618 current_learning_rate = float(
619 settings_context.get_setting(
620 "rate_limiting.learning_rate", self.learning_rate
621 )
622 )
623 else:
624 current_learning_rate = self.learning_rate
626 new_base = (
627 1 - current_learning_rate
628 ) * old_base + current_learning_rate * new_base
630 # Apply absolute cap to prevent extreme wait times
631 new_base = min(new_base, 10.0) # Cap base at 10 seconds
633 # Calculate bounds with more reasonable limits
634 min_wait = max(0.01, new_base * 0.5)
635 max_wait = min(10.0, new_base * 3.0) # Max 10 seconds absolute cap
637 # Update in memory
638 with self._cache_lock:
639 self.current_estimates[engine_type] = {
640 "base": new_base,
641 "min": min_wait,
642 "max": max_wait,
643 "confidence": min(len(attempts) / 20.0, 1.0),
644 }
646 # Persist to database
647 success_rate = len(successful_waits) / len(attempts) if attempts else 0
649 # Skip database operations in programmatic mode
650 if self.programmatic_mode:
651 logger.debug(
652 f"Skipping estimate persistence in programmatic mode for {engine_type}"
653 )
654 else:
655 # Try to get research context from search tracker
657 context = get_search_context()
658 username = None
659 password = None
660 if context is not None:
661 username = context.get("username")
662 password = context.get("user_password")
664 if username and password:
665 try:
666 from ...database.sqlcipher_utils import retry_on_db_lock
668 def _persist_estimate() -> None:
669 db_imports = _get_db_imports()
670 RateLimitEstimate = (
671 db_imports.get("RateLimitEstimate")
672 if db_imports
673 else None
674 )
675 if not RateLimitEstimate:
676 record_metric_write_drop(
677 username, "rate_limits", "database_unavailable"
678 )
679 return
681 session_inner: "Session"
682 with self._resolve_session(
683 username, password
684 ) as session_inner:
685 with self._isolated_begin_nested(session_inner):
686 # Check if estimate exists
687 estimate: Any = (
688 session_inner.query(RateLimitEstimate)
689 .filter_by(engine_type=engine_type)
690 .first()
691 )
693 if estimate:
694 # Update existing estimate
695 estimate.base_wait_seconds = new_base
696 estimate.min_wait_seconds = min_wait
697 estimate.max_wait_seconds = max_wait
698 estimate.last_updated = time.time()
699 estimate.total_attempts = len(attempts)
700 estimate.success_rate = success_rate
701 else:
702 # Create new estimate
703 estimate = RateLimitEstimate(
704 engine_type=engine_type,
705 base_wait_seconds=new_base,
706 min_wait_seconds=min_wait,
707 max_wait_seconds=max_wait,
708 last_updated=time.time(),
709 total_attempts=len(attempts),
710 success_rate=success_rate,
711 )
712 session_inner.add(estimate)
714 # Retry transient SQLITE_BUSY ("database is locked")
715 # with backoff so a short overlap with another
716 # writer does not surface as a permanent failure
717 # outside a caller write-scope. Non-lock errors
718 # propagate immediately. Each attempt re-resolves
719 # its session; outside a caller scope that is a
720 # fresh session per attempt (the common case for
721 # engine threads), while inside a scope the shared
722 # session is reused. A scope that has not written
723 # gives each attempt its own short transaction (a
724 # failed one is rolled back). Inside the caller's
725 # write transaction its open snapshot is not rolled
726 # back, so retries may fail fast and any poisoned
727 # state surfaces (with its class logged below)
728 # instead of silently failing.
729 retry_on_db_lock(_persist_estimate)
730 except Exception as exc:
731 record_metric_write_drop(
732 username, "rate_limits", "write_error"
733 )
734 # Class only, never message/traceback: the DB
735 # password is a live local in the session frames
736 # below, and rendering them would leak it
737 # (TRUST.md §5). The class name alone
738 # distinguishes lock timeouts from poisoned
739 # sessions, open failures, and constraint errors.
740 logger.warning(
741 "Failed to persist rate limit estimate "
742 f"(exc_class={type(exc).__name__})"
743 )
744 else:
745 if username:
746 record_metric_write_drop(
747 username, "rate_limits", "missing_credentials"
748 )
749 logger.debug(
750 "Skipping rate limit estimate save - no user context"
751 )
753 logger.info(
754 f"Updated rate limit for {engine_type}: {new_base:.2f}s "
755 f"(success rate: {success_rate:.1%})"
756 )
758 def _get_in_memory_stats(
759 self, engine_type: Optional[str] = None
760 ) -> List[Tuple[str, float, float, float, float, int, float]]:
761 """Return stats from in-memory caches (no DB access)."""
762 with self._cache_lock:
763 stats = []
764 engines_to_check = (
765 [engine_type]
766 if engine_type
767 else list(self.current_estimates.keys())
768 )
769 for engine in engines_to_check:
770 if engine in self.current_estimates:
771 est = self.current_estimates[engine]
772 attempts = self.recent_attempts.get(engine)
773 attempt_count = len(attempts) if attempts else 0
774 stats.append(
775 (
776 engine,
777 est["base"],
778 est["min"],
779 est["max"],
780 time.time(),
781 attempt_count,
782 est.get("confidence", 0.0),
783 )
784 )
785 return stats
787 def get_stats(
788 self, engine_type: Optional[str] = None
789 ) -> List[Tuple[str, float, float, float, float, int, float]]:
790 """
791 Get statistics for monitoring.
793 Args:
794 engine_type: Specific engine to get stats for, or None for all
796 Returns:
797 List of tuples with engine statistics
798 """
799 # Skip database operations in CI mode
800 if is_ci_environment():
801 logger.debug("Skipping database stats in CI mode")
802 return self._get_in_memory_stats(engine_type)
804 # Skip database operations in programmatic mode
805 if self.programmatic_mode:
806 return self._get_in_memory_stats(engine_type)
808 try:
809 context = get_search_context()
810 if not context:
811 return self._get_in_memory_stats(engine_type)
813 username = context.get("username")
814 password = context.get("user_password")
815 if not username or not password: 815 ↛ 816line 815 didn't jump to line 816 because the condition on line 815 was never true
816 return self._get_in_memory_stats(engine_type)
818 db_imports = _get_db_imports()
819 RateLimitEstimate = db_imports.get("RateLimitEstimate")
821 session_est: "Session"
822 with self._resolve_session(username, password) as session_est:
823 # See ``_ensure_estimates_loaded``: the production
824 # sessionmaker defaults to ``autoflush=True`` and a cold
825 # ``get_stats`` would otherwise flush caller's pending
826 # writes before the SELECT, which under a competing
827 # SQLite writer fails and invalidates the outer
828 # transaction.
829 with session_est.no_autoflush:
830 if engine_type:
831 estimates: list[Any] = (
832 session_est.query(RateLimitEstimate)
833 .filter_by(engine_type=engine_type)
834 .all()
835 )
836 else:
837 estimates = (
838 session_est.query(RateLimitEstimate)
839 .order_by(RateLimitEstimate.engine_type)
840 .all()
841 )
843 return [
844 (
845 est.engine_type,
846 est.base_wait_seconds,
847 est.min_wait_seconds,
848 est.max_wait_seconds,
849 est.last_updated,
850 est.total_attempts,
851 est.success_rate,
852 )
853 for est in estimates
854 ]
855 except Exception:
856 logger.warning("Failed to get rate limit stats from DB")
857 # Return in-memory stats as fallback
858 return self._get_in_memory_stats(engine_type)
860 def reset_engine(self, engine_type: str) -> None:
861 """
862 Reset learned values for a specific engine.
864 Args:
865 engine_type: Engine to reset
866 """
867 # Always clear from memory first
868 with self._cache_lock:
869 if engine_type in self.recent_attempts:
870 del self.recent_attempts[engine_type]
871 if engine_type in self.current_estimates:
872 del self.current_estimates[engine_type]
874 # Skip database operations in programmatic mode
875 if self.programmatic_mode:
876 logger.debug(
877 f"Reset rate limit data for {engine_type} (memory only in programmatic mode)"
878 )
879 return
881 # Skip database operations in CI mode
882 if is_ci_environment():
883 logger.debug(
884 f"Reset rate limit data for {engine_type} (memory only in CI mode)"
885 )
886 return
888 try:
889 context = get_search_context()
890 if not context: 890 ↛ 891line 890 didn't jump to line 891 because the condition on line 890 was never true
891 logger.debug(
892 f"Reset rate limit data for {engine_type} (memory only - no user context)"
893 )
894 return
896 username = context.get("username")
897 password = context.get("user_password")
898 if not username or not password: 898 ↛ 899line 898 didn't jump to line 899 because the condition on line 898 was never true
899 logger.debug(
900 f"Reset rate limit data for {engine_type} (memory only - no credentials)"
901 )
902 return
904 db_imports = _get_db_imports()
905 RateLimitAttempt = (
906 db_imports.get("RateLimitAttempt") if db_imports else None
907 )
908 RateLimitEstimate = (
909 db_imports.get("RateLimitEstimate") if db_imports else None
910 )
911 if not RateLimitAttempt or not RateLimitEstimate: 911 ↛ 912line 911 didn't jump to line 912 because the condition on line 911 was never true
912 return
914 from ...database.sqlcipher_utils import retry_on_db_lock
916 # Resolve again per attempt: owned sessions close before backoff;
917 # borrowed sessions keep main's savepoint and caller-state isolation.
918 def _reset_db():
919 session_reset: "Session"
920 with self._resolve_session(username, password) as session_reset:
921 with self._isolated_begin_nested(session_reset):
922 # Delete historical attempts
923 session_reset.query(RateLimitAttempt).filter_by(
924 engine_type=engine_type
925 ).delete()
927 # Delete estimates
928 session_reset.query(RateLimitEstimate).filter_by(
929 engine_type=engine_type
930 ).delete()
932 retry_on_db_lock(_reset_db)
934 logger.info(f"Reset rate limit data for {engine_type}")
936 except Exception:
937 logger.warning(
938 f"Failed to reset rate limit data in database for {engine_type}"
939 "In-memory data was cleared successfully."
940 )
941 # Don't re-raise in test contexts - the memory cleanup is sufficient
943 def get_search_quality_stats(
944 self, engine_type: Optional[str] = None
945 ) -> List[Dict]:
946 """
947 Get basic search quality statistics for monitoring.
949 Note: no HTTP route consumes this method any more —
950 ``/benchmark/api/search-quality`` now reports success_rate from the
951 persisted ``RateLimitEstimate`` table instead of this in-memory,
952 per-process view. It is retained as a tracker diagnostic capability
953 (with its own test coverage) for CLI / programmatic use; the
954 ``recent_avg_results`` / ``sample_size`` shape it returns is no longer
955 exposed over the API.
957 Args:
958 engine_type: Specific engine to get stats for, or None for all
960 Returns:
961 List of dictionaries with search quality metrics
962 """
963 stats = []
965 with self._cache_lock:
966 engines_to_check = (
967 [engine_type]
968 if engine_type
969 else list(self.recent_attempts.keys())
970 )
971 # Snapshot the data under lock, then process outside
972 engine_attempts = {}
973 for engine in engines_to_check:
974 if engine in self.recent_attempts:
975 engine_attempts[engine] = list(self.recent_attempts[engine])
977 for engine, recent in engine_attempts.items():
978 search_counts = [
979 attempt.get("search_result_count", 0)
980 for attempt in recent
981 if attempt.get("search_result_count") is not None
982 ]
984 if not search_counts:
985 continue
987 recent_avg = sum(search_counts) / len(search_counts)
989 stats.append(
990 {
991 "engine_type": engine,
992 "recent_avg_results": recent_avg,
993 "min_recent_results": min(search_counts),
994 "max_recent_results": max(search_counts),
995 "sample_size": len(search_counts),
996 "total_attempts": len(recent),
997 "status": self._get_quality_status(recent_avg),
998 }
999 )
1001 return stats
1003 def _get_quality_status(self, recent_avg: float) -> str:
1004 """Get quality status string based on average results."""
1005 if recent_avg < 1:
1006 return "CRITICAL"
1007 if recent_avg < 3:
1008 return "WARNING"
1009 if recent_avg < 5:
1010 return "CAUTION"
1011 if recent_avg >= 10:
1012 return "EXCELLENT"
1013 return "GOOD"
1015 def cleanup_old_data(self, days: int = 30) -> None:
1016 """
1017 Remove old retry attempt data to prevent database bloat.
1019 Args:
1020 days: Remove data older than this many days
1021 """
1022 cutoff_time = time.time() - (days * 24 * 3600)
1024 # Skip database operations in programmatic mode
1025 if self.programmatic_mode:
1026 logger.debug("Skipping database cleanup in programmatic mode")
1027 return
1029 # Skip database operations in CI mode
1030 if is_ci_environment():
1031 logger.debug("Skipping database cleanup in CI mode")
1032 return
1034 try:
1035 context = get_search_context()
1036 if not context: 1036 ↛ 1037line 1036 didn't jump to line 1037 because the condition on line 1036 was never true
1037 logger.debug("Skipping database cleanup - no user context")
1038 return
1040 username = context.get("username")
1041 password = context.get("user_password")
1042 if not username or not password: 1042 ↛ 1043line 1042 didn't jump to line 1043 because the condition on line 1042 was never true
1043 logger.debug("Skipping database cleanup - no credentials")
1044 return
1046 db_imports = _get_db_imports()
1047 RateLimitAttempt = (
1048 db_imports.get("RateLimitAttempt") if db_imports else None
1049 )
1050 if not RateLimitAttempt: 1050 ↛ 1051line 1050 didn't jump to line 1051 because the condition on line 1050 was never true
1051 return
1053 from ...database.sqlcipher_utils import retry_on_db_lock
1055 def _cleanup_db():
1056 session_clean: "Session"
1057 with self._resolve_session(username, password) as session_clean:
1058 with self._isolated_begin_nested(session_clean):
1059 old_attempts: Any = session_clean.query(
1060 RateLimitAttempt
1061 ).filter(RateLimitAttempt.timestamp < cutoff_time)
1062 deleted_count = old_attempts.count()
1063 old_attempts.delete()
1064 return deleted_count
1066 deleted_count = retry_on_db_lock(_cleanup_db)
1068 if deleted_count > 0:
1069 logger.info(f"Cleaned up {deleted_count} old retry attempts")
1071 except Exception:
1072 logger.warning("Failed to cleanup old rate limit data")
1075def get_tracker() -> AdaptiveRateLimitTracker:
1076 """Create a fresh rate limit tracker instance.
1078 Returns a new instance each call so that per-user DB credentials
1079 are never cached across requests in a multi-user Flask app.
1080 """
1081 return AdaptiveRateLimitTracker()