#!/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, Optional 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()