etm-terrain/services/pdf_indexer.py

117 lines
3.6 KiB
Python

import logging
import os
import re
log = logging.getLogger(__name__)
try:
import fitz # pymupdf
_fitz_ok = True
except ImportError:
_fitz_ok = False
log.warning("pymupdf non installé — pip install pymupdf")
PRODUITS_PDF = {
"multiplus": "Victron MultiPlus",
"quattro": "Victron Quattro",
"victron": "Victron",
"ve.bus": "Victron VE.Bus",
"ess": "Victron ESS",
"fronius": "Onduleur Fronius",
"symo": "Onduleur Fronius Symo",
"sma": "Onduleur SMA",
"huawei": "Onduleur Huawei",
"byd": "Batterie BYD",
"pylontech": "Batterie Pylontech",
"daikin": "PAC Daikin",
"viessmann": "PAC Viessmann",
"keba": "Borne IRVE Keba",
}
def detecter_produit_pdf(nom_fichier: str) -> str:
nom = nom_fichier.lower()
for mot, produit in PRODUITS_PDF.items():
if mot in nom:
return produit
return "Documentation technique"
def extraire_sections(pdf_path: str) -> list[dict]:
"""Découpe un PDF en sections logiques par titres/headings."""
if not _fitz_ok:
return []
doc = fitz.open(pdf_path)
sections = []
section_courante = {"titre": "", "contenu": "", "page": 1}
for page_num, page in enumerate(doc, 1):
blocs = page.get_text("dict")["blocks"]
for bloc in blocs:
if "lines" not in bloc:
continue
for line in bloc["lines"]:
texte = " ".join(span["text"] for span in line["spans"]).strip()
taille = max((span["size"] for span in line["spans"]), default=0)
gras = any(span["flags"] & 2**4 for span in line["spans"])
if not texte or re.match(r"^\d+$", texte): # ignorer numéros de page
continue
est_titre = (taille > 11 or gras) and len(texte) < 120
if est_titre:
if len(section_courante["contenu"].strip()) > 100:
sections.append(section_courante.copy())
section_courante = {
"titre": texte,
"contenu": texte + "\n",
"page": page_num,
}
else:
section_courante["contenu"] += texte + " "
if len(section_courante["contenu"].strip()) > 100:
sections.append(section_courante)
doc.close()
return sections
def indexer_pdf(faq_service, pdf_path: str) -> int:
"""Indexe un PDF dans ChromaDB par sections logiques."""
nom_fichier = os.path.basename(pdf_path)
produit = detecter_produit_pdf(nom_fichier)
sections = extraire_sections(pdf_path)
if not sections:
log.warning("%s : aucune section extraite", nom_fichier)
return 0
# Supprimer l'ancienne indexation de ce fichier
try:
faq_service.collection.delete(where={"source": nom_fichier})
except Exception:
pass
indexees = 0
for i, section in enumerate(sections):
contenu = section["contenu"].strip()
if len(contenu) < 80:
continue
faq_service.collection.upsert(
documents=[contenu],
ids=[f"{nom_fichier}_s{i}"],
metadatas=[{
"source": nom_fichier,
"produit": produit,
"titre": section["titre"][:200],
"page": section["page"],
"type": "documentation",
}],
)
indexees += 1
log.info("%s%d/%d sections indexées (%s)", nom_fichier, indexees, len(sections), produit)
return indexees