62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
from typing import List, Optional, Iterator
|
||
from loguru import logger
|
||
import threading
|
||
from ..config import settings
|
||
|
||
|
||
class KeyPoolManager:
|
||
"""Google API Key池管理器，用于处理限流重试"""
|
||
|
||
def __init__(self):
|
||
self._keys: List[str] = []
|
||
self._current_index = 0
|
||
self._lock = threading.Lock()
|
||
self._initialize_keys()
|
||
|
||
def _initialize_keys(self):
|
||
"""初始化key池"""
|
||
# 解析key字符串，支持单个或多个key
|
||
self._keys = [key.strip() for key in settings.google_api_keys.split(',') if key.strip()]
|
||
if not self._keys:
|
||
raise ValueError("GOOGLE_API_KEYS不能为空")
|
||
logger.info(f"初始化Google API Key池，共{len(self._keys)}个key")
|
||
|
||
def get_current_key(self) -> str:
|
||
"""获取当前key"""
|
||
with self._lock:
|
||
if not self._keys:
|
||
raise ValueError("没有可用的Google API Key")
|
||
return self._keys[self._current_index]
|
||
|
||
def switch_to_next_key(self) -> bool:
|
||
"""切换到下一个key
|
||
|
||
Returns:
|
||
bool: 如果还有下一个key返回True，否则返回False
|
||
"""
|
||
with self._lock:
|
||
if len(self._keys) <= 1:
|
||
return False
|
||
|
||
self._current_index = (self._current_index + 1) % len(self._keys)
|
||
logger.info(f"切换到下一个Google API Key，当前索引: {self._current_index}")
|
||
return True
|
||
|
||
def reset_to_first_key(self):
|
||
"""重置到第一个key"""
|
||
with self._lock:
|
||
self._current_index = 0
|
||
logger.info("重置到第一个Google API Key")
|
||
|
||
def get_all_keys(self) -> List[str]:
|
||
"""获取所有key（用于测试）"""
|
||
return self._keys.copy()
|
||
|
||
def has_multiple_keys(self) -> bool:
|
||
"""是否有多个key"""
|
||
return len(self._keys) > 1
|
||
|
||
|
||
# 全局key池管理器实例
|
||
key_pool_manager = KeyPoolManager()
|