optimize reranker

app/reranker.py

@@ -5,6 +5,7 @@ from loguru import logger
 import asyncio
 import random
 import hashlib
+import time
 from .constants import BATCH_STATUS_MESSAGES
 
 class Reranker:
@@ -21,86 +22,136 @@ class Reranker:
         else:
             raise NotImplementedError(f"Rerank provider {self.provider} not supported yet.")
         self.facebook_client = facebook_client
-        # Cache cho kết quả reranking
+        
+        # Cải thiện cache với TTL và quản lý memory
         self._rerank_cache = {}
+        self._cache_ttl = 3600  # 1 giờ
+        self._max_cache_size = 200  # Tăng cache size
+        self._cache_timestamps = {}
 
     def _get_cache_key(self, query: str, docs: List[Dict]) -> str:
         """Tạo cache key từ query và docs."""
-        # Tạo hash từ query và doc IDs
+        # 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
-        cache_content = query + "|".join(doc_ids)
+        cache_content = query_normalized + "|".join(sorted(doc_ids))
         return hashlib.md5(cache_content.encode()).hexdigest()
 
-    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 '')
-        # Tối ưu prompt ngắn gọn hơn
-        prompt = (
-            f"Luật: {content[:500]}\n"  # Giới hạn content length
-            f"Hỏi: {query}\n"
-            "Đánh giá mức độ liên quan (0-10). Chỉ trả về số."
-        )
+    def _clean_cache(self):
+        """Dọn dẹp cache cũ và quản lý memory."""
+        current_time = time.time()
         
-        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.")
+        # Xóa cache entries đã hết hạn
+        expired_keys = [
+            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]
+            )
             
-            score = float(str(score).strip().split()[0])
-            doc['rerank_score'] = score
-            return doc
+            # Xóa 20% cache entries cũ nhất
+            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]
             
-        except Exception as e:
-            logger.error(f"[RERANK] Lỗi khi tính score: {e} | doc: {doc}")
-            doc['rerank_score'] = 0
-            return doc
+            logger.info(f"[RERANK] Cleaned cache: removed {len(keys_to_remove)} old entries")
+
+    def _get_cached_result(self, cache_key: str, top_k: int) -> 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:
+                cached_result = self._rerank_cache[cache_key][:top_k]
+                logger.info(f"[RERANK] Cache hit for query, returning {len(cached_result)} cached results")
+                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()
 
     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.
+        Không cắt bớt nội dung luật.
         """
         if not docs:
             return []
         
-        # Tạo prompt batch cho tất cả documents
+        # Không giới hạn content length, giữ nguyên nội dung luật
         docs_content = []
         for i, doc in enumerate(docs):
-            content = (doc.get('tieude', '') or '') + ' ' + (doc.get('noidung', '') or '')
-            docs_content.append(f"{i+1}. {content[:300]}")  # Giới hạn length
+            tieude = doc.get('tieude', '').strip()
+            noidung = doc.get('noidung', '').strip()
+            content = f"{tieude} {noidung}".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"
             f"Các đoạn luật:\n" + "\n".join(docs_content) + "\n\n"
-            f"Đánh giá mức độ liên quan của từng đoạn (0-10). Trả về dạng: 1.8,2.5,3.0,..."
+            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':
                 loop = asyncio.get_event_loop()
-                logger.info(f"[RERANK] Sending batch prompt to Gemini")
+                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}")
                 
-                # Parse scores từ response
+                # Cải thiện parsing scores
                 scores_text = str(response).strip()
                 scores = []
-                for score_str in scores_text.split(','):
+                
+                # 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:
                     try:
-                        score = float(score_str.strip().split('.')[0])
-                        scores.append(score)
-                    except:
+                        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))
+                            scores.append(score)
+                        else:
+                            scores.append(0)
+                    except (ValueError, TypeError):
                         scores.append(0)
                 
-                # Gán scores cho documents
+                while len(scores) < len(docs):
+                    scores.append(0)
+                
                 for i, doc in enumerate(docs):
-                    doc['rerank_score'] = scores[i] if i < len(scores) else 0
+                    doc['rerank_score'] = scores[i]
                 
+                logger.info(f"[RERANK] Successfully scored {len(docs)} docs with scores: {scores}")
                 return docs
                 
             else:
@@ -108,15 +159,53 @@ class Reranker:
             
         except Exception as e:
             logger.error(f"[RERANK] Lỗi khi batch score: {e}")
-            # Fallback về individual scoring
             for doc in docs:
                 doc['rerank_score'] = 0
             return docs
 
+    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', '').strip()
+        noidung = doc.get('noidung', '').strip()
+        content = f"{tieude} {noidung}".strip()
+        prompt = (
+            f"Đánh giá mức độ liên quan:\n"
+            f"Luật: {content}\n"
+            f"Hỏi: {query}\n"
+            f"Điểm (0-10):"
+        )
+        try:
+            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_text = str(score_response).strip()
+                try:
+                    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))
+                    else:
+                        score = 0
+                except (ValueError, TypeError):
+                    score = 0
+                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
+            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 batch processing để tối ưu hiệu suất.
+        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)} | top_k: {top_k}")
         
@@ -125,9 +214,9 @@ class Reranker:
         
         # Kiểm tra cache trước
         cache_key = self._get_cache_key(query, docs)
-        if cache_key in self._rerank_cache:
-            logger.info(f"[RERANK] Cache hit for query, returning cached result")
-            cached_result = self._rerank_cache[cache_key][:top_k]
+        cached_result = self._get_cached_result(cache_key, top_k)
+        
+        if cached_result:
             return cached_result
         
         # Giới hạn số lượng docs để rerank - chỉ rerank top 15 docs có similarity cao nhất
@@ -152,7 +241,7 @@ class Reranker:
                     doc['rerank_score'] = 0
                     scored.append(doc)
         
-        # Gửi Facebook message chỉ một lần sau khi hoàn thành
+        # Gửi Facebook message sau khi hoàn thành
         if self.facebook_client:
             try:
                 message = random.choice(BATCH_STATUS_MESSAGES)
@@ -164,13 +253,8 @@ class Reranker:
         scored = sorted(scored, key=lambda x: x['rerank_score'], reverse=True)
         result = scored[:top_k]
         
-        # Cache kết quả
-        self._rerank_cache[cache_key] = scored
-        # Giới hạn cache size để tránh memory leak
-        if len(self._rerank_cache) > 100:
-            # Xóa cache cũ nhất
-            oldest_key = next(iter(self._rerank_cache))
-            del self._rerank_cache[oldest_key]
+        # Cache kết quả với system mới
+        self._set_cached_result(cache_key, scored)
         
         logger.info(f"[RERANK] Top reranked docs: {result}")
         return result