Tolan/backend/utils.py

"""
工具函数模块
"""
import time
from typing import Dict
from loguru import logger


class RateLimiter:
    """简单的速率限制器"""
    
    def __init__(self, max_requests: int = 100, window_seconds: int = 60):
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self.requests: Dict[str, list] = {}
    
    def is_allowed(self, client_id: str) -> bool:
        """检查是否允许请求"""
        now = time.time()
        
        # 清理过期记录
        if client_id in self.requests:
            self.requests[client_id] = [
                req_time for req_time in self.requests[client_id]
                if now - req_time < self.window_seconds
            ]
        else:
            self.requests[client_id] = []
        
        # 检查是否超限
        if len(self.requests[client_id]) >= self.max_requests:
            logger.warning(f"客户端 {client_id} 请求频率超限")
            return False
        
        # 记录请求
        self.requests[client_id].append(now)
        return True
    
    def get_remaining(self, client_id: str) -> int:
        """获取剩余配额"""
        if client_id not in self.requests:
            return self.max_requests
        return max(0, self.max_requests - len(self.requests[client_id]))


def validate_audio_data(audio_data: bytes, max_size: int = 1024 * 1024) -> bool:
    """验证音频数据"""
    if not audio_data:
        return False
    
    if len(audio_data) > max_size:
        logger.warning(f"音频数据过大: {len(audio_data)} bytes")
        return False
    
    return True


def format_bytes(bytes_num: int) -> str:
    """格式化字节数"""
    for unit in ['B', 'KB', 'MB', 'GB']:
        if bytes_num < 1024.0:
            return f"{bytes_num:.2f} {unit}"
        bytes_num /= 1024.0
    return f"{bytes_num:.2f} TB"


def calculate_audio_duration(audio_bytes: int, sample_rate: int = 16000, bit_depth: int = 16, channels: int = 1) -> float:
    """计算音频时长（秒）"""
    bytes_per_second = sample_rate * (bit_depth / 8) * channels
    return audio_bytes / bytes_per_second