-
-
Notifications
You must be signed in to change notification settings - Fork 17
Expand file tree
/
Copy pathtest_vector_numpy.py
More file actions
110 lines (89 loc) · 4.19 KB
/
Copy pathtest_vector_numpy.py
File metadata and controls
110 lines (89 loc) · 4.19 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
from pathlib import Path
from engraphis.backends import DeterministicEmbedder, NumpyVectorIndex
from engraphis.backends.vector_sqlitevec import _cosine_from_l2
from engraphis.core.engine import MemoryEngine
from engraphis.core.interfaces import MemoryRecord, Scope
from engraphis.core.store import Store
from scripts.repair_embed_dim import repair
def test_search_ranks_relevant_memory_first():
store = Store(":memory:")
wid = store.get_or_create_workspace("w")
rid = store.get_or_create_repo(wid, "r")
emb = DeterministicEmbedder(dim=256)
index = NumpyVectorIndex(store)
texts = {
"pm": "We standardized on pnpm as the package manager for all frontend repos.",
"sky": "The afternoon sky over the harbor was a pale shade of blue.",
}
ids = {}
for tag, text in texts.items():
vec = emb.embed([text])[0]
ids[tag] = store.add_memory(MemoryRecord(id="", content=text, scope=Scope.REPO,
workspace_id=wid, repo_id=rid, embedding=vec))
hits = index.search(emb.embed(["which package manager do we use?"])[0], k=2)
assert hits[0][0] == ids["pm"]
assert hits[0][1] >= hits[1][1]
store.close()
def test_search_skips_vectors_from_other_embedding_dimensions():
store = Store(":memory:")
wid = store.get_or_create_workspace("w")
rid = store.get_or_create_repo(wid, "r")
emb = DeterministicEmbedder(dim=384)
index = NumpyVectorIndex(store)
matching = store.add_memory(MemoryRecord(
id="", content="matching vector", workspace_id=wid, repo_id=rid,
embedding=emb.embed(["matching vector"])[0]))
store.add_memory(MemoryRecord(
id="", content="legacy vector", workspace_id=wid, repo_id=rid,
embedding=DeterministicEmbedder(dim=256).embed(["legacy vector"])[0]))
hits = index.search(emb.embed(["matching vector"])[0], k=5)
assert [memory_id for memory_id, _score in hits] == [matching]
store.close()
def test_timeline_skips_legacy_dimension_without_losing_lexical_results():
engine = MemoryEngine.create(":memory:", embed_model=None, embed_dim=384)
wid = engine.store.get_or_create_workspace("w")
rid = engine.store.get_or_create_repo(wid, "r")
mid = engine.remember(
"durable migration fact", workspace_id=wid, repo_id=rid,
resolve_conflicts=False)
engine.store.put_vector(
mid, DeterministicEmbedder(dim=256).embed(["durable migration fact"])[0])
engine.store.conn.commit()
results = engine.timeline("durable migration", workspace_id=wid, repo_id=rid)
assert [record.id for record in results] == [mid]
engine.store.close()
def test_repair_uses_active_dimension_and_creates_backup(tmp_path):
db_path = tmp_path / "mixed.db"
store = Store(str(db_path))
wid = store.get_or_create_workspace("w")
rid = store.get_or_create_repo(wid, "r")
mid = store.add_memory(MemoryRecord(
id="", content="legacy vector", workspace_id=wid, repo_id=rid,
embedding=DeterministicEmbedder(dim=256).embed(["legacy vector"])[0]))
store.close()
result = repair(str(db_path), model_name="", dim=384)
assert result["repaired"] == 1
assert result["by_dim"] == {384: 1}
assert Path(result["backup"]).is_file()
repaired = Store(str(db_path))
row = repaired.conn.execute(
"SELECT dim, model FROM mem_vectors WHERE id=?", (mid,)).fetchone()
assert (row["dim"], row["model"]) == (384, "deterministic")
repaired.close()
def test_delete_removes_from_index():
store = Store(":memory:")
wid = store.get_or_create_workspace("w")
rid = store.get_or_create_repo(wid, "r")
emb = DeterministicEmbedder(dim=128)
index = NumpyVectorIndex(store)
vec = emb.embed(["hello world"])[0]
mid = store.add_memory(MemoryRecord(id="", content="hello world", workspace_id=wid,
repo_id=rid, embedding=vec))
assert index.search(vec, k=1)[0][0] == mid
index.delete([mid])
assert index.search(vec, k=1) == []
store.close()
def test_sqlitevec_l2_distance_converts_to_cosine_similarity():
assert _cosine_from_l2(0.0) == 1.0
assert abs(_cosine_from_l2(2 ** 0.5)) < 1e-12
assert _cosine_from_l2(2.0) == -1.0