fix(rag): 强化 RagStore 并发读写验证 + close/健壮性
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user