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

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]]

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)