108 lines
3.6 KiB
Python
108 lines
3.6 KiB
Python
|
from pathlib import Path
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from gewu.adapters.sqlite_vectors import SQLiteVectorIndex, VectorIndexError
|
||
|
|
|
||
|
|
|
||
|
|
def _vectors() -> tuple[tuple[str, str, tuple[float, ...]], ...]:
|
||
|
|
return (
|
||
|
|
("dome-performance", "dome:closest", (1.0, 0.0, 0.0)),
|
||
|
|
("interactive-video", "video:closest", (0.0, 1.0, 0.0)),
|
||
|
|
("shared", "shared:closest", (0.9701425001, 0.242535625, 0.0)),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _index(path: Path) -> SQLiteVectorIndex:
|
||
|
|
pytest.importorskip("sqlite_vec")
|
||
|
|
return SQLiteVectorIndex(path)
|
||
|
|
|
||
|
|
|
||
|
|
def _build(index: SQLiteVectorIndex, generation_id: str, vectors=None) -> None:
|
||
|
|
index.build_generation(
|
||
|
|
generation_id=generation_id,
|
||
|
|
revision="a" * 40,
|
||
|
|
model_id="test-model@revision",
|
||
|
|
artifact_sha256="b" * 64,
|
||
|
|
dimensions=3,
|
||
|
|
vectors=_vectors() if vectors is None else vectors,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_vector_generation_is_scoped_and_switches_atomically(tmp_path: Path) -> None:
|
||
|
|
index = _index(tmp_path / "vectors.sqlite3")
|
||
|
|
_build(index, "generation-1")
|
||
|
|
|
||
|
|
assert index.activate("generation-1", expected_active=None) is None
|
||
|
|
active = index.active_generation()
|
||
|
|
assert active is not None
|
||
|
|
assert active.generation_id == "generation-1"
|
||
|
|
assert active.revision == "a" * 40
|
||
|
|
assert active.model_id == "test-model@revision"
|
||
|
|
assert active.artifact_sha256 == "b" * 64
|
||
|
|
assert active.dimensions == 3
|
||
|
|
assert active.node_count == 3
|
||
|
|
scoped = index.search(
|
||
|
|
query=(1.0, 0.0, 0.0),
|
||
|
|
project_ids=("dome-performance",),
|
||
|
|
include_shared=False,
|
||
|
|
limit=10,
|
||
|
|
)
|
||
|
|
assert [hit.node_id for hit in scoped] == ["dome:closest"]
|
||
|
|
with_shared = index.search(
|
||
|
|
query=(1.0, 0.0, 0.0),
|
||
|
|
project_ids=("dome-performance",),
|
||
|
|
include_shared=True,
|
||
|
|
limit=10,
|
||
|
|
)
|
||
|
|
assert {hit.node_id for hit in with_shared} == {"dome:closest", "shared:closest"}
|
||
|
|
|
||
|
|
_build(
|
||
|
|
index, "generation-2", vectors=(("dome-performance", "dome:replacement", (1.0, 0.0, 0.0)),)
|
||
|
|
)
|
||
|
|
with pytest.raises(VectorIndexError, match="changed"):
|
||
|
|
index.activate("generation-2", expected_active=None)
|
||
|
|
active = index.active_generation()
|
||
|
|
assert active is not None and active.generation_id == "generation-1"
|
||
|
|
assert index.activate("generation-2", expected_active="generation-1") == "generation-1"
|
||
|
|
assert [
|
||
|
|
hit.node_id
|
||
|
|
for hit in index.search(
|
||
|
|
query=(1.0, 0.0, 0.0),
|
||
|
|
project_ids=("dome-performance",),
|
||
|
|
include_shared=False,
|
||
|
|
limit=10,
|
||
|
|
)
|
||
|
|
] == ["dome:replacement"]
|
||
|
|
|
||
|
|
assert index.activate("generation-1", expected_active="generation-2") == "generation-2"
|
||
|
|
assert [
|
||
|
|
hit.node_id
|
||
|
|
for hit in index.search(
|
||
|
|
query=(1.0, 0.0, 0.0),
|
||
|
|
project_ids=("dome-performance",),
|
||
|
|
include_shared=False,
|
||
|
|
limit=10,
|
||
|
|
)
|
||
|
|
] == ["dome:closest"]
|
||
|
|
with pytest.raises(VectorIndexError, match="Active vector generation"):
|
||
|
|
index.delete_generation("generation-1")
|
||
|
|
index.delete_generation("generation-2")
|
||
|
|
_build(index, "generation-2")
|
||
|
|
|
||
|
|
|
||
|
|
def test_vector_generation_rejects_invalid_vectors_before_publish(tmp_path: Path) -> None:
|
||
|
|
index = _index(tmp_path / "vectors.sqlite3")
|
||
|
|
with pytest.raises(VectorIndexError, match="normalized"):
|
||
|
|
_build(index, "bad", vectors=(("dome-performance", "dome:bad", (2.0, 0.0, 0.0)),))
|
||
|
|
|
||
|
|
assert index.active_generation() is None
|
||
|
|
assert (
|
||
|
|
index.search(
|
||
|
|
query=(1.0, 0.0, 0.0),
|
||
|
|
project_ids=("dome-performance",),
|
||
|
|
include_shared=False,
|
||
|
|
limit=10,
|
||
|
|
)
|
||
|
|
== ()
|
||
|
|
)
|