486 lines
17 KiB
Python
486 lines
17 KiB
Python
"""
|
|
Tests pour la vectorisation et l'indexation des référentiels.
|
|
|
|
Ce module teste:
|
|
- Le chargement du modèle d'embeddings
|
|
- La vectorisation des chunks
|
|
- La construction de l'index HNSW avec FAISS
|
|
- La vectorisation des index alphabétiques
|
|
- La génération de hash d'index pour versionnement
|
|
"""
|
|
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import faiss
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from pipeline_mco_pmsi.rag.referentiels_manager import Chunk, ReferentielsManager, VectorIndex
|
|
|
|
|
|
class TestEmbeddingsModel:
|
|
"""Tests pour le chargement du modèle d'embeddings."""
|
|
|
|
def test_load_embeddings_model_success(self):
|
|
"""Test que le modèle d'embeddings se charge correctement."""
|
|
manager = ReferentielsManager()
|
|
model = manager._load_embeddings_model()
|
|
|
|
assert model is not None
|
|
# Vérifier que le modèle peut encoder du texte
|
|
embedding = model.encode("Test médical", convert_to_numpy=True)
|
|
assert embedding is not None
|
|
assert len(embedding.shape) == 1 # Vecteur 1D
|
|
assert embedding.shape[0] > 0 # Dimension > 0
|
|
|
|
def test_embeddings_model_produces_consistent_vectors(self):
|
|
"""Test que le modèle produit des vecteurs cohérents pour le même texte."""
|
|
manager = ReferentielsManager()
|
|
model = manager._load_embeddings_model()
|
|
|
|
text = "Gastrite aiguë avec hémorragie"
|
|
embedding1 = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
|
|
embedding2 = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
|
|
|
|
# Les embeddings doivent être identiques (ou très proches)
|
|
np.testing.assert_array_almost_equal(embedding1, embedding2, decimal=5)
|
|
|
|
def test_embeddings_model_normalized(self):
|
|
"""Test que les embeddings sont normalisés (L2 norm = 1)."""
|
|
manager = ReferentielsManager()
|
|
model = manager._load_embeddings_model()
|
|
|
|
text = "Diagnostic principal"
|
|
embedding = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
|
|
|
|
# Calculer la norme L2
|
|
norm = np.linalg.norm(embedding)
|
|
|
|
# La norme doit être proche de 1 (normalisé)
|
|
assert abs(norm - 1.0) < 0.01
|
|
|
|
|
|
class TestVectorization:
|
|
"""Tests pour la vectorisation des chunks."""
|
|
|
|
def test_vectorize_single_chunk(self):
|
|
"""Test la vectorisation d'un seul chunk."""
|
|
manager = ReferentielsManager()
|
|
|
|
chunk = Chunk(
|
|
chunk_id="test_001",
|
|
referentiel_type="cim10",
|
|
referentiel_version="2026",
|
|
content="K29.7 Gastrite, sans précision",
|
|
metadata={"chunk_type": "code_block"},
|
|
chunk_index=0,
|
|
)
|
|
|
|
# Vectoriser le chunk
|
|
model = manager._load_embeddings_model()
|
|
vector = model.encode(chunk.content, convert_to_numpy=True, normalize_embeddings=True)
|
|
|
|
assert vector is not None
|
|
assert vector.shape[0] > 0
|
|
assert isinstance(vector, np.ndarray)
|
|
|
|
def test_vectorize_multiple_chunks(self):
|
|
"""Test la vectorisation de plusieurs chunks."""
|
|
manager = ReferentielsManager()
|
|
|
|
chunks = [
|
|
Chunk(
|
|
chunk_id=f"test_{i:03d}",
|
|
referentiel_type="cim10",
|
|
referentiel_version="2026",
|
|
content=f"Code test {i}",
|
|
metadata={"chunk_type": "code_block"},
|
|
chunk_index=i,
|
|
)
|
|
for i in range(10)
|
|
]
|
|
|
|
model = manager._load_embeddings_model()
|
|
vectors = [
|
|
model.encode(chunk.content, convert_to_numpy=True, normalize_embeddings=True)
|
|
for chunk in chunks
|
|
]
|
|
|
|
assert len(vectors) == 10
|
|
# Tous les vecteurs doivent avoir la même dimension
|
|
dimensions = [v.shape[0] for v in vectors]
|
|
assert len(set(dimensions)) == 1 # Toutes les dimensions sont identiques
|
|
|
|
|
|
class TestBuildIndex:
|
|
"""Tests pour la construction de l'index HNSW."""
|
|
|
|
def test_build_index_success(self, tmp_path):
|
|
"""Test la construction réussie d'un index HNSW."""
|
|
manager = ReferentielsManager(data_dir=tmp_path)
|
|
|
|
# Créer des chunks de test
|
|
chunks = [
|
|
Chunk(
|
|
chunk_id=f"cim10_2026_{i}",
|
|
referentiel_type="cim10",
|
|
referentiel_version="2026",
|
|
content=f"K29.{i} Gastrite type {i}",
|
|
metadata={"chunk_type": "code_block"},
|
|
chunk_index=i,
|
|
)
|
|
for i in range(20)
|
|
]
|
|
|
|
# Construire l'index
|
|
vector_index = manager.build_index(chunks)
|
|
|
|
# Vérifications
|
|
assert isinstance(vector_index, VectorIndex)
|
|
assert vector_index.index_hash is not None
|
|
assert len(vector_index.index_hash) == 64 # SHA-256 hex
|
|
assert vector_index.dimension > 0
|
|
assert vector_index.num_vectors == 20
|
|
assert vector_index.index_type == "HNSW"
|
|
assert vector_index.created_at is not None
|
|
|
|
def test_build_index_saves_to_disk(self, tmp_path):
|
|
"""Test que l'index est sauvegardé sur disque."""
|
|
manager = ReferentielsManager(data_dir=tmp_path)
|
|
|
|
chunks = [
|
|
Chunk(
|
|
chunk_id=f"cim10_2026_{i}",
|
|
referentiel_type="cim10",
|
|
referentiel_version="2026",
|
|
content=f"Code {i}",
|
|
metadata={},
|
|
chunk_index=i,
|
|
)
|
|
for i in range(10)
|
|
]
|
|
|
|
manager.build_index(chunks)
|
|
|
|
# Vérifier que les fichiers sont créés
|
|
index_file = tmp_path / "cim10_2026_index.faiss"
|
|
chunks_file = tmp_path / "cim10_2026_chunks.json"
|
|
|
|
assert index_file.exists()
|
|
assert chunks_file.exists()
|
|
|
|
# Vérifier que l'index peut être rechargé
|
|
loaded_index = faiss.read_index(str(index_file))
|
|
assert loaded_index.ntotal == 10
|
|
|
|
def test_build_index_chunks_json_valid(self, tmp_path):
|
|
"""Test que le fichier JSON des chunks est valide."""
|
|
manager = ReferentielsManager(data_dir=tmp_path)
|
|
|
|
chunks = [
|
|
Chunk(
|
|
chunk_id=f"ccam_2025_{i}",
|
|
referentiel_type="ccam",
|
|
referentiel_version="2025",
|
|
content=f"YYYY00{i} Acte {i}",
|
|
metadata={"section": "Section A"},
|
|
chunk_index=i,
|
|
)
|
|
for i in range(5)
|
|
]
|
|
|
|
manager.build_index(chunks)
|
|
|
|
# Charger et vérifier le JSON
|
|
chunks_file = tmp_path / "ccam_2025_chunks.json"
|
|
with open(chunks_file, "r", encoding="utf-8") as f:
|
|
loaded_chunks = json.load(f)
|
|
|
|
assert len(loaded_chunks) == 5
|
|
assert loaded_chunks[0]["chunk_id"] == "ccam_2025_0"
|
|
assert loaded_chunks[0]["referentiel_type"] == "ccam"
|
|
assert loaded_chunks[0]["content"] == "YYYY000 Acte 0"
|
|
|
|
def test_build_index_empty_chunks_raises_error(self):
|
|
"""Test qu'une liste vide de chunks lève une erreur."""
|
|
manager = ReferentielsManager()
|
|
|
|
with pytest.raises(ValueError, match="ne peut pas être vide"):
|
|
manager.build_index([])
|
|
|
|
def test_build_index_hash_consistency(self, tmp_path):
|
|
"""Test que le hash de l'index est cohérent."""
|
|
manager = ReferentielsManager(data_dir=tmp_path)
|
|
|
|
chunks = [
|
|
Chunk(
|
|
chunk_id=f"test_{i}",
|
|
referentiel_type="cim10",
|
|
referentiel_version="2026",
|
|
content=f"Content {i}",
|
|
metadata={},
|
|
chunk_index=i,
|
|
)
|
|
for i in range(10)
|
|
]
|
|
|
|
# Construire l'index deux fois
|
|
index1 = manager.build_index(chunks)
|
|
|
|
# Le hash devrait être le même (même paramètres, même nombre de vecteurs)
|
|
# Note: En pratique, le hash dépend des paramètres, pas du contenu exact
|
|
assert len(index1.index_hash) == 64
|
|
|
|
|
|
class TestAlphabeticalIndexVectorization:
|
|
"""Tests pour la vectorisation des index alphabétiques."""
|
|
|
|
def test_vectorize_alphabetical_index_cim10(self):
|
|
"""Test la vectorisation d'un index alphabétique CIM-10."""
|
|
manager = ReferentielsManager()
|
|
|
|
# Simuler un index alphabétique CIM-10
|
|
alpha_text = """A
|
|
Abcès - voir A00.1
|
|
Abdomen aigu - voir R10.0
|
|
|
|
B
|
|
Bronchite - voir J20.9
|
|
Bronchopneumonie - voir J18.0
|
|
|
|
C
|
|
Cataracte - voir H26.9
|
|
Cholécystite - voir K81.9
|
|
"""
|
|
|
|
chunks = manager.vectorize_alphabetical_indexes(
|
|
alpha_text, "cim10", "2026"
|
|
)
|
|
|
|
assert len(chunks) > 0
|
|
# Vérifier que les chunks ont le bon type
|
|
for chunk in chunks:
|
|
assert chunk.referentiel_type == "cim10"
|
|
assert chunk.referentiel_version == "2026"
|
|
assert chunk.metadata["chunk_type"] == "alphabetical_index"
|
|
assert "letter" in chunk.metadata
|
|
|
|
def test_vectorize_alphabetical_index_ccam(self):
|
|
"""Test la vectorisation d'un index alphabétique CCAM."""
|
|
manager = ReferentielsManager()
|
|
|
|
# Simuler un index alphabétique CCAM
|
|
alpha_text = """A
|
|
Ablation - YYYY001
|
|
Amputation - YYYY002
|
|
|
|
B
|
|
Biopsie - ZZZZ001
|
|
Bronchoscopie - ZZZZ002
|
|
"""
|
|
|
|
chunks = manager.vectorize_alphabetical_indexes(
|
|
alpha_text, "ccam", "2025"
|
|
)
|
|
|
|
assert len(chunks) > 0
|
|
for chunk in chunks:
|
|
assert chunk.referentiel_type == "ccam"
|
|
assert chunk.referentiel_version == "2025"
|
|
assert chunk.metadata["chunk_type"] == "alphabetical_index"
|
|
|
|
def test_alphabetical_index_extracts_codes(self):
|
|
"""Test que les codes sont extraits dans les métadonnées."""
|
|
manager = ReferentielsManager()
|
|
|
|
alpha_text = """Gastrite - voir K29.7
|
|
Gastro-entérite - voir A09
|
|
Glaucome - voir H40.9
|
|
"""
|
|
|
|
chunks = manager.vectorize_alphabetical_indexes(
|
|
alpha_text, "cim10", "2026"
|
|
)
|
|
|
|
# Au moins un chunk devrait contenir des codes dans les métadonnées
|
|
codes_found = False
|
|
for chunk in chunks:
|
|
if chunk.metadata.get("codes"):
|
|
codes_found = True
|
|
# Vérifier que les codes sont au format CIM-10
|
|
codes = chunk.metadata["codes"].split(",")
|
|
for code in codes:
|
|
assert len(code) >= 3 # Au moins A00
|
|
|
|
assert codes_found
|
|
|
|
def test_alphabetical_index_chunk_size(self):
|
|
"""Test que les chunks d'index alphabétique respectent la taille cible."""
|
|
manager = ReferentielsManager()
|
|
|
|
# Créer un long index alphabétique
|
|
entries = []
|
|
for i in range(100):
|
|
entries.append(f"Terme médical {i} - voir K{i:02d}.{i % 10}")
|
|
|
|
alpha_text = "\n".join(entries)
|
|
|
|
chunks = manager.vectorize_alphabetical_indexes(
|
|
alpha_text, "cim10", "2026"
|
|
)
|
|
|
|
# Vérifier que les chunks ne sont pas trop grands
|
|
for chunk in chunks:
|
|
assert len(chunk.content) <= 2048 # max_chunk_size
|
|
|
|
|
|
class TestIntegrationVectorizationIndexation:
|
|
"""Tests d'intégration pour le workflow complet."""
|
|
|
|
def test_full_workflow_import_chunk_vectorize_index(self, tmp_path):
|
|
"""Test le workflow complet: import → chunk → vectorize → index."""
|
|
# Créer un fichier PDF de test
|
|
pdf_path = tmp_path / "test_ref.pdf"
|
|
|
|
# Pour ce test, on va créer un fichier texte au lieu d'un PDF
|
|
# car créer un PDF valide est complexe
|
|
text_content = """CHAPITRE I - Maladies infectieuses
|
|
|
|
A00-A09 Maladies intestinales infectieuses
|
|
|
|
A00 Choléra
|
|
A00.0 Choléra dû à Vibrio cholerae 01, biovar cholerae
|
|
A00.1 Choléra dû à Vibrio cholerae 01, biovar El Tor
|
|
A00.9 Choléra, sans précision
|
|
|
|
Inclus: infection à Vibrio cholerae
|
|
Exclut: intoxication alimentaire (A05.-)
|
|
"""
|
|
|
|
# Sauvegarder le texte directement (simuler l'extraction PDF)
|
|
manager = ReferentielsManager(data_dir=tmp_path)
|
|
text_file = tmp_path / "cim10_2026_text.txt"
|
|
with open(text_file, "w", encoding="utf-8") as f:
|
|
f.write(text_content)
|
|
|
|
# Créer une version de référentiel manuellement
|
|
from pipeline_mco_pmsi.models.metadata import ReferentielVersion
|
|
from datetime import datetime
|
|
import hashlib
|
|
|
|
ref_version = ReferentielVersion(
|
|
type="cim10",
|
|
version="2026",
|
|
import_date=datetime.now(),
|
|
file_hash=hashlib.sha256(text_content.encode()).hexdigest(),
|
|
chunk_count=0,
|
|
index_hash="0" * 64,
|
|
)
|
|
|
|
# Chunker
|
|
chunks = manager.chunk_referentiel(ref_version)
|
|
assert len(chunks) > 0
|
|
|
|
# Construire l'index
|
|
vector_index = manager.build_index(chunks)
|
|
assert vector_index.num_vectors == len(chunks)
|
|
|
|
# Vérifier que les fichiers sont créés
|
|
assert (tmp_path / "cim10_2026_index.faiss").exists()
|
|
assert (tmp_path / "cim10_2026_chunks.json").exists()
|
|
|
|
|
|
class TestIndexSearch:
|
|
"""Tests pour la recherche dans l'index vectoriel."""
|
|
|
|
def test_index_search_basic(self, tmp_path):
|
|
"""Test une recherche basique dans l'index."""
|
|
manager = ReferentielsManager(data_dir=tmp_path)
|
|
|
|
# Créer des chunks avec du contenu médical
|
|
chunks = [
|
|
Chunk(
|
|
chunk_id=f"cim10_2026_{i}",
|
|
referentiel_type="cim10",
|
|
referentiel_version="2026",
|
|
content=content,
|
|
metadata={},
|
|
chunk_index=i,
|
|
)
|
|
for i, content in enumerate([
|
|
"K29.7 Gastrite, sans précision",
|
|
"K29.0 Gastrite hémorragique aiguë",
|
|
"J18.0 Bronchopneumonie, sans précision",
|
|
"I10 Hypertension essentielle",
|
|
"E11 Diabète sucré non insulino-dépendant",
|
|
])
|
|
]
|
|
|
|
# Construire l'index
|
|
manager.build_index(chunks)
|
|
|
|
# Charger l'index
|
|
index_file = tmp_path / "cim10_2026_index.faiss"
|
|
index = faiss.read_index(str(index_file))
|
|
|
|
# Effectuer une recherche
|
|
model = manager._load_embeddings_model()
|
|
query = "gastrite"
|
|
query_vector = model.encode(query, convert_to_numpy=True, normalize_embeddings=True)
|
|
query_vector = query_vector.reshape(1, -1).astype(np.float32)
|
|
|
|
# Rechercher les 2 plus proches voisins
|
|
distances, indices = index.search(query_vector, 2)
|
|
|
|
# Les résultats devraient inclure les chunks sur la gastrite
|
|
assert len(indices[0]) == 2
|
|
# Les indices devraient être 0 ou 1 (les deux chunks sur la gastrite)
|
|
assert indices[0][0] in [0, 1]
|
|
assert indices[0][1] in [0, 1]
|
|
|
|
def test_index_search_similarity_scores(self, tmp_path):
|
|
"""Test que les scores de similarité sont cohérents."""
|
|
manager = ReferentielsManager(data_dir=tmp_path)
|
|
|
|
chunks = [
|
|
Chunk(
|
|
chunk_id=f"test_{i}",
|
|
referentiel_type="cim10",
|
|
referentiel_version="2026",
|
|
content=content,
|
|
metadata={},
|
|
chunk_index=i,
|
|
)
|
|
for i, content in enumerate([
|
|
"Gastrite aiguë avec hémorragie",
|
|
"Gastrite chronique",
|
|
"Pneumonie bactérienne",
|
|
])
|
|
]
|
|
|
|
manager.build_index(chunks)
|
|
|
|
# Charger l'index
|
|
index_file = tmp_path / "cim10_2026_index.faiss"
|
|
index = faiss.read_index(str(index_file))
|
|
|
|
# Rechercher avec une requête proche du premier chunk
|
|
model = manager._load_embeddings_model()
|
|
query = "gastrite hémorragique"
|
|
query_vector = model.encode(query, convert_to_numpy=True, normalize_embeddings=True)
|
|
query_vector = query_vector.reshape(1, -1).astype(np.float32)
|
|
|
|
distances, indices = index.search(query_vector, 3)
|
|
|
|
# Le premier résultat devrait être le plus proche (distance la plus petite)
|
|
# Avec des embeddings normalisés, la distance est 1 - cosine_similarity
|
|
# Donc une distance plus petite = plus similaire
|
|
assert distances[0][0] <= distances[0][1]
|
|
assert distances[0][1] <= distances[0][2]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|