refactor request limiter
Browse files- app/gemini_client.py +13 -5
- app/request_limit_manager.py +3 -3
app/gemini_client.py
CHANGED
|
@@ -30,6 +30,16 @@ class GeminiClient:
|
|
| 30 |
logger.info(f"[GEMINI] Created new model instance for key={key[:5]}...{key[-5:]} model={model}")
|
| 31 |
return self._cached_model_instance
|
| 32 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
def generate_text(self, prompt: str, **kwargs) -> str:
|
| 34 |
last_error = None
|
| 35 |
max_retries = 3
|
|
@@ -67,13 +77,11 @@ class GeminiClient:
|
|
| 67 |
if m:
|
| 68 |
retry_delay = int(m.group(1))
|
| 69 |
|
| 70 |
-
# Log failure
|
| 71 |
self.limit_manager.log_request(key, model, success=False, retry_delay=retry_delay)
|
| 72 |
|
| 73 |
-
#
|
| 74 |
-
|
| 75 |
-
self._cached_key = None
|
| 76 |
-
self._cached_model = None
|
| 77 |
|
| 78 |
logger.warning(f"[GEMINI] Rate limit hit, will retry with new key/model (attempt {attempt + 1}/{max_retries})")
|
| 79 |
last_error = e
|
|
|
|
| 30 |
logger.info(f"[GEMINI] Created new model instance for key={key[:5]}...{key[-5:]} model={model}")
|
| 31 |
return self._cached_model_instance
|
| 32 |
|
| 33 |
+
def _clear_cache_if_needed(self, new_key: str, new_model: str):
|
| 34 |
+
"""
|
| 35 |
+
Chỉ clear cache khi key/model thực sự thay đổi.
|
| 36 |
+
"""
|
| 37 |
+
if (self._cached_key != new_key or self._cached_model != new_model):
|
| 38 |
+
logger.info(f"[GEMINI] Clearing cache due to key/model change: {self._cached_key}->{new_key}, {self._cached_model}->{new_model}")
|
| 39 |
+
self._cached_model_instance = None
|
| 40 |
+
self._cached_key = None
|
| 41 |
+
self._cached_model = None
|
| 42 |
+
|
| 43 |
def generate_text(self, prompt: str, **kwargs) -> str:
|
| 44 |
last_error = None
|
| 45 |
max_retries = 3
|
|
|
|
| 77 |
if m:
|
| 78 |
retry_delay = int(m.group(1))
|
| 79 |
|
| 80 |
+
# Log failure với key/model thực tế đang được sử dụng
|
| 81 |
self.limit_manager.log_request(key, model, success=False, retry_delay=retry_delay)
|
| 82 |
|
| 83 |
+
# Chỉ clear cache nếu key/model thay đổi
|
| 84 |
+
# Không clear cache ngay lập tức để tránh recreate không cần thiết
|
|
|
|
|
|
|
| 85 |
|
| 86 |
logger.warning(f"[GEMINI] Rate limit hit, will retry with new key/model (attempt {attempt + 1}/{max_retries})")
|
| 87 |
last_error = e
|
app/request_limit_manager.py
CHANGED
|
@@ -120,11 +120,11 @@ class RequestLimitManager:
|
|
| 120 |
self.status[key][model]["status"] = "blocked"
|
| 121 |
self.status[key][model]["timestamp"] = now + (retry_delay or 60)
|
| 122 |
|
| 123 |
-
#
|
|
|
|
| 124 |
if key == self.current_key and model == self.current_model:
|
| 125 |
logger.warning(f"[LIMIT] Current pair blocked, will scan for new pair on next request")
|
| 126 |
-
|
| 127 |
-
self.current_model = None
|
| 128 |
|
| 129 |
def iterate_key_model(self) -> Iterator[Tuple[str, str]]:
|
| 130 |
"""
|
|
|
|
| 120 |
self.status[key][model]["status"] = "blocked"
|
| 121 |
self.status[key][model]["timestamp"] = now + (retry_delay or 60)
|
| 122 |
|
| 123 |
+
# Chỉ clear current pair nếu chính xác là pair đang được sử dụng
|
| 124 |
+
# Không clear ngay lập tức để tránh trigger scan không cần thiết
|
| 125 |
if key == self.current_key and model == self.current_model:
|
| 126 |
logger.warning(f"[LIMIT] Current pair blocked, will scan for new pair on next request")
|
| 127 |
+
# Không clear ngay lập tức, để get_current_key_model() tự xử lý
|
|
|
|
| 128 |
|
| 129 |
def iterate_key_model(self) -> Iterator[Tuple[str, str]]:
|
| 130 |
"""
|