feat(rag): 新增 RagStore(SQLite 向量存储 + 余弦检索,线程安全)
This commit is contained in:
@@ -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]]
|
||||
Reference in New Issue
Block a user