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 |
"""
|