fix(rag): 强化 RagStore 并发读写验证 + close/健壮性
This commit is contained in:
@@ -31,7 +31,9 @@ class RagStore:
|
|||||||
self.conn.commit()
|
self.conn.commit()
|
||||||
|
|
||||||
def add(self, scope: str, chunks: List[str], embeddings: List[List[float]]) -> None:
|
def add(self, scope: str, chunks: List[str], embeddings: List[List[float]]) -> None:
|
||||||
# 先清后写(增量索引),整段加锁串行化
|
# 按 scope 全量替换(先清后写),非追加;整段加锁串行化
|
||||||
|
if len(chunks) != len(embeddings):
|
||||||
|
raise ValueError("chunks 与 embeddings 长度不一致")
|
||||||
with self._lock:
|
with self._lock:
|
||||||
self.conn.execute("DELETE FROM rag_chunks WHERE scope=?", (scope,))
|
self.conn.execute("DELETE FROM rag_chunks WHERE scope=?", (scope,))
|
||||||
self.conn.executemany(
|
self.conn.executemany(
|
||||||
@@ -41,6 +43,10 @@ class RagStore:
|
|||||||
self.conn.commit()
|
self.conn.commit()
|
||||||
|
|
||||||
def search(self, scope: str, query_vec: List[float], k: int = 5) -> List[str]:
|
def search(self, scope: str, query_vec: List[float], k: int = 5) -> List[str]:
|
||||||
|
k = max(0, int(k))
|
||||||
|
if k <= 0:
|
||||||
|
return []
|
||||||
|
# 读也加锁,保证跨线程读写串行化
|
||||||
with self._lock:
|
with self._lock:
|
||||||
rows = self.conn.execute(
|
rows = self.conn.execute(
|
||||||
"SELECT chunk, embedding FROM rag_chunks WHERE scope=?", (scope,)
|
"SELECT chunk, embedding FROM rag_chunks WHERE scope=?", (scope,)
|
||||||
@@ -50,3 +56,6 @@ class RagStore:
|
|||||||
scored.append((_cosine(query_vec, json.loads(emb)), chunk))
|
scored.append((_cosine(query_vec, json.loads(emb)), chunk))
|
||||||
scored.sort(key=lambda x: x[0], reverse=True)
|
scored.sort(key=lambda x: x[0], reverse=True)
|
||||||
return [c for _, c in scored[:k]]
|
return [c for _, c in scored[:k]]
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
self.conn.close()
|
||||||
|
|||||||
+26
-9
@@ -16,6 +16,9 @@ def test_search_returns_most_similar_chunk():
|
|||||||
s.add("p1", ["订单模块处理创建", "用户认证登录"], [_vec(0, 1), _vec(4, 5)])
|
s.add("p1", ["订单模块处理创建", "用户认证登录"], [_vec(0, 1), _vec(4, 5)])
|
||||||
res = s.search("p1", _vec(0, 1), k=1)
|
res = s.search("p1", _vec(0, 1), k=1)
|
||||||
assert res == ["订单模块处理创建"]
|
assert res == ["订单模块处理创建"]
|
||||||
|
# k<=0 健壮性
|
||||||
|
assert s.search("p1", _vec(0, 1), k=0) == []
|
||||||
|
s.close()
|
||||||
|
|
||||||
|
|
||||||
def test_reset_scope_clears():
|
def test_reset_scope_clears():
|
||||||
@@ -24,6 +27,13 @@ def test_reset_scope_clears():
|
|||||||
s.add("p1", ["a"], [_vec(0)])
|
s.add("p1", ["a"], [_vec(0)])
|
||||||
s.reset_scope("p1")
|
s.reset_scope("p1")
|
||||||
assert s.search("p1", _vec(0), k=3) == []
|
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():
|
def test_concurrent_add_and_search_no_crash():
|
||||||
@@ -32,13 +42,20 @@ def test_concurrent_add_and_search_no_crash():
|
|||||||
s.reset_scope("p1")
|
s.reset_scope("p1")
|
||||||
|
|
||||||
def worker(i):
|
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)]
|
try:
|
||||||
for t in threads:
|
threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)]
|
||||||
t.start()
|
for t in threads:
|
||||||
for t in threads:
|
t.start()
|
||||||
t.join()
|
for t in threads:
|
||||||
res = s.search("p1", _vec(0), k=3)
|
t.join()
|
||||||
assert len(res) <= 3
|
res = s.search("p1", _vec(0), k=3)
|
||||||
assert all(isinstance(r, str) for r in res)
|
assert len(res) <= 3
|
||||||
|
assert all(isinstance(r, str) for r in res)
|
||||||
|
finally:
|
||||||
|
s.close()
|
||||||
|
|||||||
Reference in New Issue
Block a user