Code/batch_ia/nettoyer_pgpt.py
2025-05-26 17:27:59 +02:00

101 lines
2.8 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
"""
import json
import re
import requests
import time
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 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 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(pattern) -> int:
"""
Supprime des documents selon différents critères
Retourne le nombre de documents supprimés
"""
documents = list_documents()
if not documents or not pattern:
return 0
# Comptage des suppressions réussies
success_count = 0
# Filtrer les documents à supprimer
docs_to_delete = []
try:
regex = re.compile(pattern)
docs_to_delete = [doc for doc in documents if regex.search(doc["filename"])]
except re.error as e:
return 0
# Supprimer les documents
for doc in docs_to_delete:
delete_document(doc["id"])
# Petite pause pour éviter de surcharger l'API
time.sleep(0.1)
return success_count