Coverage for src/local_deep_research/vector_stores/base.py: 93%

38 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-10-22 21:20 +0000

1"""Abstract base class for vector stores. 

2 

3Mirrors the embeddings/LLM provider pattern (``BaseEmbeddingProvider``, 

4``BaseLLM`` + factory): a small typed interface plus a registry so a future 

5release can add backends (Qdrant, pgvector, Milvus, ...) selectable by a 

6setting. **Only FAISS is implemented today** — there is intentionally no 

7settings selector wired yet. 

8 

9Scope / honest boundary 

10----------------------- 

11This abstraction covers the *vector-storage operations* — create / add / 

12search / delete / reconstruct / persist / count. Local-file constructors accept 

13injected persistence resources, but this interface deliberately does **not** 

14prescribe their per-user paths, integrity checksums, or reload/merge locking 

15model that :class:`LibraryRAGService` layers on top. Those are *local-file* 

16concerns (see :attr:`is_local_file`): a future server-backed store (e.g. 

17Qdrant) would set ``is_local_file = False``, no-op :meth:`persist` / 

18:meth:`load`, implement :meth:`drop_collection` (the deletion path cannot 

19``unlink()`` a remote collection), and the service layer would skip the file 

20lock + integrity machinery for it. So this base class makes adding a backend *easier* (a clean, 

21verified query surface), not *free* — a real server backend still requires 

22refactoring the file/lock/integrity layer in the service. 

23 

24Identity model 

25-------------- 

26Vectors are keyed by an application-supplied **int64 id**. In this codebase 

27that id is ``DocumentChunk.id`` (the encrypted-DB primary key); the store holds 

28only vectors + ids, and all text/metadata is rehydrated from the DB by id. 

29 

30============================================================================ 

31SECURITY INVARIANT — NEVER pass document text to a vector store 

32============================================================================ 

33Every method here accepts ONLY an integer id and its embedding vector. 

34Callers MUST NOT pass document/chunk text (or any other user content) into a 

35vector-store method, and implementations MUST NOT persist such text. 

36 

37The per-user database is encrypted; a vector index is not (it is a local file 

38for FAISS, or an external service for a future backend). Keeping user content 

39out of the vector store keeps the encrypted database the sole home of that 

40content. Because the interface has no text parameter, the protection is 

41STRUCTURAL rather than a matter of discipline — no backend (FAISS today, or a 

42swapped-in Qdrant/pgvector tomorrow) can receive or store the text. The 

43authoritative text lives solely in the encrypted DB 

44(``DocumentChunk.chunk_text``), keyed by the same int id; retrieval rehydrates 

45snippets from there by id (see ``LibraryRAGService`` / the search engines). 

46============================================================================ 

47""" 

48 

49from abc import ABC, abstractmethod 

50from pathlib import Path 

51from types import TracebackType 

52from typing import Callable, List, Optional, Protocol, Self, Sequence, Tuple 

53 

54import numpy as np 

55 

56 

57IntegrityRecord = Callable[[Path], None] 

58IntegrityVerify = Callable[[Path], Tuple[bool, Optional[str]]] 

59 

60 

61class WriteLock(Protocol): 

62 """Minimal lock capability used to guard a vector store's file writes. 

63 

64 Implementations MUST be reentrant. Write paths nest acquisition on a 

65 single thread — the service layer holds the lock across a document 

66 operation while the store's own reload/apply/persist choreography 

67 re-acquires it — so a plain ``threading.Lock`` satisfies this Protocol 

68 structurally but self-deadlocks at runtime. Use ``threading.RLock`` (or 

69 a wrapper around one, as the production tracked lock is). 

70 """ 

71 

72 def __enter__(self) -> bool | Self: ... 72 ↛ anywhereline 72 didn't jump anywhere: it always raised an exception.

73 

74 def __exit__( 74 ↛ exitline 74 didn't return from function '__exit__' because

75 self, 

76 exception_type: type[BaseException] | None, 

77 exception: BaseException | None, 

78 traceback: TracebackType | None, 

79 ) -> None: ... 

80 

81 

82class BaseVectorStore(ABC): 

83 """Abstract interface for a per-collection vector index. 

84 

85 An instance wraps a single collection's index (stateful, in-memory, backed 

86 by a persisted file for local-file stores). Construct via :meth:`create` 

87 (new, empty) or :meth:`load` (from disk). 

88 """ 

89 

90 # Override in subclasses. 

91 provider_key: str = "base" # unique id; matches the (future) setting value 

92 provider_name: str = "Base" # display name for logs/UI 

93 # True when the store persists to a local file this process owns end to end 

94 # — i.e. the service's file write-lock + integrity checksum + reload/merge 

95 # model applies. A server-backed store sets this False. 

96 is_local_file: bool = True 

97 # Settings key holding this store's endpoint URI, for a store that is NOT 

98 # a local file. The egress policy reads it exactly as it reads a search 

99 # engine's ``url_setting`` (``security/egress/policy.py``): a 

100 # server-backed store whose configured host resolves to a PUBLIC address 

101 # ships the collection's embeddings and every query off the machine, so 

102 # the library/collection engines fail up to exposing/public and 

103 # PRIVATE_ONLY denies them. A non-local-file store that declares no key 

104 # cannot be shown to be on-box and is therefore treated as remote — 

105 # declaring it is how a backend proves the opposite. Unused for a 

106 # local-file store (FAISS), which is contained by construction. 

107 uri_setting: Optional[str] = None 

108 # True when the store can invert id -> vector (needed for the in-place 

109 # format migration, which re-keys vectors without re-embedding). 

110 supports_reconstruct: bool = True 

111 # Embedding dimension this instance was constructed/loaded for. Concrete 

112 # subclasses set this in __init__/create()/load(); declared here (no 

113 # default — there is no sensible universal value) so callers like the 

114 # facade can read it via the abstract interface. 

115 dimension: int 

116 

117 # ------------------------------------------------------------------ # 

118 # Construction 

119 # ------------------------------------------------------------------ # 

120 @classmethod 

121 @abstractmethod 

122 def create( 

123 cls, 

124 *, 

125 dimension: int, 

126 index_type: str, 

127 metric: str, 

128 normalize: bool, 

129 path: Optional[Path] = None, 

130 lock: Optional[WriteLock] = None, 

131 integrity_record: Optional[IntegrityRecord] = None, 

132 integrity_verify: Optional[IntegrityVerify] = None, 

133 ) -> "BaseVectorStore": 

134 """Create a new, empty store for vectors of ``dimension``. 

135 

136 Args: 

137 dimension: Embedding dimension. 

138 index_type: Backend index family (e.g. ``"flat"``, ``"hnsw"``). 

139 metric: Distance metric (``"l2"``, ``"cosine"``, ``"dot_product"``). 

140 normalize: Whether query/doc vectors are L2-normalized before use 

141 (cosine similarity via inner product on normalized vectors). 

142 path: Optional local persistence path. 

143 lock: Optional reentrant lock for local persistence (see 

144 :class:`WriteLock`). 

145 integrity_record: Optional persisted-file integrity callback. 

146 integrity_verify: Optional persisted-file integrity check. 

147 """ 

148 

149 @classmethod 

150 @abstractmethod 

151 def load( 

152 cls, 

153 path: Path, 

154 *, 

155 dimension: int, 

156 index_type: str, 

157 metric: str, 

158 normalize: bool, 

159 lock: Optional[WriteLock] = None, 

160 integrity_record: Optional[IntegrityRecord] = None, 

161 integrity_verify: Optional[IntegrityVerify] = None, 

162 ) -> "BaseVectorStore": 

163 """Load a persisted store from ``path`` (local-file stores only). 

164 

165 ``dimension`` / ``index_type`` / ``metric`` / ``normalize`` describe how 

166 the index was built (some backends persist this, some do not); callers 

167 pass the values from settings so search-time normalization matches 

168 build-time normalization. ``lock`` / ``integrity_record`` / 

169 ``integrity_verify`` are the same injected persistence resources as on 

170 :meth:`create`. 

171 """ 

172 

173 # ------------------------------------------------------------------ # 

174 # Vector operations 

175 # ------------------------------------------------------------------ # 

176 @abstractmethod 

177 def add(self, ids: List[int], vectors: np.ndarray) -> None: 

178 """Add vectors under the given int64 ids (aligned 1:1, same length). 

179 

180 ids + vectors ONLY — never text (see the module-level "SECURITY 

181 INVARIANT" block). The text stays in the encrypted DB. 

182 """ 

183 

184 @abstractmethod 

185 def search( 

186 self, query_vector: np.ndarray, k: int 

187 ) -> List[Tuple[int, float]]: 

188 """Return up to ``k`` ``(id, distance)`` pairs, nearest first. 

189 

190 ``distance`` is the backend's raw score (L2 distance or inner product); 

191 callers map it to a relevance score. Absent/empty slots are filtered. 

192 """ 

193 

194 @abstractmethod 

195 def delete(self, ids: List[int]) -> int: 

196 """Remove the given ids. Returns the number actually removed. 

197 

198 May raise for index families that do not support removal (e.g. FAISS 

199 HNSW) — callers handle that the same way they did pre-abstraction. 

200 """ 

201 

202 @abstractmethod 

203 def live_ids(self) -> List[int]: 

204 """Return every id currently stored (the authoritative membership).""" 

205 

206 @abstractmethod 

207 def count(self) -> int: 

208 """Return the number of vectors currently stored.""" 

209 

210 @abstractmethod 

211 def apply( 

212 self, 

213 *, 

214 add_ids: List[int], 

215 add_vectors: Optional[np.ndarray], 

216 remove_ids: Sequence[int] = (), 

217 dedup: bool = True, 

218 ) -> dict: 

219 """Durably apply a batch of removals + additions as one atomic unit. 

220 

221 This is the single write primitive. The caller computes *which* ids to 

222 remove and add (from its own authoritative source — for LDR, the 

223 encrypted DB); the store applies them and persists so that a concurrent 

224 writer to the same logical index cannot lose either party's changes. 

225 *How* that atomicity/isolation is achieved is backend-specific (a 

226 local-file backend takes a per-index lock and reload-merges; a 

227 server-backed store issues upserts/deletes) and is not part of this 

228 contract. 

229 

230 Contract: removals are applied before additions. A "replace" is 

231 expressed as ``remove_ids`` + ``add_ids`` of the same id. 

232 

233 ``dedup`` de-duplicates ids WITHIN one call (a repeated id in 

234 ``add_ids`` is written once). It does NOT mean "skip an add whose id 

235 is already live in the store" — an already-live ``add_id`` MUST be 

236 replaced, and the NEW vector must win. 

237 

238 That is not a nicety. An ``add_id`` that is already live is either an 

239 idempotent re-add or — critically — a REUSED ``DocumentChunk.id`` 

240 whose prior row was rolled back: SQLite recycles AUTOINCREMENT ids on 

241 ROLLBACK, so a failed ``index()`` commit can leave a stale orphan 

242 vector under an id that a later, unrelated chunk then reuses. Keeping 

243 the old vector makes search rank on one document's embedding and then 

244 rehydrate a DIFFERENT document's text by id. Nothing downstream 

245 detects it: a reconciler that checks ``live_ids()`` membership sees a 

246 healthy index, because the skipped add left the id live. 

247 

248 On failure the store must not be left durably half-applied. 

249 

250 Returns a stats dict (e.g. ``{"added": n, "removed": m}``). ids + 

251 vectors ONLY — never text (see the module "SECURITY INVARIANT"). 

252 """ 

253 

254 # ------------------------------------------------------------------ # 

255 # Optional (local-file / reconstructable stores) 

256 # ------------------------------------------------------------------ # 

257 def reconstruct(self, id: int) -> Optional[np.ndarray]: 

258 """Return the stored vector for ``id`` (migration helper). 

259 

260 Only meaningful when :attr:`supports_reconstruct` is True. 

261 """ 

262 raise NotImplementedError( 

263 f"{type(self).__name__} does not support reconstruct()" 

264 ) 

265 

266 def persist(self, path: Path) -> None: 

267 """Durably write the store to ``path`` (local-file stores only).""" 

268 raise NotImplementedError( 

269 f"{type(self).__name__} does not support persist()" 

270 ) 

271 

272 # ------------------------------------------------------------------ # 

273 # Lifecycle (server-backed stores) 

274 # ------------------------------------------------------------------ # 

275 @classmethod 

276 def drop_collection( 

277 cls, 

278 *, 

279 username: str, 

280 index_hash: str, 

281 settings_snapshot: Optional[dict] = None, 

282 ) -> bool: 

283 """Permanently remove the backing collection for one RAG index. 

284 

285 The deletion counterpart of :meth:`create`, so the library deletion 

286 cascade can dispose of vectors it cannot simply ``unlink()``. It is a 

287 **classmethod on purpose**: a deletion path must be able to drop a 

288 collection without first loading it (loading a possibly-corrupt or 

289 already-absent index just to delete it is both wasteful and 

290 failure-prone). 

291 

292 Identity 

293 -------- 

294 The collection is addressed by ``(username, index_hash)`` rather than 

295 by an instance, because those two values are the only stable handle 

296 available at deletion time: ``index_hash`` lives on the ``RAGIndex`` 

297 row, and that row is deleted by the same cascade. An implementation 

298 MUST derive the backend collection name from these same two inputs at 

299 BOTH creation and deletion, or the name it computes here will not 

300 match the one it created and drop_collection will silently miss. 

301 

302 Nothing enforces that yet: :meth:`create` does not currently receive 

303 ``username`` or ``index_hash`` (it takes ``dimension``/``index_type``/ 

304 ``metric``/``normalize``/``path``/``lock``/``integrity_record``/ 

305 ``integrity_verify``), so there is no shared derivation to reuse. The 

306 first server-backed provider MUST extend the creation path to accept 

307 both, thread them through into whatever it uses to name the remote 

308 collection, and use that exact derivation here too (see the follow-up 

309 design issue, #6769). 

310 

311 **Ordering contract**: callers MUST invoke this while the ``RAGIndex`` 

312 row still exists. Once the row is gone ``index_hash`` is unrecoverable 

313 and the remote collection becomes *unaddressable* -- orphaned vectors 

314 that no longer have a name (see issue #5765). 

315 

316 Contract for implementations: 

317 

318 * Idempotent -- return ``False`` (not raise) when the collection is 

319 already absent, ``True`` when one was actually dropped. 

320 * Raise on *transport / permission* failure. Returning ``False`` for an 

321 unreachable backend would let the caller record a successful deletion 

322 while the user's vectors survive on third-party infrastructure. 

323 * Make the failure classifiable. The cascade tells the user to retry 

324 only for a *transient* failure: a ``ConnectionError`` / 

325 ``TimeoutError`` (or subclass), or any exception carrying a boolean 

326 ``retryable`` attribute set to True. Everything else -- auth or 

327 permission errors, missing connection config, unrecognised client 

328 exceptions -- is reported as permanent, so translate the backend 

329 client's own "unreachable" / "timed out" exceptions accordingly. 

330 * Bound each call with a client timeout: the cascade calls this inside 

331 the user's database transaction. 

332 * Read connection config from ``settings_snapshot`` -- the deletion 

333 callers pass the user's snapshot. Only the exception's type name is 

334 surfaced to users and logs, never its message. 

335 

336 Local-file stores (:attr:`is_local_file` True) do not implement this -- 

337 their deletion path unlinks the index files instead -- so the base 

338 implementation raises :class:`NotImplementedError`. Server-backed 

339 stores MUST override it. 

340 """ 

341 raise NotImplementedError( 

342 f"{cls.__name__} does not support drop_collection()" 

343 )