dedup_store.py

"""
Déduplication de la base ChromaDB
=================================
Supprime les chunks redondants : quand une même paire (source, chunk_index)
existe plusieurs fois (= fichier ingéré plusieurs fois), on garde 1 exemplaire
et on supprime les autres. Lecture/écriture locale, AUCUN appel Azure.

Usage : ./.venv/Scripts/python.exe dedup_store.py
"""

import os
from pathlib import Path
from collections import defaultdict

os.environ.setdefault("ANONYMIZED_TELEMETRY", "False")
import chromadb
from chromadb.config import Settings

PROJECT_ROOT = Path(__file__).parent
DATA_DIR = Path("/data") if Path("/data").is_dir() else PROJECT_ROOT / "data"
COLLECTION_NAME = "rag_documents"

client = chromadb.PersistentClient(path=str(DATA_DIR / "chroma_db"), settings=Settings(anonymized_telemetry=False))
col = client.get_or_create_collection(name=COLLECTION_NAME, metadata={"hnsw:space": "cosine"})

total = col.count()
print(f"Total chunks avant : {total}")

# (source, chunk_index) -> liste d'ids
pair_to_ids = defaultdict(list)
off, step = 0, 2000
while off < total:
    got = col.get(include=["metadatas"], limit=step, offset=off)
    for _id, m in zip(got["ids"], got["metadatas"]):
        pair_to_ids[(m.get("source"), m.get("chunk_index"))].append(_id)
    off += step

# pour chaque paire en double, on garde le 1er id, on supprime les autres
to_delete = []
affected = defaultdict(int)
for (src, idx), ids in pair_to_ids.items():
    if len(ids) > 1:
        extra = ids[1:]            # on garde ids[0]
        to_delete.extend(extra)
        affected[src] += len(extra)

print(f"Chunks redondants à supprimer : {len(to_delete)}")
print("Mémoires concernés :")
for src, n in sorted(affected.items(), key=lambda x: -x[1]):
    print(f"   -{n}  {src}")

if not to_delete:
    print("Aucun doublon. Rien à faire.")
else:
    # suppression par lots (limite SQLite sur le nombre de variables)
    B = 500
    for i in range(0, len(to_delete), B):
        col.delete(ids=to_delete[i:i + B])
        print(f"  supprimé {min(i + B, len(to_delete))}/{len(to_delete)}")
    print(f"Total chunks après : {col.count()}")