add reranker

app/config.py

@@ -42,6 +42,9 @@ class Settings(BaseSettings):
     embedding_provider: str = os.getenv("EMBEDDING_PROVIDER", "gemini") or ""
     embedding_model: str = os.getenv("EMBEDDING_MODEL", "models/embedding-001") or ""
 
+    rerank_provider: str = os.getenv("RERANK_PROVIDER", "") or llm_provider
+    rerank_model: str = os.getenv("RERANK_MODEL", "") or llm_model
+
     class Config:
         env_file = ".env"
 

app/main.py

@@ -18,6 +18,7 @@ from .utils import setup_logging, extract_command, extract_keywords, timing_deco
 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
 
 app = FastAPI(title="WeBot Facebook Messenger API")
 
@@ -67,6 +68,8 @@ llm_client = create_llm_client(
     model=settings.gemini_model
 )
 
+reranker = Reranker()
+
 logger.info("[STARTUP] Mount health router...")
 app.include_router(health_router)
 
@@ -395,6 +398,13 @@ async def process_business_logic(log_kwargs: Dict[str, Any], page_token: str) ->
 async def format_search_results(question: str, matches: List[Dict[str, Any]]) -> str:
     if not matches:
         return "Không tìm thấy kết quả phù hợp."
+    # Rerank matches trước khi format cho LLM
+    try:
+        reranked = await reranker.rerank(question, matches, top_k=5)
+        if reranked:
+            matches = reranked
+    except Exception as e:
+        logger.error(f"[RERANK] Lỗi khi rerank: {e}")
     # Tìm item có similarity cao nhất
     top = None
     top_result_text = ""

app/reranker.py

@@ -0,0 +1,47 @@
+from typing import List, Dict
+from .config import get_settings
+from .gemini_client import GeminiClient
+from loguru import logger
+import asyncio
+
+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.client = GeminiClient(settings.gemini_api_key, model=self.model)
+        # 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.")
+
+    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.
+        """
+        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()
+                    score = await loop.run_in_executor(None, self.client.generate_text, prompt)
+                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)
+        scored = sorted(scored, key=lambda x: x['rerank_score'], reverse=True)
+        return scored[:top_k] 

app/supabase_db.py

@@ -31,7 +31,7 @@ class SupabaseClient:
             return None
 
     @timing_decorator_sync
-    def match_documents(self, embedding: List[float], match_count: int = 10, vehicle_keywords: Optional[List[str]] = None):
+    def match_documents(self, embedding: List[float], match_count: int = 20, vehicle_keywords: Optional[List[str]] = None):
         """
         Truy vấn vector similarity search qua RPC match_documents.
         Input: embedding (list[float]), match_count (int), vehicle_keywords (list[str] hoặc None)