Funk-lab/ansible/roles/rag/files/rag-ingest

#!/usr/bin/env python3
"""
Indexe les fichiers Markdown de la doc Funk dans Qdrant.
Usage: rag-ingest [docs_dir]
       docs_dir : répertoire contenant les .md (défaut: /srv/data/rag/docs)
"""
import os
import sys
import json
import hashlib
import re
import urllib.request
import urllib.error

QDRANT_URL  = os.environ.get("QDRANT_URL",  "http://127.0.0.1:6333")
EMBED_URL   = os.environ.get("EMBED_URL",   "http://192.168.10.20:1234/v1/embeddings")
EMBED_MODEL = os.environ.get("EMBED_MODEL", "qwen3-8b")
COLLECTION  = os.environ.get("RAG_COLLECTION", "funk-docs")
VECTOR_DIM  = 4096
CHUNK_MAX   = 2000


def _request(method, url, data=None):
    body = json.dumps(data).encode() if data is not None else None
    req = urllib.request.Request(
        url, data=body,
        headers={"Content-Type": "application/json"} if body else {},
        method=method
    )
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.loads(r.read())


def ensure_collection():
    try:
        _request("GET", f"{QDRANT_URL}/collections/{COLLECTION}")
        print(f"Collection '{COLLECTION}' existante")
    except urllib.error.HTTPError as e:
        if e.code == 404:
            print(f"Création de la collection '{COLLECTION}'...")
            _request("PUT", f"{QDRANT_URL}/collections/{COLLECTION}", {
                "vectors": {"size": VECTOR_DIM, "distance": "Cosine"}
            })
            print("Collection créée")
        else:
            raise


def embed(text):
    result = _request("POST", EMBED_URL, {"model": EMBED_MODEL, "input": text})
    return result["data"][0]["embedding"]


def chunk_markdown(rel_path, content):
    """Découpe un fichier Markdown par sections H2, puis H3 si trop grand."""
    chunks = []

    # Séparer par H2
    parts = re.split(r'\n(?=## )', content)

    # Intro avant le premier H2
    if parts and not parts[0].strip().startswith('## '):
        intro = parts[0].strip()
        if len(intro) > 80:
            title = intro.split('\n')[0].lstrip('#').strip() or rel_path
            chunks.append({"file": rel_path, "section": title, "text": intro[:CHUNK_MAX]})
        parts = parts[1:]

    for part in parts:
        if not part.strip():
            continue
        h2_title = part.split('\n')[0].lstrip('#').strip()

        if len(part) <= CHUNK_MAX:
            chunks.append({"file": rel_path, "section": h2_title, "text": part.strip()})
        else:
            # Découper par H3
            subs = re.split(r'\n(?=### )', part)
            for sub in subs:
                if not sub.strip():
                    continue
                h3_title = sub.split('\n')[0].lstrip('#').strip()
                label = f"{h2_title} — {h3_title}" if h3_title != h2_title else h2_title
                chunks.append({"file": rel_path, "section": label, "text": sub.strip()[:CHUNK_MAX]})

    return chunks


def point_id(file_path, section):
    """ID stable uint64 = MD5(file::section)."""
    h = hashlib.md5(f"{file_path}::{section}".encode()).hexdigest()
    return int(h[:16], 16)


def ingest(docs_dir):
    ensure_collection()

    md_files = []
    for root, dirs, files in os.walk(docs_dir):
        dirs[:] = sorted(d for d in dirs if not d.startswith('.'))
        for f in sorted(files):
            if f.endswith('.md'):
                md_files.append(os.path.join(root, f))

    print(f"{len(md_files)} fichiers Markdown trouvés dans {docs_dir}")

    total = 0
    errors = 0

    for filepath in md_files:
        rel = os.path.relpath(filepath, docs_dir)
        try:
            with open(filepath, encoding='utf-8') as f:
                content = f.read()
        except OSError as e:
            print(f"  SKIP {rel}: {e}")
            continue

        chunks = chunk_markdown(rel, content)
        if not chunks:
            continue

        points = []
        for chunk in chunks:
            print(f"  embed  {rel} § {chunk['section'][:55]}")
            try:
                vector = embed(chunk['text'])
            except Exception as e:
                print(f"    ERREUR embedding: {e}")
                errors += 1
                continue

            points.append({
                "id": point_id(chunk['file'], chunk['section']),
                "vector": vector,
                "payload": {
                    "file": chunk['file'],
                    "section": chunk['section'],
                    "text": chunk['text'],
                }
            })

        if points:
            _request("PUT", f"{QDRANT_URL}/collections/{COLLECTION}/points?wait=true", {"points": points})
            total += len(points)
            print(f"  → {len(points)} chunks indexés  ({rel})")

    print(f"\nTerminé : {total} chunks dans '{COLLECTION}'" + (f" ({errors} erreurs)" if errors else ""))


if __name__ == "__main__":
    docs_dir = sys.argv[1] if len(sys.argv) > 1 else "/srv/data/rag/docs"
    if not os.path.isdir(docs_dir):
        print(f"ERREUR : {docs_dir} n'existe pas", file=sys.stderr)
        sys.exit(1)
    ingest(docs_dir)