diff --git a/ingest/fallout_ingest_v2.py b/ingest/fallout_ingest_v2.py new file mode 100644 index 0000000..8be73df --- /dev/null +++ b/ingest/fallout_ingest_v2.py @@ -0,0 +1,308 @@ +""" +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()