feat(rag): 新增 RagStore(SQLite 向量存储 + 余弦检索,线程安全)
This commit is contained in:
@@ -0,0 +1,44 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user