feat(rag): 新增 RagStore(SQLite 向量存储 + 余弦检索,线程安全)

This commit is contained in:
lhl
2026-08-29 22:39:37 +08:00
parent afc09fc908
commit 2eab0183ee
2 changed files with 96 additions and 0 deletions
+52
View File
@@ -0,0 +1,52 @@
from __future__ import annotations
import json
import math
import sqlite3
import threading
from typing import List
def _cosine(a, b):
dot = sum(x * y for x, y in zip(a, b))
na = math.sqrt(sum(x * x for x in a))
nb = math.sqrt(sum(y * y for y in b))
return dot / (na * nb or 1.0)
class RagStore:
def __init__(self, db_path: str):
# D2:服务端在 worker 线程复用连接,须关闭同线程检查
self.conn = sqlite3.connect(db_path, check_same_thread=False)
self._lock = threading.Lock()
self.conn.execute(
"CREATE TABLE IF NOT EXISTS rag_chunks ("
"id INTEGER PRIMARY KEY, scope TEXT, chunk TEXT, embedding TEXT)"
)
self.conn.commit()
def reset_scope(self, scope: str) -> None:
with self._lock:
self.conn.execute("DELETE FROM rag_chunks WHERE scope=?", (scope,))
self.conn.commit()
def add(self, scope: str, chunks: List[str], embeddings: List[List[float]]) -> None:
# 先清后写(增量索引),整段加锁串行化
with self._lock:
self.conn.execute("DELETE FROM rag_chunks WHERE scope=?", (scope,))
self.conn.executemany(
"INSERT INTO rag_chunks(scope, chunk, embedding) VALUES(?,?,?)",
[(scope, c, json.dumps(e)) for c, e in zip(chunks, embeddings)],
)
self.conn.commit()
def search(self, scope: str, query_vec: List[float], k: int = 5) -> List[str]:
with self._lock:
rows = self.conn.execute(
"SELECT chunk, embedding FROM rag_chunks WHERE scope=?", (scope,)
).fetchall()
scored = []
for chunk, emb in rows:
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]]