Add embed_fallout v3 — structural chunking + multi-collection routing

pipelines/embed_fallout.py

@@ -0,0 +1,328 @@
+"""
+embed_fallout.py — Ingestion PDFs Fallout JdR → ChromaDB (v3)
+
+Fusion :
+  - Chunking structurel par paragraphes (fallout_ingest_v2, 15/06)
+  - Routing multi-collections selon README priorités (16/06)
+
+Collections :
+  fallout_lore_canon    ← lore_post_guerre/       (canonique, priorité max)
+  fallout_regles        ← regles/core|supplements|fiches_de_jeu/
+  fallout_lore_contexte ← lore_pre_guerre/ + aventures/ + ambiance/radio/
+
+Chemins container : PDF_DIR=/pdfs/jdr  STATE=/pdfs/fallout_embed_state.json
+Env vars : CHROMA_URL, OLLAMA_URL
+"""
+import json, os, re, subprocess, time, urllib.request
+from pathlib import Path
+
+CHROMA_URL  = os.getenv("CHROMA_URL", "http://chromadb:8000")
+CHROMA_BASE = CHROMA_URL + "/api/v2/tenants/default_tenant/databases/default_database"
+OLLAMA_URL  = os.getenv("OLLAMA_URL", "http://172.18.0.1:11434")
+EMBED_MODEL = "nomic-embed-text"
+
+PDF_DIR    = Path(os.getenv("PDF_DIR", "/pdfs/jdr"))
+STATE_FILE = Path(os.getenv("STATE_FILE", "/pdfs/fallout_embed_state.json"))
+BATCH_SIZE = 50
+
+ROUTING = [
+    ("lore_post_guerre",     "fallout_lore_canon",    "lore_post_guerre",  450, 80, True),
+    ("regles/core",          "fallout_regles",         "regles_core",       350, 60, False),
+    ("regles/supplements",   "fallout_regles",         "regles_supplement", 350, 60, False),
+    ("regles/fiches_de_jeu", "fallout_regles",         "fiches_jeu",        350, 60, False),
+    ("lore_pre_guerre",      "fallout_lore_contexte",  "lore_pre_guerre",   450, 80, False),
+    ("aventures",            "fallout_lore_contexte",  "aventures",         450, 80, False),
+    ("ambiance/radio",       "fallout_lore_contexte",  "ambiance_radio",    200, 40, False),
+]
+
+SKIP_DIRS = {
+    "assets_graphiques", "battlemaps", "fonts", "photoshop_psd",
+    "fonds_de_page", "images_ref", "croquis", "maps", "logos",
+    "couvertures", ".claude", "archive",
+}
+SKIP_SLUGS = {"archive__meta_data__compressed"}
+MIN_WORDS  = 40
+
+
+# ── State ─────────────────────────────────────────────────────────────────────
+
+def load_state():
+    try:
+        return json.loads(STATE_FILE.read_text())
+    except Exception:
+        return {"done": [], "failed": []}
+
+def save_state(s):
+    STATE_FILE.write_text(json.dumps(s, indent=2))
+
+
+# ── ChromaDB ──────────────────────────────────────────────────────────────────
+
+def _http(method, url, data=None, timeout=30):
+    body = json.dumps(data).encode() if data is not None else None
+    req = urllib.request.Request(url, data=body, method=method,
+                                 headers={"Content-Type": "application/json"})
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        return json.loads(r.read())
+
+def get_or_create_collection(name):
+    try:
+        r = _http("POST", f"{CHROMA_BASE}/collections",
+                  {"name": name, "metadata": {"hnsw:space": "cosine"}})
+        return r["id"]
+    except urllib.error.HTTPError as e:
+        if e.code == 409:
+            cols = _http("GET", f"{CHROMA_BASE}/collections")
+            return next(c["id"] for c in cols if c["name"] == name)
+        raise
+
+def upsert_batch(col_id, ids, embeddings, documents, metadatas):
+    _http("POST", f"{CHROMA_BASE}/collections/{col_id}/upsert", {
+        "ids": ids, "embeddings": embeddings,
+        "documents": documents, "metadatas": metadatas,
+    }, timeout=60)
+
+
+# ── Embeddings ────────────────────────────────────────────────────────────────
+
+def embed(text):
+    safe = " ".join(text.split()[:280])
+    r = _http("POST", f"{OLLAMA_URL}/api/embeddings",
+              {"model": EMBED_MODEL, "prompt": safe}, timeout=60)
+    return r["embedding"]
+
+
+# ── PDF extraction ────────────────────────────────────────────────────────────
+
+def extract_text(pdf_path):
+    """pdftotext -layout en priorité, fallback pypdf, fallback fitz."""
+    text = ""
+    try:
+        r = subprocess.run(
+            ["pdftotext", "-layout", str(pdf_path), "-"],
+            capture_output=True, timeout=180,
+        )
+        if r.returncode == 0:
+            text = r.stdout.decode("utf-8", errors="replace")
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        pass
+
+    if len(text.strip()) < 200:
+        try:
+            from pypdf import PdfReader
+            reader = PdfReader(str(pdf_path))
+            text = "\n".join(p.extract_text() or "" for p in reader.pages)
+        except Exception:
+            try:
+                import fitz
+                doc = fitz.open(str(pdf_path))
+                text = "\n".join(p.get_text() for p in doc)
+                doc.close()
+            except Exception as e:
+                raise RuntimeError(f"Aucune méthode d'extraction n'a fonctionné : {e}")
+
+    return text.strip()
+
+
+# ── Chunking structurel ───────────────────────────────────────────────────────
+
+def is_section_header(line):
+    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_structural(text, chunk_words=450, overlap_words=80):
+    raw_blocks = re.split(r'\n{2,}', text)
+    blocks = []
+    for b in raw_blocks:
+        b = b.strip()
+        if not b:
+            continue
+        if len(b.split()) > 800:
+            sub = [s.strip() for s in b.split('\n') if s.strip()]
+            blocks.extend(sub)
+        else:
+            blocks.append(b)
+
+    chunks = []
+    current_blocks, current_words = [], 0
+
+    for block in blocks:
+        bwords = len(block.split())
+        if bwords < MIN_WORDS and not is_section_header(block):
+            current_blocks.append(block)
+            current_words += bwords
+            continue
+
+        if current_words >= chunk_words:
+            chunk = "\n\n".join(current_blocks).strip()
+            if len(chunk.split()) >= MIN_WORDS:
+                chunks.append(chunk)
+            overlap_b, overlap_w = [], 0
+            for ob in reversed(current_blocks):
+                ow = len(ob.split())
+                if overlap_w + ow <= overlap_words:
+                    overlap_b.insert(0, ob)
+                    overlap_w += ow
+                else:
+                    break
+            current_blocks = overlap_b + [block]
+            current_words = overlap_w + bwords
+        else:
+            current_blocks.append(block)
+            current_words += bwords
+
+    if current_blocks:
+        chunk = "\n\n".join(current_blocks).strip()
+        if len(chunk.split()) >= MIN_WORDS:
+            chunks.append(chunk)
+
+    return chunks
+
+
+# ── Routing & métadonnées ─────────────────────────────────────────────────────
+
+def route_pdf(rel_path):
+    rel = rel_path.replace("\\", "/")
+    for part in rel.split("/")[:-1]:
+        if part.lower() in SKIP_DIRS:
+            return None
+    rel_low = rel.lower()
+    for prefix, col, cat, cw, ov, canon in ROUTING:
+        if rel_low.startswith(prefix.lower()):
+            return col, cat, cw, ov, canon
+    return None
+
+def detect_langue(filename):
+    n = filename.lower()
+    if re.search(r'(_fr[_\-\.])|(_fr$)|(manuel.*officiel)|(bestiaire)|(armes)|'
+                 r'(special.*comp)|(regles_maison)|(frequence)|(pizza)', n):
+        return "fr"
+    if re.search(r'(_en[_\-\.])|(_en$)|(quickstart)|(booklet)|(toolkit)|(settlers)|'
+                 r'(enclave)|(rust_devil)|(winter_of_atom)|(cascadia)|(hunted)|'
+                 r'(showdown)|(fully_operational)|(last_boat)', n):
+        return "en"
+    return "fr"
+
+def detect_version(filename):
+    m = re.search(r'v(\d+[\.\d]*)', filename.lower())
+    return m.group(1) if m else ""
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+
+def main():
+    print("=== embed_fallout v3 — chunking structurel + routing multi-collections ===\n")
+
+    all_pdfs = []
+    for root, dirs, files in os.walk(PDF_DIR):
+        dirs[:] = [d for d in dirs if d.lower() not in SKIP_DIRS]
+        for f in sorted(files):
+            if f.lower().endswith(".pdf"):
+                all_pdfs.append(Path(root) / f)
+    all_pdfs.sort()
+
+    to_process, skipped = [], []
+    for pdf_path in all_pdfs:
+        rel  = str(pdf_path.relative_to(PDF_DIR))
+        slug = rel.replace(os.sep, "__").replace(" ", "_").lower()[:-4]
+        r    = route_pdf(rel)
+        if r is None or slug in SKIP_SLUGS:
+            skipped.append(rel)
+            continue
+        col, cat, cw, ov, canon = r
+        to_process.append({
+            "path": pdf_path, "rel": rel, "slug": slug,
+            "col": col, "cw": cw, "ov": ov,
+            "meta": {
+                "source":    pdf_path.name,
+                "categorie": cat,
+                "langue":    detect_langue(pdf_path.name),
+                "canonique": "true" if canon else "false",
+                "version":   detect_version(pdf_path.name),
+            },
+        })
+
+    print(f"PDFs trouvés : {len(all_pdfs)}")
+    print(f"  A ingérer  : {len(to_process)}")
+    print(f"  Ignorés    : {len(skipped)}")
+    for s in skipped:
+        print(f"    [SKIP] {s}")
+
+    col_ids = {}
+    print()
+    for name in sorted(set(p["col"] for p in to_process)):
+        col_ids[name] = get_or_create_collection(name)
+        print(f"  Collection '{name}' : {col_ids[name][:8]}...")
+
+    state = load_state()
+    # Reset state pour repartir propre (collections wipées)
+    state = {"done": [], "failed": []}
+    save_state(state)
+
+    total_chunks = 0
+
+    for item in to_process:
+        slug = item["slug"]
+        if slug in state["done"]:
+            print(f"[résumé] {item['rel']}")
+            continue
+
+        print(f"\n[->] {item['rel']}")
+        print(f"     col={item['col']}  cat={item['meta']['categorie']}"
+              f"  lang={item['meta']['langue']}  canon={item['meta']['canonique']}")
+
+        try:
+            text = extract_text(item["path"])
+            print(f"     {len(text)} chars extraits")
+
+            chunks = chunk_structural(text, item["cw"], item["ov"])
+            print(f"     {len(chunks)} chunks (~{item['cw']}w cible, overlap {item['ov']}w)")
+
+            ids, embeddings, documents, metadatas = [], [], [], []
+            for i, chunk in enumerate(chunks):
+                t0 = time.time()
+                emb = embed(chunk)
+                wc  = len(chunk.split())
+                ids.append(f"{slug}__{i:04d}")
+                embeddings.append(emb)
+                documents.append(chunk)
+                metadatas.append({**item["meta"],
+                                  "chunk": i, "total_chunks": len(chunks),
+                                  "word_count": wc})
+                print(f"    chunk {i+1}/{len(chunks)} ({wc}w) {time.time()-t0:.1f}s")
+
+                if len(ids) >= BATCH_SIZE:
+                    upsert_batch(col_ids[item["col"]], ids, embeddings, documents, metadatas)
+                    ids, embeddings, documents, metadatas = [], [], [], []
+
+            if ids:
+                upsert_batch(col_ids[item["col"]], ids, embeddings, documents, metadatas)
+
+            total_chunks += len(chunks)
+            state["done"].append(slug)
+            state["failed"] = [f for f in state["failed"] if f.get("slug") != slug]
+            print(f"  OK — {len(chunks)} chunks")
+
+        except Exception as e:
+            state["failed"].append({"slug": slug, "error": str(e)})
+            print(f"  ERREUR : {e}")
+
+        save_state(state)
+
+    print(f"\n{'='*60}")
+    print(f"TERMINÉ — {len(state['done'])}/{len(to_process)} PDFs")
+    print(f"Chunks totaux : {total_chunks}")
+    if state.get("failed"):
+        print("Échecs :")
+        for f in state["failed"]:
+            print(f"  - {f['slug']}: {f['error']}")
+
+
+if __name__ == "__main__":
+    main()