diff --git a/src/genesis/rag/store.py b/src/genesis/rag/store.py index 2a05a72..4c4ac09 100644 --- a/src/genesis/rag/store.py +++ b/src/genesis/rag/store.py @@ -31,7 +31,9 @@ class RagStore: self.conn.commit() 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: self.conn.execute("DELETE FROM rag_chunks WHERE scope=?", (scope,)) self.conn.executemany( @@ -41,6 +43,10 @@ class RagStore: self.conn.commit() 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: rows = self.conn.execute( "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.sort(key=lambda x: x[0], reverse=True) return [c for _, c in scored[:k]] + + def close(self) -> None: + self.conn.close() diff --git a/tests/test_rag_store.py b/tests/test_rag_store.py index c96315d..e350888 100644 --- a/tests/test_rag_store.py +++ b/tests/test_rag_store.py @@ -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()