""" 非同期レート制限管理クラス 1日1回制限とAPI呼び出し制限を管理し、 デバッグモード時の制限緩和機能を提供します。 """ import asyncio import os from datetime import datetime, timedelta from typing import Dict, Any, Optional, Tuple import logging # ログ設定 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class RateLimitError(Exception): """レート制限エラー""" pass class AsyncRateLimitManager: """非同期レート制限管理クラス""" def __init__(self, storage_manager, max_requests: int = 1): self.storage = storage_manager # 設定値（環境変数から取得、デフォルト値あり） self.max_daily_requests = int(os.getenv("MAX_DAILY_REQUESTS", "1")) self.max_api_calls_per_day = int(os.getenv("MAX_API_CALLS_PER_DAY", "10")) self.debug_mode = os.getenv("DEBUG_MODE", "false").lower() == "true" # デバッグモード時の制限緩和 if self.debug_mode: self.max_daily_requests = int(os.getenv("DEBUG_MAX_DAILY_REQUESTS", "10")) self.max_api_calls_per_day = int(os.getenv("DEBUG_MAX_API_CALLS", "100")) logger.info("デバッグモードが有効です。制限が緩和されています") logger.info(f"レート制限設定 - 1日のリクエスト上限: {self.max_daily_requests}, API呼び出し上限: {self.max_api_calls_per_day}") async def check_daily_request_limit(self, user_id: str) -> Tuple[bool, Dict[str, Any]]: """1日のリクエスト制限をチェック""" try: user_data = await self.storage.get_user_data(user_id) today = datetime.now().strftime("%Y-%m-%d") # 今日のリクエスト数を取得 daily_requests = user_data["rate_limits"]["daily_requests"] today_requests = daily_requests.get(today, 0) # 制限チェック is_allowed = today_requests < self.max_daily_requests limit_info = { "today_requests": today_requests, "max_requests": self.max_daily_requests, "remaining": max(0, self.max_daily_requests - today_requests), "reset_time": self._get_next_reset_time(), "debug_mode": self.debug_mode } if not is_allowed: logger.warning(f"ユーザー {user_id} の1日のリクエスト制限に達しました ({today_requests}/{self.max_daily_requests})") return is_allowed, limit_info except Exception as e: logger.error(f"1日のリクエスト制限チェックエラー: {e}") # エラー時は制限を適用 return False, {"error": str(e)} async def check_api_call_limit(self, user_id: str) -> Tuple[bool, Dict[str, Any]]: """API呼び出し制限をチェック""" try: user_data = await self.storage.get_user_data(user_id) today = datetime.now().strftime("%Y-%m-%d") # 今日のAPI呼び出し数を取得 api_calls = user_data["rate_limits"]["api_calls"] today_calls = api_calls.get(today, 0) # 制限チェック is_allowed = today_calls < self.max_api_calls_per_day limit_info = { "today_calls": today_calls, "max_calls": self.max_api_calls_per_day, "remaining": max(0, self.max_api_calls_per_day - today_calls), "reset_time": self._get_next_reset_time(), "debug_mode": self.debug_mode } if not is_allowed: logger.warning(f"ユーザー {user_id} のAPI呼び出し制限に達しました ({today_calls}/{self.max_api_calls_per_day})") return is_allowed, limit_info except Exception as e: logger.error(f"API呼び出し制限チェックエラー: {e}") # エラー時は制限を適用 return False, {"error": str(e)} async def record_request(self, user_id: str) -> None: """リクエストを記録""" try: user_data = await self.storage.get_user_data(user_id) today = datetime.now().strftime("%Y-%m-%d") # 今日のリクエスト数を増加 if "daily_requests" not in user_data["rate_limits"]: user_data["rate_limits"]["daily_requests"] = {} user_data["rate_limits"]["daily_requests"][today] = \ user_data["rate_limits"]["daily_requests"].get(today, 0) + 1 # プロファイルの最終リクエスト日を更新 user_data["profile"]["last_request"] = today await self.storage.update_user_data(user_id, user_data) logger.info(f"ユーザー {user_id} のリクエストを記録しました") except Exception as e: logger.error(f"リクエスト記録エラー: {e}") raise RateLimitError(f"リクエストの記録に失敗しました: {e}") async def record_api_call(self, user_id: str, api_type: str = "general") -> None: """API呼び出しを記録""" try: user_data = await self.storage.get_user_data(user_id) today = datetime.now().strftime("%Y-%m-%d") # 今日のAPI呼び出し数を増加 if "api_calls" not in user_data["rate_limits"]: user_data["rate_limits"]["api_calls"] = {} user_data["rate_limits"]["api_calls"][today] = \ user_data["rate_limits"]["api_calls"].get(today, 0) + 1 await self.storage.update_user_data(user_id, user_data) logger.info(f"ユーザー {user_id} のAPI呼び出し ({api_type}) を記録しました") except Exception as e: logger.error(f"API呼び出し記録エラー: {e}") raise RateLimitError(f"API呼び出しの記録に失敗しました: {e}") async def get_user_limits_status(self, user_id: str) -> Dict[str, Any]: """ユーザーの制限状況を取得""" try: # リクエスト制限の確認 request_allowed, request_info = await self.check_daily_request_limit(user_id) # API呼び出し制限の確認 api_allowed, api_info = await self.check_api_call_limit(user_id) # 次回リクエスト可能時刻の計算 next_request_time = None if not request_allowed: next_request_time = self._get_next_reset_time() return { "request_limit": { "allowed": request_allowed, "info": request_info }, "api_limit": { "allowed": api_allowed, "info": api_info }, "next_request_time": next_request_time, "debug_mode": self.debug_mode } except Exception as e: logger.error(f"制限状況取得エラー: {e}") return {"error": str(e)} async def reset_daily_counters(self) -> int: """1日のカウンターをリセット（古いデータを削除）""" try: # 7日以上前のデータを削除 cutoff_date = datetime.now() - timedelta(days=7) cutoff_str = cutoff_date.strftime("%Y-%m-%d") all_users = await self.storage.get_all_users() reset_count = 0 for user_id in all_users: user_data = await self.storage.get_user_data(user_id) # 古い1日のリクエストデータを削除 daily_requests = user_data["rate_limits"]["daily_requests"] dates_to_delete = [date for date in daily_requests.keys() if date < cutoff_str] for date in dates_to_delete: del daily_requests[date] reset_count += 1 # 古いAPI呼び出しデータを削除 api_calls = user_data["rate_limits"]["api_calls"] dates_to_delete = [date for date in api_calls.keys() if date < cutoff_str] for date in dates_to_delete: del api_calls[date] reset_count += 1 if reset_count > 0: await self.storage.update_user_data(user_id, user_data) if reset_count > 0: logger.info(f"{reset_count}件の古い制限データをリセットしました") return reset_count except Exception as e: logger.error(f"カウンターリセットエラー: {e}") return 0 def _get_next_reset_time(self) -> str: """次のリセット時刻を取得（翌日の0時）""" tomorrow = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) return tomorrow.isoformat() async def is_request_allowed(self, user_id: str) -> Tuple[bool, str]: """リクエストが許可されているかチェック（統合チェック）""" try: # 1日のリクエスト制限チェック request_allowed, request_info = await self.check_daily_request_limit(user_id) if not request_allowed: remaining_time = self._calculate_remaining_time() return False, f"1日のリクエスト制限に達しています。次回リクエスト可能時刻: {remaining_time}" # API呼び出し制限チェック api_allowed, api_info = await self.check_api_call_limit(user_id) if not api_allowed: remaining_time = self._calculate_remaining_time() return False, f"API呼び出し制限に達しています。次回リクエスト可能時刻: {remaining_time}" return True, "リクエスト可能です" except Exception as e: logger.error(f"リクエスト許可チェックエラー: {e}") return False, f"制限チェック中にエラーが発生しました: {e}" def _calculate_remaining_time(self) -> str: """次回リクエスト可能までの残り時間を計算""" now = datetime.now() tomorrow = now.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) remaining = tomorrow - now hours = remaining.seconds // 3600 minutes = (remaining.seconds % 3600) // 60 return f"{hours}時間{minutes}分後" async def get_rate_limit_stats(self) -> Dict[str, Any]: """レート制限の統計情報を取得""" try: all_users = await self.storage.get_all_users() today = datetime.now().strftime("%Y-%m-%d") total_requests_today = 0 total_api_calls_today = 0 active_users_today = 0 for user_id in all_users: user_data = await self.storage.get_user_data(user_id) # 今日のリクエスト数 daily_requests = user_data["rate_limits"]["daily_requests"] user_requests_today = daily_requests.get(today, 0) total_requests_today += user_requests_today # 今日のAPI呼び出し数 api_calls = user_data["rate_limits"]["api_calls"] user_api_calls_today = api_calls.get(today, 0) total_api_calls_today += user_api_calls_today # アクティブユーザー数 if user_requests_today > 0 or user_api_calls_today > 0: active_users_today += 1 return { "total_users": len(all_users), "active_users_today": active_users_today, "total_requests_today": total_requests_today, "total_api_calls_today": total_api_calls_today, "max_daily_requests": self.max_daily_requests, "max_api_calls_per_day": self.max_api_calls_per_day, "debug_mode": self.debug_mode, "date": today } except Exception as e: logger.error(f"統計情報取得エラー: {e}") return {"error": str(e)} def is_debug_mode(self) -> bool: """デバッグモードかどうかを確認""" return self.debug_mode async def set_debug_mode(self, enabled: bool) -> None: """デバッグモードの設定（動的変更）""" self.debug_mode = enabled if enabled: self.max_daily_requests = int(os.getenv("DEBUG_MAX_DAILY_REQUESTS", "10")) self.max_api_calls_per_day = int(os.getenv("DEBUG_MAX_API_CALLS", "100")) logger.info("デバッグモードを有効にしました") else: self.max_daily_requests = int(os.getenv("MAX_DAILY_REQUESTS", "1")) self.max_api_calls_per_day = int(os.getenv("MAX_API_CALLS_PER_DAY", "10")) logger.info("デバッグモードを無効にしました") async def force_reset_user_limits(self, user_id: str) -> None: """特定ユーザーの制限を強制リセット（デバッグ用）""" if not self.debug_mode: raise RateLimitError("デバッグモードでのみ利用可能です") try: user_data = await self.storage.get_user_data(user_id) today = datetime.now().strftime("%Y-%m-%d") # 今日の制限をリセット user_data["rate_limits"]["daily_requests"][today] = 0 user_data["rate_limits"]["api_calls"][today] = 0 await self.storage.update_user_data(user_id, user_data) logger.info(f"ユーザー {user_id} の制限を強制リセットしました") except Exception as e: logger.error(f"強制リセットエラー: {e}") raise RateLimitError(f"制限のリセットに失敗しました: {e}") # テスト用の関数 async def test_rate_limit_manager(): """RateLimitManagerのテスト""" import tempfile import uuid from async_storage_manager import AsyncStorageManager # 一時ディレクトリでテスト with tempfile.TemporaryDirectory() as temp_dir: test_file = os.path.join(temp_dir, "test_letters.json") storage = AsyncStorageManager(test_file) rate_limiter = AsyncRateLimitManager(storage) print("=== RateLimitManagerテスト開始 ===") user_id = str(uuid.uuid4()) # 初回リクエストチェック allowed, message = await rate_limiter.is_request_allowed(user_id) print(f"✓ 初回リクエストチェック: {allowed} - {message}") # リクエスト記録 await rate_limiter.record_request(user_id) print("✓ リクエスト記録成功") # API呼び出し記録 await rate_limiter.record_api_call(user_id, "groq") print("✓ API呼び出し記録成功") # 制限状況確認 status = await rate_limiter.get_user_limits_status(user_id) print(f"✓ 制限状況確認: {status}") # 統計情報取得 stats = await rate_limiter.get_rate_limit_stats() print(f"✓ 統計情報取得: {stats}") # デバッグモードテスト await rate_limiter.set_debug_mode(True) print("✓ デバッグモード有効化成功") # 強制リセットテスト（デバッグモード時のみ） await rate_limiter.force_reset_user_limits(user_id) print("✓ 強制リセット成功") print("=== 全てのテストが完了しました！ ===") if __name__ == "__main__": asyncio.run(test_rate_limit_manager())