Spaces:
Sleeping
Sleeping
| 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) |