diff --git a/src/genesis/rag/store.py b/src/genesis/rag/store.py new file mode 100644 index 0000000..2a05a72 --- /dev/null +++ b/src/genesis/rag/store.py @@ -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]] diff --git a/tests/test_rag_store.py b/tests/test_rag_store.py new file mode 100644 index 0000000..c96315d --- /dev/null +++ b/tests/test_rag_store.py @@ -0,0 +1,44 @@ +import math +from genesis.rag.store import RagStore + + +def _vec(*ones): + v = [0.0] * 8 + for i in ones: + v[i] = 1.0 + n = math.sqrt(sum(x * x for x in v)) + return [x / n for x in v] + + +def test_search_returns_most_similar_chunk(): + s = RagStore(":memory:") + s.reset_scope("p1") + s.add("p1", ["订单模块处理创建", "用户认证登录"], [_vec(0, 1), _vec(4, 5)]) + res = s.search("p1", _vec(0, 1), k=1) + assert res == ["订单模块处理创建"] + + +def test_reset_scope_clears(): + s = RagStore(":memory:") + s.reset_scope("p1") + s.add("p1", ["a"], [_vec(0)]) + s.reset_scope("p1") + assert s.search("p1", _vec(0), k=3) == [] + + +def test_concurrent_add_and_search_no_crash(): + import threading + s = RagStore(":memory:") + s.reset_scope("p1") + + def worker(i): + s.add("p1", [f"chunk-{i}"], [_vec(i % 8)]) + + 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)