62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
import math
|
|
from genesis.rag.store import RagStore
|
|
|
|
|
|
def _vec(*ones):
|
|
v = [0.0] * 8
|
|
for i in ones:
|
|
v[i] = 1.0
|
|
n = math.sqrt(sum(x * x for x in v))
|
|
return [x / n for x in v]
|
|
|
|
|
|
def test_search_returns_most_similar_chunk():
|
|
s = RagStore(":memory:")
|
|
s.reset_scope("p1")
|
|
s.add("p1", ["订单模块处理创建", "用户认证登录"], [_vec(0, 1), _vec(4, 5)])
|
|
res = s.search("p1", _vec(0, 1), k=1)
|
|
assert res == ["订单模块处理创建"]
|
|
# k<=0 健壮性
|
|
assert s.search("p1", _vec(0, 1), k=0) == []
|
|
s.close()
|
|
|
|
|
|
def test_reset_scope_clears():
|
|
s = RagStore(":memory:")
|
|
s.reset_scope("p1")
|
|
s.add("p1", ["a"], [_vec(0)])
|
|
s.reset_scope("p1")
|
|
assert s.search("p1", _vec(0), k=3) == []
|
|
# chunks 与 embeddings 长度不一致应抛 ValueError
|
|
try:
|
|
import pytest
|
|
with pytest.raises(ValueError):
|
|
s.add("p1", ["a", "b"], [_vec(0)])
|
|
finally:
|
|
s.close()
|
|
|
|
|
|
def test_concurrent_add_and_search_no_crash():
|
|
import threading
|
|
s = RagStore(":memory:")
|
|
s.reset_scope("p1")
|
|
|
|
def worker(i):
|
|
# 每个 worker 线程内交替执行 add 与 search，让读写真正并发跨线程运行
|
|
for j in range(5):
|
|
s.add("p1", [f"chunk-{i}-{j}"], [_vec(i % 8)])
|
|
res = s.search("p1", _vec(i % 8), k=3)
|
|
assert all(isinstance(r, str) for r in res)
|
|
|
|
try:
|
|
threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
res = s.search("p1", _vec(0), k=3)
|
|
assert len(res) <= 3
|
|
assert all(isinstance(r, str) for r in res)
|
|
finally:
|
|
s.close()
|