Spaces:
Sleeping
Sleeping
File size: 5,573 Bytes
4211d95 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """Pipeline RAG complet : document → chunks → embeddings → retrieval → génération."""
from __future__ import annotations
import logging
from pathlib import Path
from src.config import SYSTEM_PROMPT, TOP_K
from src.document_processor import DocumentProcessor
from src.llm_provider import LLMProvider, get_llm_provider
from src.vector_store import VectorStore
logger = logging.getLogger(__name__)
class RAGEngine:
"""Orchestre le pipeline RAG de bout en bout."""
def __init__(self):
self.processor = DocumentProcessor()
self.store = VectorStore()
self.llm: LLMProvider = get_llm_provider()
# Tenter de charger un index existant
self.store.load("default")
# ── Ingestion ────────────────────────────────────────────────────────
def ingest_pdf(self, pdf_path: str | Path) -> str:
"""Ingère un PDF : parse → chunk → embed → index."""
parsed = self.processor.process_pdf(pdf_path)
added = self.store.add_document(parsed)
self.store.save("default")
if added == 0:
return f"ℹ️ «{parsed.filename}» était déjà indexé."
return parsed.summary
def ingest_directory(self, dir_path: str | Path) -> list[str]:
"""Ingère tous les PDFs d'un répertoire."""
dir_path = Path(dir_path)
pdfs = sorted(dir_path.glob("*.pdf"))
if not pdfs:
return ["⚠️ Aucun PDF trouvé dans ce répertoire."]
results = []
for pdf in pdfs:
try:
results.append(self.ingest_pdf(pdf))
except Exception as e:
results.append(f"❌ Erreur sur {pdf.name} : {e}")
return results
# ── Question-Réponse ─────────────────────────────────────────────────
def ask(self, question: str, top_k: int = TOP_K) -> dict:
"""Pose une question et retourne la réponse avec les sources."""
if not question.strip():
return {"answer": "⚠️ Veuillez poser une question.", "sources": []}
if self.store.index.ntotal == 0:
return {
"answer": (
"⚠️ Aucun document n'est indexé. "
"Veuillez d'abord uploader un PDF."
),
"sources": [],
}
# 1. Retrieval
results = self.store.search(question, top_k=top_k)
# 2. Construire le contexte
context = self._build_context(results)
# 3. Génération via le provider LLM (local ou cloud)
user_message = (
f"Contexte documentaire :\n\n{context}\n\n"
f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n\n"
f"Question de l'utilisateur : {question}"
)
answer = self.llm.generate(SYSTEM_PROMPT, user_message)
# 4. Formater les sources
sources = self._format_sources(results)
return {"answer": answer, "sources": sources}
# ── Construction du contexte ─────────────────────────────────────────
@staticmethod
def _build_context(results: list[dict]) -> str:
"""Formate les chunks retrouvés en contexte pour le LLM."""
if not results:
return "Aucun passage pertinent trouvé."
sections = []
for i, r in enumerate(results, 1):
meta = r["metadata"]
source = meta.get("source", "inconnu")
page = meta.get("page", "?")
score = r["score"]
sections.append(
f"━━━ Extrait {i} [Document: «{source}», Page {page}] "
f"(pertinence: {score:.2f}) ━━━\n{r['text']}"
)
return "\n\n".join(sections)
# ── Formatage des sources ────────────────────────────────────────────
@staticmethod
def _format_sources(results: list[dict]) -> list[dict]:
"""Formate les sources pour l'affichage dans l'UI."""
sources = []
for r in results:
meta = r["metadata"]
sources.append(
{
"document": meta.get("source", "inconnu"),
"page": meta.get("page", "?"),
"score": round(r["score"], 3),
"extrait": r["text"][:200] + "..."
if len(r["text"]) > 200
else r["text"],
}
)
return sources
# ── Info ─────────────────────────────────────────────────────────────
@property
def status(self) -> dict:
"""État actuel du pipeline."""
store_stats = self.store.stats
return {
"documents_indexés": len(store_stats["indexed_documents"]),
"liste_documents": store_stats["indexed_documents"],
"chunks_total": store_stats["total_chunks"],
"vecteurs_total": store_stats["total_vectors"],
"llm_disponible": self.llm.is_available(),
"modèle_llm": self.llm.name,
}
|