add quota manager

app/config.py

@@ -31,9 +31,11 @@ class Settings(BaseSettings):
     log_level: str = os.getenv("LOG_LEVEL", "INFO") or "INFO"
 
     # Gemini Configuration
-    gemini_api_key: str = os.getenv("GEMINI_API_KEY") or ""
+    # Hỗ trợ nhiều API key và model cho Gemini
+    # Định nghĩa biến môi trường: GEMINI_API_KEYS="key1,key2,..."; GEMINI_MODELS="model1,model2,..."
+    gemini_api_keys: list[str] = os.getenv("GEMINI_API_KEYS", "").split(",") if os.getenv("GEMINI_API_KEYS") else []
+    gemini_models: list[str] = os.getenv("GEMINI_MODELS", "").split(",") if os.getenv("GEMINI_MODELS") else []
     gemini_base_url: str = os.getenv("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com/v1/models/gemini-2.5-flash:generateContent") or ""
-    gemini_model: str = os.getenv("GEMINI_MODEL", "gemini-2.5-flash") or ""
 
     # LLM (chat/completion) provider/model
     llm_provider: str = os.getenv("LLM_PROVIDER", "gemini") or ""

app/gemini_client.py

@@ -2,48 +2,75 @@ from google.generativeai.embedding import embed_content
 from google.generativeai.client import configure
 from google.generativeai.generative_models import GenerativeModel
 from loguru import logger
+from .request_limit_manager import RequestLimitManager
 
 class GeminiClient:
-    def __init__(self, api_key: str, model: str = "gemini-1.5-flash-latest"):
-        self.api_key = api_key
-        self.model = model
-        configure(api_key=api_key)
-        self._model = GenerativeModel(model)
+    def __init__(self):
+        self.limit_manager = RequestLimitManager("gemini")
 
     def generate_text(self, prompt: str, **kwargs) -> str:
-        try:
-            response = self._model.generate_content(prompt, **kwargs)
-            if hasattr(response, 'usage_metadata'):
-                logger.info(f"[GEMINI][USAGE] Prompt Token Count: {response.usage_metadata.prompt_token_count} - Candidate Token Count: {response.usage_metadata.candidates_token_count} - Total Token Count: {response.usage_metadata.total_token_count}")
-            # response có thể là GenerativeModelResponse, lấy text hoặc trả về str
-            if hasattr(response, 'text'):
-                logger.info(f"[GEMINI][TEXT_RESPONSE] {response.text}")
-                return response.text
-            elif hasattr(response, 'candidates') and response.candidates:
-                logger.info(f"[GEMINI][CANDIDATES_RESPONSE] {response.candidates[0].content.parts[0].text}")
-                return response.candidates[0].content.parts[0].text
-            logger.info(f"[GEMINI][RAW_RESPONSE] {response}")
-            return str(response)
-        except Exception as e:
-            logger.error(f"[GEMINI] Error: {e}")
-            raise
+        last_error = None
+        for key, model in self.limit_manager.iterate_key_model():
+            try:
+                configure(api_key=key)
+                _model = GenerativeModel(model)
+                response = _model.generate_content(prompt, **kwargs)
+                self.limit_manager.log_request(key, model, success=True)
+                if hasattr(response, 'usage_metadata'):
+                    logger.info(f"[GEMINI][USAGE] Prompt Token Count: {response.usage_metadata.prompt_token_count} - Candidate Token Count: {response.usage_metadata.candidates_token_count} - Total Token Count: {response.usage_metadata.total_token_count}")
+                if hasattr(response, 'text'):
+                    logger.info(f"[GEMINI][TEXT_RESPONSE] {response.text}")
+                    return response.text
+                elif hasattr(response, 'candidates') and response.candidates:
+                    logger.info(f"[GEMINI][CANDIDATES_RESPONSE] {response.candidates[0].content.parts[0].text}")
+                    return response.candidates[0].content.parts[0].text
+                logger.info(f"[GEMINI][RAW_RESPONSE] {response}")
+                return str(response)
+            except Exception as e:
+                import re
+                msg = str(e)
+                if "429" in msg or "rate limit" in msg.lower():
+                    retry_delay = 60
+                    m = re.search(r'retry_delay.*?seconds: (\d+)', msg)
+                    if m:
+                        retry_delay = int(m.group(1))
+                    self.limit_manager.log_request(key, model, success=False, retry_delay=retry_delay)
+                last_error = e
+                continue
+        raise last_error or RuntimeError("No available Gemini API key/model")
 
     def count_tokens(self, prompt: str) -> int:
-        try:
-            return self._model.count_tokens(prompt).total_tokens
-        except Exception as e:
-            logger.error(f"[GEMINI] Token count error: {e}")
-            return 0
+        for key, model in self.limit_manager.iterate_key_model():
+            try:
+                configure(api_key=key)
+                _model = GenerativeModel(model)
+                return _model.count_tokens(prompt).total_tokens
+            except Exception:
+                continue
+        return 0
 
     def create_embedding(self, text: str) -> list:
-        try:
-            response = embed_content(
-                model=self.model,
-                content=text,
-                task_type="retrieval_query"
-            )
-            logger.info(f"[GEMINI][EMBEDDING][RAW_RESPONSE] {response['embedding'][:10]} ..... {response['embedding'][-10:]}")
-            return response['embedding']
-        except Exception as e:
-            logger.error(f"[GEMINI][EMBEDDING] Error: {e}")
-            raise 
+        last_error = None
+        for key, model in self.limit_manager.iterate_key_model():
+            try:
+                configure(api_key=key)
+                response = embed_content(
+                    model=model,
+                    content=text,
+                    task_type="retrieval_query"
+                )
+                self.limit_manager.log_request(key, model, success=True)
+                logger.info(f"[GEMINI][EMBEDDING][RAW_RESPONSE] {response['embedding'][:10]} ..... {response['embedding'][-10:]}")
+                return response['embedding']
+            except Exception as e:
+                import re
+                msg = str(e)
+                if "429" in msg or "rate limit" in msg.lower():
+                    retry_delay = 60
+                    m = re.search(r'retry_delay.*?seconds: (\d+)', msg)
+                    if m:
+                        retry_delay = int(m.group(1))
+                    self.limit_manager.log_request(key, model, success=False, retry_delay=retry_delay)
+                last_error = e
+                continue
+        raise last_error or RuntimeError("No available Gemini API key/model") 

app/llm.py

@@ -88,9 +88,7 @@ class LLMClient:
     
     def _setup_gemini(self, config: Dict[str, Any]):
         """Cấu hình cho Gemini."""
-        self.api_key = config.get("api_key", "")
-        self.model = config.get("model", "gemini-1.5-flash-latest")
-        self.gemini_client = GeminiClient(self.api_key, self.model)
+        self.gemini_client = GeminiClient()
     
     @timing_decorator_async
     async def generate_text(
@@ -479,7 +477,6 @@ if __name__ == "__main__":
         llm_client = create_llm_client(
             provider=settings.llm_provider,
             model=settings.llm_model,
-            api_key=settings.gemini_api_key,
             # ... các config khác nếu cần ...
         )
         

app/main.py

@@ -63,9 +63,8 @@ VEHICLE_KEYWORDS = ["xe máy", "ô tô", "xe đạp", "xe hơi"]
 # Khởi tạo LLM client Gemini
 llm_client = create_llm_client(
     provider="gemini",
-    api_key=settings.gemini_api_key,
     base_url=settings.gemini_base_url,
-    model=settings.gemini_model
+    model=settings.gemini_models[0] if settings.gemini_models else "gemini-2.5-flash"
 )
 
 reranker = Reranker()

app/request_limit_manager.py

@@ -0,0 +1,62 @@
+import time
+import threading
+from typing import Dict, List, Tuple, Optional, Iterator, Union
+from app.config import get_settings
+
+class RequestLimitManager:
+    def __init__(self, provider: str):
+        self.provider = provider
+        self.lock = threading.Lock()
+        self._init_keys_models()
+
+    def _init_keys_models(self):
+        settings = get_settings()
+        if self.provider == "gemini":
+            self.api_keys: List[str] = [k.strip() for k in getattr(settings, 'gemini_api_keys', []) if k.strip()]
+            self.models: List[str] = [m.strip() for m in getattr(settings, 'gemini_models', []) if m.strip()]
+        # Có thể mở rộng cho provider khác ở đây
+        self.status: Dict[str, Dict[str, Dict[str, Union[str, float]]]] = {}
+        now = time.time()
+        for key in self.api_keys:
+            self.status[key] = {}
+            for model in self.models:
+                self.status[key][model] = {"status": "active", "timestamp": now}
+        self.default_key: Optional[str] = self.api_keys[0] if self.api_keys else None
+        self.default_model: Optional[str] = self.models[0] if self.models else None
+
+    def log_request(self, key: str, model: str, success: bool, retry_delay: Optional[int] = None):
+        with self.lock:
+            now = time.time()
+            if key not in self.status:
+                self.status[key] = {}
+            if model not in self.status[key]:
+                self.status[key][model] = {"status": "active", "timestamp": now}
+            if success:
+                self.status[key][model]["status"] = "active"
+                self.status[key][model]["timestamp"] = now
+            else:
+                self.status[key][model]["status"] = "blocked"
+                self.status[key][model]["timestamp"] = now + (retry_delay or 60)
+
+    def iterate_key_model(self) -> Iterator[Tuple[str, str]]:
+        now = time.time()
+        keys = self.api_keys[:]
+        models = self.models[:]
+        # Ưu tiên default key/model nếu có
+        if self.default_key and self.default_key in keys:
+            keys.remove(self.default_key)
+            keys = [self.default_key] + keys
+        if self.default_model and self.default_model in models:
+            models.remove(self.default_model)
+            models = [self.default_model] + models
+        for key in keys:
+            for model in models:
+                info = self.status.get(key, {}).get(model, {"status": "active", "timestamp": 0.0})
+                status = info.get("status", "active")
+                ts = float(info.get("timestamp", 0.0))
+                if status == "active":
+                    yield key, model
+                elif status == "blocked" and now > ts:
+                    # Đã hết thời gian block, cho thử lại
+                    yield key, model
+        # Nếu không có key/model nào hợp lệ, không yield gì 

app/reranker.py

@@ -10,7 +10,7 @@ class Reranker:
         self.provider = getattr(settings, 'rerank_provider', settings.llm_provider)
         self.model = getattr(settings, 'rerank_model', settings.llm_model)
         if self.provider == 'gemini':
-            self.client = GeminiClient(settings.gemini_api_key, model=self.model)
+            self.client = GeminiClient()
         # elif self.provider == 'openai':
         #     self.client = OpenAIClient(settings.openai_api_key, model=self.model)
         # elif self.provider == 'cohere':