Spaces:
Sleeping
Sleeping
File size: 6,421 Bytes
fe4792e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
import asyncio
import pytz
from datetime import datetime, timedelta
import time
from typing import Dict, Optional
import logging
from utils.logging_utils import log_to_file
from utils.admin import supabase # Ajoutez cet import avec les autres
from utils.cache_manager import CacheManager
cache_manager = CacheManager() # optimiser les statistiques d'utilisation avec le cache
class RateLimiter:
def __init__(self, requests_per_second: int = 10, burst_limit: int = 20): # 10 requetes par seconde, 20 requetes en rafale
"""
Initialise le rate limiter
requests_per_second: Nombre maximal de requêtes par seconde
burst_limit: Nombre maximal de requêtes en rafale
"""
self.rate = requests_per_second # Nombre maximal de requetes pas seconde (aujourd'hui)
self.burst_limit = burst_limit # Nombre maximal de requetes en rafale (aujourd'hui)
self.tokens = burst_limit # Nombre de jetons restants en rafale (aujourd'hui)
self.last_update = time.monotonic() # Derniere mise à jour
self.lock = asyncio.Lock() # Prise en charge du verrouillage
self._requests_count: Dict[str, int] = {} # Compteur de requetes
self.last_cleanup = time.time() # Date du dernier nettoyage
self.cleanup_interval = 3600 # 1 heure en secondes pour le nettoyage au bout d'une heure
self.cache_manager = CacheManager() # utilisation du cache pour stocker les statisqtiques d'utilisation (aujourd'hui)
async def acquire(self, user_id: Optional[str] = None) -> bool:
"""Tente d'acquérir un token pour faire une requête"""
async with self.lock:
now = time.monotonic()
time_passed = now - self.last_update
self.tokens = min(
self.burst_limit,
self.tokens + time_passed * self.rate
)
self.last_update = now
if self.tokens >= 1:
self.tokens -= 1
if user_id:
await self._track_request(user_id)
return True
log_to_file(f"Rate limit atteint pour l'utilisateur {user_id}", level=logging.WARNING)
return False
async def _get_stats_from_cache(self, user_id: str) -> Optional[Dict]:
"""Récupère les statistiques depuis le cache"""
cache_key = f"usage_stats:{user_id}"
return await cache_manager.get(cache_key)
async def _update_stats_cache(self, user_id: str, stats: Dict):
"""Met à jour les statistiques dans le cache"""
cache_key = f"usage_stats:{user_id}"
await cache_manager.set(cache_key, stats, expiry=timedelta(hours=1))
async def _track_request(self, user_id: str) -> None:
"""Suit les requêtes par utilisateur avec cache"""
try:
current_time = datetime.now(pytz.UTC).isoformat()
stats_data = None # Initialiser stats_data
# Vérifier d'abord le cache
cached_stats = await self._get_stats_from_cache(user_id)
if cached_stats:
# Mettre à jour les stats en cache
cached_stats['request_count'] += 1
cached_stats['last_request'] = current_time
await self._update_stats_cache(user_id, cached_stats)
stats_data = cached_stats
# Mettre à jour Supabase de manière asynchrone
supabase.table("usage_statistics").update({
"request_count": cached_stats['request_count'],
"last_request": current_time
}).eq("user_id", user_id).execute()
else:
# Vérifier Supabase
existing = supabase.table("usage_statistics").select("*").eq("user_id", user_id).execute()
if existing.data:
# Mettre à jour les stats existantes
stats = existing.data[0]
new_count = stats.get('request_count', 0) + 1
stats_data = {
"request_count": new_count,
"last_request": current_time
}
supabase.table("usage_statistics").update(stats_data).eq("user_id", user_id).execute()
await self._update_stats_cache(user_id, stats_data)
else:
# Créer nouvelles stats
stats_data = {
"user_id": user_id,
"request_count": 1,
"last_request": current_time,
"last_cleanup": current_time
}
supabase.table("usage_statistics").insert(stats_data).execute()
await self._update_stats_cache(user_id, stats_data)
if stats_data: # Vérifier que stats_data existe
self._requests_count[user_id] = stats_data['request_count']
log_to_file(f"Nouvelle requête pour l'utilisateur {user_id}. Total: {stats_data['request_count']}", level=logging.INFO)
except Exception as e:
log_to_file(f"Erreur lors du suivi des requêtes: {str(e)}", level=logging.ERROR)
# Continuer avec le comptage en mémoire en cas d'erreur
if user_id not in self._requests_count:
self._requests_count[user_id] = 1
else:
self._requests_count[user_id] += 1
async def get_user_requests_count(self, user_id: str) -> int:
"""Retourne le nombre de requêtes avec cache"""
try:
cached_stats = await self._get_stats_from_cache(user_id)
if cached_stats:
return cached_stats.get('request_count', 0)
response = supabase.table("usage_statistics").select("request_count").eq("user_id", user_id).execute()
if response.data:
count = response.data[0].get('request_count', 0)
await self._update_stats_cache(user_id, {"request_count": count})
return count
return 0
except Exception as e:
log_to_file(f"Erreur lors de la récupération du compte des requêtes: {str(e)}", level=logging.ERROR)
return self._requests_count.get(user_id, 0) |