src/knowlang/utils/rate_limiter.py

from dataclasses import dataclass, field
from collections import defaultdict
import time
from typing import Dict, List
from threading import Lock


@dataclass
class RateLimiter:
    """Rate limiter implementation for chat interface"""
    requests_per_minute: int = 2
    window_size: int = 60  # seconds
    
    # Track request timestamps per user IP
    _requests: Dict[str, List[float]] = field(default_factory=lambda: defaultdict(list))
    _lock: Lock = field(default_factory=Lock)

    def _clean_old_requests(self, user_ip: str) -> None:
        """Remove requests older than the window size"""
        current_time = time.time()
        with self._lock:
            self._requests[user_ip] = [
                timestamp for timestamp in self._requests[user_ip]
                if current_time - timestamp < self.window_size
            ]

    def check_rate_limit(self, user_ip: str) -> bool:
        """
        Check if user has exceeded rate limit
        Returns True if rate limit exceeded, False otherwise
        """
        self._clean_old_requests(user_ip)
        
        with self._lock:
            if len(self._requests[user_ip]) >= self.requests_per_minute:
                return True
                
            self._requests[user_ip].append(time.time())
            return False

    def get_remaining_time(self, user_ip: str) -> float:
        """Get remaining time until next request is allowed"""
        self._clean_old_requests(user_ip)
        
        with self._lock:
            if not self._requests[user_ip]:
                return 0
                
            oldest_request = min(self._requests[user_ip])
            current_time = time.time()
            time_until_reset = self.window_size - (current_time - oldest_request)
            
            return max(0, time_until_reset)