fix(rag): 强化 RagStore 并发读写验证 + close/健壮性

This commit is contained in:
lhl
2026-08-29 22:45:17 +08:00
parent 2eab0183ee
commit 69a7ba7226
2 changed files with 36 additions and 10 deletions
+26 -9
View File
@@ -16,6 +16,9 @@ def test_search_returns_most_similar_chunk():
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():
@@ -24,6 +27,13 @@ def test_reset_scope_clears():
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():
@@ -32,13 +42,20 @@ def test_concurrent_add_and_search_no_crash():
s.reset_scope("p1")
def worker(i):
s.add("p1", [f"chunk-{i}"], [_vec(i % 8)])
# 每个 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)
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)
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()