Corback
309 lines
12 KiB
Python
309 lines
12 KiB
Python
"""
|
|
Ingestion Fallout PDFs → ChromaDB v2
|
|
Améliorations v2 :
|
|
- Chunking structurel (paragraphes/sections) au lieu de token-count fixe
|
|
- Métadonnées enrichies : systeme (2d20 | tardy | none), category, priority
|
|
- Évite la collision mécanique 2D20 vs Tardy en taggant chaque chunk
|
|
- Chunks 400-600 mots, overlap 80 mots, respecte les frontières de paragraphes
|
|
"""
|
|
import json, os, re, subprocess, tempfile, time, urllib.request, urllib.parse
|
|
from pathlib import Path
|
|
|
|
PDFFETCH_URL = "http://localhost:15001/files"
|
|
OLLAMA_EMBED = "http://localhost:11434/api/embeddings"
|
|
CHROMA_URL = "http://localhost:8800"
|
|
CHROMA_BASE = "/api/v2/tenants/default_tenant/databases/default_database"
|
|
COLLECTION = "fallout_lore"
|
|
EMBED_MODEL = "nomic-embed-text"
|
|
STATE_FILE = Path("/home/ubuntu/fallout/ingest_state.json")
|
|
FALLOUT_PREFIX = "fallout/"
|
|
|
|
# Chunking params (en mots, pas tokens)
|
|
CHUNK_WORDS = 450
|
|
OVERLAP_WORDS = 80
|
|
MIN_WORDS = 40
|
|
|
|
|
|
# ── État ──────────────────────────────────────────────────────────────────────
|
|
|
|
def load_state():
|
|
if STATE_FILE.exists():
|
|
return json.loads(STATE_FILE.read_text())
|
|
return {"ingested": [], "failed": [], "version": 2}
|
|
|
|
def save_state(s):
|
|
STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
STATE_FILE.write_text(json.dumps(s, indent=2))
|
|
|
|
|
|
# ── ChromaDB ──────────────────────────────────────────────────────────────────
|
|
|
|
def chroma_request(method, path, data=None):
|
|
url = CHROMA_URL + path
|
|
body = json.dumps(data).encode() if data else None
|
|
req = urllib.request.Request(url, data=body, method=method,
|
|
headers={"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(req, timeout=30) as r:
|
|
return json.loads(r.read())
|
|
|
|
def get_or_create_collection():
|
|
try:
|
|
r = chroma_request("POST", f"{CHROMA_BASE}/collections", {
|
|
"name": COLLECTION,
|
|
"metadata": {"hnsw:space": "cosine"}
|
|
})
|
|
print(f"Collection créée : {COLLECTION} ({r['id']})")
|
|
return r["id"]
|
|
except urllib.error.HTTPError as e:
|
|
if e.code == 409:
|
|
cols = chroma_request("GET", f"{CHROMA_BASE}/collections")
|
|
col_id = next(c["id"] for c in cols if c["name"] == COLLECTION)
|
|
print(f"Collection existante : {COLLECTION} ({col_id})")
|
|
return col_id
|
|
raise
|
|
|
|
def add_documents(col_id, ids, embeddings, documents, metadatas):
|
|
chroma_request("POST", f"{CHROMA_BASE}/collections/{col_id}/add", {
|
|
"ids": ids, "embeddings": embeddings,
|
|
"documents": documents, "metadatas": metadatas
|
|
})
|
|
|
|
|
|
# ── Embeddings ────────────────────────────────────────────────────────────────
|
|
|
|
def embed(text):
|
|
payload = json.dumps({"model": EMBED_MODEL, "prompt": text}).encode()
|
|
req = urllib.request.Request(OLLAMA_EMBED, data=payload,
|
|
headers={"Content-Type": "application/json"}, method="POST")
|
|
with urllib.request.urlopen(req, timeout=60) as r:
|
|
return json.loads(r.read())["embedding"]
|
|
|
|
|
|
# ── PDF helpers ───────────────────────────────────────────────────────────────
|
|
|
|
def list_fallout_pdfs():
|
|
req = urllib.request.Request(PDFFETCH_URL,
|
|
headers={"User-Agent": "fallout-ingest/2.0"})
|
|
with urllib.request.urlopen(req, timeout=15) as r:
|
|
all_files = json.loads(r.read())
|
|
return [f for f in all_files if f.startswith(FALLOUT_PREFIX)]
|
|
|
|
def download_pdf(filename):
|
|
url = f"{PDFFETCH_URL}/{urllib.parse.quote(filename)}"
|
|
req = urllib.request.Request(url, headers={"User-Agent": "fallout-ingest/2.0"})
|
|
with urllib.request.urlopen(req, timeout=300) as r:
|
|
return r.read()
|
|
|
|
def extract_text(pdf_bytes):
|
|
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
|
|
f.write(pdf_bytes); tmp = f.name
|
|
text = ""
|
|
try:
|
|
r = subprocess.run(["pdftotext", "-layout", tmp, "-"],
|
|
capture_output=True, timeout=180)
|
|
if r.returncode == 0:
|
|
text = r.stdout.decode("utf-8", errors="replace")
|
|
finally:
|
|
os.unlink(tmp)
|
|
if len(text.strip()) < 200:
|
|
try:
|
|
import io
|
|
from pypdf import PdfReader
|
|
reader = PdfReader(io.BytesIO(pdf_bytes))
|
|
text = "\n".join(p.extract_text() or "" for p in reader.pages)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Extraction impossible: {e}")
|
|
return text
|
|
|
|
|
|
# ── Chunking structurel ───────────────────────────────────────────────────────
|
|
|
|
def is_section_header(line):
|
|
"""Détecte les titres de section : ligne courte, capitalisée, pas de ponctuation finale."""
|
|
l = line.strip()
|
|
if not l or len(l) > 80:
|
|
return False
|
|
if l.isupper() and len(l.split()) <= 8:
|
|
return True
|
|
if re.match(r'^(CHAPITRE|CHAPTER|PARTIE|PART|SECTION|\d+[\.\)])\s+\S', l, re.I):
|
|
return True
|
|
return False
|
|
|
|
def chunk_by_structure(text, filename):
|
|
"""
|
|
Chunking basé sur les paragraphes et sections, pas sur le token count fixe.
|
|
Respecte les frontières naturelles du texte.
|
|
"""
|
|
# Nettoyage
|
|
text = re.sub(r'[ \t]+', ' ', text)
|
|
text = re.sub(r'\n{3,}', '\n\n', text)
|
|
|
|
# Découpe en blocs (paragraphes ou sections)
|
|
raw_blocks = re.split(r'\n{2,}', text)
|
|
|
|
# Regroupe les headers avec le paragraphe suivant
|
|
blocks = []
|
|
i = 0
|
|
while i < len(raw_blocks):
|
|
b = raw_blocks[i].strip()
|
|
if not b:
|
|
i += 1
|
|
continue
|
|
if is_section_header(b) and i + 1 < len(raw_blocks):
|
|
merged = b + "\n\n" + raw_blocks[i+1].strip()
|
|
blocks.append(merged)
|
|
i += 2
|
|
else:
|
|
blocks.append(b)
|
|
i += 1
|
|
|
|
# Accumulation en chunks de ~CHUNK_WORDS mots avec overlap
|
|
chunks = []
|
|
current_blocks = []
|
|
current_words = 0
|
|
|
|
for block in blocks:
|
|
bwords = len(block.split())
|
|
if bwords < 5:
|
|
continue
|
|
|
|
if current_words + bwords > CHUNK_WORDS and current_blocks:
|
|
chunk = "\n\n".join(current_blocks).strip()
|
|
if len(chunk.split()) >= MIN_WORDS:
|
|
chunks.append(chunk)
|
|
|
|
# Overlap : garder les derniers blocs jusqu'à OVERLAP_WORDS mots
|
|
overlap_blocks = []
|
|
overlap_words = 0
|
|
for ob in reversed(current_blocks):
|
|
ow = len(ob.split())
|
|
if overlap_words + ow <= OVERLAP_WORDS:
|
|
overlap_blocks.insert(0, ob)
|
|
overlap_words += ow
|
|
else:
|
|
break
|
|
current_blocks = overlap_blocks + [block]
|
|
current_words = overlap_words + bwords
|
|
else:
|
|
current_blocks.append(block)
|
|
current_words += bwords
|
|
|
|
# Flush final
|
|
if current_blocks:
|
|
chunk = "\n\n".join(current_blocks).strip()
|
|
if len(chunk.split()) >= MIN_WORDS:
|
|
chunks.append(chunk)
|
|
|
|
return chunks
|
|
|
|
|
|
# ── Métadonnées ───────────────────────────────────────────────────────────────
|
|
|
|
def infer_metadata(filename):
|
|
f = filename.lower()
|
|
name = Path(filename).name
|
|
meta = {"source": name, "path": filename}
|
|
|
|
# Catégorie
|
|
if "lore_post" in f: meta["category"] = "lore_canon"
|
|
elif "lore_pre" in f: meta["category"] = "lore_inspiration"
|
|
elif "regles/core" in f: meta["category"] = "regles_core"
|
|
elif "regles/suppl" in f: meta["category"] = "regles_supplement"
|
|
elif "aventures" in f: meta["category"] = "aventure"
|
|
elif "ambiance" in f: meta["category"] = "ambiance"
|
|
else: meta["category"] = "divers"
|
|
|
|
# Système de règles — CRITIQUE pour éviter la collision 2D20 vs Tardy
|
|
if "tardy" in f:
|
|
meta["systeme"] = "tardy"
|
|
elif any(k in f for k in ["manuel_officiel", "quickstart", "rules_booklet",
|
|
"gm_toolkit", "bestiaire", "armes", "special",
|
|
"competences", "pizza", "frequence"]):
|
|
meta["systeme"] = "2d20"
|
|
elif "regles" in f:
|
|
meta["systeme"] = "2d20"
|
|
else:
|
|
meta["systeme"] = "none"
|
|
|
|
# Priorité lore : la bible canon prime
|
|
if "bible_v2" in f: meta["priority"] = "haute"
|
|
elif "bible_v1" in f: meta["priority"] = "normale"
|
|
elif "lore_canon" == meta.get("category"): meta["priority"] = "normale"
|
|
else: meta["priority"] = "basse"
|
|
|
|
return meta
|
|
|
|
|
|
# ── Main ──────────────────────────────────────────────────────────────────────
|
|
|
|
def main():
|
|
print("=== Ingestion Fallout v2 → ChromaDB ===")
|
|
print(f"Chunking structurel, ~{CHUNK_WORDS} mots/chunk, overlap {OVERLAP_WORDS} mots")
|
|
print(f"Modèle embedding : {EMBED_MODEL}\n")
|
|
|
|
state = load_state()
|
|
col_id = get_or_create_collection()
|
|
|
|
pdfs = list_fallout_pdfs()
|
|
print(f"{len(pdfs)} PDFs Fallout sur Guardian\n")
|
|
|
|
for filename in pdfs:
|
|
if filename in state["ingested"]:
|
|
print(f"[SKIP] {filename}")
|
|
continue
|
|
|
|
print(f"\n[{filename}]")
|
|
try:
|
|
pdf_bytes = download_pdf(filename)
|
|
print(f" {len(pdf_bytes)/1024:.0f} Ko téléchargés")
|
|
|
|
text = extract_text(pdf_bytes)
|
|
print(f" {len(text)} chars extraits")
|
|
|
|
chunks = chunk_by_structure(text, filename)
|
|
meta_base = infer_metadata(filename)
|
|
print(f" {len(chunks)} chunks | catégorie: {meta_base['category']} | système: {meta_base['systeme']}")
|
|
|
|
ids, embeddings, documents, metadatas = [], [], [], []
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
chunk_id = f"{meta_base['source']}_{i:04d}"
|
|
t0 = time.time()
|
|
emb = embed(chunk)
|
|
word_count = len(chunk.split())
|
|
ids.append(chunk_id)
|
|
embeddings.append(emb)
|
|
documents.append(chunk)
|
|
metadatas.append({
|
|
**meta_base,
|
|
"chunk_index": i,
|
|
"total_chunks": len(chunks),
|
|
"word_count": word_count,
|
|
})
|
|
print(f" chunk {i+1}/{len(chunks)} ({word_count}w) OK ({time.time()-t0:.1f}s)")
|
|
|
|
if len(ids) >= 50:
|
|
add_documents(col_id, ids, embeddings, documents, metadatas)
|
|
ids, embeddings, documents, metadatas = [], [], [], []
|
|
|
|
if ids:
|
|
add_documents(col_id, ids, embeddings, documents, metadatas)
|
|
|
|
print(f"\n Indexé — {len(chunks)} chunks")
|
|
state["ingested"].append(filename)
|
|
# Retire des failed si présent
|
|
state["failed"] = [f for f in state["failed"] if f.get("file") != filename]
|
|
save_state(state)
|
|
|
|
except Exception as e:
|
|
print(f" ERREUR: {e}")
|
|
state["failed"].append({"file": filename, "error": str(e)})
|
|
save_state(state)
|
|
|
|
print(f"\n=== Terminé : {len(state['ingested'])} PDFs indexés, {len(state['failed'])} échecs ===")
|
|
for f in state.get("failed", []):
|
|
print(f" ECHEC: {f['file']} — {f['error']}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|