replypier/backend/app/utils/rate_limit.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import functools
from fastapi import HTTPException, Request
from app.core.config import settings
from app.core.logger import get_logger
from app.core.redis import get_redis
from typing import Optional, Callable

logger = get_logger(__name__)

class RateLimitConfig:
    """Configuration class for rate limiting settings"""
    
    # Rate limits for instruction generation
    INSTRUCTION_GEN_LIMIT = 10  # requests
    INSTRUCTION_GEN_WINDOW = 86400  # 24 hours in seconds
    INSTRUCTION_GEN_PREFIX = "gen_instruct"
    
    # Add more rate limit configurations here as needed
    # Example:
    # CHAT_MESSAGE_LIMIT = 100
    # CHAT_MESSAGE_WINDOW = 3600  # 1 hour
    # CHAT_MESSAGE_PREFIX = "chat_msg"

def format_time_remaining(seconds: int) -> str:
    """Format remaining time into a human-readable string"""
    hours = seconds // 3600
    minutes = (seconds % 3600) // 60
    seconds = seconds % 60
    
    if hours > 0:
        time_msg = f"{hours} hours"
        if minutes > 0:
            time_msg += f" and {minutes} minutes"
    elif minutes > 0:
        time_msg = f"{minutes} minutes"
        if seconds > 0:
            time_msg += f" and {seconds} seconds"
    else:
        time_msg = f"{seconds} seconds"
    
    return time_msg

def check_rate_limit(
    redis_client: Optional[object],
    key: str,
    limit: int,
    window: int
) -> None:
    """
    Check rate limit for a given key
    
    Args:
        redis_client: Redis client instance
        key: Redis key for rate limiting
        limit: Maximum number of requests allowed
        window: Time window in seconds
        
    Raises:
        HTTPException: If rate limit is exceeded
    """
    if not redis_client:
        return
        
    # Get current count
    current = redis_client.get(key)
    
    if current is None:
        # First request, set initial count
        redis_client.setex(key, window, 1)
    else:
        current = int(current)
        if current >= limit:
            # Calculate time until reset
            ttl = redis_client.ttl(key)
            time_msg = format_time_remaining(ttl)
            
            logger.warning(f"Rate limit exceeded for key: {key}")
            raise HTTPException(
                status_code=429,
                detail=f"Rate limit exceeded. You can try again in {time_msg}."
            )
        
        # Increment counter
        redis_client.incr(key)

def rate_limit(limit: int, window: int, key_prefix: str):
    """
    Rate limiting decorator that uses Redis to track request counts.
    
    Args:
        limit (int): Maximum number of requests allowed within the window
        window (int): Time window in seconds
        key_prefix (str): Prefix for the Redis key (e.g., 'gen_instruct', 'api_call')
    
    Returns:
        function: Decorated function with rate limiting
    """
    def decorator(func: Callable):
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            redis_client = get_redis()
            if not settings.REDIS_ENABLED or not redis_client:
                return await func(*args, **kwargs)

            # Get request object from kwargs
            request = None
            for arg in args:
                if isinstance(arg, Request):
                    request = arg
                    break
            if not request and 'request' in kwargs:
                request = kwargs['request']

            # Get client IP
            client_ip = request.client.host if request and request.client else "unknown"
            
            # Get organization ID from kwargs
            org_id = kwargs.get('current_user', {}).organization_id
            if not org_id:
                logger.warning("No organization ID found for rate limiting")
                return await func(*args, **kwargs)

            # Create a key specific to this organization and IP
            rate_key = f"{key_prefix}:{org_id}:{client_ip}"
            
            # Check rate limit
            check_rate_limit(redis_client, rate_key, limit, window)

            return await func(*args, **kwargs)
        return wrapper
    return decorator

def limit_instruction_generation(func: Callable):
    """
    Specific rate limiting decorator for instruction generation.
    Limits to 10 requests per day per organization/IP.
    
    Args:
        func: The function to be decorated
        
    Returns:
        function: Decorated function with rate limiting
    """
    return rate_limit(
        limit=RateLimitConfig.INSTRUCTION_GEN_LIMIT,
        window=RateLimitConfig.INSTRUCTION_GEN_WINDOW,
        key_prefix=RateLimitConfig.INSTRUCTION_GEN_PREFIX
    )(func)

# Add more specific rate limiting decorators here as needed
# Example:
# def limit_chat_messages(func: Callable):
#     """Rate limiting decorator for chat messages"""
#     return rate_limit(
#         limit=RateLimitConfig.CHAT_MESSAGE_LIMIT,
#         window=RateLimitConfig.CHAT_MESSAGE_WINDOW,
#         key_prefix=RateLimitConfig.CHAT_MESSAGE_PREFIX
#     )(func)