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

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 threading 

19import time 

20from contextlib import contextmanager 

21from pathlib import Path 

22from typing import Callable, Iterator, List, Optional, Sequence, Tuple, cast 

23 

24import numpy as np 

25from faiss import ( 

26 METRIC_INNER_PRODUCT, 

27 METRIC_L2, 

28 IDSelectorBatch, 

29 IndexFlatIP, 

30 IndexFlatL2, 

31 IndexHNSWFlat, 

32 IndexIDMap2, 

33 downcast_index, 

34 normalize_L2, 

35 read_index, 

36 vector_to_array, 

37 write_index, 

38) 

39from loguru import logger 

40 

41from ..base import BaseVectorStore 

42 

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

44_HNSW_M = 32 

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

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

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

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

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

50_STALE_TMP_AGE_SECONDS = 3600 

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

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

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

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

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

56# pathological/overflowing magnitudes. 

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

58 

59 

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

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

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

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

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

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

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

67 

68 

69class FaissVectorStore(BaseVectorStore): 

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

71 

72 Concurrency: the injected write lock is a ``threading.Lock`` — it serializes 

73 writers (and, via ``_read_guard``, readers) within ONE process only. This 

74 assumes a 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[threading.Lock] = None 

108 self._integrity_record: Optional[Callable[[Path], None]] = None 

109 self._integrity_verify: Optional[ 

110 Callable[[Path], Tuple[bool, Optional[str]]] 

111 ] = None 

112 

113 def _bind( 

114 self, 

115 path: Optional[Path], 

116 lock: Optional[threading.Lock], 

117 integrity_record: Optional[Callable[[Path], None]], 

118 integrity_verify: Optional[ 

119 Callable[[Path], Tuple[bool, Optional[str]]] 

120 ], 

121 ) -> None: 

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

123 

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

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

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

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

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

129 """ 

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

131 self._lock = lock 

132 self._integrity_record = integrity_record 

133 self._integrity_verify = integrity_verify 

134 if ( 

135 self._path is not None 

136 and lock is not None 

137 and (integrity_record is None or integrity_verify is None) 

138 ): 

139 logger.warning( 

140 "FaissVectorStore bound for writes without integrity hooks; " 

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

142 ) 

143 

144 # ------------------------------------------------------------------ # 

145 # Construction 

146 # ------------------------------------------------------------------ # 

147 @staticmethod 

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

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

150 

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

152 """ 

153 index_type = _canon(index_type, "flat") 

154 metric = _canon(metric, "cosine") 

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

156 metric_type = ( 

157 METRIC_INNER_PRODUCT if metric in _IP_METRICS else METRIC_L2 

158 ) 

159 if index_type == "hnsw": 

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

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

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

163 return IndexHNSWFlat(dimension, _HNSW_M, metric_type) 

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

165 if metric in _IP_METRICS: 

166 return IndexFlatIP(dimension) 

167 return IndexFlatL2(dimension) 

168 

169 @staticmethod 

170 def _physical_config_reason( 

171 index: IndexIDMap2, index_type: str, metric: str 

172 ) -> Optional[str]: 

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

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

175 

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

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

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

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

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

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

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

183 """ 

184 base = downcast_index(index.index) 

185 want_type = _canon(index_type, "flat") 

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

187 return ( 

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

189 f"{want_type!r}" 

190 ) 

191 want_metric_type = ( 

192 METRIC_INNER_PRODUCT 

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

194 else METRIC_L2 

195 ) 

196 if base.metric_type != want_metric_type: 

197 return ( 

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

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

200 ) 

201 return None 

202 

203 @classmethod 

204 def create( 

205 cls, 

206 *, 

207 dimension: int, 

208 index_type: str, 

209 metric: str, 

210 normalize: bool, 

211 path: Optional[Path] = None, 

212 lock: Optional[threading.Lock] = None, 

213 integrity_record: Optional[Callable[[Path], None]] = None, 

214 integrity_verify: Optional[ 

215 Callable[[Path], Tuple[bool, Optional[str]]] 

216 ] = None, 

217 **kwargs, 

218 ) -> "FaissVectorStore": 

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

220 store = cls( 

221 IndexIDMap2(base), 

222 dimension=dimension, 

223 index_type=index_type, 

224 metric=metric, 

225 normalize=normalize, 

226 ) 

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

228 return store 

229 

230 @classmethod 

231 def load( 

232 cls, 

233 path: Path, 

234 *, 

235 dimension: int, 

236 index_type: str, 

237 metric: str, 

238 normalize: bool, 

239 lock: Optional[threading.Lock] = None, 

240 integrity_record: Optional[Callable[[Path], None]] = None, 

241 integrity_verify: Optional[ 

242 Callable[[Path], Tuple[bool, Optional[str]]] 

243 ] = None, 

244 **kwargs, 

245 ) -> "FaissVectorStore": 

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

247 

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

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

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

251 surfaced rather than silently coerced. 

252 

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

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

255 """ 

256 index = read_index(str(path)) 

257 if not isinstance(index, IndexIDMap2): 

258 raise ValueError( 

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

260 "(pre-migration or corrupt format)" 

261 ) 

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

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

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

265 if index.d != dimension: 

266 raise ValueError( 

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

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

269 ) 

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

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

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

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

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

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

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

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

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

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

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

281 # IP-vs-L2.) 

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

283 if reason: 

284 raise ValueError( 

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

286 "(reindex required)" 

287 ) 

288 store = cls( 

289 index, 

290 dimension=dimension, 

291 index_type=index_type, 

292 metric=metric, 

293 normalize=normalize, 

294 ) 

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

296 return store 

297 

298 # ------------------------------------------------------------------ # 

299 # Internal helpers 

300 # ------------------------------------------------------------------ # 

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

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

303 

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

305 touches the caller's array. 

306 """ 

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

308 if vecs.ndim == 1: 

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

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

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

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

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

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

315 raise ValueError( 

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

317 f"{self.dimension}" 

318 ) 

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

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

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

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

323 raise ValueError( 

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

325 ) 

326 if self.normalize: 

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

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

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

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

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

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

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

334 raise ValueError( 

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

336 "that L2 normalization cannot represent — bad embedding" 

337 ) 

338 normalize_L2(vecs) 

339 else: 

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

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

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

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

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

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

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

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

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

349 raise ValueError( 

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

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

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

353 ) 

354 return vecs 

355 

356 # ------------------------------------------------------------------ # 

357 # Vector operations 

358 # ------------------------------------------------------------------ # 

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

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

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

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

363 if len(id_arr) == 0: 

364 return 

365 vecs = self._prepare(vectors) 

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

367 raise ValueError( 

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

369 ) 

370 self._index.add_with_ids(vecs, id_arr) 

371 

372 @contextmanager 

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

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

375 poisoned instance. 

376 

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

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

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

380 that already hold the lock (e.g. ``_dedup_new``) must use the 

381 ``*_unlocked`` helpers, not these public methods, to avoid re-entrancy 

382 deadlock (``threading.Lock`` is not reentrant). 

383 """ 

384 if self._lock is not None: 

385 with self._lock: 

386 if self._poisoned: 

387 raise RuntimeError( 

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

389 "write; reload the index before reading" 

390 ) 

391 yield 

392 else: 

393 if self._poisoned: 

394 raise RuntimeError( 

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

396 "write; reload the index before reading" 

397 ) 

398 yield 

399 

400 def search( 

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

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

403 if k <= 0: 

404 return [] 

405 with self._read_guard(): 

406 if self._index.ntotal == 0: 

407 return [] 

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

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

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

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

412 # this never changes the returned hits. 

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

414 q = self._prepare(query_vector) 

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

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

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

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

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

420 continue 

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

422 return results 

423 

424 @property 

425 def supports_delete(self) -> bool: 

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

427 return self.index_type != "hnsw" 

428 

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

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

431 if not id_list: 

432 return 0 

433 if not self.supports_delete: 

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

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

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

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

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

439 raise RuntimeError( 

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

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

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

443 ) 

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

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

446 

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

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

449 rebuilding it from its own reconstructed vectors. 

450 

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

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

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

454 

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

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

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

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

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

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

461 """ 

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

463 live = self._live_ids_unlocked() 

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

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

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

467 return 0 

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

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

470 for row, vid in enumerate(keep): 

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

472 base = self._build_base_index( 

473 self.dimension, self.index_type, self.metric 

474 ) 

475 new_index = IndexIDMap2(base) 

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

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

478 self._index = new_index 

479 return removed 

480 

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

482 if self._index.ntotal == 0: 

483 return [] 

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

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

486 return cast( 

487 List[int], 

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

489 ) 

490 

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

492 with self._read_guard(): 

493 return self._live_ids_unlocked() 

494 

495 def count(self) -> int: 

496 with self._read_guard(): 

497 return int(self._index.ntotal) 

498 

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

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

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

502 with self._read_guard(): 

503 try: 

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

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

506 except RuntimeError: 

507 return None 

508 

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

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

511 

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

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

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

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

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

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

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

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

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

521 """ 

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

523 return None 

524 try: 

525 st = os.stat(self._path) 

526 except OSError: 

527 return None 

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

529 

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

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

532 

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

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

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

536 """ 

537 path = Path(path) 

538 path.parent.mkdir(parents=True, exist_ok=True) 

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

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

541 # left to the caller. 

542 try: 

543 path.parent.chmod(0o700) 

544 except OSError: 

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

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

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

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

549 # per-process (a threading.Lock in a per-process registry), so under 

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

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

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

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

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

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

556 now = time.time() 

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

558 try: 

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

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

561 stale.unlink() 

562 except OSError: 

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

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

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

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

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

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

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

570 fd, tmp_name = tempfile.mkstemp( 

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

572 ) 

573 os.close(fd) 

574 tmp = Path(tmp_name) 

575 try: 

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

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

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

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

580 with open(tmp, "rb") as f: 

581 os.fsync(f.fileno()) 

582 os.replace(tmp, path) 

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

584 try: 

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

586 try: 

587 os.fsync(dir_fd) 

588 finally: 

589 os.close(dir_fd) 

590 except OSError: 

591 pass # directory fsync is unsupported on some platforms 

592 finally: 

593 if tmp.exists(): 

594 try: 

595 tmp.unlink() 

596 except OSError: 

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

598 

599 # ------------------------------------------------------------------ # 

600 # Durable batch write (apply) 

601 # ------------------------------------------------------------------ # 

602 def _reload_under_lock(self) -> None: 

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

604 

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

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

607 

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

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

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

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

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

613 guarantee and integrity detection. The caller decides whether to 

614 quarantine + rebuild. 

615 """ 

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

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

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

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

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

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

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

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

624 if self._poisoned: 

625 base = self._build_base_index( 

626 self.dimension, self.index_type, self.metric 

627 ) 

628 self._index = IndexIDMap2(base) 

629 self._poisoned = False 

630 logger.warning( 

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

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

633 ) 

634 return 

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

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

637 if not ok: 

638 raise RuntimeError( 

639 "Refusing to overwrite index whose integrity check failed " 

640 f"on reload ({reason})" 

641 ) 

642 try: 

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

644 except Exception as exc: 

645 raise RuntimeError( 

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

647 ) from exc 

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

649 raise RuntimeError( 

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

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

652 ) 

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

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

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

656 # corrupting the index. Fail closed instead. 

657 if index.d != self.dimension: 

658 raise RuntimeError( 

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

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

661 ) 

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

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

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

665 # silently rebuilds the file as HNSW). 

666 reason = self._physical_config_reason( 

667 index, self.index_type, self.metric 

668 ) 

669 if reason: 

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

671 self._index = index 

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

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

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

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

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

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

678 self._poisoned = False 

679 

680 def _restore_from_disk(self) -> None: 

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

682 

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

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

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

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

687 """ 

688 try: 

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

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

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

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

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

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

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

696 isinstance(index, IndexIDMap2) 

697 and index.d == self.dimension 

698 and self._physical_config_reason( 

699 index, self.index_type, self.metric 

700 ) 

701 is None 

702 ): 

703 self._index = index 

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

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

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

707 self._poisoned = False 

708 return 

709 except Exception: 

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

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

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

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

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

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

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

717 self._poisoned = True 

718 logger.warning( 

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

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

721 ) 

722 

723 @staticmethod 

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

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

726 

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

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

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

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

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

732 """ 

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

734 raise ValueError("add_ids given without add_vectors") 

735 return add_vectors 

736 

737 def _dedup_new( 

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

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

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

741 

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

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

744 avoid re-entrant self-deadlock. 

745 """ 

746 live = set(self._live_ids_unlocked()) 

747 keep_rows: List[int] = [] 

748 seen: set = set() 

749 for row, vid in enumerate(ids): 

750 if vid in live or vid in seen: 

751 continue 

752 seen.add(vid) 

753 keep_rows.append(row) 

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

755 return list(ids), vectors 

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

757 kept_vecs = ( 

758 np.asarray(vectors)[keep_rows] 

759 if keep_rows 

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

761 ) 

762 return kept_ids, kept_vecs 

763 

764 def apply( 

765 self, 

766 *, 

767 add_ids: List[int], 

768 add_vectors: Optional[np.ndarray], 

769 remove_ids: Sequence[int] = (), 

770 dedup: bool = True, 

771 ) -> dict: 

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

773 raise RuntimeError( 

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

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

776 ) 

777 add_ids = list(add_ids or []) 

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

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

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

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

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

783 if add_ids and add_vectors is None: 

784 raise ValueError("add_ids given without add_vectors") 

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

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

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

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

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

790 if not self.supports_delete: 

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

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

793 

794 def _apply_locked( 

795 self, 

796 add_ids: List[int], 

797 add_vectors: Optional[np.ndarray], 

798 remove_ids: List[int], 

799 dedup: bool, 

800 ) -> dict: 

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

802 

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

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

805 """ 

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

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

808 # so mypy sees the same invariant. 

809 lock = self._lock 

810 path = self._path 

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

812 raise RuntimeError( 

813 "FaissVectorStore._apply_locked() requires a persistence " 

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

815 "validated this before dispatching here" 

816 ) 

817 with lock: 

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

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

820 self._reload_under_lock() 

821 try: 

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

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

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

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

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

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

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

829 effective_remove = list(remove_ids) 

830 if add_ids: 

831 live_now = set(self._live_ids_unlocked()) 

832 effective_remove.extend( 

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

834 ) 

835 

836 if not effective_remove: 

837 removed = 0 

838 elif self.supports_delete: 

839 removed = self.delete(effective_remove) 

840 else: 

841 removed = self._rebuild_dropping(effective_remove) 

842 

843 added = 0 

844 if add_ids: 

845 checked_vectors = self._require_add_vectors(add_vectors) 

846 ids, vecs = ( 

847 self._dedup_new(add_ids, checked_vectors) 

848 if dedup 

849 else (add_ids, checked_vectors) 

850 ) 

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

852 self.add(ids, vecs) 

853 added = len(ids) 

854 

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

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

857 self.persist(path) 

858 self._record_integrity() 

859 except Exception: 

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

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

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

863 self._restore_from_disk() 

864 raise 

865 self._poisoned = False 

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

867 

868 def _apply_hnsw( 

869 self, 

870 add_ids: List[int], 

871 add_vectors: Optional[np.ndarray], 

872 remove_ids: List[int], 

873 dedup: bool, 

874 _attempts: int = 3, 

875 ) -> dict: 

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

877 

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

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

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

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

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

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

884 not the whole O(collection) graph build. 

885 """ 

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

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

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

889 lock = self._lock 

890 path = self._path 

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

892 raise RuntimeError( 

893 "FaissVectorStore._apply_hnsw() requires a persistence " 

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

895 "validated this before dispatching here" 

896 ) 

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

898 with lock: