Update rag.py

rag.py

@@ -17,7 +17,7 @@ from FlagEmbedding import BGEM3FlagModel
 from rich import print
 
 ROOT = Path("Fiches")                # dossier racine des fiches
-K    = 4                               # nombre de passages remis au LLM
+K    = 10                               # nombre de passages remis au LLM
 
 # ------------------ utilitaires de découpe identiques à l'index -------------
 CHUNK, OVERLAP = 800, 100              # garder cohérent avec index.py
@@ -80,8 +80,10 @@ try:
 
         D, I = idx.search(q_emb.astype("float32").reshape(1, -1), K)
         hits = I[0]
+        # réordonne pour mettre en tête les passages contenant “Seuil”
+        hits = sorted(hits, key=lambda i: "Seuil" not in docs[int(i)])
 
-        context = "\n\n".join(docs[int(i)] for i in hits)
+        context = "\n\n".join(docs[int(i)] for i in hits[:K])
         prompt = (
             "<system>Réponds en français, de façon précise, et uniquement à partir du contexte fourni. Si l'information n'est pas dans le contexte, réponds : 'Je ne sais pas'.</system>\n"
             f"<context>{context}</context>\n<user>{q}</user>"