""" Tests pour l'optimisation FAISS IVF Valide: - Migration automatique Flat → IVF - Entraînement automatique de l'index IVF - Optimisation périodique - Calcul de nlist optimal """ import pytest import numpy as np from pathlib import Path import tempfile import shutil from core.embedding.faiss_manager import FAISSManager class TestFAISSIVFOptimization: """Tests pour l'optimisation IVF""" def setup_method(self): """Setup avant chaque test""" self.dimensions = 512 self.temp_dir = Path(tempfile.mkdtemp()) def teardown_method(self): """Cleanup après chaque test""" if self.temp_dir.exists(): shutil.rmtree(self.temp_dir) def test_ivf_training(self): """Test entraînement automatique de l'index IVF""" manager = FAISSManager( dimensions=self.dimensions, index_type="IVF", metric="cosine", auto_optimize=False ) # Au début, pas entraîné assert not manager.is_trained # Ajouter des vecteurs (moins de 100) for i in range(50): vector = np.random.randn(self.dimensions).astype(np.float32) manager.add_embedding(f"emb_{i}", vector) # Toujours pas entraîné assert not manager.is_trained # Ajouter plus de vecteurs pour déclencher l'entraînement for i in range(50, 150): vector = np.random.randn(self.dimensions).astype(np.float32) manager.add_embedding(f"emb_{i}", vector) # Maintenant entraîné assert manager.is_trained assert manager.index.ntotal == 150 def test_nlist_calculation(self): """Test calcul de nlist optimal""" manager = FAISSManager( dimensions=self.dimensions, index_type="Flat", auto_optimize=False ) # Test différentes tailles assert manager._calculate_nlist(100) == 100 # min assert manager._calculate_nlist(10000) == 100 # sqrt(10000) = 100 assert manager._calculate_nlist(40000) == 200 # sqrt(40000) = 200 assert manager._calculate_nlist(1000000) == 1000 # sqrt(1000000) = 1000 def test_auto_migration_flat_to_ivf(self): """Test migration automatique Flat → IVF""" # Créer index Flat avec seuil bas pour test manager = FAISSManager( dimensions=self.dimensions, index_type="Flat", metric="cosine", auto_optimize=True ) # Réduire le seuil pour test manager.migration_threshold = 100 # Vérifier qu'on commence avec Flat assert manager.index_type == "Flat" # Ajouter des vecteurs jusqu'au seuil for i in range(110): vector = np.random.randn(self.dimensions).astype(np.float32) manager.add_embedding(f"emb_{i}", vector) # Devrait avoir migré vers IVF assert manager.index_type == "IVF" assert manager.is_trained assert manager.index.ntotal == 110 def test_ivf_search_quality(self): """Test qualité de recherche avec IVF""" # Créer index IVF manager = FAISSManager( dimensions=self.dimensions, index_type="IVF", metric="cosine", nlist=50, nprobe=20 # Augmenter nprobe pour meilleure qualité ) # Ajouter des vecteurs vectors = [] for i in range(200): vector = np.random.randn(self.dimensions).astype(np.float32) vectors.append(vector) manager.add_embedding(f"emb_{i}", vector) # Vérifier que l'index contient bien 200 vecteurs assert manager.index.ntotal == 200, f"Expected 200 vectors, got {manager.index.ntotal}" # Rechercher avec un vecteur après l'entraînement (index 150) # Les 100 premiers sont utilisés pour l'entraînement query = vectors[150] results = manager.search_similar(query, k=5) # Le premier résultat devrait être le vecteur lui-même assert len(results) > 0, "No results returned" # Pour IVF, la recherche est approximative, donc on vérifie juste # que le vecteur est dans les résultats (pas forcément en premier) embedding_ids = [r.embedding_id for r in results] assert "emb_150" in embedding_ids, f"emb_150 not found in results: {embedding_ids}" # Vérifier qu'au moins un résultat a une bonne similarité max_similarity = max(r.similarity for r in results) assert max_similarity > 0.8, f"Max similarity too low: {max_similarity}" def test_ivf_nprobe_effect(self): """Test effet de nprobe sur la qualité de recherche""" # Créer index avec nprobe faible manager_low = FAISSManager( dimensions=self.dimensions, index_type="IVF", metric="cosine", nlist=50, nprobe=1 # Très faible ) # Créer index avec nprobe élevé manager_high = FAISSManager( dimensions=self.dimensions, index_type="IVF", metric="cosine", nlist=50, nprobe=20 # Élevé ) # Ajouter les mêmes vecteurs aux deux vectors = [] for i in range(200): vector = np.random.randn(self.dimensions).astype(np.float32) vectors.append(vector) manager_low.add_embedding(f"emb_{i}", vector) manager_high.add_embedding(f"emb_{i}", vector) # Rechercher avec un vecteur query = vectors[0] results_low = manager_low.search_similar(query, k=10) results_high = manager_high.search_similar(query, k=10) # nprobe élevé devrait donner de meilleurs résultats # (plus de résultats ou meilleure similarité moyenne) assert len(results_high) >= len(results_low) def test_optimize_index(self): """Test optimisation périodique de l'index""" manager = FAISSManager( dimensions=self.dimensions, index_type="IVF", metric="cosine", nlist=50, # Petit nlist initial auto_optimize=False ) # Ajouter beaucoup de vecteurs for i in range(500): vector = np.random.randn(self.dimensions).astype(np.float32) manager.add_embedding(f"emb_{i}", vector) # Vérifier nlist initial initial_nlist = manager.index.nlist assert initial_nlist == 50 # Optimiser manager.optimize_index() # nlist devrait avoir changé (optimal pour 500 vecteurs ≈ 22) # Mais on garde le nlist actuel car pas assez différent # Testons avec un cas plus extrême # Ajouter encore plus de vecteurs for i in range(500, 2000): vector = np.random.randn(self.dimensions).astype(np.float32) manager.add_embedding(f"emb_{i}", vector) # Maintenant optimiser devrait changer nlist manager.optimize_index() # Pour 2000 vecteurs, optimal ≈ sqrt(2000) ≈ 45 # Différence avec 50 n'est pas assez grande (< 50%) # Donc nlist ne change pas dans ce cas # Test avec un nlist vraiment sous-optimal manager.index.nlist = 10 # Forcer un nlist très bas manager.optimize_index() # Devrait avoir augmenté assert manager.index.nlist > 10 def test_save_load_ivf(self): """Test sauvegarde/chargement d'index IVF""" # Créer et peupler index IVF manager = FAISSManager( dimensions=self.dimensions, index_type="IVF", metric="cosine", nlist=50, nprobe=8 ) # Ajouter des vecteurs for i in range(200): vector = np.random.randn(self.dimensions).astype(np.float32) manager.add_embedding(f"emb_{i}", vector, metadata={"index": i}) # Sauvegarder index_path = self.temp_dir / "test_ivf.index" metadata_path = self.temp_dir / "test_ivf.meta" manager.save(index_path, metadata_path) # Charger loaded_manager = FAISSManager.load(index_path, metadata_path) # Vérifier assert loaded_manager.dimensions == self.dimensions assert loaded_manager.index_type == "IVF" assert loaded_manager.metric == "cosine" assert loaded_manager.index.ntotal == 200 assert loaded_manager.is_trained assert loaded_manager.index.nlist == 50 assert loaded_manager.index.nprobe == 8 # Vérifier métadonnées assert len(loaded_manager.metadata_store) == 200 assert loaded_manager.metadata_store[0]["metadata"]["index"] == 0 def test_stats_with_ivf(self): """Test statistiques avec index IVF""" manager = FAISSManager( dimensions=self.dimensions, index_type="IVF", metric="cosine", nlist=50 ) # Ajouter des vecteurs for i in range(200): vector = np.random.randn(self.dimensions).astype(np.float32) manager.add_embedding(f"emb_{i}", vector) # Obtenir stats stats = manager.get_stats() # Vérifier assert stats["index_type"] == "IVF" assert stats["total_vectors"] == 200 assert stats["is_trained"] == True assert stats["nlist"] == 50 assert stats["nprobe"] == 8 assert "optimal_nlist" in stats assert "nlist_efficiency" in stats if __name__ == "__main__": pytest.main([__file__, "-v"])