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 q = self._prepare(query_vector) 

399 if self._index.ntotal == 0: 

400 return [] 

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

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

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

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

405 # this never changes the returned hits. 

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

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 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Index persistence intentionally targets its service-supplied local index path; glob is limited to that basename's temp files 

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

557 try: 

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

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

560 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Only aged temp files from this index basename's parent glob are unlinked 

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 try: 

581 # nosemgrep: semgrep.rules.path-traversal-risk, reason: tmp is the random O_EXCL mkstemp file created above inside the private index directory 

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

583 os.fsync(f.fileno()) 

584 except OSError: 

585 pass # best-effort fsync 

586 # nosemgrep: semgrep.rules.path-traversal-risk, reason: tmp is this persist call's mkstemp file in the destination directory; path is the explicit service-supplied index output already used by this persistence API 

587 os.replace(tmp, path) 

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

589 try: 

590 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Open the service-supplied index parent directory only to fsync the completed atomic rename 

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

592 try: 

593 os.fsync(dir_fd) 

594 finally: 

595 os.close(dir_fd) 

596 except OSError: 

597 pass # directory fsync is unsupported on some platforms 

598 finally: 

599 if tmp.exists(): 

600 try: 

601 # nosemgrep: semgrep.rules.path-traversal-risk, reason: tmp is this persist call's mkstemp file; final cleanup never uses the index payload as a path 

602 tmp.unlink() 

603 except OSError: 

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

605 

606 # ------------------------------------------------------------------ # 

607 # Durable batch write (apply) 

608 # ------------------------------------------------------------------ # 

609 def _reload_under_lock(self) -> None: 

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

611 

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

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

614 

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

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

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

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

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

620 guarantee and integrity detection. The caller decides whether to 

621 quarantine + rebuild. 

622 """ 

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

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

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

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

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

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

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

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

631 if self._poisoned: 

632 base = self._build_base_index( 

633 self.dimension, self.index_type, self.metric 

634 ) 

635 self._index = IndexIDMap2(base) 

636 self._poisoned = False 

637 logger.warning( 

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

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

640 ) 

641 return 

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

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

644 if not ok: 

645 raise RuntimeError( 

646 "Refusing to overwrite index whose integrity check failed " 

647 f"on reload ({reason})" 

648 ) 

649 try: 

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

651 except Exception as exc: 

652 raise RuntimeError( 

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

654 ) from exc 

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

656 raise RuntimeError( 

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

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

659 ) 

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

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

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

663 # corrupting the index. Fail closed instead. 

664 if index.d != self.dimension: 

665 raise RuntimeError( 

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

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

668 ) 

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

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

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

672 # silently rebuilds the file as HNSW). 

673 reason = self._physical_config_reason( 

674 index, self.index_type, self.metric 

675 ) 

676 if reason: 

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

678 self._index = index 

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

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

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

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

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

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

685 self._poisoned = False 

686 

687 def _restore_from_disk(self) -> None: 

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

689 

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

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

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

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

694 """ 

695 try: 

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

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

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

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

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

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

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

703 isinstance(index, IndexIDMap2) 

704 and index.d == self.dimension 

705 and self._physical_config_reason( 

706 index, self.index_type, self.metric 

707 ) 

708 is None 

709 ): 

710 self._index = index 

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

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

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

714 self._poisoned = False 

715 return 

716 except Exception: 

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

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

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

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

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

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

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

724 self._poisoned = True 

725 logger.warning( 

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

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

728 ) 

729 

730 @staticmethod 

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

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

733 

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

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

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

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

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

739 """ 

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

741 raise ValueError("add_ids given without add_vectors") 

742 return add_vectors 

743 

744 def _dedup_new( 

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

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

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

748 

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

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

751 avoid re-entrant self-deadlock. 

752 """ 

753 live = set(self._live_ids_unlocked()) 

754 keep_rows: List[int] = [] 

755 seen: set = set() 

756 for row, vid in enumerate(ids): 

757 if vid in live or vid in seen: 

758 continue 

759 seen.add(vid) 

760 keep_rows.append(row) 

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

762 return list(ids), vectors 

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

764 kept_vecs = ( 

765 np.asarray(vectors)[keep_rows] 

766 if keep_rows 

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

768 ) 

769 return kept_ids, kept_vecs 

770 

771 def apply( 

772 self, 

773 *, 

774 add_ids: List[int], 

775 add_vectors: Optional[np.ndarray], 

776 remove_ids: Sequence[int] = (), 

777 dedup: bool = True, 

778 ) -> dict: 

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

780 raise RuntimeError( 

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

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

783 ) 

784 add_ids = list(add_ids or []) 

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

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

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

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

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

790 if add_ids and add_vectors is None: 

791 raise ValueError("add_ids given without add_vectors") 

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

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

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

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

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

797 if not self.supports_delete: 

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

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

800 

801 def _apply_locked( 

802 self, 

803 add_ids: List[int], 

804 add_vectors: Optional[np.ndarray], 

805 remove_ids: List[int], 

806 dedup: bool, 

807 ) -> dict: 

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

809 

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

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

812 """ 

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

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

815 # so mypy sees the same invariant. 

816 lock = self._lock 

817 path = self._path 

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

819 raise RuntimeError( 

820 "FaissVectorStore._apply_locked() requires a persistence " 

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

822 "validated this before dispatching here" 

823 ) 

824 with lock: 

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

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

827 self._reload_under_lock() 

828 try: 

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

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

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

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

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

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

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

836 effective_remove = list(remove_ids) 

837 if add_ids: 

838 live_now = set(self._live_ids_unlocked()) 

839 effective_remove.extend( 

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

841 ) 

842 

843 if not effective_remove: 

844 removed = 0 

845 elif self.supports_delete: 

846 removed = self.delete(effective_remove) 

847 else: 

848 removed = self._rebuild_dropping(effective_remove) 

849 

850 added = 0 

851 if add_ids: 

852 checked_vectors = self._require_add_vectors(add_vectors) 

853 ids, vecs = ( 

854 self._dedup_new(add_ids, checked_vectors) 

855 if dedup 

856 else (add_ids, checked_vectors) 

857 ) 

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

859 self.add(ids, vecs) 

860 added = len(ids) 

861 

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

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

864 self.persist(path) 

865 self._record_integrity() 

866 except Exception: 

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

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

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

870 self._restore_from_disk() 

871 raise 

872 self._poisoned = False 

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

874 

875 def _apply_hnsw( 

876 self, 

877 add_ids: List[int], 

878 add_vectors: Optional[np.ndarray], 

879 remove_ids: List[int], 

880 dedup: bool, 

881 _attempts: int = 3, 

882 ) -> dict: 

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

884 

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

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

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

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

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

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

891 not the whole O(collection) graph build. 

892 """ 

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

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

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

896 lock = self._lock 

897 path = self._path 

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

899 raise RuntimeError( 

900 "FaissVectorStore._apply_hnsw() requires a persistence " 

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

902 "validated this before dispatching here" 

903 ) 

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

905 with lock: 

906 self._reload_under_lock() 

907 live = set(self._live_ids_unlocked()) 

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

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