Spaces:
Paused
Paused
File size: 1,793 Bytes
8c1b9fe | 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 | from __future__ import annotations
from auralynq.embeddings import get_embedder
from auralynq.retrieval.models import Filter
from auralynq.vectorstore.memory_store import MemoryStore
def _populate(tmp_path, chunks):
emb = get_embedder()
batch = emb.embed([c.text for c in chunks])
store = MemoryStore(path=tmp_path / "ms")
store.upsert(chunks, batch)
return store, emb
def test_upsert_and_count(tmp_path, sample_chunks):
store, _ = _populate(tmp_path, sample_chunks)
assert store.count() == len(sample_chunks)
def test_hybrid_search_returns_relevant(tmp_path, sample_chunks):
store, emb = _populate(tmp_path, sample_chunks)
q = emb.embed_query("flow based pruning of relational paths")
res = store.search(q, k=3)
assert res
assert "pruning" in res[0].chunk.text.lower() or "path" in res[0].chunk.text.lower()
assert res[0].method == "hybrid"
def test_dense_and_sparse_primitives(tmp_path, sample_chunks):
store, emb = _populate(tmp_path, sample_chunks)
q = emb.embed_query("reciprocal rank fusion")
dense = store.search_dense(q.dense, 3)
sparse = store.search_sparse(q.sparse, 3)
assert dense and dense[0].method == "dense"
assert sparse and sparse[0].method == "sparse"
def test_metadata_filter(tmp_path, sample_chunks):
store, emb = _populate(tmp_path, sample_chunks)
q = emb.embed_query("anything")
res = store.search_dense(q.dense, 10, filt=Filter(doc_ids=["nope"]))
assert res == []
def test_persistence_roundtrip(tmp_path, sample_chunks):
store, emb = _populate(tmp_path, sample_chunks)
store.save()
reloaded = MemoryStore(path=tmp_path / "ms")
assert reloaded.count() == len(sample_chunks)
q = emb.embed_query("paris france")
assert reloaded.search(q, k=1)
|