"""Tests for the vector index.""" from __future__ import annotations from pathlib import Path import pytest from mathvision_explorer.index import VectorIndex def test_vector_index_searches_by_cosine_similarity() -> None: """Nearest-neighbor search ranks vectors by cosine similarity.""" index = VectorIndex() index.add("same", (1.0, 0.0)) index.add("opposite", (-1.0, 0.0)) index.add("orthogonal", (0.0, 1.0)) neighbors = index.search((1.0, 0.0), limit=2) assert [neighbor.item_id for neighbor in neighbors] == ["same", "orthogonal"] assert neighbors[0].score == pytest.approx(1.0) def test_vector_index_round_trips_tsv(tmp_path: Path) -> None: """Saved indexes can be loaded back from disk.""" path = tmp_path / "index.tsv" index = VectorIndex() index.add("a", (0.25, 0.75)) index.save_tsv(path) loaded = VectorIndex.load_tsv(path) assert len(loaded) == 1 assert loaded.search((0.25, 0.75))[0].item_id == "a" def test_vector_dimensions_must_match() -> None: """Search vectors must use the same dimensionality as stored vectors.""" index = VectorIndex() index.add("a", (1.0, 2.0)) with pytest.raises(ValueError, match="same dimensions"): index.search((1.0,))