update query and rerank in parallel

app/gemini_client.py

@@ -6,29 +6,40 @@ from typing import Dict, List, Optional
 from google.generativeai.types import GenerationConfig, HarmCategory, HarmBlockThreshold
 
 from .request_limit_manager import RequestLimitManager
-from .utils import (
-    _safe_truncate
-)
+from .utils import _safe_truncate
 from .config import get_settings
 
+
 class GeminiResponseError(Exception):
     """Custom exception for non-retriable Gemini response issues like safety or token limits."""
+
     def __init__(self, message, finish_reason=None, usage_metadata=None):
         super().__init__(message)
         self.finish_reason = finish_reason
         self.usage_metadata = usage_metadata
 
     def __str__(self):
-        usage_str = f"Prompt: {self.usage_metadata.prompt_token_count}, Candidates: {self.usage_metadata.candidates_token_count}, Total: {self.usage_metadata.total_token_count}" if self.usage_metadata else "N/A"
+        usage_str = (
+            f"Prompt: {self.usage_metadata.prompt_token_count}, Candidates: {self.usage_metadata.candidates_token_count}, Total: {self.usage_metadata.total_token_count}"
+            if self.usage_metadata
+            else "N/A"
+        )
         return f"{super().__str__()} (Finish Reason: {self.finish_reason}, Usage: {usage_str})"
 
+
 class GeminiClient:
     def __init__(self):
         self.limit_manager = RequestLimitManager("gemini")
         settings = get_settings()
-        num_keys = len(settings.gemini_api_keys.split(',')) if settings.gemini_api_keys else 0
-        num_models = len(settings.gemini_models.split(',')) if settings.gemini_models else 0
-        logger.info(f"[GEMINI_INIT] Limiter is considering {num_keys} API keys and {num_models} models.")
+        num_keys = (
+            len(settings.gemini_api_keys.split(",")) if settings.gemini_api_keys else 0
+        )
+        num_models = (
+            len(settings.gemini_models.split(",")) if settings.gemini_models else 0
+        )
+        logger.info(
+            f"[GEMINI_INIT] Limiter is considering {num_keys} API keys and {num_models} models."
+        )
         self._cached_model = None
         self._cached_key = None
         self._cached_model_instance = None
@@ -44,26 +55,32 @@ class GeminiClient:
         """
         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):
+        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}")
+
+        logger.info(
+            f"[GEMINI] Created new model instance for key={key[:5]}...{key[-5:]} model={model}"
+        )
         return self._cached_model_instance
 
     def _clear_cache_if_needed(self, new_key: str, new_model: str):
         """
         Chỉ clear cache khi key/model thực sự thay đổi.
         """
-        if (self._cached_key != new_key or self._cached_model != new_model):
-            logger.info(f"[GEMINI] Clearing cache due to key/model change: {self._cached_key}->{new_key}, {self._cached_model}->{new_model}")
+        if self._cached_key != new_key or self._cached_model != new_model:
+            logger.info(
+                f"[GEMINI] Clearing cache due to key/model change: {self._cached_key}->{new_key}, {self._cached_model}->{new_model}"
+            )
             self._cached_model_instance = None
             self._cached_key = None
             self._cached_model = None
@@ -71,21 +88,19 @@ class GeminiClient:
     def generate_text(self, prompt: str, **kwargs) -> str:
         last_error = None
         max_retries = 3
-        
+
         for attempt in range(max_retries):
             try:
                 # 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,
-                    safety_settings=self.safety_settings,
-                    **kwargs
+                    prompt, safety_settings=self.safety_settings, **kwargs
                 )
-                
+
                 # Log toàn bộ nội dung response ở mức INFO để tiện gỡ lỗi
                 logger.info(f"[GEMINI][RAW_RESPONSE] {response}")
 
@@ -93,47 +108,81 @@ class GeminiClient:
                 # 1. Kiểm tra response có hợp lệ không
                 if not response.candidates:
                     # Lỗi này nên được coi là lỗi tạm thời, thử lại với key/model khác
-                    raise ValueError("Gemini response is missing 'candidates' field. Retrying...")
+                    raise ValueError(
+                        "Gemini response is missing 'candidates' field. Retrying..."
+                    )
 
-                candidate = response.candidates[0] 
-                finish_reason_name = getattr(getattr(candidate, 'finish_reason', None), 'name', 'UNKNOWN') 
+                candidate = response.candidates[0]
+                finish_reason_name = getattr(
+                    getattr(candidate, "finish_reason", None), "name", "UNKNOWN"
+                )
                 # Kiểm tra xem có nội dung thực sự không
                 # Sửa: Dùng getattr để tránh AttributeError nếu 'parts' không tồn tại
-                has_content = bool(candidate.content and getattr(candidate.content, 'parts', None))
+                has_content = bool(
+                    candidate.content and getattr(candidate.content, "parts", None)
+                )
 
                 # 2. Phân loại lỗi và xử lý
                 # Case 1: Lỗi nội dung không thể thử lại (SAFETY, MAX_TOKENS, etc.)
                 if finish_reason_name != "STOP":
-                    usage_metadata = response.usage_metadata if hasattr(response, 'usage_metadata') else None
+                    usage_metadata = (
+                        response.usage_metadata
+                        if hasattr(response, "usage_metadata")
+                        else None
+                    )
                     error_message = f"Gemini response finished with non-OK reason: {finish_reason_name}."
                     raise GeminiResponseError(
-                        error_message, finish_reason=finish_reason_name, usage_metadata=usage_metadata
+                        error_message,
+                        finish_reason=finish_reason_name,
+                        usage_metadata=usage_metadata,
                     )
 
                 # Case 2: Lỗi có thể thử lại (STOP nhưng không có nội dung)
-                if not has_content: # Tại đây, ta biết chắc chắn finish_reason_name là "STOP"
-                    usage_metadata = response.usage_metadata if hasattr(response, 'usage_metadata') else None
-                    last_error = GeminiResponseError("Gemini response finished with STOP but has no content parts.", finish_reason='STOP_NO_CONTENT', usage_metadata=usage_metadata)
-                    logger.warning(f"[GEMINI] Model returned STOP with no content. Retrying with another key/model... (Attempt {attempt + 1}/{max_retries})")
-                    self.limit_manager.log_request(key, model, success=False, retry_delay=5)
-                    continue # Thử lại vòng lặp với key/model mới
+                if (
+                    not has_content
+                ):  # Tại đây, ta biết chắc chắn finish_reason_name là "STOP"
+                    usage_metadata = (
+                        response.usage_metadata
+                        if hasattr(response, "usage_metadata")
+                        else None
+                    )
+                    last_error = GeminiResponseError(
+                        "Gemini response finished with STOP but has no content parts.",
+                        finish_reason="STOP_NO_CONTENT",
+                        usage_metadata=usage_metadata,
+                    )
+                    logger.warning(
+                        f"[GEMINI] Model returned STOP with no content. Retrying with another key/model... (Attempt {attempt + 1}/{max_retries})"
+                    )
+                    self.limit_manager.log_request(
+                        key, model, success=False, retry_delay=5
+                    )
+                    continue  # Thử lại vòng lặp với key/model mới
 
                 # Case 3: Thành công (STOP và có nội dung)
                 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, "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}"
+                    )
+
                 try:
-                    logger.info(f"[GEMINI][TEXT_RESPONSE] {_safe_truncate(response.text)}")
+                    logger.info(
+                        f"[GEMINI][TEXT_RESPONSE] {_safe_truncate(response.text)}"
+                    )
                     return response.text
                 except ValueError as ve:
                     # Safety net: Nếu truy cập .text thất bại dù các kiểm tra trước đó đã qua,
                     # coi như đây là lỗi STOP_NO_CONTENT và ném ra để tầng trên xử lý.
-                    usage_metadata = response.usage_metadata if hasattr(response, 'usage_metadata') else None
+                    usage_metadata = (
+                        response.usage_metadata
+                        if hasattr(response, "usage_metadata")
+                        else None
+                    )
                     raise GeminiResponseError(
                         f"Gemini response has no valid content part. Original error: {ve}",
-                        finish_reason='STOP_NO_CONTENT',
-                        usage_metadata=usage_metadata
+                        finish_reason="STOP_NO_CONTENT",
+                        usage_metadata=usage_metadata,
                     ) from ve
                 # --- END: Cải tiến logic xử lý response ---
             except GeminiResponseError as e:
@@ -142,28 +191,41 @@ class GeminiClient:
                 raise e
             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))
-                    
-                    # Log failure với key/model thực tế đang được sử dụng
-                    self.limit_manager.log_request(key, model, success=False, retry_delay=retry_delay)
-                    
-                    # Chỉ clear cache nếu key/model thay đổi
-                    # Không clear cache ngay lập tức để tránh recreate không cần thiết
-                    
-                    logger.warning(f"[GEMINI] Rate limit hit, will retry with new key/model (attempt {attempt + 1}/{max_retries})")
+                # Kiểm tra lỗi rate limit hoặc lỗi server (5xx)
+                is_rate_limit = "429" in msg or "rate limit" in msg.lower()
+                is_server_error = any(
+                    code in msg for code in ["500", "502", "503", "504"]
+                )
+
+                if is_rate_limit or is_server_error:
+                    retry_delay = 60  # Mặc định cho lỗi server
+                    if is_rate_limit:
+                        m = re.search(r"retry_delay.*?seconds: (\d+)", msg)
+                        if m:
+                            retry_delay = int(m.group(1))
+
+                    # Log lỗi và chặn cặp key/model hiện tại trong một khoảng thời gian
+                    self.limit_manager.log_request(
+                        key, model, success=False, retry_delay=retry_delay
+                    )
+
+                    error_type = "Rate limit" if is_rate_limit else "Server"
+                    logger.warning(
+                        f"[GEMINI] {error_type} error hit, will retry with new key/model "
+                        f"(attempt {attempt + 1}/{max_retries}). Error: {e}"
+                    )
                     last_error = e
-                    continue
+                    continue  # Tiếp tục vòng lặp để thử key/model mới
                 else:
-                    # Lỗi khác không phải rate limit (vd: timeout, server error)
-                    # sẽ được propagate lên để lớp llm.py xử lý retry với backoff.
-                    logger.error(f"[GEMINI] Error generating text: {e}")
+                    # Các lỗi khác không phải rate limit hoặc server error (vd: network timeout, invalid argument)
+                    # sẽ được propagate lên để lớp llm.py/reranker.py xử lý retry với backoff.
+                    logger.error(
+                        f"[GEMINI] Unhandled error generating text, propagating up: {e}"
+                    )
                     raise e
-        
+
         raise last_error or RuntimeError("No available Gemini API key/model")
 
     def count_tokens(self, prompt: str) -> int:
@@ -175,51 +237,60 @@ class GeminiClient:
             logger.error(f"[GEMINI] Error counting tokens: {e}")
             return 0
 
-    def create_embedding(self, text: str, model: Optional[str] = None, task_type: str = "retrieval_query") -> list:
+    def create_embedding(
+        self, text: str, model: Optional[str] = None, task_type: str = "retrieval_query"
+    ) -> list:
         last_error = None
         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}), task_type={task_type}")
-                
+
+                logger.info(
+                    f"[GEMINI][EMBEDDING] Using model={use_model} (requested={model}, default={default_model}), task_type={task_type}"
+                )
+
                 configure(api_key=key)
                 response = embed_content(
-                    model=use_model,
-                    content=text,
-                    task_type=task_type
+                    model=use_model, content=text, task_type=task_type
                 )
-                
+
                 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']
-                
+                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_retry = re.search(r'retry_delay.*?seconds: (\d+)', msg)
+                    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)
-                    
-                    logger.warning(f"[GEMINI] Rate limit hit in embedding, will retry with new key/model (attempt {attempt + 1}/{max_retries})")
+                    self.limit_manager.log_request(
+                        key, use_model, success=False, retry_delay=retry_delay
+                    )
+
+                    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") 
+
+        raise last_error or RuntimeError("No available Gemini API key/model")

app/reranker.py

@@ -2,27 +2,47 @@ from typing import List, Dict
 
 from app.utils import timing_decorator_async
 from .config import get_settings
-from .gemini_client import GeminiClient
+from .gemini_client import GeminiClient, GeminiResponseError
 from loguru import logger
+import re
 import asyncio
 import hashlib
 import time
+from tenacity import retry, stop_after_attempt, wait_exponential
+
 # from .constants import BATCH_STATUS_MESSAGES
 # from .utils import get_random_message
 
+# --- Retry decorator cho các lỗi tạm thời của Reranker (network, server-side) ---
+retry_on_rerank_transient_error = retry(
+    stop=stop_after_attempt(4),  # 1 lần gọi gốc + 3 lần thử lại
+    wait=wait_exponential(multiplier=5, min=10, max=60),  # Chờ 10s, 20s, 40s
+    retry=lambda retry_state: (
+        retry_state.outcome.failed
+        and not isinstance(retry_state.outcome.exception(), GeminiResponseError)
+    ),
+    before_sleep=lambda retry_state: logger.warning(
+        f"[RERANK][RETRY] Rerank call failed with transient error, retrying... "
+        f"Attempt: {retry_state.attempt_number}, Error: {retry_state.outcome.exception()}"
+    ),
+)
+
+
 class Reranker:
     def __init__(self):
         settings = get_settings()
-        self.provider = getattr(settings, 'rerank_provider', settings.llm_provider)
-        self.model = getattr(settings, 'rerank_model', settings.llm_model)
-        if self.provider == 'gemini':
+        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()
         # elif self.provider == 'openai':
         #     self.client = OpenAIClient(settings.openai_api_key, model=self.model)
         # elif self.provider == 'cohere':
         #     self.client = CohereClient(settings.cohere_api_key, model=self.model)
         else:
-            raise NotImplementedError(f"Rerank provider {self.provider} not supported yet.")
+            raise NotImplementedError(
+                f"Rerank provider {self.provider} not supported yet."
+            )
         # Cải thiện cache với TTL và quản lý memory
         self._rerank_cache = {}
         self._cache_ttl = 3600  # 1 giờ
@@ -35,65 +55,77 @@ class Reranker:
         """Tạo cache key từ query và docs."""
         # Tối ưu hóa cache key generation
         query_normalized = query.lower().strip()
-        doc_ids = [str(doc.get('id', '')) for doc in docs[:15]]  # Chỉ cache top 15 docs
+        doc_ids = [str(doc.get("id", "")) for doc in docs[:15]]  # Chỉ cache top 15 docs
         cache_content = query_normalized + "|".join(sorted(doc_ids))
         return hashlib.md5(cache_content.encode()).hexdigest()
 
     def _clean_cache(self):
         """Dọn dẹp cache cũ và quản lý memory."""
         current_time = time.time()
-        
+
         # Xóa cache entries đã hết hạn
         expired_keys = [
-            key for key, timestamp in self._cache_timestamps.items()
+            key
+            for key, timestamp in self._cache_timestamps.items()
             if current_time - timestamp > self._cache_ttl
         ]
-        
+
         for key in expired_keys:
             del self._rerank_cache[key]
             del self._cache_timestamps[key]
-        
+
         # Nếu cache vẫn quá lớn, xóa entries cũ nhất
         if len(self._rerank_cache) > self._max_cache_size:
             sorted_keys = sorted(
-                self._cache_timestamps.keys(),
-                key=lambda k: self._cache_timestamps[k]
+                self._cache_timestamps.keys(), key=lambda k: self._cache_timestamps[k]
             )
-            
+
             # Xóa 20% cache entries cũ nhất
-            keys_to_remove = sorted_keys[:len(sorted_keys) // 5]
+            keys_to_remove = sorted_keys[: len(sorted_keys) // 5]
             for key in keys_to_remove:
                 del self._rerank_cache[key]
                 del self._cache_timestamps[key]
-            
-            logger.info(f"[RERANK] Cleaned cache: removed {len(keys_to_remove)} old entries")
+
+            logger.info(
+                f"[RERANK] Cleaned cache: removed {len(keys_to_remove)} old entries"
+            )
 
     def _get_cached_result(self, cache_key: str, min_score: float) -> List[Dict]:
         """Lấy kết quả từ cache nếu có và còn hợp lệ."""
         if cache_key in self._rerank_cache:
             current_time = time.time()
-            if current_time - self._cache_timestamps.get(cache_key, 0) <= self._cache_ttl:
+            if (
+                current_time - self._cache_timestamps.get(cache_key, 0)
+                <= self._cache_ttl
+            ):
                 # Lọc theo điểm thay vì lấy top_k
                 cached_docs = self._rerank_cache[cache_key]
-                cached_result = [doc for doc in cached_docs if doc.get('rerank_score', 0) >= min_score]
-                logger.info(f"[RERANK] Cache hit for query, returning {len(cached_result)} cached results with score >= {min_score}")
+                cached_result = [
+                    doc
+                    for doc in cached_docs
+                    if doc.get("rerank_score", 0) >= min_score
+                ]
+                logger.info(
+                    f"[RERANK] Cache hit for query, returning {len(cached_result)} cached results with score >= {min_score}"
+                )
                 return cached_result
             else:
                 # Cache đã hết hạn, xóa
                 del self._rerank_cache[cache_key]
                 del self._cache_timestamps[cache_key]
-        
+
         return []
 
     def _set_cached_result(self, cache_key: str, scored_docs: List[Dict]):
         """Lưu kết quả vào cache."""
         self._rerank_cache[cache_key] = scored_docs
         self._cache_timestamps[cache_key] = time.time()
-        
+
         # Dọn dẹp cache nếu cần
         if len(self._rerank_cache) > self._max_cache_size:
             self._clean_cache()
 
+    @retry_on_rerank_transient_error
     async def _batch_score_docs(self, query: str, docs: List[Dict]) -> List[Dict]:
         """
         Score nhiều documents cùng lúc bằng một prompt duy nhất.
@@ -101,16 +133,16 @@ class Reranker:
         """
         if not docs:
             return []
-        
+
         # Không giới hạn content length, giữ nguyên nội dung luật
         docs_content = []
         for i, doc in enumerate(docs):
             # tieude = (doc.get('tieude') or '').strip()
             # noidung = (doc.get('noidung') or '').strip()
             # content = f"{tieude} {noidung}".strip()
-            content = (doc.get('fullcontent') or '').strip()
+            content = (doc.get("fullcontent") or "").strip()
             docs_content.append(f"{i+1}. {content}")
-        
+
         batch_prompt = (
             f"Đánh giá mức độ liên quan giữa câu hỏi và các đoạn luật sau:\n\n"
             f"Câu hỏi: {query}\n\n"
@@ -118,63 +150,70 @@ class Reranker:
             f"Trả về điểm số từ 0-10 cho từng đoạn, phân cách bằng dấu phẩy.\n"
             f"Ví dụ: 8,5,7,3,9"
         )
-        
+
         try:
-            if self.provider == 'gemini':
+            if self.provider == "gemini":
                 loop = asyncio.get_event_loop()
-                logger.info(f"[RERANK] Sending batch prompt to Gemini for {len(docs)} docs")
-                response = await loop.run_in_executor(None, self.client.generate_text, batch_prompt)
+                logger.info(
+                    f"[RERANK] Sending batch prompt to Gemini for {len(docs)} docs"
+                )
+                response = await loop.run_in_executor(
+                    None, self.client.generate_text, batch_prompt
+                )
                 logger.info(f"[RERANK] Got batch scores from Gemini: {response}")
-                
-                # Cải thiện parsing scores
+
+                # Cải thiện parsing scores bằng regex để chỉ lấy các số hợp lệ
                 scores_text = str(response).strip()
+                # Tìm tất cả các chuỗi số (integer hoặc float) trong văn bản trả về
+                score_strings = re.findall(r"\b\d+(?:\.\d+)?\b", scores_text)
+
                 scores = []
-                
-                # Xử lý nhiều format response có thể có
-                if ',' in scores_text:
-                    score_parts = scores_text.split(',')
-                elif ' ' in scores_text:
-                    score_parts = scores_text.split()
-                else:
-                    score_parts = scores_text.replace('.', ',').split(',')
-                
-                for score_str in score_parts:
+                for s in score_strings:
                     try:
-                        clean_score = ''.join(c for c in score_str.strip() if c.isdigit() or c == '.')
-                        if clean_score:
-                            score = float(clean_score)
-                            score = max(0, min(10, score))
+                        score = float(s)
+                        # Chỉ chấp nhận các điểm số trong khoảng 0-10 để tăng độ chính xác
+                        if 0 <= score <= 10:
                             scores.append(score)
-                        else:
-                            scores.append(0)
                     except (ValueError, TypeError):
-                        scores.append(0)
-                
-                while len(scores) < len(docs):
-                    scores.append(0)
-                
+                        # Bỏ qua các giá trị không phải là số hợp lệ
+                        continue
+
+                # Đảm bảo số lượng điểm khớp với số lượng văn bản
+                # Nếu thiếu, thêm điểm 0. Nếu thừa, cắt bớt.
+                if len(scores) < len(docs):
+                    scores.extend([0.0] * (len(docs) - len(scores)))
+                else:
+                    scores = scores[: len(docs)]
+
                 for i, doc in enumerate(docs):
-                    doc['rerank_score'] = scores[i]
-                
-                logger.info(f"[RERANK] Successfully scored {len(docs)} docs with scores: {scores}")
+                    doc["rerank_score"] = scores[i]
+
+                logger.info(
+                    f"[RERANK] Successfully scored {len(docs)} docs with scores: {scores}"
+                )
                 return docs
-                
+
             else:
-                raise NotImplementedError(f"Rerank provider {self.provider} not supported yet in batch method.")
-            
-        except Exception as e:
-            logger.error(f"[RERANK] Lỗi khi batch score: {e}")
+                raise NotImplementedError(
+                    f"Rerank provider {self.provider} not supported yet in batch method."
+                )
+
+        except GeminiResponseError as e:
+            # Lỗi nội dung không thể retry (safety, max_tokens), gán điểm 0 và trả về.
+            # Các lỗi khác (network, 500) sẽ được decorator retry.
+            logger.error(f"[RERANK] Lỗi nội dung không thể retry khi batch score: {e}")
             for doc in docs:
-                doc['rerank_score'] = 0
+                doc["rerank_score"] = 0
             return docs
 
+    @retry_on_rerank_transient_error
     async def _score_doc(self, query: str, doc: Dict) -> Dict:
         """
         Score một document với query.
         Không cắt bớt nội dung luật.
         """
-        tieude = (doc.get('tieude') or '').strip()
-        noidung = (doc.get('noidung') or '').strip()
+        tieude = (doc.get("tieude") or "").strip()
+        noidung = (doc.get("noidung") or "").strip()
         content = f"{tieude} {noidung}".strip()
         prompt = (
             f"Đánh giá mức độ liên quan:\n"
@@ -183,14 +222,20 @@ class Reranker:
             f"Điểm (0-10):"
         )
         try:
-            if self.provider == 'gemini':
+            if self.provider == "gemini":
                 loop = asyncio.get_event_loop()
                 logger.info(f"[RERANK] Sending individual prompt to Gemini")
-                score_response = await loop.run_in_executor(None, self.client.generate_text, prompt)
-                logger.info(f"[RERANK] Got individual score from Gemini: {score_response}")
+                score_response = await loop.run_in_executor(
+                    None, self.client.generate_text, prompt
+                )
+                logger.info(
+                    f"[RERANK] Got individual score from Gemini: {score_response}"
+                )
                 score_text = str(score_response).strip()
                 try:
-                    clean_score = ''.join(c for c in score_text if c.isdigit() or c == '.')
+                    clean_score = "".join(
+                        c for c in score_text if c.isdigit() or c == "."
+                    )
                     if clean_score:
                         score = float(clean_score)
                         score = max(0, min(10, score))
@@ -198,44 +243,59 @@ class Reranker:
                         score = 0
                 except (ValueError, TypeError):
                     score = 0
-                doc['rerank_score'] = score
+                doc["rerank_score"] = score
                 return doc
             else:
-                raise NotImplementedError(f"Rerank provider {self.provider} not supported yet in rerank method.")
-        except Exception as e:
-            logger.error(f"[RERANK] Lỗi khi tính score: {e} | doc: {doc}")
-            doc['rerank_score'] = 0
+                raise NotImplementedError(
+                    f"Rerank provider {self.provider} not supported yet in rerank method."
+                )
+        except GeminiResponseError as e:
+            # Lỗi nội dung không thể retry (safety, max_tokens), gán điểm 0 và trả về.
+            logger.error(
+                f"[RERANK] Lỗi nội dung không thể retry khi tính score: {e} | doc: {doc}"
+            )
+            doc["rerank_score"] = 0
             return doc
 
     @timing_decorator_async
-    async def rerank(self, query: str, docs: List[Dict], min_score: float = 7.0) -> List[Dict]:
+    async def rerank(
+        self, query: str, docs: List[Dict], min_score: float = 7.0
+    ) -> List[Dict]:
         """
         Rerank docs theo độ liên quan với query, trả về các docs có điểm >= min_score.
         Sử dụng batch processing và caching để tối ưu hiệu suất.
         """
-        logger.info(f"[RERANK] Start rerank for query: {query} | docs: {len(docs)} | min_score: {min_score}")
-        
+        logger.info(
+            f"[RERANK] Start rerank for query: {query} | docs: {len(docs)} | min_score: {min_score}"
+        )
+
         if not docs:
             return []
-        
+
         # Kiểm tra cache trước
         cache_key = self._get_cache_key(query, docs)
         cached_result = self._get_cached_result(cache_key, min_score)
-        
+
         if cached_result:
             return cached_result
-        
+
         # Giới hạn số lượng docs để rerank - chỉ rerank top N docs có similarity cao nhất
         max_docs_to_rerank = self.max_docs_to_rerank
         docs_to_rerank = docs[:max_docs_to_rerank]
-        logger.info(f"[RERANK] Will rerank {len(docs_to_rerank)} docs (limited to top {max_docs_to_rerank})")
-        
+        logger.info(
+            f"[RERANK] Will rerank {len(docs_to_rerank)} docs (limited to top {max_docs_to_rerank})"
+        )
+
         # Sử dụng batch processing thay vì individual scoring
         try:
             scored = await self._batch_score_docs(query, docs_to_rerank)
-            logger.info(f"[RERANK] Batch processing completed, scored {len(scored)} docs")
+            logger.info(
+                f"[RERANK] Batch processing completed, scored {len(scored)} docs"
+            )
         except Exception as e:
-            logger.error(f"[RERANK] Batch processing failed, falling back to individual scoring: {e}")
+            logger.error(
+                f"[RERANK] Batch processing failed, falling back to individual scoring: {e}"
+            )
             # Fallback về individual scoring nếu batch processing thất bại
             scored = []
             for doc in docs_to_rerank:
@@ -244,17 +304,19 @@ class Reranker:
                     scored.append(scored_doc)
                 except Exception as e:
                     logger.error(f"[RERANK] Error scoring individual doc: {e}")
-                    doc['rerank_score'] = 0
+                    doc["rerank_score"] = 0
                     scored.append(doc)
-        
+
         # Sort theo score
-        scored = sorted(scored, key=lambda x: x.get('rerank_score', 0), reverse=True)
-        
+        scored = sorted(scored, key=lambda x: x.get("rerank_score", 0), reverse=True)
+
         # Lọc theo min_score
-        result = [doc for doc in scored if doc.get('rerank_score', 0) >= min_score]
-        
+        result = [doc for doc in scored if doc.get("rerank_score", 0) >= min_score]
+
         # Cache kết quả đã được chấm điểm (toàn bộ, trước khi lọc)
         self._set_cached_result(cache_key, scored)
-        
-        logger.info(f"[RERANK] Found {len(result)} docs with score >= {min_score}. Top results: {result[:2]}...{result[-2:] if len(result) > 2 else ''}")
-        return result
+
+        logger.info(
+            f"[RERANK] Found {len(result)} docs with score >= {min_score}. Top results: {result[:2]}...{result[-2:] if len(result) > 2 else ''}"
+        )
+        return result