Coverage for src/local_deep_research/vector_stores/implementations/faiss_store.py: 92%

1"""FAISS-backed vector store. 

2 

3Drives raw ``faiss`` directly via an :class:`IndexIDMap2` keyed by 

4application-supplied int64 ids (``DocumentChunk.id``). The store holds **only** 

5vectors + ids — no document text, no metadata (see the "SECURITY INVARIANT" 

6block in :mod:`..base`: text must never reach a vector store). Persistence is 

7``faiss.write_index`` / ``faiss.read_index`` (pure binary), with no companion 

8docstore file. 

9 

10``IndexIDMap2`` keeps id management in the faiss C++ layer (no Python-side 

11position map, no position renumbering on delete), which lets the caller key 

12vectors directly by ``DocumentChunk.id`` and rehydrate all text/metadata from 

13the encrypted DB by that id. 

14""" 

15 

16import os 

17import tempfile 

18import time 

19from contextlib import contextmanager 

20from pathlib import Path 

21from typing import Iterator, List, Optional, Sequence, Tuple, cast 

22 

23import numpy as np 

24from faiss import ( 

25 METRIC_INNER_PRODUCT, 

26 METRIC_L2, 

27 IDSelectorBatch, 

28 IndexFlatIP, 

29 IndexFlatL2, 

30 IndexHNSWFlat, 

31 IndexIDMap2, 

32 downcast_index, 

33 normalize_L2, 

34 read_index, 

35 vector_to_array, 

36 write_index, 

37) 

38from loguru import logger 

39 

40from ..base import BaseVectorStore, IntegrityRecord, IntegrityVerify, WriteLock 

41 

42# HNSW connections-per-layer (parity with the previous wrapper's default). 

43_HNSW_M = 32 

44# Metrics that use inner product on L2-normalized vectors (cosine similarity). 

45_IP_METRICS = ("cosine", "dot_product") 

46# persist() only sweeps ``.tmp`` files older than this — long enough that no 

47# single index write could still be in flight, so a concurrent (multi-process) 

48# writer's just-created temp is never mistaken for an orphan and deleted. 

49_STALE_TMP_AGE_SECONDS = 3600 

50# Max squared L2 norm allowed for a raw (normalize=False) vector. faiss's L2/IP 

51# distance math (e.g. ||q-v||^2 ~ 4*norm^2 worst case) is done in float32; a row 

52# whose squared norm approaches float32's max overflows internally and faiss then 

53# returns the reserved -1 sentinel for a REAL vector, which search() drops. Real 

54# embeddings sit many orders of magnitude below this; this only rejects 

55# pathological/overflowing magnitudes. 

56_MAX_SAFE_SQ_NORM = float(np.finfo(np.float32).max) / 8.0 

57 

58 

59def _canon(value: Optional[str], default: str) -> str: 

60 """Canonicalize an index_type / metric string (strip + lower) so every 

61 construction path — the facade, the legacy rekey builder, direct create() — 

62 agrees. A non-canonical 'HNSW'/'Cosine' would otherwise build one structure 

63 but be interpreted as another (wrong index type, wrong metric, corrupted 

64 relevance) depending on which case-sensitive comparison ran.""" 

65 return (value or default).strip().lower() 

66 

67 

68class FaissVectorStore(BaseVectorStore): 

69 """A single collection's FAISS index (``IndexIDMap2`` over a base index). 

70 

71 Concurrency: the injected write lock (a reentrant :class:`WriteLock`; in 

72 production a tracked ``threading.RLock`` wrapper) serializes writers (and, 

73 via ``_read_guard``, readers) within ONE process only. This assumes a 

74 single-process deployment. Under a multi-process server (e.g. 

75 ``gunicorn -w N``) two workers hold independent locks and a lost update is 

76 possible; a cross-process (file) lock would be required for that deployment. 

77 """ 

78 

79 provider_key = "faiss" 

80 provider_name = "FAISS" 

81 is_local_file = True 

82 supports_reconstruct = True 

83 

84 def __init__( 

85 self, 

86 index: IndexIDMap2, 

87 *, 

88 dimension: int, 

89 index_type: str, 

90 metric: str, 

91 normalize: bool, 

92 ) -> None: 

93 self._index = index 

94 self.dimension = dimension 

95 self.index_type = _canon(index_type, "flat") 

96 self.metric = _canon(metric, "cosine") 

97 self.normalize = normalize 

98 # Set if a failed apply() could not resync in-memory state to disk; a 

99 # poisoned instance must refuse reads rather than serve state that 

100 # disagrees with the durable file (a fresh load/apply clears it). 

101 self._poisoned = False 

102 # Persistence binding (set by _bind via create()/load()). Required only 

103 # for apply(); a search-only load can leave these unset. Injected by the 

104 # caller so this store never imports a lock registry or DB integrity 

105 # model — it only owns the choreography, not the resources. 

106 self._path: Optional[Path] = None 

107 self._lock: Optional[WriteLock] = None 

108 self._integrity_record: Optional[IntegrityRecord] = None 

109 self._integrity_verify: Optional[IntegrityVerify] = None 

110 

111 def _bind( 

112 self, 

113 path: Optional[Path], 

114 lock: Optional[WriteLock], 

115 integrity_record: Optional[IntegrityRecord], 

116 integrity_verify: Optional[IntegrityVerify], 

117 ) -> None: 

118 """Attach the persistence resources apply() needs (injected, not imported). 

119 

120 ``path``: where this index's ``.faiss`` lives (caller resolves the 

121 per-user path). ``lock``: the per-index write lock (caller owns the 

122 registry + its lifecycle). ``integrity_record`` / ``integrity_verify``: 

123 closures over the caller's file-integrity manager, so save+record stay 

124 under the same lock without this class knowing about the DB. 

125 """ 

126 self._path = Path(path) if path is not None else None 

127 self._lock = lock 

128 self._integrity_record = integrity_record 

129 self._integrity_verify = integrity_verify 

130 if ( 

131 self._path is not None 

132 and lock is not None 

133 and (integrity_record is None or integrity_verify is None) 

134 ): 

135 logger.warning( 

136 "FaissVectorStore bound for writes without integrity hooks; " 

137 "torn-write/corruption detection is disabled for this index" 

138 ) 

139 

140 # ------------------------------------------------------------------ # 

141 # Construction 

142 # ------------------------------------------------------------------ # 

143 @staticmethod 

144 def _build_base_index(dimension: int, index_type: str, metric: str): 

145 """Build the base faiss index (parity with the prior wrapper). 

146 

147 ``ivf`` is not implemented and falls back to flat, exactly as before. 

148 """ 

149 index_type = _canon(index_type, "flat") 

150 metric = _canon(metric, "cosine") 

151 # Inner product for cosine/dot-product (on normalized vectors), else L2. 

152 metric_type = ( 

153 METRIC_INNER_PRODUCT if metric in _IP_METRICS else METRIC_L2 

154 ) 

155 if index_type == "hnsw": 

156 logger.info(f"Created HNSW index with M={_HNSW_M} connections") 

157 # Pass the metric explicitly — IndexHNSWFlat defaults to L2, which 

158 # would return L2 distances even for a cosine collection. 

159 return IndexHNSWFlat(dimension, _HNSW_M, metric_type) 

160 # "flat" (default), "ivf" (falls back to flat), or anything else. 

161 if metric in _IP_METRICS: 

162 return IndexFlatIP(dimension) 

163 return IndexFlatL2(dimension) 

164 

165 @staticmethod 

166 def _physical_config_reason( 

167 index: IndexIDMap2, index_type: str, metric: str 

168 ) -> Optional[str]: 

169 """Why an on-disk IndexIDMap2's PHYSICAL base index / metric does NOT 

170 match ``index_type``/``metric`` (or ``None`` if it matches). 

171 

172 A stale/foreign file of the wrong type must never be adopted: an HNSW 

173 file treated as flat makes ``supports_delete`` wrongly True so 

174 ``delete()`` raises, and a flat file treated as hnsw makes 

175 ``apply(remove_ids)`` silently REBUILD the file as HNSW. cosine and 

176 dot_product are physically identical (both inner-product), so the metric 

177 FAMILY is compared, not the exact string. Shared by load() and the 

178 reload/restore paths so all three enforce the same invariant. 

179 """ 

180 base = downcast_index(index.index) 

181 want_type = _canon(index_type, "flat") 

182 if isinstance(base, IndexHNSWFlat) != (want_type == "hnsw"): 

183 return ( 

184 f"physical index {type(base).__name__} != index_type " 

185 f"{want_type!r}" 

186 ) 

187 want_metric_type = ( 

188 METRIC_INNER_PRODUCT 

189 if _canon(metric, "cosine") in _IP_METRICS 

190 else METRIC_L2 

191 ) 

192 if base.metric_type != want_metric_type: 

193 return ( 

194 f"physical metric_type {base.metric_type} != metric " 

195 f"{_canon(metric, 'cosine')!r}" 

196 ) 

197 return None 

198 

199 @classmethod 

200 def create( 

201 cls, 

202 *, 

203 dimension: int, 

204 index_type: str, 

205 metric: str, 

206 normalize: bool, 

207 path: Optional[Path] = None, 

208 lock: Optional[WriteLock] = None, 

209 integrity_record: Optional[IntegrityRecord] = None, 

210 integrity_verify: Optional[IntegrityVerify] = None, 

211 ) -> "FaissVectorStore": 

212 base = cls._build_base_index(dimension, index_type, metric) 

213 store = cls( 

214 IndexIDMap2(base), 

215 dimension=dimension, 

216 index_type=index_type, 

217 metric=metric, 

218 normalize=normalize, 

219 ) 

220 store._bind(path, lock, integrity_record, integrity_verify) 

221 return store 

222 

223 @classmethod 

224 def load( 

225 cls, 

226 path: Path, 

227 *, 

228 dimension: int, 

229 index_type: str, 

230 metric: str, 

231 normalize: bool, 

232 lock: Optional[WriteLock] = None, 

233 integrity_record: Optional[IntegrityRecord] = None, 

234 integrity_verify: Optional[IntegrityVerify] = None, 

235 ) -> "FaissVectorStore": 

236 """Load a persisted ``IndexIDMap2`` from ``path`` (a ``.faiss`` file). 

237 

238 ``read_index`` is a pure binary read — no pickle. Post-migration every 

239 on-disk index is an ``IndexIDMap2``; a raw base index here means either 

240 a pre-migration file (should be converted first) or corruption, and is 

241 surfaced rather than silently coerced. 

242 

243 ``lock`` / ``integrity_*`` are only needed if the caller will later 

244 write via :meth:`apply`; a search-only load can omit them. 

245 """ 

246 index = read_index(str(path)) 

247 if not isinstance(index, IndexIDMap2): 

248 raise ValueError( 

249 f"Index at {path} is {type(index).__name__}, not IndexIDMap2 " 

250 "(pre-migration or corrupt format)" 

251 ) 

252 # Catch embedding-model drift early: a persisted index whose dimension 

253 # no longer matches the configured model would otherwise fail with an 

254 # opaque faiss error deep inside a later add()/search(). 

255 if index.d != dimension: 

256 raise ValueError( 

257 f"Index at {path} has dimension {index.d}, expected {dimension} " 

258 "(embedding model/config changed — reindex required)" 

259 ) 

260 # Verify the PHYSICAL index matches the caller-declared index_type/metric 

261 # instead of trusting the label. A stale file left at this path (e.g. a 

262 # force-reindex whose unlink failed, or a foreign/pre-canonicalization 

263 # file) could otherwise be adopted under a mismatched config: an HNSW 

264 # file loaded as "flat" makes supports_delete wrongly True so .delete() 

265 # raises a C++ RuntimeError, while a Flat file loaded as "hnsw" makes 

266 # .apply(remove_ids) silently REBUILD the file as HNSW. Introspect the 

267 # real base index + metric and refuse a mismatch so the caller 

268 # quarantines + rebuilds. (cosine/dot_product are physically identical — 

269 # both METRIC_INNER_PRODUCT — so we compare the metric FAMILY, not the 

270 # exact string; _build_base_index only ever encodes HNSW-vs-Flat and 

271 # IP-vs-L2.) 

272 reason = cls._physical_config_reason(index, index_type, metric) 

273 if reason: 

274 raise ValueError( 

275 f"Index at {path}: {reason} — refusing a mismatched index " 

276 "(reindex required)" 

277 ) 

278 store = cls( 

279 index, 

280 dimension=dimension, 

281 index_type=index_type, 

282 metric=metric, 

283 normalize=normalize, 

284 ) 

285 store._bind(path, lock, integrity_record, integrity_verify) 

286 return store 

287 

288 # ------------------------------------------------------------------ # 

289 # Internal helpers 

290 # ------------------------------------------------------------------ # 

291 def _prepare(self, vectors: np.ndarray) -> np.ndarray: 

292 """Return a contiguous float32 2-D copy, L2-normalized if configured. 

293 

294 A copy is always made so ``normalize_L2`` (which mutates in place) never 

295 touches the caller's array. 

296 """ 

297 vecs = np.array(vectors, dtype="float32", copy=True, order="C") 

298 if vecs.ndim == 1: 

299 vecs = vecs.reshape(1, -1) 

300 # Validate dimension explicitly. FAISS only guards this with an internal 

301 # ``assert`` (no message), which is STRIPPED under ``python -O`` — a 

302 # wrong-dimension array would then be added/searched silently, corrupting 

303 # the index with no error. Fail loudly instead. 

304 if vecs.shape[1] != self.dimension: 

305 raise ValueError( 

306 f"vector dimension {vecs.shape[1]} != index dimension " 

307 f"{self.dimension}" 

308 ) 

309 # Reject non-finite INPUT first: NaN/Inf from a bad embedding poison the 

310 # index — the vector becomes permanently unsearchable and NaN distances 

311 # corrupt ranking — with no error. Fail loudly so the caller fixes it. 

312 if not np.isfinite(vecs).all(): 

313 raise ValueError( 

314 "vectors contain non-finite values (NaN/Inf) — bad embedding" 

315 ) 

316 if self.normalize: 

317 # normalize_L2 silently maps a zero-magnitude row (and one whose 

318 # squared norm overflows float32) to ALL-ZEROS, not NaN, so the 

319 # finite input passes the check above yet becomes a degenerate, 

320 # effectively-unsearchable vector. Reject those explicitly (norm in 

321 # float32 to match faiss: a real overflow is a non-finite norm here). 

322 norms = np.sqrt((vecs**2).sum(axis=1)) 

323 if not np.isfinite(norms).all() or (norms == 0).any(): 

324 raise ValueError( 

325 "vectors contain a zero-magnitude or overflow-magnitude row " 

326 "that L2 normalization cannot represent — bad embedding" 

327 ) 

328 normalize_L2(vecs) 

329 else: 

330 # normalize=False: raw vectors go straight into faiss's L2/IP 

331 # distance math. A large-but-finite magnitude (which passes the 

332 # isfinite check above) overflows float32 INSIDE that computation, 

333 # and faiss then returns the reserved -1 sentinel for a REAL vector — 

334 # search() drops it as an "empty slot", silently losing the hit. 

335 # Reject rows whose squared L2 norm approaches float32's max (computed 

336 # in float64 so the check itself can't overflow). 

337 sq_norms = (vecs.astype("float64") ** 2).sum(axis=1) 

338 if (sq_norms > _MAX_SAFE_SQ_NORM).any(): 338 ↛ 339line 338 didn't jump to line 339 because the condition on line 338 was never true

339 raise ValueError( 

340 "vectors contain a magnitude too large for float32 distance " 

341 "computation (would overflow to a dropped/mis-ranked search " 

342 "hit) — normalize the embeddings or use a bounded model" 

343 ) 

344 return vecs 

345 

346 # ------------------------------------------------------------------ # 

347 # Vector operations 

348 # ------------------------------------------------------------------ # 

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

350 id_arr = np.asarray(list(ids), dtype="int64") 

351 # Guard empty BEFORE _prepare: a length-0 1-D array would otherwise 

352 # reshape to one row of dimension 0 and raise a spurious mismatch. 

353 if len(id_arr) == 0: 

354 return 

355 vecs = self._prepare(vectors) 

356 if len(id_arr) != len(vecs): 

357 raise ValueError( 

358 f"ids/vectors length mismatch: {len(id_arr)} != {len(vecs)}" 

359 ) 

360 self._index.add_with_ids(vecs, id_arr) 

361 

362 @contextmanager 

363 def _read_guard(self) -> Iterator[None]: 

364 """Serialize a read against a concurrent ``apply()`` and refuse a 

365 poisoned instance. 

366 

367 faiss is not safe for concurrent search-vs-add/remove on the *same* 

368 index object; when this instance is bound for writes (a shared, cached 

369 instance), reads take the same lock ``apply()`` uses. Internal callers 

370 that already hold the lock (e.g. ``_dedup_new``) use the ``*_unlocked`` 

371 helpers, not these public methods — single-acquisition discipline, so 

372 this store's own choreography never depends on the injected lock's 

373 reentrancy (which the :class:`WriteLock` contract reserves for the 

374 service layer's nested holds). 

375 """ 

376 if self._lock is not None: 

377 with self._lock: 

378 if self._poisoned: 

379 raise RuntimeError( 

380 "vector store in-memory state is poisoned by a failed " 

381 "write; reload the index before reading" 

382 ) 

383 yield 

384 else: 

385 if self._poisoned: 

386 raise RuntimeError( 

387 "vector store in-memory state is poisoned by a failed " 

388 "write; reload the index before reading" 

389 ) 

390 yield 

391 

392 def search( 

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

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

395 if k <= 0: 

396 return [] 

397 with self._read_guard(): 

398 if self._index.ntotal == 0: 

399 return [] 

400 # Clamp k to the number of stored vectors: faiss allocates k*4 bytes 

401 # of result buffer per query, so an unbounded k (a buggy/hostile 

402 # caller passing e.g. 10**9) would OOM. Searching for more neighbours 

403 # than exist only pads the result with -1 slots we skip anyway, so 

404 # this never changes the returned hits. 

405 k = min(k, self._index.ntotal) 

406 q = self._prepare(query_vector) 

407 distances, ids = self._index.search(q, k) 

408 results: List[Tuple[int, float]] = [] 

409 for dist, vid in zip(distances[0], ids[0]): 

410 # faiss returns -1 for empty slots when fewer than k neighbors. 

411 if vid == -1: 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true

412 continue 

413 results.append((int(vid), float(dist))) 

414 return results 

415 

416 @property 

417 def supports_delete(self) -> bool: 

418 """False for HNSW — faiss ``remove_ids`` is unimplemented for it.""" 

419 return self.index_type != "hnsw" 

420 

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

422 id_list = [i for i in ids if i is not None] 

423 if not id_list: 

424 return 0 

425 if not self.supports_delete: 

426 # faiss HNSW has no remove_ids; surface a clear, actionable error 

427 # rather than the opaque C++ "remove_ids not implemented" RuntimeError 

428 # (or a silent no-op). apply() never reaches this — it routes HNSW 

429 # removals through _rebuild_dropping() (a full index rebuild) — so 

430 # this guards only a direct delete() call on an HNSW store. 

431 raise RuntimeError( 

432 f"delete/replace is unsupported for index_type=" 

433 f"'{self.index_type}' (faiss HNSW cannot remove vectors); " 

434 "removals go through apply(), which rebuilds the index" 

435 ) 

436 id_arr = np.asarray(id_list, dtype="int64") 

437 return int(self._index.remove_ids(IDSelectorBatch(id_arr))) 

438 

439 def _rebuild_dropping(self, remove_ids: List[int]) -> int: 

440 """Remove ids from an index that has no ``remove_ids`` (HNSW) by 

441 rebuilding it from its own reconstructed vectors. 

442 

443 The only way to delete from a faiss HNSW index is to build a fresh one 

444 holding every surviving vector. Runs under ``apply()``'s write lock and 

445 uses the raw/unlocked helpers to avoid re-entrant self-deadlock. 

446 

447 No re-embedding: ``IndexIDMap2`` preserves each vector for 

448 reconstruction by id, so survivors are reconstructed and re-added 

449 VERBATIM via ``add_with_ids`` — deliberately bypassing ``_prepare`` so 

450 the stored vectors (already L2-normalized if ``normalize`` is set) are 

451 not normalized a second time. Returns the count actually removed (ids 

452 absent from the index are ignored, matching ``delete``'s leniency). 

453 """ 

454 remove_set = {int(i) for i in remove_ids} 

455 live = self._live_ids_unlocked() 

456 keep = [i for i in live if i not in remove_set] 

457 removed = len(live) - len(keep) 

458 if removed == 0: 458 ↛ 459line 458 didn't jump to line 459 because the condition on line 458 was never true

459 return 0 

460 if keep: 460 ↛ 464line 460 didn't jump to line 464 because the condition on line 460 was always true

461 vecs = np.empty((len(keep), self.dimension), dtype="float32") 

462 for row, vid in enumerate(keep): 

463 vecs[row] = self._index.reconstruct(int(vid)) 

464 base = self._build_base_index( 

465 self.dimension, self.index_type, self.metric 

466 ) 

467 new_index = IndexIDMap2(base) 

468 if keep: 468 ↛ 470line 468 didn't jump to line 470 because the condition on line 468 was always true

469 new_index.add_with_ids(vecs, np.asarray(keep, dtype="int64")) 

470 self._index = new_index 

471 return removed 

472 

473 def _live_ids_unlocked(self) -> List[int]: 

474 if self._index.ntotal == 0: 

475 return [] 

476 # faiss is untyped (ignore_missing_imports); tell mypy what the raw 

477 # C++ id-map array becomes after the numpy round-trip. 

478 return cast( 

479 List[int], 

480 vector_to_array(self._index.id_map).astype("int64").tolist(), 

481 ) 

482 

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

484 with self._read_guard(): 

485 return self._live_ids_unlocked() 

486 

487 def count(self) -> int: 

488 with self._read_guard(): 

489 return int(self._index.ntotal) 

490 

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

492 # Honor the Optional contract: faiss raises RuntimeError for an id not 

493 # in the index; return None instead of propagating. 

494 with self._read_guard(): 

495 try: 

496 # faiss is untyped; reconstruct() returns a raw ndarray here. 

497 return cast(np.ndarray, self._index.reconstruct(int(id))) 

498 except RuntimeError: 

499 return None 

500 

501 def _file_fingerprint(self) -> Optional[tuple]: 

502 """A cheap staleness token for the on-disk index file (stat, no read). 

503 

504 Shared across FaissVectorStore INSTANCES via the file itself. This is 

505 why the off-lock rebuild uses it and NOT a per-instance counter: the 

506 normal caller builds a FRESH store per operation (only the write lock is 

507 shared via the registry), so a concurrent writer is a different Python 

508 object whose mutation would never bump this instance's counter. The 

509 on-disk file, in contrast, is written by every writer. ``os.replace`` 

510 gives the target a fresh inode on every persist (mkstemp temp + rename), 

511 and mtime_ns/size change with content, so any persist by any writer 

512 changes this tuple. Returns None if the file does not exist yet. 

513 """ 

514 if self._path is None: 514 ↛ 515line 514 didn't jump to line 515 because the condition on line 514 was never true

515 return None 

516 try: 

517 st = os.stat(self._path) 

518 except OSError: 

519 return None 

520 return (st.st_ino, st.st_mtime_ns, st.st_size) 

521 

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

523 """Atomically write the index to ``path`` (temp file + ``os.replace``). 

524 

525 The temp+rename closes the torn-write window (#4197): a concurrent 

526 reader/verifier never observes a half-written ``.faiss`` — it sees 

527 either the old bytes or the new bytes, never a truncated mix. 

528 """ 

529 # Lazy import to avoid pulling security into this module's 

530 # top-level import graph (vector store implementations load early). 

531 from ...security.directory_creation import create_directory 

532 

533 path = Path(path) 

534 create_directory( 

535 path.parent, context="FAISS vector store index directory" 

536 ) 

537 # Keep the index dir private: the raw vectors are weakly invertible, so 

538 # other local accounts must not be able to read them. Structural, not 

539 # left to the caller. 

540 try: 

541 path.parent.chmod(0o700) 

542 except OSError: 

543 logger.warning(f"Could not chmod {path.parent} to 0o700") 

544 # Sweep stale temp files left by a previously hard-killed write. Within 

545 # one process persist() runs under this index's write lock, so no 

546 # in-process writer owns a ``<name>.*.tmp`` here. But that lock is 

547 # per-process (an in-process lock in a per-process registry), so under 

548 # ``gunicorn -w N`` a SECOND worker can be mid-persist for the SAME path 

549 # with its temp not yet renamed — unlinking it would crash that write 

550 # with FileNotFoundError at fsync. Only sweep temps OLDER than a generous 

551 # threshold: a genuinely orphaned temp (from a past hard kill) is 

552 # minutes/hours old, while an in-flight one was just created, and no 

553 # single write approaches this age — so a live temp is never swept. 

554 now = time.time() 

555 for stale in path.parent.glob(f"{path.name}.*.tmp"): 

556 try: 

557 if now - stale.stat().st_mtime < _STALE_TMP_AGE_SECONDS: 

558 continue # too recent — may be a concurrent worker's temp 

559 stale.unlink() 

560 except OSError: 

561 logger.warning(f"Could not remove stale temp file {stale}") 

562 # mkstemp creates the temp file O_EXCL with a random name inside the 

563 # (private) index dir — no predictable-name symlink-swap window. 

564 # Safe: this writes only the faiss index — vectors + integer 

565 # DocumentChunk.ids, never chunk text/metadata (see base.py's 

566 # "SECURITY INVARIANT" — the interface has no text parameter, so no 

567 # backend can persist it). Not sensitive data at rest. 

568 fd, tmp_name = tempfile.mkstemp( 

569 dir=str(path.parent), prefix=f"{path.name}.", suffix=".tmp" 

570 ) 

571 os.close(fd) 

572 tmp = Path(tmp_name) 

573 try: 

574 write_index(self._index, str(tmp)) 

575 # fsync the temp bytes BEFORE the rename so a crash/power-loss can't 

576 # leave a renamed-but-unflushed (torn) index — the rename is only 

577 # atomic w.r.t. what has actually reached disk. 

578 try: 

579 with open(tmp, "r+b") as f: 

580 os.fsync(f.fileno()) 

581 except OSError: 

582 pass # best-effort fsync 

583 os.replace(tmp, path) 

584 # fsync the directory so the rename itself is durable across a crash. 

585 try: 

586 dir_fd = os.open(str(path.parent), os.O_RDONLY) 

587 try: 

588 os.fsync(dir_fd) 

589 finally: 

590 os.close(dir_fd) 

591 except OSError: 

592 pass # directory fsync is unsupported on some platforms 

593 finally: 

594 if tmp.exists(): 

595 try: 

596 tmp.unlink() 

597 except OSError: 

598 logger.warning(f"Could not remove temp index file {tmp}") 

599 

600 # ------------------------------------------------------------------ # 

601 # Durable batch write (apply) 

602 # ------------------------------------------------------------------ # 

603 def _reload_under_lock(self) -> None: 

604 """Re-read the on-disk index so this write absorbs other writers' saves. 

605 

606 Called while holding the write lock. If the file does not exist yet 

607 (first write), keep the current fresh in-memory index and proceed. 

608 

609 But if the file EXISTS and cannot be trusted — integrity verify fails, 

610 the read fails, or it is a foreign format — FAIL CLOSED: raise rather 

611 than proceed to overwrite it with our (possibly stale) in-memory copy. 

612 Overwriting would clobber a concurrent writer's committed save and 

613 silently "heal" a corrupt/tampered file, defeating both the reload-merge 

614 guarantee and integrity detection. The caller decides whether to 

615 quarantine + rebuild. 

616 """ 

617 if self._path is None or not self._path.exists(): 

618 # No durable file to merge against. Normally keep the current 

619 # in-memory index (a genuine first write). BUT if a prior apply() 

620 # failed to persist and poisoned this instance, the in-memory index 

621 # holds a stale, never-persisted mutation — and _read_guard blocks 

622 # reads but apply() doesn't check _poisoned, so this now-succeeding 

623 # write would silently persist that earlier failed change. Discard it 

624 # (rebuild an empty index) so only THIS operation's mutation lands. 

625 if self._poisoned: 

626 base = self._build_base_index( 

627 self.dimension, self.index_type, self.metric 

628 ) 

629 self._index = IndexIDMap2(base) 

630 self._poisoned = False 

631 logger.warning( 

632 "Discarded a stale in-memory mutation from a prior failed " 

633 "apply() (no durable file to reconcile against)" 

634 ) 

635 return 

636 if self._integrity_verify is not None: 636 ↛ 643line 636 didn't jump to line 643 because the condition on line 636 was always true

637 ok, reason = self._integrity_verify(self._path) 

638 if not ok: 

639 raise RuntimeError( 

640 "Refusing to overwrite index whose integrity check failed " 

641 f"on reload ({reason})" 

642 ) 

643 try: 

644 index = read_index(str(self._path)) 

645 except Exception as exc: 

646 raise RuntimeError( 

647 f"Refusing to overwrite index that failed to reload: {exc}" 

648 ) from exc 

649 if not isinstance(index, IndexIDMap2): 649 ↛ 650line 649 didn't jump to line 650 because the condition on line 649 was never true

650 raise RuntimeError( 

651 f"Refusing to overwrite: on-disk index is " 

652 f"{type(index).__name__}, not IndexIDMap2 (foreign/corrupt)" 

653 ) 

654 # Same dimension guard load() enforces: adding a self.dimension-wide 

655 # vector into a differently-sized index would otherwise trip only 

656 # faiss's internal assert (stripped under `python -O`), silently 

657 # corrupting the index. Fail closed instead. 

658 if index.d != self.dimension: 

659 raise RuntimeError( 

660 f"Refusing to overwrite: on-disk index dimension {index.d} " 

661 f"!= configured {self.dimension} (embedding model/config changed)" 

662 ) 

663 # Same physical-type/metric guard load() enforces: a stale/foreign file 

664 # of the wrong base type or metric must not be adopted here either, or 

665 # apply() would misroute its delete (HNSW-as-flat raises; flat-as-hnsw 

666 # silently rebuilds the file as HNSW). 

667 reason = self._physical_config_reason( 

668 index, self.index_type, self.metric 

669 ) 

670 if reason: 

671 raise RuntimeError(f"Refusing to overwrite: {reason}") 

672 self._index = index 

673 # In-memory state now equals the verified durable file (integrity, 

674 # type, dimension and physical-config all checked above), so any stale 

675 # never-persisted mutation that poisoned this instance is resolved. 

676 # Clear the flag — mirroring the no-durable-file branch above — so a 

677 # concurrent read during _apply_hnsw's off-lock rebuild window is not 

678 # spuriously refused by _read_guard while self._index is actually valid. 

679 self._poisoned = False 

680 

681 def _restore_from_disk(self) -> None: 

682 """Best-effort resync of the in-memory index to the durable file. 

683 

684 Used after a failed ``apply()`` so a partial in-memory mutation (e.g. a 

685 removal that landed before an add/persist raised) never leaves this 

686 instance diverged from disk. If the file can't be read, drop nothing — 

687 the next ``apply()`` reloads under the lock anyway. 

688 """ 

689 try: 

690 if self._path is not None and self._path.exists(): 

691 index = read_index(str(self._path)) 

692 # Only trust a resync to a same-format, same-dimension AND 

693 # same physical-type/metric index; anything else means we cannot 

694 # safely represent disk state in memory, so fall through to 

695 # poisoning (rather than silently adopting a wrong-typed file). 

696 if ( 696 ↛ 718line 696 didn't jump to line 718 because the condition on line 696 was always true

697 isinstance(index, IndexIDMap2) 

698 and index.d == self.dimension 

699 and self._physical_config_reason( 

700 index, self.index_type, self.metric 

701 ) 

702 is None 

703 ): 

704 self._index = index 

705 # In-memory now matches the durable file again, so this 

706 # instance is consistent — clear any prior poisoning so reads 

707 # aren't blocked forever after a recovered failure. 

708 self._poisoned = False 

709 return 

710 except Exception: 

711 # Re-read/verify of the durable index failed; the poison flag and 

712 # warning just below handle it (reads are refused until the next 

713 # clean apply() reloads under the lock). 

714 logger.debug("Resync re-read of the durable index failed") 

715 # Could not resync in-memory state to disk. Mark poisoned so reads 

716 # refuse to serve a possibly half-mutated index; the next successful 

717 # apply() (which reloads under the lock) clears it. 

718 self._poisoned = True 

719 logger.warning( 

720 "Could not restore in-memory index from disk after a failed " 

721 "apply(); marking store poisoned until the next successful write" 

722 ) 

723 

724 @staticmethod 

725 def _require_add_vectors(add_vectors: Optional[np.ndarray]) -> np.ndarray: 

726 """Narrow ``add_vectors`` to non-Optional for the type checker. 

727 

728 ``apply()`` already raises this identical ``ValueError`` before 

729 dispatching to ``_apply_locked``/``_apply_hnsw`` whenever ``add_ids`` 

730 is non-empty and ``add_vectors`` is ``None``; this re-checks (rather 

731 than just asserting) so mypy can see the narrowed, non-Optional type 

732 at each call site without an inline ``raise`` inside a ``try`` block. 

733 """ 

734 if add_vectors is None: 734 ↛ 735line 734 didn't jump to line 735 because the condition on line 734 was never true

735 raise ValueError("add_ids given without add_vectors") 

736 return add_vectors 

737 

738 def _dedup_new( 

739 self, ids: List[int], vectors: np.ndarray 

740 ) -> Tuple[List[int], np.ndarray]: 

741 """Drop ids already present in the index and within-batch duplicates. 

742 

743 Keeps the first occurrence of each id. Returns aligned (ids, vectors). 

744 Runs under ``apply()``'s lock — uses the unlocked live-ids helper to 

745 avoid re-entrant self-deadlock. 

746 """ 

747 live = set(self._live_ids_unlocked()) 

748 keep_rows: List[int] = [] 

749 seen: set = set() 

750 for row, vid in enumerate(ids): 

751 if vid in live or vid in seen: 

752 continue 

753 seen.add(vid) 

754 keep_rows.append(row) 

755 if len(keep_rows) == len(ids): 

756 return list(ids), vectors 

757 kept_ids = [ids[r] for r in keep_rows] 

758 kept_vecs = ( 

759 np.asarray(vectors)[keep_rows] 

760 if keep_rows 

761 else np.empty((0, self.dimension), dtype="float32") 

762 ) 

763 return kept_ids, kept_vecs 

764 

765 def apply( 

766 self, 

767 *, 

768 add_ids: List[int], 

769 add_vectors: Optional[np.ndarray], 

770 remove_ids: Sequence[int] = (), 

771 dedup: bool = True, 

772 ) -> dict: 

773 if self._lock is None or self._path is None: 

774 raise RuntimeError( 

775 "FaissVectorStore.apply() requires a persistence binding " 

776 "(path + lock); construct via create()/load() with them set." 

777 ) 

778 add_ids = list(add_ids or []) 

779 # Drop None remove-ids up front so BOTH dispatch paths are safe: the 

780 # flat/IVF delete() already filters None, but the HNSW rebuild path 

781 # (_rebuild_dropping / IDSelectorBatch) does not and would raise an 

782 # unhandled TypeError. A None remove-id matches nothing anyway. 

783 remove_ids = [i for i in (remove_ids or []) if i is not None] 

784 if add_ids and add_vectors is None: 

785 raise ValueError("add_ids given without add_vectors") 

786 # HNSW has no in-place remove_ids: any removal or id-collision forces a 

787 # full O(collection) graph rebuild. Route that through the off-lock path 

788 # so the (expensive) rebuild does NOT hold the write lock that also 

789 # gates every search() — only the snapshot and the final swap are 

790 # locked. flat/IVF (real remove_ids) stay on the fully-locked path. 

791 if not self.supports_delete: 

792 return self._apply_hnsw(add_ids, add_vectors, remove_ids, dedup) 

793 return self._apply_locked(add_ids, add_vectors, remove_ids, dedup) 

794 

795 def _apply_locked( 

796 self, 

797 add_ids: List[int], 

798 add_vectors: Optional[np.ndarray], 

799 remove_ids: List[int], 

800 dedup: bool, 

801 ) -> dict: 

802 """Fully-locked apply: reload -> remove(+collision) -> add -> persist. 

803 

804 Used for flat/IVF (in-place ``remove_ids``) and as the race fallback for 

805 the off-lock HNSW path. Holds the write lock for the whole operation. 

806 """ 

807 # apply() already required a persistence binding before dispatching 

808 # here (path + lock not None) -- narrow the Optionals once, locally, 

809 # so mypy sees the same invariant. 

810 lock = self._lock 

811 path = self._path 

812 if lock is None or path is None: 812 ↛ 813line 812 didn't jump to line 813 because the condition on line 812 was never true

813 raise RuntimeError( 

814 "FaissVectorStore._apply_locked() requires a persistence " 

815 "binding (path + lock); apply() should have already " 

816 "validated this before dispatching here" 

817 ) 

818 with lock: 

819 # #4200: reload under the lock so we merge onto other writers' saves. 

820 # Fails closed if the on-disk file exists but can't be trusted. 

821 self._reload_under_lock() 

822 try: 

823 # Replace-on-collision. An add_id that is ALREADY live is either 

824 # an idempotent re-add OR — critically — a REUSED DocumentChunk.id 

825 # whose prior row was rolled back: SQLite recycles AUTOINCREMENT 

826 # ids on ROLLBACK, so a failed index() commit can leave a stale 

827 # orphan vector under an id that a later, unrelated chunk then 

828 # reuses. Either way the NEW vector must win, so every colliding 

829 # live id is removed along with remove_ids BEFORE the add. 

830 effective_remove = list(remove_ids) 

831 if add_ids: 

832 live_now = set(self._live_ids_unlocked()) 

833 effective_remove.extend( 

834 i for i in dict.fromkeys(add_ids) if i in live_now 

835 ) 

836 

837 if not effective_remove: 

838 removed = 0 

839 elif self.supports_delete: 

840 removed = self.delete(effective_remove) 

841 else: 

842 removed = self._rebuild_dropping(effective_remove) 

843 

844 added = 0 

845 if add_ids: 

846 checked_vectors = self._require_add_vectors(add_vectors) 

847 ids, vecs = ( 

848 self._dedup_new(add_ids, checked_vectors) 

849 if dedup 

850 else (add_ids, checked_vectors) 

851 ) 

852 if ids: 852 ↛ 858line 852 didn't jump to line 858 because the condition on line 852 was always true

853 self.add(ids, vecs) 

854 added = len(ids) 

855 

856 # Atomic save + integrity record, both under the same lock so a 

857 # concurrent writer can't slip bytes between them (#4197). 

858 self.persist(path) 

859 self._record_integrity() 

860 except Exception: 

861 # delete()/add() mutate self._index in place before persist; if 

862 # anything after the reload raised, resync in-memory state to the 

863 # durable file so search()/count() can't disagree with disk. 

864 self._restore_from_disk() 

865 raise 

866 self._poisoned = False 

867 return {"added": added, "removed": removed} 

868 

869 def _apply_hnsw( 

870 self, 

871 add_ids: List[int], 

872 add_vectors: Optional[np.ndarray], 

873 remove_ids: List[int], 

874 dedup: bool, 

875 _attempts: int = 3, 

876 ) -> dict: 

877 """Apply for HNSW (which has no in-place ``remove_ids``). 

878 

879 A pure add goes in-place under the lock (cheap — O(added)). A removal or 

880 id-collision forces a full graph rebuild; that expensive rebuild runs 

881 OFF the write lock from a locked snapshot, then is swapped in under a 

882 short lock only if no concurrent write happened meanwhile (version 

883 check). On a race it retries, then falls back to a fully-locked rebuild. 

884 So concurrent search() blocks only for the fast snapshot + the swap, 

885 not the whole O(collection) graph build. 

886 """ 

887 # apply() already required a persistence binding before dispatching 

888 # here (path + lock not None) -- narrow the Optionals once, locally, 

889 # so mypy sees the same invariant across both phases below. 

890 lock = self._lock 

891 path = self._path 

892 if lock is None or path is None: 892 ↛ 893line 892 didn't jump to line 893 because the condition on line 892 was never true

893 raise RuntimeError( 

894 "FaissVectorStore._apply_hnsw() requires a persistence " 

895 "binding (path + lock); apply() should have already " 

896 "validated this before dispatching here" 

897 ) 

898 # ---- Phase 1: snapshot under the lock (reconstruct is O(N) but fast) -- 

899 with lock: 

900 self._reload_under_lock() 

901 live = set(self._live_ids_unlocked()) 

902 # A live add id is a collision that must be REPLACED, so treat it as 

903 # a removal — the rebuild then re-adds it with the NEW vector. 

904 remove_set = ( 

905 {int(i) for i in remove_ids} | {i for i in add_ids if i in live} 

906 ) & live