jurix-api / tests /test_scripts /test_regenerate_embeddings.py
rahim237cm's picture
Claude Opus 5.5
feat(embeddings): voyage-4-large remplace EmbeddingGemma (API Voyage, halfvec 1024)
57c8117
Raw History Blame Contribute Delete
18.1 kB
"""
Tests de scripts/regenerate_embeddings.py : ce qui est a refaire, et la trace
laissee.
La selection repose sur la PROVENANCE des vecteurs (articles.embedding_model),
plus seulement sur leur absence. C'est ce qui rend un changement de fournisseur
sur : chaque vecteur de l'ancien modele est reconnu et refait, au lieu de
cohabiter en silence avec ceux du nouveau — deux espaces dont les cosinus
croises ne veulent rien dire.
Usage:
pytest tests/test_scripts/test_regenerate_embeddings.py -v
"""
import json
import numpy as np
import pytest
from sqlalchemy import text
from app.models.law import Article, Law
from app.services.embedding_service import EmbeddingService, VoyageProvider
from scripts import regenerate_embeddings as regen
EMPREINTE = "doublure|actuelle"
DIM = EmbeddingService.EMBEDDING_DIM
def _vecteur(graine: int) -> np.ndarray:
v = np.random.default_rng(graine).normal(size=DIM)
return (v / np.linalg.norm(v)).astype(np.float32)
@pytest.fixture
def base(sync_db_session):
"""Un article dans chacun des etats possibles."""
session = sync_db_session
session.add(Law(id=1, reference="LOI-REGEN", title="Loi", content="Contenu.",
type="loi", language="fr", status="published"))
session.flush()
etats = {
1: (None, None, True), # jamais vectorise
2: (_vecteur(2), "gemini|ancien", True), # autre fournisseur
3: (_vecteur(3), None, True), # origine inconnue
4: (_vecteur(4), EMPREINTE, True), # a jour
5: (None, None, False), # ecarte par le raffineur
}
for article_id, (vecteur, empreinte, embed) in etats.items():
session.add(Article(
id=article_id, law_id=1, number=str(article_id),
content=f"Article {article_id}, contenu de test.", order=article_id,
embedding=None if vecteur is None else vecteur.tolist(),
embedding_model=empreinte, embed=embed,
))
session.commit()
return session
def _ids(lot):
return [r["id"] for r in lot]
class TestSelection:
def test_refait_ce_qui_manque_ou_vient_d_ailleurs(self, base):
"""
L'article 3 est le cas qui justifie IS DISTINCT FROM : avec `<>`, son
empreinte NULL rendrait la condition NULL, et ce vecteur d'origine
inconnue passerait pour bon.
"""
lot = regen._fetch_batch(base, None, False, EMPREINTE, 0, 100)
assert _ids(lot) == [1, 2, 3]
def test_le_compte_annonce_suit_la_meme_selection(self, base):
assert regen._count_remaining(base, None, False, EMPREINTE) == 3
def test_force_reprend_tout_ce_qui_est_vectorisable(self, base):
lot = regen._fetch_batch(base, None, True, EMPREINTE, 0, 100)
assert _ids(lot) == [1, 2, 3, 4]
assert regen._count_remaining(base, None, True, EMPREINTE) == 4
def test_un_autre_fournisseur_voit_tout_a_refaire(self, base):
"""Le changement de fournisseur, vu du script : rien n'est a jour."""
assert regen._count_remaining(base, None, False, "voyage|voyage-4-large") == 4
class TestEcriture:
def test_le_vecteur_part_avec_son_empreinte(self, base):
vecteurs = [_vecteur(10 + i) for i in range(3)]
regen._write_batch(base, [1, 2, 3], vecteurs, EMPREINTE)
base.commit()
rows = base.execute(text(
"SELECT id, embedding_model FROM articles WHERE embed ORDER BY id"
)).all()
assert [(r.id, r.embedding_model) for r in rows] == [
(i, EMPREINTE) for i in (1, 2, 3, 4)
]
relu = np.array(base.get(Article, 1).embedding, dtype=np.float32)
# halfvec = fp16, environ 3 chiffres significatifs : ecart maximal
# mesure de 3e-5 sur un vecteur unitaire de 1024. 1e-6 echouerait.
assert np.allclose(relu, vecteurs[0], atol=1e-4)
assert float(relu @ vecteurs[0]) > 0.9999
assert regen._count_remaining(base, None, False, EMPREINTE) == 0
def test_un_vecteur_d_une_autre_dimension_est_refuse_par_la_colonne(self, base):
"""Le typmod de la colonne garde la dimension, au caractere pres."""
court = np.ones(768, dtype=np.float32) / np.sqrt(768)
with pytest.raises(Exception, match="expected 1024 dimensions"):
regen._write_batch(base, [1], [court], EMPREINTE)
base.commit()
base.rollback()
class _Fournisseur:
label = "doublure"
fingerprint = EMPREINTE
inter_batch_delay_s = 0.0
class _Service:
"""Doublure d'EmbeddingService : aucun modele, aucun reseau."""
MAX_TEXT_LENGTH = EmbeddingService.MAX_TEXT_LENGTH
EMBEDDING_DIM = DIM
encodes = []
def __init__(self, *args, **kwargs):
self.provider = _Fournisseur()
def generate_batch_embeddings(self, texts, batch_size=None, normalize=True):
_Service.encodes.extend(texts)
return [_vecteur(len(t)) for t in texts]
class TestDeBoutEnBout:
@pytest.fixture(autouse=True)
def _doublure(self, monkeypatch):
_Service.encodes = []
monkeypatch.setattr(regen, "EmbeddingService", _Service)
def test_une_passe_suffit_et_la_suivante_ne_trouve_rien(self, base):
assert regen.main(["--all", "--batch-size", "2"]) == 0
assert len(_Service.encodes) == 3
assert regen._count_remaining(base, None, False, EMPREINTE) == 0
_Service.encodes = []
assert regen.main(["--all"]) == 0
assert _Service.encodes == []
def test_le_compte_a_blanc_n_ecrit_rien(self, base):
assert regen.main(["--all", "--dry-run"]) == 0
assert _Service.encodes == []
assert regen._count_remaining(base, None, False, EMPREINTE) == 3
def test_une_cle_refusee_arrete_tout_au_premier_lot(self, base, monkeypatch):
"""
Sans cet arret, une cle revoquee envoyait des milliers de requetes
vouees au meme refus, un lot apres l'autre.
"""
from app.services.embedding_service import EmbeddingDefinitiveError
appels = []
def _refus(self, texts, **kwargs):
appels.append(len(texts))
raise EmbeddingDefinitiveError("Voyage HTTP 401 : Provided API key is invalid.", code=401)
monkeypatch.setattr(_Service, "generate_batch_embeddings", _refus)
assert regen.main(["--all", "--batch-size", "1"]) == 2
assert len(appels) == 1
def test_un_arret_apres_des_lots_commites_vide_quand_meme_le_cache(self, base, monkeypatch):
from app.services.embedding_service import EmbeddingDefinitiveError
appels = []
def _refus_au_second(self, texts, **kwargs):
appels.append(len(texts))
if len(appels) == 2:
raise EmbeddingDefinitiveError("Voyage HTTP 401 : revoquee", code=401)
return [_vecteur(len(t)) for t in texts]
monkeypatch.setattr(_Service, "generate_batch_embeddings", _refus_au_second)
base.execute(text(
"INSERT INTO query_cache (cache_key, response_json, expires_at) "
"VALUES ('k', '{}', now() + interval '1 hour')"
))
base.commit()
assert regen.main(["--all", "--batch-size", "1"]) == 2
assert base.execute(text("SELECT count(*) FROM query_cache")).scalar() == 0
def test_par_defaut_les_lots_sont_ceux_du_fournisseur(self, base, monkeypatch):
"""Des lots de 16 enverraient environ 2 000 requetes a Voyage pour le corpus."""
tailles = []
def _compter(self, texts, batch_size=None, **kwargs):
tailles.append((len(texts), batch_size))
return [_vecteur(len(t)) for t in texts]
monkeypatch.setattr(_Fournisseur, "TEXTES_PAR_REQUETE", 2, raising=False)
monkeypatch.setattr(_Service, "generate_batch_embeddings", _compter)
assert regen.main(["--all"]) == 0
assert tailles == [(2, 2), (1, 2)]
def test_une_colonne_d_une_autre_dimension_arrete_avant_tout_appel(self, base, monkeypatch):
"""Sur une base non migree, chaque lot serait paye a Voyage, puis refuse."""
monkeypatch.setattr(_Service, "EMBEDDING_DIM", 768)
assert regen.main(["--all"]) == 2
assert _Service.encodes == []
def test_un_lot_refuse_pour_son_texte_est_saute(self, base, monkeypatch):
"""Un 400 vise un texte : les autres lots continuent."""
from app.services.embedding_service import EmbeddingDefinitiveError
appels = []
def _refus_du_premier(self, texts, **kwargs):
appels.append(len(texts))
if len(appels) == 1:
raise EmbeddingDefinitiveError("Voyage HTTP 400 : bad input", code=400)
return [_vecteur(len(t)) for t in texts]
monkeypatch.setattr(_Service, "generate_batch_embeddings", _refus_du_premier)
assert regen.main(["--all", "--batch-size", "1"]) == 1
assert len(appels) == 3
assert regen._count_remaining(base, None, False, EMPREINTE) == 1
def test_une_passe_vide_le_cache_des_recherches(self, base):
"""Les reponses gardees pendant que les vecteurs manquaient etaient du plein texte seul."""
base.execute(text(
"INSERT INTO query_cache (cache_key, response_json, expires_at) "
"VALUES ('k', '{}', now() + interval '1 hour')"
))
base.commit()
assert regen.main(["--all"]) == 0
assert base.execute(text("SELECT count(*) FROM query_cache")).scalar() == 0
def _index_embedding(session):
return session.execute(text(
"SELECT indexrelid::regclass::text, indisvalid FROM pg_index "
"WHERE indrelid = 'articles'::regclass "
"AND indexrelid::regclass::text LIKE 'idx_articles_embedding_hnsw%'"
)).all()
class TestReconstructionDeLIndex:
def test_le_nom_est_celui_de_la_migration(self):
assert regen.INDEX_EMBEDDING == "idx_articles_embedding_hnsw_halfvec1024"
def test_index_reconstruit_et_valide(self, base):
regen.reindex(base)
assert _index_embedding(base) == [(regen.INDEX_EMBEDDING, True)]
def test_echec_ne_laisse_pas_de_copie_invalide(self, base, monkeypatch):
"""
Un REINDEX CONCURRENTLY interrompu laisse une copie _ccnew invalide,
maintenue a chaque ecriture. On la simule, puis on fait echouer.
"""
with regen.sync_engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
conn.execute(text(
f"CREATE INDEX {regen.INDEX_EMBEDDING}_ccnew ON articles (id)"
))
monkeypatch.setattr(
regen, "text",
lambda sql: text(sql.replace("REINDEX INDEX CONCURRENTLY", "REINDEX INDEX CONCURRENTLY inexistant_")
if sql.startswith("REINDEX") else sql),
)
with pytest.raises(Exception):
regen.reindex(base)
assert _index_embedding(base) == [(regen.INDEX_EMBEDDING, True)]
# ==================== CHARGEMENT DES LOTS DEJA CALCULES ====================
def _texte(article_id):
"""Le texte vectorise de l'article : coalesce(embed_text, content)."""
return f"Article {article_id}, contenu de test."
def _instantane(dossier, textes):
chemin = dossier / "morceaux.jsonl"
with open(chemin, "w", encoding="utf-8") as f:
for article_id, texte in textes.items():
f.write(json.dumps({"id": article_id, "law_id": 1, "texte": texte}, ensure_ascii=False) + "\n")
return chemin
def _lots(dossier, vecteurs, nom="lot"):
dossier.mkdir(parents=True, exist_ok=True)
ids = sorted(vecteurs)
np.savez(dossier / f"{nom}.npz", ids=np.array(ids, dtype=np.int64),
v=np.array([vecteurs[i] for i in ids], dtype=np.float32))
return dossier
@pytest.fixture
def lots(base, tmp_path):
"""
Un cas par motif : 1 a charger ; 2 vecteur non unitaire ; 3 texte modifie
depuis l'instantane ; 4 deja a jour ; 5 non vectorisable ; 99 absent de la
base ; 77 absent de l'instantane.
"""
textes = {i: _texte(i) for i in (1, 2, 4, 5, 99)}
textes[3] = "Ancienne version de l'article 3."
vecteurs = {i: _vecteur(100 + i) for i in (1, 3, 4, 5, 77, 99)}
vecteurs[2] = _vecteur(102) * 0.5
dossier = _lots(tmp_path / "voyage-4-large", vecteurs)
return dossier, _instantane(tmp_path, textes), vecteurs
class TestChargementDesLots:
def test_chaque_motif_est_compte_et_seul_le_bon_est_charge(self, base, lots):
dossier, instantane, vecteurs = lots
rapport = regen.charger_les_lots(base, dossier, instantane, EMPREINTE, DIM)
assert dict(rapport) == {
"charges": 1,
"vecteur_invalide": 1,
"texte_modifie_depuis_l_instantane": 1,
"deja_a_jour": 1,
"non_vectorisable": 1,
"absent_de_la_base": 1,
"absent_de_l_instantane": 1,
}
etat = {r.id: r.embedding_model for r in base.execute(text(
"SELECT id, embedding_model FROM articles ORDER BY id"
))}
assert etat == {1: EMPREINTE, 2: "gemini|ancien", 3: None, 4: EMPREINTE, 5: None}
relu = np.array(base.get(Article, 1).embedding, dtype=np.float32)
assert np.allclose(relu, vecteurs[1], atol=1e-4)
def test_le_compte_a_blanc_n_ecrit_rien(self, base, lots):
dossier, instantane, _ = lots
rapport = regen.charger_les_lots(base, dossier, instantane, EMPREINTE, DIM, dry_run=True)
assert rapport["a_charger"] == 1
assert rapport["a_charger_qui_manquaient"] == 1
assert rapport["charges"] == 0
assert base.execute(text("SELECT embedding IS NULL FROM articles WHERE id = 1")).scalar()
def test_le_compte_a_blanc_avec_force_distingue_ce_qui_manquait(self, base, lots):
"""Avec --force, l'article 4, deja a jour, est recharge mais ne manquait pas."""
dossier, instantane, _ = lots
rapport = regen.charger_les_lots(
base, dossier, instantane, EMPREINTE, DIM, force=True, dry_run=True,
)
assert rapport["a_charger"] == 2
assert rapport["a_charger_qui_manquaient"] == 1
def test_force_reecrit_aussi_ce_qui_est_a_jour(self, base, lots):
dossier, instantane, vecteurs = lots
rapport = regen.charger_les_lots(base, dossier, instantane, EMPREINTE, DIM, force=True)
assert rapport["charges"] == 2
relu = np.array(base.get(Article, 4).embedding, dtype=np.float32)
assert np.allclose(relu, vecteurs[4], atol=1e-4)
def test_une_reprise_ne_recharge_rien(self, base, lots):
dossier, instantane, _ = lots
regen.charger_les_lots(base, dossier, instantane, EMPREINTE, DIM)
rapport = regen.charger_les_lots(base, dossier, instantane, EMPREINTE, DIM)
assert rapport["charges"] == 0
assert rapport["deja_a_jour"] == 2
def test_un_id_present_dans_deux_lots_est_ecarte(self, base, tmp_path):
dossier = tmp_path / "voyage-4-large"
_lots(dossier, {1: _vecteur(1)}, nom="a")
_lots(dossier, {1: _vecteur(2)}, nom="b")
instantane = _instantane(tmp_path, {1: _texte(1)})
rapport = regen.charger_les_lots(base, dossier, instantane, EMPREINTE, DIM)
assert rapport["id_en_double_dans_les_lots"] == 1
assert rapport["charges"] == 0
def test_un_lot_d_une_autre_dimension_arrete_tout(self, base, tmp_path):
dossier = tmp_path / "voyage-4-large"
dossier.mkdir()
np.savez(dossier / "lot.npz", ids=np.array([1], dtype=np.int64),
v=np.ones((1, 768), dtype=np.float32))
instantane = _instantane(tmp_path, {1: _texte(1)})
with pytest.raises(SystemExit, match="768"):
regen.charger_les_lots(base, dossier, instantane, EMPREINTE, DIM)
class TestChargementParLaLigneDeCommande:
def test_l_empreinte_ecrite_est_celle_du_fournisseur(self, base, lots):
"""
Calculee par le code, jamais recopiee : un caractere d'ecart, et les
vecteurs charges passeraient pour etrangers.
"""
dossier, instantane, _ = lots
assert regen.main(["--depuis-lots", str(dossier), "--instantane", str(instantane)]) == 0
empreinte = base.execute(text("SELECT embedding_model FROM articles WHERE id = 1")).scalar()
assert empreinte == VoyageProvider(api_key="factice").fingerprint == "voyage|voyage-4-large"
def test_le_chargement_vide_le_cache_des_recherches(self, base, lots):
dossier, instantane, _ = lots
base.execute(text(
"INSERT INTO query_cache (cache_key, response_json, expires_at) "
"VALUES ('k', '{}', now() + interval '1 hour')"
))
base.commit()
assert regen.main(["--depuis-lots", str(dossier), "--instantane", str(instantane)]) == 0
assert base.execute(text("SELECT count(*) FROM query_cache")).scalar() == 0
def test_l_instantane_est_exige(self, base, lots):
dossier, _, _ = lots
assert regen.main(["--depuis-lots", str(dossier)]) == 2
def test_refuse_si_le_fournisseur_n_est_pas_voyage(self, base, lots, monkeypatch):
dossier, instantane, _ = lots
monkeypatch.setattr(regen.settings, "EMBEDDING_PROVIDER", "gemini")
assert regen.main(["--depuis-lots", str(dossier), "--instantane", str(instantane)]) == 2
def test_refuse_un_dossier_d_un_autre_modele(self, base, lots, tmp_path):
"""Des vecteurs voyage-context-4 ne doivent pas passer pour voyage-4-large."""
dossier, instantane, _ = lots
autre = dossier.rename(tmp_path / "voyage-context-4")
assert regen.main(["--depuis-lots", str(autre), "--instantane", str(instantane)]) == 2
assert base.execute(text("SELECT embedding IS NULL FROM articles WHERE id = 1")).scalar()
def test_instantane_sans_depuis_lots_refuse(self, base, tmp_path):
instantane = _instantane(tmp_path, {1: _texte(1)})
assert regen.main(["--all", "--instantane", str(instantane)]) == 2