45 lines
1.1 KiB
Python
45 lines
1.1 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 == ["订单模块处理创建"]
|
|
|
|
|
|
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) == []
|
|
|
|
|
|
def test_concurrent_add_and_search_no_crash():
|
|
import threading
|
|
s = RagStore(":memory:")
|
|
s.reset_scope("p1")
|
|
|
|
def worker(i):
|
|
s.add("p1", [f"chunk-{i}"], [_vec(i % 8)])
|
|
|
|
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)
|