local-deep-research 0.5.9__py3-none-any.whl → 0.6.0__py3-none-any.whl

local_deep_research/web_search_engines/rate_limiting/tracker.py

@@ -0,0 +1,506 @@
+"""
+Adaptive rate limit tracker that learns optimal retry wait times for each search engine.
+"""
+
+import time
+import random
+import logging
+from collections import deque
+from typing import Dict, Optional, Tuple, List
+
+
+from ...utilities.db_utils import get_db_session
+from ...web.database.models import RateLimitAttempt, RateLimitEstimate
+
+logger = logging.getLogger(__name__)
+
+
+class AdaptiveRateLimitTracker:
+    """
+    Tracks and learns optimal retry wait times for each search engine.
+    Persists learned patterns to the main application database using SQLAlchemy.
+    """
+
+    def __init__(self):
+        # Load configuration from database settings
+        from ...utilities.db_utils import get_db_setting
+
+        self.memory_window = int(
+            get_db_setting("rate_limiting.memory_window", 100)
+        )
+        self.exploration_rate = float(
+            get_db_setting("rate_limiting.exploration_rate", 0.1)
+        )
+        self.learning_rate = float(
+            get_db_setting("rate_limiting.learning_rate", 0.3)
+        )
+        self.decay_per_day = float(
+            get_db_setting("rate_limiting.decay_per_day", 0.95)
+        )
+        self.enabled = bool(get_db_setting("rate_limiting.enabled", True))
+
+        # Apply rate limiting profile
+        self._apply_profile(get_db_setting("rate_limiting.profile", "balanced"))
+
+        # In-memory cache for fast access
+        self.recent_attempts: Dict[str, deque] = {}
+        self.current_estimates: Dict[str, Dict[str, float]] = {}
+
+        # Load estimates from database
+        self._load_estimates()
+
+        logger.info(
+            f"AdaptiveRateLimitTracker initialized: enabled={self.enabled}, profile={get_db_setting('rate_limiting.profile', 'balanced')}"
+        )
+
+    def _apply_profile(self, profile: str) -> None:
+        """Apply rate limiting profile settings."""
+        if profile == "conservative":
+            # More conservative: lower exploration, slower learning
+            self.exploration_rate = min(
+                self.exploration_rate * 0.5, 0.05
+            )  # 5% max exploration
+            self.learning_rate = min(
+                self.learning_rate * 0.7, 0.2
+            )  # Slower learning
+            logger.info("Applied conservative rate limiting profile")
+        elif profile == "aggressive":
+            # More aggressive: higher exploration, faster learning
+            self.exploration_rate = min(
+                self.exploration_rate * 1.5, 0.2
+            )  # Up to 20% exploration
+            self.learning_rate = min(
+                self.learning_rate * 1.3, 0.5
+            )  # Faster learning
+            logger.info("Applied aggressive rate limiting profile")
+        else:  # balanced
+            # Use settings as-is
+            logger.info("Applied balanced rate limiting profile")
+
+    def _load_estimates(self) -> None:
+        """Load estimates from database into memory."""
+        try:
+            session = get_db_session()
+            estimates = session.query(RateLimitEstimate).all()
+
+            for estimate in estimates:
+                # Apply decay for old estimates
+                age_hours = (time.time() - estimate.last_updated) / 3600
+                decay = self.decay_per_day ** (age_hours / 24)
+
+                self.current_estimates[estimate.engine_type] = {
+                    "base": estimate.base_wait_seconds,
+                    "min": estimate.min_wait_seconds,
+                    "max": estimate.max_wait_seconds,
+                    "confidence": decay,
+                }
+
+                logger.debug(
+                    f"Loaded estimate for {estimate.engine_type}: base={estimate.base_wait_seconds:.2f}s, confidence={decay:.2f}"
+                )
+
+        except Exception as e:
+            logger.warning(f"Could not load rate limit estimates: {e}")
+            # Continue with empty estimates - they'll be learned
+
+    def get_wait_time(self, engine_type: str) -> float:
+        """
+        Get adaptive wait time for a search engine.
+        Includes exploration to discover better rates.
+
+        Args:
+            engine_type: Name of the search engine
+
+        Returns:
+            Wait time in seconds
+        """
+        # If rate limiting is disabled, return minimal wait time
+        if not self.enabled:
+            return 0.1
+
+        if engine_type not in self.current_estimates:
+            # First time seeing this engine - start optimistic and learn from real responses
+            # Use engine-specific optimistic defaults only for what we know for sure
+            optimistic_defaults = {
+                "LocalSearchEngine": 0.0,  # No network calls
+                "SearXNGSearchEngine": 0.1,  # Self-hosted default engine
+            }
+
+            wait_time = optimistic_defaults.get(
+                engine_type, 0.5
+            )  # Default optimistic for others
+            logger.info(
+                f"No rate limit data for {engine_type}, starting optimistic with {wait_time}s"
+            )
+            return wait_time
+
+        estimate = self.current_estimates[engine_type]
+        base_wait = estimate["base"]
+
+        # Exploration vs exploitation
+        if random.random() < self.exploration_rate:
+            # Explore: try a faster rate to see if API limits have relaxed
+            wait_time = base_wait * random.uniform(0.5, 0.9)
+            logger.debug(
+                f"Exploring faster rate for {engine_type}: {wait_time:.2f}s"
+            )
+        else:
+            # Exploit: use learned estimate with jitter
+            wait_time = base_wait * random.uniform(0.9, 1.1)
+
+        # Enforce bounds
+        wait_time = max(estimate["min"], min(wait_time, estimate["max"]))
+        return wait_time
+
+    def record_outcome(
+        self,
+        engine_type: str,
+        wait_time: float,
+        success: bool,
+        retry_count: int,
+        error_type: Optional[str] = None,
+        search_result_count: Optional[int] = None,
+    ) -> None:
+        """
+        Record the outcome of a retry attempt.
+
+        Args:
+            engine_type: Name of the search engine
+            wait_time: How long we waited before this attempt
+            success: Whether the attempt succeeded
+            retry_count: Which retry attempt this was (1, 2, 3, etc.)
+            error_type: Type of error if failed
+            search_result_count: Number of search results returned (for quality monitoring)
+        """
+        # If rate limiting is disabled, don't record outcomes
+        if not self.enabled:
+            return
+        timestamp = time.time()
+
+        try:
+            # Save to database
+            session = get_db_session()
+            attempt = RateLimitAttempt(
+                engine_type=engine_type,
+                timestamp=timestamp,
+                wait_time=wait_time,
+                retry_count=retry_count,
+                success=success,
+                error_type=error_type,
+            )
+            session.add(attempt)
+            session.commit()
+        except Exception as e:
+            logger.error(f"Failed to record rate limit outcome: {e}")
+
+        # Update in-memory tracking
+        if engine_type not in self.recent_attempts:
+            # Get current memory window setting
+            from ...utilities.db_utils import get_db_setting
+
+            current_memory_window = int(
+                get_db_setting(
+                    "rate_limiting.memory_window", self.memory_window
+                )
+            )
+            self.recent_attempts[engine_type] = deque(
+                maxlen=current_memory_window
+            )
+
+        self.recent_attempts[engine_type].append(
+            {
+                "wait_time": wait_time,
+                "success": success,
+                "timestamp": timestamp,
+                "retry_count": retry_count,
+                "search_result_count": search_result_count,
+            }
+        )
+
+        # Update estimates
+        self._update_estimate(engine_type)
+
+    def _update_estimate(self, engine_type: str) -> None:
+        """Update wait time estimate based on recent attempts."""
+        if (
+            engine_type not in self.recent_attempts
+            or len(self.recent_attempts[engine_type]) < 3
+        ):
+            return
+
+        attempts = list(self.recent_attempts[engine_type])
+
+        # Calculate success rate and optimal wait time
+        successful_waits = [a["wait_time"] for a in attempts if a["success"]]
+        failed_waits = [a["wait_time"] for a in attempts if not a["success"]]
+
+        if not successful_waits:
+            # All attempts failed - increase wait time with a cap
+            new_base = max(failed_waits) * 1.5 if failed_waits else 10.0
+            # Cap the base wait time to prevent runaway growth
+            new_base = min(new_base, 10.0)  # Max 10 seconds base when all fail
+        else:
+            # Use 75th percentile of successful waits
+            successful_waits.sort()
+            percentile_75 = successful_waits[int(len(successful_waits) * 0.75)]
+            new_base = percentile_75
+
+        # Update estimate with learning rate (exponential moving average)
+        if engine_type in self.current_estimates:
+            old_base = self.current_estimates[engine_type]["base"]
+            # Get current learning rate from settings
+            from ...utilities.db_utils import get_db_setting
+
+            current_learning_rate = float(
+                get_db_setting(
+                    "rate_limiting.learning_rate", self.learning_rate
+                )
+            )
+            new_base = (
+                1 - current_learning_rate
+            ) * old_base + current_learning_rate * new_base
+
+        # Apply absolute cap to prevent extreme wait times
+        new_base = min(new_base, 10.0)  # Cap base at 10 seconds
+
+        # Calculate bounds with more reasonable limits
+        min_wait = max(0.5, new_base * 0.5)
+        max_wait = min(10.0, new_base * 3.0)  # Max 10 seconds absolute cap
+
+        # Update in memory
+        self.current_estimates[engine_type] = {
+            "base": new_base,
+            "min": min_wait,
+            "max": max_wait,
+            "confidence": min(len(attempts) / 20.0, 1.0),
+        }
+
+        # Persist to database
+        success_rate = len(successful_waits) / len(attempts) if attempts else 0
+
+        try:
+            session = get_db_session()
+
+            # Check if estimate exists
+            estimate = (
+                session.query(RateLimitEstimate)
+                .filter_by(engine_type=engine_type)
+                .first()
+            )
+
+            if estimate:
+                # Update existing estimate
+                estimate.base_wait_seconds = new_base
+                estimate.min_wait_seconds = min_wait
+                estimate.max_wait_seconds = max_wait
+                estimate.last_updated = time.time()
+                estimate.total_attempts = len(attempts)
+                estimate.success_rate = success_rate
+            else:
+                # Create new estimate
+                estimate = RateLimitEstimate(
+                    engine_type=engine_type,
+                    base_wait_seconds=new_base,
+                    min_wait_seconds=min_wait,
+                    max_wait_seconds=max_wait,
+                    last_updated=time.time(),
+                    total_attempts=len(attempts),
+                    success_rate=success_rate,
+                )
+                session.add(estimate)
+
+            session.commit()
+
+        except Exception as e:
+            logger.error(f"Failed to persist rate limit estimate: {e}")
+
+        logger.info(
+            f"Updated rate limit for {engine_type}: {new_base:.2f}s "
+            f"(success rate: {success_rate:.1%})"
+        )
+
+    def get_stats(
+        self, engine_type: Optional[str] = None
+    ) -> List[Tuple[str, float, float, float, float, int, float]]:
+        """
+        Get statistics for monitoring.
+
+        Args:
+            engine_type: Specific engine to get stats for, or None for all
+
+        Returns:
+            List of tuples with engine statistics
+        """
+        try:
+            session = get_db_session()
+
+            if engine_type:
+                estimates = (
+                    session.query(RateLimitEstimate)
+                    .filter_by(engine_type=engine_type)
+                    .all()
+                )
+            else:
+                estimates = (
+                    session.query(RateLimitEstimate)
+                    .order_by(RateLimitEstimate.engine_type)
+                    .all()
+                )
+
+            return [
+                (
+                    est.engine_type,
+                    est.base_wait_seconds,
+                    est.min_wait_seconds,
+                    est.max_wait_seconds,
+                    est.last_updated,
+                    est.total_attempts,
+                    est.success_rate,
+                )
+                for est in estimates
+            ]
+        except Exception as e:
+            logger.error(f"Failed to get rate limit stats: {e}")
+            return []
+
+    def reset_engine(self, engine_type: str) -> None:
+        """
+        Reset learned values for a specific engine.
+
+        Args:
+            engine_type: Engine to reset
+        """
+        try:
+            session = get_db_session()
+
+            # Delete historical attempts
+            session.query(RateLimitAttempt).filter_by(
+                engine_type=engine_type
+            ).delete()
+
+            # Delete estimates
+            session.query(RateLimitEstimate).filter_by(
+                engine_type=engine_type
+            ).delete()
+
+            session.commit()
+
+            # Clear from memory
+            if engine_type in self.recent_attempts:
+                del self.recent_attempts[engine_type]
+            if engine_type in self.current_estimates:
+                del self.current_estimates[engine_type]
+
+            logger.info(f"Reset rate limit data for {engine_type}")
+
+        except Exception as e:
+            logger.error(
+                f"Failed to reset rate limit data for {engine_type}: {e}"
+            )
+            # Still try to clear from memory even if database operation failed
+            if engine_type in self.recent_attempts:
+                del self.recent_attempts[engine_type]
+            if engine_type in self.current_estimates:
+                del self.current_estimates[engine_type]
+            # Re-raise the exception so callers know it failed
+            raise
+
+    def get_search_quality_stats(
+        self, engine_type: Optional[str] = None
+    ) -> List[Dict]:
+        """
+        Get basic search quality statistics for monitoring.
+
+        Args:
+            engine_type: Specific engine to get stats for, or None for all
+
+        Returns:
+            List of dictionaries with search quality metrics
+        """
+        stats = []
+
+        engines_to_check = (
+            [engine_type] if engine_type else list(self.recent_attempts.keys())
+        )
+
+        for engine in engines_to_check:
+            if engine not in self.recent_attempts:
+                continue
+
+            recent = list(self.recent_attempts[engine])
+            search_counts = [
+                attempt.get("search_result_count", 0)
+                for attempt in recent
+                if attempt.get("search_result_count") is not None
+            ]
+
+            if not search_counts:
+                continue
+
+            recent_avg = sum(search_counts) / len(search_counts)
+
+            stats.append(
+                {
+                    "engine_type": engine,
+                    "recent_avg_results": recent_avg,
+                    "min_recent_results": min(search_counts),
+                    "max_recent_results": max(search_counts),
+                    "sample_size": len(search_counts),
+                    "total_attempts": len(recent),
+                    "status": self._get_quality_status(recent_avg),
+                }
+            )
+
+        return stats
+
+    def _get_quality_status(self, recent_avg: float) -> str:
+        """Get quality status string based on average results."""
+        if recent_avg < 1:
+            return "CRITICAL"
+        elif recent_avg < 3:
+            return "WARNING"
+        elif recent_avg < 5:
+            return "CAUTION"
+        elif recent_avg >= 10:
+            return "EXCELLENT"
+        else:
+            return "GOOD"
+
+    def cleanup_old_data(self, days: int = 30) -> None:
+        """
+        Remove old retry attempt data to prevent database bloat.
+
+        Args:
+            days: Remove data older than this many days
+        """
+        cutoff_time = time.time() - (days * 24 * 3600)
+
+        try:
+            session = get_db_session()
+
+            # Count and delete old attempts
+            old_attempts = session.query(RateLimitAttempt).filter(
+                RateLimitAttempt.timestamp < cutoff_time
+            )
+            deleted_count = old_attempts.count()
+            old_attempts.delete()
+
+            session.commit()
+
+            if deleted_count > 0:
+                logger.info(f"Cleaned up {deleted_count} old retry attempts")
+
+        except Exception as e:
+            logger.error(f"Failed to cleanup old rate limit data: {e}")
+
+
+# Create a singleton instance
+_tracker_instance: Optional[AdaptiveRateLimitTracker] = None
+
+
+def get_tracker() -> AdaptiveRateLimitTracker:
+    """Get the global rate limit tracker instance."""
+    global _tracker_instance
+    if _tracker_instance is None:
+        _tracker_instance = AdaptiveRateLimitTracker()
+    return _tracker_instance

local_deep_research/web_search_engines/retriever_registry.py

@@ -0,0 +1,108 @@
+"""
+Registry for dynamically registering LangChain retrievers as search engines.
+"""
+
+from typing import Dict, Optional
+from threading import Lock
+from langchain.schema import BaseRetriever
+from loguru import logger
+
+
+class RetrieverRegistry:
+    """
+    Thread-safe registry for LangChain retrievers.
+
+    This allows users to register retrievers programmatically and use them
+    as search engines within LDR.
+    """
+
+    def __init__(self):
+        self._retrievers: Dict[str, BaseRetriever] = {}
+        self._lock = Lock()
+
+    def register(self, name: str, retriever: BaseRetriever) -> None:
+        """
+        Register a retriever with a given name.
+
+        Args:
+            name: Name to register the retriever under
+            retriever: LangChain BaseRetriever instance
+        """
+        with self._lock:
+            self._retrievers[name] = retriever
+            logger.info(
+                f"Registered retriever '{name}' of type {type(retriever).__name__}"
+            )
+
+    def register_multiple(self, retrievers: Dict[str, BaseRetriever]) -> None:
+        """
+        Register multiple retrievers at once.
+
+        Args:
+            retrievers: Dictionary of {name: retriever} pairs
+        """
+        with self._lock:
+            for name, retriever in retrievers.items():
+                self._retrievers[name] = retriever
+                logger.info(
+                    f"Registered retriever '{name}' of type {type(retriever).__name__}"
+                )
+
+    def get(self, name: str) -> Optional[BaseRetriever]:
+        """
+        Get a registered retriever by name.
+
+        Args:
+            name: Name of the retriever
+
+        Returns:
+            The retriever if found, None otherwise
+        """
+        with self._lock:
+            return self._retrievers.get(name)
+
+    def unregister(self, name: str) -> None:
+        """
+        Remove a registered retriever.
+
+        Args:
+            name: Name of the retriever to remove
+        """
+        with self._lock:
+            if name in self._retrievers:
+                del self._retrievers[name]
+                logger.info(f"Unregistered retriever '{name}'")
+
+    def clear(self) -> None:
+        """Clear all registered retrievers."""
+        with self._lock:
+            count = len(self._retrievers)
+            self._retrievers.clear()
+            logger.info(f"Cleared {count} registered retrievers")
+
+    def is_registered(self, name: str) -> bool:
+        """
+        Check if a retriever is registered.
+
+        Args:
+            name: Name of the retriever
+
+        Returns:
+            True if registered, False otherwise
+        """
+        with self._lock:
+            return name in self._retrievers
+
+    def list_registered(self) -> list[str]:
+        """
+        Get list of all registered retriever names.
+
+        Returns:
+            List of retriever names
+        """
+        with self._lock:
+            return list(self._retrievers.keys())
+
+
+# Global registry instance
+retriever_registry = RetrieverRegistry()