281 lines
10 KiB
Python
281 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
Script pour l'injection automatique de fichiers dans Private GPT.
|
|
Ce script scanne un répertoire source et injecte les nouveaux fichiers via l'API de Private GPT.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import json
|
|
import argparse
|
|
import logging
|
|
import requests
|
|
from pathlib import Path
|
|
from typing import List, Set, Dict, Any
|
|
from datetime import datetime
|
|
|
|
# Configuration du logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
|
handlers=[
|
|
logging.StreamHandler(sys.stdout),
|
|
logging.FileHandler('auto_ingest.log')
|
|
]
|
|
)
|
|
logger = logging.getLogger("auto_ingest")
|
|
|
|
# Extensions de fichiers supportées par défaut
|
|
DEFAULT_SUPPORTED_EXTENSIONS = {
|
|
'.txt', '.pdf', '.csv', '.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx',
|
|
'.md', '.html', '.htm', '.json', '.xml', '.rtf', '.odt', '.ods', '.odp'
|
|
}
|
|
|
|
class PrivateGPTIngestor:
|
|
"""Classe pour gérer l'ingestion de fichiers dans Private GPT."""
|
|
|
|
def __init__(self, api_url: str = "http://localhost:8001",
|
|
processed_file: str = "processed_files.json"):
|
|
"""
|
|
Initialise l'ingesteur.
|
|
|
|
Args:
|
|
api_url: URL de l'API Private GPT
|
|
processed_file: Fichier pour stocker les fichiers déjà traités
|
|
"""
|
|
self.api_url = api_url
|
|
self.processed_file = processed_file
|
|
self.processed_files = self._load_processed_files()
|
|
|
|
def _load_processed_files(self) -> Set[str]:
|
|
"""Charge la liste des fichiers déjà traités."""
|
|
try:
|
|
if os.path.exists(self.processed_file):
|
|
with open(self.processed_file, 'r', encoding='utf-8') as f:
|
|
return set(json.load(f))
|
|
return set()
|
|
except Exception as e:
|
|
logger.error(f"Erreur lors du chargement des fichiers traités: {e}")
|
|
return set()
|
|
|
|
def _save_processed_files(self) -> None:
|
|
"""Sauvegarde la liste des fichiers déjà traités."""
|
|
try:
|
|
with open(self.processed_file, 'w', encoding='utf-8') as f:
|
|
json.dump(list(self.processed_files), f, ensure_ascii=False, indent=2)
|
|
except Exception as e:
|
|
logger.error(f"Erreur lors de la sauvegarde des fichiers traités: {e}")
|
|
|
|
def scan_directory(self, directory: str, extensions: Set[str] = None,
|
|
recursive: bool = True) -> List[str]:
|
|
"""
|
|
Scanne un répertoire pour trouver des fichiers à injecter.
|
|
|
|
Args:
|
|
directory: Le répertoire à scanner
|
|
extensions: Extensions de fichiers à prendre en compte
|
|
recursive: Si True, scanne les sous-répertoires
|
|
|
|
Returns:
|
|
Liste des chemins des fichiers à injecter
|
|
"""
|
|
if extensions is None:
|
|
extensions = DEFAULT_SUPPORTED_EXTENSIONS
|
|
|
|
files_to_ingest = []
|
|
directory_path = Path(directory)
|
|
|
|
if not directory_path.exists():
|
|
logger.error(f"Le répertoire {directory} n'existe pas")
|
|
return []
|
|
|
|
logger.info(f"Scan du répertoire {directory}")
|
|
|
|
# Fonction de scan
|
|
def scan_dir(path: Path):
|
|
for item in path.iterdir():
|
|
if item.is_file() and any(item.name.lower().endswith(ext) for ext in extensions):
|
|
abs_path = str(item.absolute())
|
|
if abs_path not in self.processed_files:
|
|
files_to_ingest.append(abs_path)
|
|
elif item.is_dir() and recursive:
|
|
scan_dir(item)
|
|
|
|
scan_dir(directory_path)
|
|
logger.info(f"Trouvé {len(files_to_ingest)} fichiers à injecter")
|
|
return files_to_ingest
|
|
|
|
def ingest_file(self, file_path: str) -> bool:
|
|
"""
|
|
Injecte un fichier dans Private GPT via l'API.
|
|
|
|
Args:
|
|
file_path: Chemin du fichier à injecter
|
|
|
|
Returns:
|
|
True si l'injection a réussi, False sinon
|
|
"""
|
|
logger.info(f"Injection du fichier: {file_path}")
|
|
|
|
try:
|
|
with open(file_path, 'rb') as f:
|
|
files = {'file': (os.path.basename(file_path), f)}
|
|
response = requests.post(f"{self.api_url}/v1/ingest/file", files=files)
|
|
|
|
if response.status_code == 200:
|
|
logger.info(f"Injection réussie pour {file_path}")
|
|
self.processed_files.add(file_path)
|
|
self._save_processed_files()
|
|
return True
|
|
else:
|
|
logger.error(f"Échec de l'injection pour {file_path}: {response.status_code} - {response.text}")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"Erreur lors de l'injection de {file_path}: {e}")
|
|
return False
|
|
|
|
def list_documents(self) -> List[Dict[str, Any]]:
|
|
"""
|
|
Liste les documents déjà injectés dans Private GPT.
|
|
|
|
Returns:
|
|
Liste des documents injectés
|
|
"""
|
|
try:
|
|
response = requests.get(f"{self.api_url}/v1/ingest/list")
|
|
if response.status_code == 200:
|
|
return response.json()
|
|
else:
|
|
logger.error(f"Échec de la récupération des documents: {response.status_code} - {response.text}")
|
|
return []
|
|
except Exception as e:
|
|
logger.error(f"Erreur lors de la récupération des documents: {e}")
|
|
return []
|
|
|
|
def run_ingestion(self, directory: str, extensions: Set[str] = None,
|
|
recursive: bool = True, batch_size: int = 5,
|
|
delay: float = 2.0) -> None:
|
|
"""
|
|
Exécute l'ingestion des fichiers d'un répertoire.
|
|
|
|
Args:
|
|
directory: Répertoire à scanner
|
|
extensions: Extensions à prendre en compte
|
|
recursive: Si True, scanne les sous-répertoires
|
|
batch_size: Nombre de fichiers à injecter par lot
|
|
delay: Délai entre chaque lot (en secondes)
|
|
"""
|
|
files_to_ingest = self.scan_directory(directory, extensions, recursive)
|
|
|
|
if not files_to_ingest:
|
|
logger.info("Aucun nouveau fichier à injecter")
|
|
return
|
|
|
|
total_files = len(files_to_ingest)
|
|
successful = 0
|
|
failed = 0
|
|
|
|
for i, file_path in enumerate(files_to_ingest):
|
|
logger.info(f"Progression: {i+1}/{total_files}")
|
|
|
|
if self.ingest_file(file_path):
|
|
successful += 1
|
|
else:
|
|
failed += 1
|
|
|
|
# Pause après chaque lot
|
|
if (i + 1) % batch_size == 0 and i < total_files - 1:
|
|
logger.info(f"Pause de {delay} secondes après le lot de {batch_size} fichiers")
|
|
time.sleep(delay)
|
|
|
|
logger.info(f"Ingestion terminée: {successful} réussis, {failed} échoués sur {total_files} fichiers")
|
|
|
|
def parse_args():
|
|
"""Parse les arguments de ligne de commande."""
|
|
parser = argparse.ArgumentParser(description="Outil d'injection automatique pour Private GPT")
|
|
|
|
parser.add_argument("--directory", "-d", type=str, required=True,
|
|
help="Répertoire contenant les fichiers à injecter")
|
|
parser.add_argument("--api-url", type=str, default="http://localhost:8001",
|
|
help="URL de l'API Private GPT (défaut: http://localhost:8001)")
|
|
parser.add_argument("--recursive", "-r", action="store_true", default=True,
|
|
help="Scanner récursivement les sous-répertoires")
|
|
parser.add_argument("--extensions", "-e", type=str, nargs="+",
|
|
help="Extensions de fichiers à prendre en compte (ex: .pdf .txt)")
|
|
parser.add_argument("--batch-size", "-b", type=int, default=5,
|
|
help="Nombre de fichiers à injecter par lot (défaut: 5)")
|
|
parser.add_argument("--delay", type=float, default=2.0,
|
|
help="Délai entre les lots en secondes (défaut: 2.0)")
|
|
parser.add_argument("--list", "-l", action="store_true",
|
|
help="Lister les documents déjà injectés et quitter")
|
|
parser.add_argument("--watch", "-w", action="store_true",
|
|
help="Mode surveillance: vérifier périodiquement les nouveaux fichiers")
|
|
parser.add_argument("--watch-interval", type=int, default=300,
|
|
help="Intervalle de surveillance en secondes (défaut: 300)")
|
|
|
|
return parser.parse_args()
|
|
|
|
def main():
|
|
"""Fonction principale."""
|
|
args = parse_args()
|
|
|
|
ingestor = PrivateGPTIngestor(api_url=args.api_url)
|
|
|
|
# Option pour lister les documents
|
|
if args.list:
|
|
documents = ingestor.list_documents()
|
|
print(f"Documents déjà injectés ({len(documents)}):")
|
|
for doc in documents:
|
|
print(f"- {doc.get('doc_id')}: {doc.get('doc_metadata', {}).get('file_name', 'Inconnu')}")
|
|
return
|
|
|
|
# Conversion des extensions
|
|
extensions = None
|
|
if args.extensions:
|
|
extensions = set()
|
|
for ext in args.extensions:
|
|
if not ext.startswith('.'):
|
|
ext = '.' + ext
|
|
extensions.add(ext.lower())
|
|
|
|
# Mode surveillance
|
|
if args.watch:
|
|
logger.info(f"Mode surveillance activé. Vérification toutes les {args.watch_interval} secondes")
|
|
try:
|
|
while True:
|
|
start_time = datetime.now()
|
|
logger.info(f"Démarrage d'un scan à {start_time.strftime('%H:%M:%S')}")
|
|
|
|
ingestor.run_ingestion(
|
|
directory=args.directory,
|
|
extensions=extensions,
|
|
recursive=args.recursive,
|
|
batch_size=args.batch_size,
|
|
delay=args.delay
|
|
)
|
|
|
|
# Calcul du temps à attendre
|
|
elapsed = (datetime.now() - start_time).total_seconds()
|
|
wait_time = max(0, args.watch_interval - elapsed)
|
|
|
|
if wait_time > 0:
|
|
logger.info(f"En attente pendant {wait_time:.1f} secondes jusqu'au prochain scan")
|
|
time.sleep(wait_time)
|
|
except KeyboardInterrupt:
|
|
logger.info("Arrêt du mode surveillance")
|
|
else:
|
|
# Exécution unique
|
|
ingestor.run_ingestion(
|
|
directory=args.directory,
|
|
extensions=extensions,
|
|
recursive=args.recursive,
|
|
batch_size=args.batch_size,
|
|
delay=args.delay
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|