224 lines
8.1 KiB
Python
224 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Script de nettoyage pour PrivateGPT
|
|
|
|
Ce script permet de lister et supprimer les documents ingérés dans PrivateGPT.
|
|
Options:
|
|
- Lister tous les documents
|
|
- Supprimer des documents par préfixe (ex: "temp_section_")
|
|
- Supprimer des documents par motif
|
|
- Supprimer tous les documents
|
|
|
|
Usage:
|
|
python nettoyer_pgpt.py --list
|
|
python nettoyer_pgpt.py --delete-prefix "temp_section_"
|
|
python nettoyer_pgpt.py --delete-pattern "rapport_.*\.md"
|
|
python nettoyer_pgpt.py --delete-all
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
import requests
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from typing import List, Dict, Any, Optional
|
|
|
|
# Configuration de l'API PrivateGPT
|
|
PGPT_URL = "http://127.0.0.1:8001"
|
|
API_URL = f"{PGPT_URL}/v1"
|
|
|
|
|
|
def check_api_availability() -> bool:
|
|
"""Vérifie si l'API PrivateGPT est disponible"""
|
|
try:
|
|
response = requests.get(f"{PGPT_URL}/health")
|
|
if response.status_code == 200:
|
|
print("✅ API PrivateGPT disponible")
|
|
return True
|
|
else:
|
|
print(f"❌ L'API PrivateGPT a retourné le code d'état {response.status_code}")
|
|
return False
|
|
except requests.RequestException as e:
|
|
print(f"❌ Erreur de connexion à l'API PrivateGPT: {e}")
|
|
return False
|
|
|
|
|
|
def list_documents() -> List[Dict[str, Any]]:
|
|
"""Liste tous les documents ingérés et renvoie la liste des métadonnées"""
|
|
try:
|
|
# Récupérer la liste des documents
|
|
response = requests.get(f"{API_URL}/ingest/list")
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
# Format de réponse OpenAI
|
|
if "data" in data:
|
|
documents = data.get("data", [])
|
|
# Format alternatif
|
|
else:
|
|
documents = data.get("documents", [])
|
|
|
|
# Construire une liste normalisée des documents
|
|
normalized_docs = []
|
|
for doc in documents:
|
|
doc_id = doc.get("doc_id") or doc.get("id")
|
|
metadata = doc.get("doc_metadata", {})
|
|
filename = metadata.get("file_name") or metadata.get("filename", "Inconnu")
|
|
|
|
normalized_docs.append({
|
|
"id": doc_id,
|
|
"filename": filename,
|
|
"metadata": metadata
|
|
})
|
|
|
|
return normalized_docs
|
|
|
|
except Exception as e:
|
|
print(f"❌ Erreur lors de la récupération des documents: {e}")
|
|
return []
|
|
|
|
|
|
def print_documents(documents: List[Dict[str, Any]]) -> None:
|
|
"""Affiche la liste des documents de façon lisible"""
|
|
if not documents:
|
|
print("📋 Aucun document trouvé dans PrivateGPT")
|
|
return
|
|
|
|
print(f"📋 {len(documents)} documents trouvés dans PrivateGPT:")
|
|
|
|
# Regrouper par nom de fichier pour un affichage plus compact
|
|
files_grouped = {}
|
|
for doc in documents:
|
|
filename = doc["filename"]
|
|
if filename not in files_grouped:
|
|
files_grouped[filename] = []
|
|
files_grouped[filename].append(doc["id"])
|
|
|
|
# Afficher les résultats groupés
|
|
for i, (filename, ids) in enumerate(files_grouped.items(), 1):
|
|
print(f"{i}. {filename} ({len(ids)} chunks)")
|
|
if args.verbose:
|
|
for j, doc_id in enumerate(ids, 1):
|
|
print(f" {j}. ID: {doc_id}")
|
|
|
|
|
|
def delete_document(doc_id: str) -> bool:
|
|
"""Supprime un document par son ID"""
|
|
try:
|
|
response = requests.delete(f"{API_URL}/ingest/{doc_id}")
|
|
if response.status_code == 200:
|
|
return True
|
|
else:
|
|
print(f"⚠️ Échec de la suppression de l'ID {doc_id}: Code {response.status_code}")
|
|
return False
|
|
except Exception as e:
|
|
print(f"❌ Erreur lors de la suppression de l'ID {doc_id}: {e}")
|
|
return False
|
|
|
|
|
|
def delete_documents_by_criteria(documents: List[Dict[str, Any]],
|
|
prefix: Optional[str] = None,
|
|
pattern: Optional[str] = None,
|
|
delete_all: bool = False) -> int:
|
|
"""
|
|
Supprime des documents selon différents critères
|
|
Retourne le nombre de documents supprimés
|
|
"""
|
|
if not documents:
|
|
print("❌ Aucun document à supprimer")
|
|
return 0
|
|
|
|
if not (prefix or pattern or delete_all):
|
|
print("❌ Aucun critère de suppression spécifié")
|
|
return 0
|
|
|
|
# Comptage des suppressions réussies
|
|
success_count = 0
|
|
|
|
# Filtrer les documents à supprimer
|
|
docs_to_delete = []
|
|
|
|
if delete_all:
|
|
docs_to_delete = documents
|
|
print(f"🗑️ Suppression de tous les documents ({len(documents)} chunks)...")
|
|
elif prefix:
|
|
docs_to_delete = [doc for doc in documents if doc["filename"].startswith(prefix)]
|
|
print(f"🗑️ Suppression des documents dont le nom commence par '{prefix}' ({len(docs_to_delete)} chunks)...")
|
|
elif pattern:
|
|
try:
|
|
regex = re.compile(pattern)
|
|
docs_to_delete = [doc for doc in documents if regex.search(doc["filename"])]
|
|
print(f"🗑️ Suppression des documents correspondant au motif '{pattern}' ({len(docs_to_delete)} chunks)...")
|
|
except re.error as e:
|
|
print(f"❌ Expression régulière invalide: {e}")
|
|
return 0
|
|
|
|
# Demander confirmation si beaucoup de documents
|
|
if len(docs_to_delete) > 5 and not args.force:
|
|
confirm = input(f"⚠️ Vous êtes sur le point de supprimer {len(docs_to_delete)} chunks. Confirmer ? (o/N) ")
|
|
if confirm.lower() != 'o':
|
|
print("❌ Opération annulée")
|
|
return 0
|
|
|
|
# Supprimer les documents
|
|
for doc in docs_to_delete:
|
|
if delete_document(doc["id"]):
|
|
success_count += 1
|
|
if args.verbose:
|
|
print(f"✅ Document supprimé: {doc['filename']} (ID: {doc['id']})")
|
|
|
|
# Petite pause pour éviter de surcharger l'API
|
|
time.sleep(0.1)
|
|
|
|
print(f"✅ {success_count}/{len(docs_to_delete)} documents supprimés avec succès")
|
|
return success_count
|
|
|
|
|
|
def generate_unique_prefix() -> str:
|
|
"""Génère un préfixe unique basé sur un UUID pour différencier les fichiers temporaires"""
|
|
unique_id = str(uuid.uuid4())[:8] # Prendre les 8 premiers caractères de l'UUID
|
|
return f"temp_{unique_id}_"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Utilitaire de nettoyage pour PrivateGPT")
|
|
|
|
# Options principales
|
|
group = parser.add_mutually_exclusive_group(required=True)
|
|
group.add_argument("--list", action="store_true", help="Lister tous les documents ingérés")
|
|
group.add_argument("--delete-prefix", type=str, help="Supprimer les documents dont le nom commence par PREFIX")
|
|
group.add_argument("--delete-pattern", type=str, help="Supprimer les documents dont le nom correspond au motif PATTERN (regex)")
|
|
group.add_argument("--delete-all", action="store_true", help="Supprimer tous les documents (⚠️ DANGER)")
|
|
group.add_argument("--generate-prefix", action="store_true", help="Générer un préfixe unique pour les fichiers temporaires")
|
|
|
|
# Options additionnelles
|
|
parser.add_argument("--force", action="store_true", help="Ne pas demander de confirmation")
|
|
parser.add_argument("--verbose", action="store_true", help="Afficher plus de détails")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Vérifier la disponibilité de l'API
|
|
if not check_api_availability():
|
|
sys.exit(1)
|
|
|
|
# Générer un préfixe unique
|
|
if args.generate_prefix:
|
|
unique_prefix = generate_unique_prefix()
|
|
print(f"🔑 Préfixe unique généré: {unique_prefix}")
|
|
print(f"Utilisez ce préfixe pour les fichiers temporaires de votre script.")
|
|
sys.exit(0)
|
|
|
|
# Récupérer la liste des documents
|
|
documents = list_documents()
|
|
|
|
# Traiter selon l'option choisie
|
|
if args.list:
|
|
print_documents(documents)
|
|
elif args.delete_prefix:
|
|
delete_documents_by_criteria(documents, prefix=args.delete_prefix)
|
|
elif args.delete_pattern:
|
|
delete_documents_by_criteria(documents, pattern=args.delete_pattern)
|
|
elif args.delete_all:
|
|
delete_documents_by_criteria(documents, delete_all=True) |