refactor request limiter

app/embedding.py

@@ -19,6 +19,8 @@ class EmbeddingClient:
         self.provider = getattr(settings, 'embedding_provider', 'default')
         self.model = getattr(settings, 'embedding_model', 'models/embedding-001')
         self.gemini_client: Optional[GeminiClient] = GeminiClient() if self.provider == 'gemini' else None
+        
+        logger.info(f"[EMBEDDING] Initialized with provider={self.provider}, model={self.model}")
 
     @timing_decorator_async
     async def create_embedding(self, text: str) -> List[float]:
@@ -35,31 +37,37 @@ class EmbeddingClient:
                 import asyncio
                 loop = asyncio.get_event_loop()
                 gemini_client = self.gemini_client  # type: ignore
+                
+                # Luôn sử dụng model từ config, không phụ thuộc vào key/model từ RequestLimitManager
+                logger.info(f"[EMBEDDING] Creating embedding with model={self.model}")
                 embedding = await loop.run_in_executor(None, lambda: gemini_client.create_embedding(text, model=self.model))
+                
                 # Kiểm tra kiểu dữ liệu trả về
                 if isinstance(embedding, list):
                     preview = f"{embedding[:10]}...{embedding[-10:]}" if len(embedding) > 20 else str(embedding)
-                    logger.info(f"[DEBUG] Embedding API response: {preview}")
+                    logger.info(f"[EMBEDDING] API response: {preview}")
                     return embedding
                 else:
-                    logger.error(f"[DEBUG] Unknown embedding type: {type(embedding)} - value: {embedding}")
+                    logger.error(f"[EMBEDDING] Unknown embedding type: {type(embedding)} - value: {embedding}")
                     raise RuntimeError(f"Embedding returned unexpected type: {type(embedding)}")
             except Exception as e:
-                logger.error(f"Error creating embedding with Gemini: {e}")
+                logger.error(f"[EMBEDDING] Error creating embedding with Gemini: {e}")
                 raise
+        
+        # Fallback to HuggingFace embedding
         url = "https://vietcat-vietnameseembeddingv2.hf.space/embed"
         payload = {"text": text}
         try:
             response = await call_endpoint_with_retry(self._client, url, payload)
             if response is not None:
                 data = response.json()
-                logger.info(f"[DEBUG] Embedding API response: {data['embedding'][:10]}...{data['embedding'][-10:]}")
+                logger.info(f"[EMBEDDING] HuggingFace API response: {data['embedding'][:10]}...{data['embedding'][-10:]}")
                 return data["embedding"]
             else:
-                logger.error("Embedding API response is None")
-                raise RuntimeError("Embedding API response is None")
+                logger.error("[EMBEDDING] HuggingFace API response is None")
+                raise RuntimeError("HuggingFace API response is None")
         except Exception as e:
-            logger.error(f"Error creating embedding: {e}")
+            logger.error(f"[EMBEDDING] Error creating embedding with HuggingFace: {e}")
             raise
 
     def cosine_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
@@ -73,5 +81,12 @@ class EmbeddingClient:
             b = np.array(embedding2)
             return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
         except Exception as e:
-            logger.error(f"Error calculating similarity: {e}")
-            return 0.0 
+            logger.error(f"[EMBEDDING] Error calculating similarity: {e}")
+            return 0.0
+
+    def get_embedding_model(self) -> str:
+        """
+        Trả về model được config cho embedding.
+        Dùng để verify rằng model đúng được sử dụng.
+        """
+        return self.model 

app/gemini_client.py

@@ -8,25 +8,56 @@ from typing import List, Optional
 class GeminiClient:
     def __init__(self):
         self.limit_manager = RequestLimitManager("gemini")
+        self._cached_model = None
+        self._cached_key = None
+        self._cached_model_instance = None
+
+    def _get_model_instance(self, key: str, model: str):
+        """
+        Cache model instance để tránh recreate mỗi lần.
+        """
+        if (self._cached_key == key and 
+            self._cached_model == model and 
+            self._cached_model_instance is not None):
+            return self._cached_model_instance
+        
+        # Configure và tạo model instance mới
+        configure(api_key=key)
+        self._cached_model_instance = GenerativeModel(model)
+        self._cached_key = key
+        self._cached_model = model
+        
+        logger.info(f"[GEMINI] Created new model instance for key={key[:5]}...{key[-5:]} model={model}")
+        return self._cached_model_instance
 
     def generate_text(self, prompt: str, **kwargs) -> str:
         last_error = None
-        for key, model in self.limit_manager.iterate_key_model():
+        max_retries = 3
+        
+        for attempt in range(max_retries):
             try:
-                configure(api_key=key)
-                _model = GenerativeModel(model)
+                # Lấy current key/model từ manager
+                key, model = self.limit_manager.get_current_key_model()
+                
+                # Sử dụng cached model instance
+                _model = self._get_model_instance(key, 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)
@@ -35,39 +66,62 @@ class GeminiClient:
                     m = re.search(r'retry_delay.*?seconds: (\d+)', msg)
                     if m:
                         retry_delay = int(m.group(1))
+                    
+                    # Log failure và trigger scan cho key/model mới
                     self.limit_manager.log_request(key, model, success=False, retry_delay=retry_delay)
-                last_error = e
-                continue
+                    
+                    # Clear cache để force tạo model instance mới với key/model mới
+                    self._cached_model_instance = None
+                    self._cached_key = None
+                    self._cached_model = None
+                    
+                    logger.warning(f"[GEMINI] Rate limit hit, will retry with new key/model (attempt {attempt + 1}/{max_retries})")
+                    last_error = e
+                    continue
+                else:
+                    # Lỗi khác không phải rate limit
+                    logger.error(f"[GEMINI] Error generating text: {e}")
+                    last_error = e
+                    break
+        
         raise last_error or RuntimeError("No available Gemini API key/model")
 
     def count_tokens(self, prompt: str) -> int:
-        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
+        try:
+            key, model = self.limit_manager.get_current_key_model()
+            _model = self._get_model_instance(key, model)
+            return _model.count_tokens(prompt).total_tokens
+        except Exception as e:
+            logger.error(f"[GEMINI] Error counting tokens: {e}")
+            return 0
 
     def create_embedding(self, text: str, model: Optional[str] = None) -> list:
         last_error = None
-        for key, m in self.limit_manager.iterate_key_model():
-            m = m or ""
-            use_model = model if model not in (None, "") else m
-            if not use_model:
-                continue
-            use_model = str(use_model)
+        max_retries = 3
+        
+        for attempt in range(max_retries):
             try:
+                key, default_model = self.limit_manager.get_current_key_model()
+                
+                # Ưu tiên model được truyền vào parameter, chỉ fallback về default_model nếu không có
+                use_model = model if model and model.strip() else default_model
+                
+                if not use_model:
+                    raise ValueError("No model specified for embedding")
+                
+                logger.info(f"[GEMINI][EMBEDDING] Using model={use_model} (requested={model}, default={default_model})")
+                
                 configure(api_key=key)
                 response = embed_content(
                     model=use_model,
                     content=text,
                     task_type="retrieval_query"
                 )
+                
                 self.limit_manager.log_request(key, use_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)
@@ -76,7 +130,16 @@ class GeminiClient:
                     m_retry = re.search(r'retry_delay.*?seconds: (\d+)', msg)
                     if m_retry:
                         retry_delay = int(m_retry.group(1))
+                    
+                    # Log failure và trigger scan cho key/model mới
                     self.limit_manager.log_request(key, use_model, success=False, retry_delay=retry_delay)
-                last_error = e
-                continue
+                    
+                    logger.warning(f"[GEMINI] Rate limit hit in embedding, will retry with new key/model (attempt {attempt + 1}/{max_retries})")
+                    last_error = e
+                    continue
+                else:
+                    logger.error(f"[GEMINI] Error creating embedding: {e}")
+                    last_error = e
+                    break
+        
         raise last_error or RuntimeError("No available Gemini API key/model") 

app/main.py

@@ -19,6 +19,7 @@ from .constants import VEHICLE_KEYWORDS, SHEET_RANGE, VEHICLE_KEYWORD_TO_COLUMN
 from .health import router as health_router
 from .llm import create_llm_client
 from .reranker import Reranker
+from .request_limit_manager import RequestLimitManager
 
 app = FastAPI(title="WeBot Facebook Messenger API")
 
@@ -39,6 +40,10 @@ logger.info("[STARTUP] Đang lấy PORT từ biến môi trường hoặc config
 port = int(os.environ.get("PORT", settings.port if hasattr(settings, 'port') else 7860))
 logger.info(f"[STARTUP] PORT sử dụng: {port}")
 
+logger.info("[STARTUP] Khởi tạo global RequestLimitManager...")
+# Global RequestLimitManager instance - singleton
+request_limit_manager = RequestLimitManager("gemini")
+
 logger.info("[STARTUP] Khởi tạo FacebookClient...")
 facebook_client = FacebookClient(settings.facebook_app_secret)
 logger.info("[STARTUP] Khởi tạo SheetsClient...")

app/request_limit_manager.py

@@ -5,10 +5,23 @@ from app.config import get_settings
 from loguru import logger
 
 class RequestLimitManager:
+    _instance = None
+    _lock = threading.Lock()
+    
+    def __new__(cls, provider: str):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+    
     def __init__(self, provider: str):
+        if hasattr(self, 'initialized'):
+            return
         self.provider = provider
         self.lock = threading.Lock()
         self._init_keys_models()
+        self.initialized = True
 
     def _init_keys_models(self):
         settings = get_settings()
@@ -22,16 +35,82 @@ class RequestLimitManager:
             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
+        self.current_key: Optional[str] = self.api_keys[0] if self.api_keys else None
+        self.current_model: Optional[str] = self.models[0] if self.models else None
+        key_display = f"{self.current_key[:5]}...{self.current_key[-5:]}" if self.current_key else "None"
+        logger.info(f"[LIMIT] Initialized with current key={key_display} model={self.current_model}")
+
+    def get_current_key_model(self) -> Tuple[str, str]:
+        """
+        Trả về cặp key/model hiện tại đang active.
+        Chỉ scan tìm key/model mới khi current pair bị blocked.
+        """
+        with self.lock:
+            now = time.time()
+            
+            # Check if current pair is still available
+            if self.current_key and self.current_model:
+                info = self.status.get(self.current_key, {}).get(self.current_model, {})
+                status = info.get("status", "active")
+                ts = float(info.get("timestamp", 0.0))
+                
+                if status == "active" or (status == "blocked" and now > ts):
+                    logger.info(f"[LIMIT] Using current key={self.current_key[:5]}...{self.current_key[-5:]} model={self.current_model}")
+                    return self.current_key, self.current_model
+            
+            # Current pair not available, scan for new one
+            logger.warning(f"[LIMIT] Current pair not available, scanning for new key/model...")
+            new_key, new_model = self._find_available_key_model()
+            
+            if new_key and new_model:
+                self.current_key = new_key
+                self.current_model = new_model
+                logger.info(f"[LIMIT] Switched to new key={self.current_key[:5]}...{self.current_key[-5:]} model={self.current_model}")
+                return self.current_key, self.current_model
+            else:
+                logger.error(f"[LIMIT] No available key/model found for provider {self.provider}")
+                raise RuntimeError(f"No available key/model for provider {self.provider}")
+
+    def _find_available_key_model(self) -> Tuple[Optional[str], Optional[str]]:
+        """
+        Tìm cặp key/model khả dụng gần nhất.
+        """
+        now = time.time()
+        keys = self.api_keys[:]
+        models = self.models[:]
+        
+        # Ưu tiên default key/model nếu có
+        if self.current_key and self.current_key in keys:
+            keys.remove(self.current_key)
+            keys = [self.current_key] + keys
+        if self.current_model and self.current_model in models:
+            models.remove(self.current_model)
+            models = [self.current_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" or (status == "blocked" and now > ts):
+                    logger.info(f"[LIMIT] Found available key={key[:5]}...{key[-5:]} model={model}")
+                    return key, model
+        
+        return None, None
 
     def log_request(self, key: str, model: str, success: bool, retry_delay: Optional[int] = None):
+        """
+        Log kết quả request và cập nhật status.
+        Nếu request fail với 429, trigger scan cho key/model mới.
+        """
         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:
                 logger.info(f"[LIMIT] Mark key={key[:5]}...{key[-5:]} - model={model} as active at {now}")
                 self.status[key][model]["status"] = "active"
@@ -40,34 +119,17 @@ class RequestLimitManager:
                 logger.warning(f"[LIMIT] Mark key={key[:5]}...{key[-5:]} - model={model} as blocked until {now + (retry_delay or 60)} (retry_delay={retry_delay})")
                 self.status[key][model]["status"] = "blocked"
                 self.status[key][model]["timestamp"] = now + (retry_delay or 60)
+                
+                # Nếu current pair bị blocked, trigger scan cho pair mới
+                if key == self.current_key and model == self.current_model:
+                    logger.warning(f"[LIMIT] Current pair blocked, will scan for new pair on next request")
+                    self.current_key = None
+                    self.current_model = None
 
     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
-        logger.info(f"[LIMIT] Trying key/model candidates: {[(k[:6]+'...', m) for k in keys for m in models]}")
-        found = False
-        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":
-                    logger.info(f"[LIMIT] Use key={key[:5]}...{key[-5:]} - model={model} (active)")
-                    found = True
-                    yield key, model
-                elif status == "blocked" and now > ts:
-                    logger.info(f"[LIMIT] Use key={key[:5]}...{key[-5:]} - model={model} (was blocked, now retry)")
-                    found = True
-                    yield key, model
-        if not found:
-            logger.warning(f"[LIMIT] No available key/model for provider {self.provider}")
-            pass
-        # Nếu không có key/model nào hợp lệ, không yield gì 
+        """
+        Legacy method - chỉ trả về current pair.
+        Để tương thích với code cũ.
+        """
+        key, model = self.get_current_key_model()
+        yield key, model 

app/reranker.py

@@ -18,34 +18,76 @@ class Reranker:
         else:
             raise NotImplementedError(f"Rerank provider {self.provider} not supported yet.")
 
+    async def _score_doc(self, query: str, doc: Dict) -> Dict:
+        """
+        Score một document với query.
+        """
+        content = (doc.get('tieude', '') or '') + ' ' + (doc.get('noidung', '') or '')
+        prompt = (
+            f"Đoạn luật: {content}\n"
+            f"Câu hỏi: {query}\n"
+            "Hãy đánh giá mức độ liên quan giữa đoạn luật và câu hỏi trên thang điểm 0-10. "
+            "Chỉ trả về một số duy nhất."
+        )
+        
+        try:
+            if self.provider == 'gemini':
+                loop = asyncio.get_event_loop()
+                logger.info(f"[RERANK] Sending prompt to Gemini: {prompt}")
+                score = await loop.run_in_executor(None, self.client.generate_text, prompt)
+                logger.info(f"[RERANK] Got score from Gemini: {score}")
+            else:
+                raise NotImplementedError(f"Rerank provider {self.provider} not supported yet in rerank method.")
+            
+            score = float(str(score).strip().split()[0])
+            doc['rerank_score'] = score
+            return doc
+            
+        except Exception as e:
+            logger.error(f"[RERANK] Lỗi khi tính score: {e} | doc: {doc}")
+            doc['rerank_score'] = 0
+            return doc
+
     async def rerank(self, query: str, docs: List[Dict], top_k: int = 5) -> List[Dict]:
         """
         Rerank docs theo độ liên quan với query, trả về top_k docs.
+        Sử dụng concurrency để process nhiều docs cùng lúc.
         """
         logger.info(f"[RERANK] Start rerank for query: {query} | docs: {len(docs)} | top_k: {top_k}")
+        
+        if not docs:
+            return []
+        
+        # Giới hạn số docs để rerank (tối đa 10 docs)
+        docs_to_rerank = docs[:10] if len(docs) > 10 else docs
+        logger.info(f"[RERANK] Will rerank {len(docs_to_rerank)} docs (limited from {len(docs)})")
+        
+        # Process docs với concurrency
+        batch_size = 5  # Process 5 docs cùng lúc
         scored = []
-        for doc in docs:
-            content = (doc.get('tieude', '') or '') + ' ' + (doc.get('noidung', '') or '')
-            prompt = (
-                f"Đoạn luật: {content}\n"
-                f"Câu hỏi: {query}\n"
-                "Hãy đánh giá mức độ liên quan giữa đoạn luật và câu hỏi trên thang điểm 0-10. "
-                "Chỉ trả về một số duy nhất."
-            )
-            try:
-                if self.provider == 'gemini':
-                    loop = asyncio.get_event_loop()
-                    logger.info(f"[RERANK] Sending prompt to Gemini: {prompt}")
-                    score = await loop.run_in_executor(None, self.client.generate_text, prompt)
-                    logger.info(f"[RERANK] Got score from Gemini: {score}")
-                else:
-                    raise NotImplementedError(f"Rerank provider {self.provider} not supported yet in rerank method.")
-                score = float(str(score).strip().split()[0])
-            except Exception as e:
-                logger.error(f"[RERANK] Lỗi khi tính score: {e} | doc: {doc}")
-                score = 0
-            doc['rerank_score'] = score
-            scored.append(doc)
+        
+        for i in range(0, len(docs_to_rerank), batch_size):
+            batch = docs_to_rerank[i:i + batch_size]
+            logger.info(f"[RERANK] Processing batch {i//batch_size + 1}: {len(batch)} docs")
+            
+            # Tạo tasks cho batch hiện tại
+            tasks = [self._score_doc(query, doc) for doc in batch]
+            
+            # Chạy batch concurrently
+            batch_results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            # Xử lý kết quả
+            for result in batch_results:
+                if isinstance(result, Exception):
+                    logger.error(f"[RERANK] Batch processing error: {result}")
+                    continue
+                scored.append(result)
+            
+            logger.info(f"[RERANK] Completed batch {i//batch_size + 1}, processed {len(scored)} docs so far")
+        
+        # Sort theo score và trả về top_k
         scored = sorted(scored, key=lambda x: x['rerank_score'], reverse=True)
-        logger.info(f"[RERANK] Top reranked docs: {scored[:top_k]}")
-        return scored[:top_k] 
+        result = scored[:top_k]
+        
+        logger.info(f"[RERANK] Top reranked docs: {result}")
+        return result