Coverage for src/local_deep_research/research_library/services/library_rag_service.py: 87%
930 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""
2Library RAG Service
4Handles indexing and searching library documents using RAG:
5- Index text documents into vector database
6- Chunk documents for semantic search
7- Generate embeddings using local models
8- Manage FAISS indices per research
9- Track RAG status in library
10"""
12import threading
13import time
14import json
15from contextlib import contextmanager
16from dataclasses import dataclass
17from datetime import datetime, UTC
19import numpy as np
20from pathlib import Path
21from typing import Any, Callable, Dict, List, Optional, Set, Tuple
22from concurrent.futures import ThreadPoolExecutor, as_completed, CancelledError
24from langchain_core.documents import Document as LangchainDocument
25from loguru import logger
26from sqlalchemy import func
27from sqlalchemy.exc import IntegrityError
28from sqlalchemy.orm import Session
30from ...config.paths import get_cache_directory
31from ...constants import (
32 DEFAULT_LOCAL_SEARCH_CHUNK_OVERLAP,
33 DEFAULT_LOCAL_SEARCH_CHUNK_SIZE,
34 DEFAULT_LOCAL_SEARCH_DISTANCE_METRIC,
35 DEFAULT_LOCAL_SEARCH_INDEX_TYPE,
36 DEFAULT_LOCAL_SEARCH_MODEL,
37 DEFAULT_LOCAL_SEARCH_NORMALIZE_VECTORS,
38 DEFAULT_LOCAL_SEARCH_PROVIDER,
39 DEFAULT_LOCAL_SEARCH_SPLITTER_TYPE,
40 DEFAULT_LOCAL_SEARCH_TEXT_SEPARATORS,
41)
42from ...database.models.library import (
43 Document,
44 DocumentChunk,
45 DocumentCollection,
46 Collection,
47 RAGIndex,
48 RAGIndexStatus,
49 RagDocumentStatus,
50 EmbeddingProvider,
51)
52from ...database.session_context import get_user_db_session, safe_rollback
53from ...utilities.type_utils import to_bool
54from ..utils import ensure_in_collection
55from ...embeddings.splitters import get_text_splitter
56from ...web_search_engines.engines.local_embedding_manager import (
57 LocalEmbeddingManager,
58)
59from ...security.directory_creation import create_directory
60from ...security.file_integrity import FileIntegrityManager, FAISSIndexVerifier
61from ...security.file_integrity.integrity_manager import NO_INTEGRITY_RECORD
62from ...vector_stores.facade import VectorIndex, ChunkInput, SearchResult
63from langchain_core.embeddings import Embeddings
64import hashlib
67@dataclass
68class _PreparedDocument:
69 document_id: str
70 collection_id: str
71 chunk_inputs: List[ChunkInput]
72 vectors: np.ndarray
75class _LazyEmbeddings(Embeddings):
76 """Embeddings proxy that defers backend construction until first use.
78 A ``VectorIndex`` built purely to ``delete()`` never touches ``embeddings``
79 (delete resolves rows by id and calls ``store.apply(remove_ids=...)``).
80 ``LocalEmbeddingManager.embeddings`` is a lazy property that, on first
81 access, synchronously loads the real backend (a SentenceTransformer costs
82 seconds and hundreds of MB of RSS, an Ollama client opens live httpx
83 connections). Handing that property straight to the delete path forces that
84 cost once per (deleted-document x collection). This proxy pulls the real
85 embeddings from the manager only if ``embed_*`` is genuinely called, so the
86 delete path pays nothing while index/search paths still work unchanged.
87 """
89 def __init__(self, manager) -> None:
90 self._manager = manager
92 def embed_documents(self, texts):
93 return self._manager.embeddings.embed_documents(texts)
95 def embed_query(self, text):
96 return self._manager.embeddings.embed_query(text)
99class _TrackedRLock:
100 """RLock proxy that records active keys even for direct acquire() calls.
102 The recursion depth is tracked locally (not via ``threading.RLock._is_owned``,
103 which is a CPython-private API and is not part of the threading contract —
104 absent on PyPy/other implementations). The key set membership is mutated
105 atomically with respect to ``acquire`` and ``pop_faiss_locks_for_user``: the
106 per-(user, index_path) entry is added when ownership transitions 0 → 1 and
107 discarded when it transitions 1 → 0, and both transitions happen under
108 ``_faiss_write_locks_lock``.
110 Without holding ``_faiss_write_locks_lock`` across release, there would be a
111 race window: thread A could release → thread B (blocked in ``acquire``) unblock
112 and add its key → thread A then observe depth 0 and discard B's key. With
113 ``pop_faiss_locks_for_user`` evicting a "free" lock, the next writer would
114 create a fresh lock for the same ``(username, index_path)`` and race B on the
115 FAISS file. Holding ``_faiss_write_locks_lock`` across inner release + depth
116 check + discard ensures no concurrent acquire's key-add can interleave with
117 key-discard.
118 """
120 def __init__(self, key: Tuple[str, str]):
121 self._key = key
122 self._lock = threading.RLock()
123 # Per-wrapper recursion depth (independent of the inner RLock so we
124 # never need the CPython-private ``_is_owned`` accessor). ``_owner``
125 # records the outer holder's thread id, or ``None`` when free.
126 # ``_tracks_active_key`` is True only when this wrapper added
127 # ``self._key`` to ``_faiss_active_lock_keys`` on its outer acquire
128 # (canonical wrappers only) so a non-canonical release cannot erase
129 # another holder's active-key entry.
130 self._depth = 0
131 self._owner = None
132 self._tracks_active_key = False
134 def acquire(self, *args, **kwargs):
135 first_acquire = self._lock.acquire(*args, **kwargs)
136 if first_acquire:
137 # ``RLock.acquire`` returns True for the outer acquisition AND for
138 # inner re-acquires. We detect outer-vs-inner via the wrapper's
139 # tracked ``_owner``: if the current thread already holds the lock,
140 # this is a re-entry; otherwise it's a fresh outer acquisition.
141 #
142 # Only the *canonical* wrapper for ``self._key`` (the one currently
143 # stored in ``_faiss_write_locks``) may mark the key active. A
144 # non-canonical wrapper can still be acquired briefly during the
145 # pop-during-acquire TOCTOU window handled by
146 # ``_hold_faiss_write_lock``; if it also added the key, its matching
147 # ``release()`` would discard the key while the *canonical* holder
148 # still owned the file — reopening the concurrent-FAISS-writer race.
149 was_owner = self._owner == threading.get_ident()
150 if not was_owner:
151 with _faiss_write_locks_lock:
152 if _faiss_write_locks.get(self._key) is self:
153 _faiss_active_lock_keys.add(self._key)
154 self._tracks_active_key = True
155 else:
156 self._tracks_active_key = False
157 self._owner = threading.get_ident()
158 self._depth += 1
159 return first_acquire
161 def release(self):
162 # Atomic with respect to acquire's add: hold the key-set lock across
163 # the inner release and the depth check so a concurrent acquire's
164 # key-add cannot interleave with our key-discard. ``_faiss_write_locks_lock``
165 # is taken BEFORE the inner release (the old bug was the reverse order).
166 #
167 # Only discard the active key if *this* wrapper registered it on
168 # acquire. Non-canonical wrappers never register, so their release
169 # cannot erase a concurrent canonical holder's active-key entry.
170 with _faiss_write_locks_lock:
171 self._lock.release()
172 self._depth -= 1
173 if self._depth == 0:
174 if self._tracks_active_key:
175 _faiss_active_lock_keys.discard(self._key)
176 self._tracks_active_key = False
177 self._owner = None
179 def locked(self) -> bool:
180 """True iff the current thread holds the wrapper at depth >= 1."""
181 return self._owner == threading.get_ident() and self._depth > 0
183 def __enter__(self):
184 self.acquire()
185 return self
187 def __exit__(self, *args):
188 self.release()
191# load_or_create verify→quarantine→build sequence per (username, index_path).
192# MUST be module-level: each auto-index / scheduler / search worker constructs
193# its own LibraryRAGService, so an instance-scoped lock would coordinate
194# nothing. Pattern is adapted from web/queue/processor_v2._user_critical_locks
195# — instance-scoped→module-scoped, username-only-key→(username, path)-key.
196# See #4197 for the race this guards (concurrent save_local interleaves bytes,
197# producing checksum_mismatch → destructive unlink, lost data).
198_faiss_write_locks: Dict[Tuple[str, str], _TrackedRLock] = {}
199_faiss_write_locks_lock = threading.Lock()
200_faiss_active_lock_keys: Set[Tuple[str, str]] = set()
202# Hard cap on suffix-increment retries when generating the .corrupt-<ns> path.
203# Normal case is one attempt — same-ns collisions only happen if the user
204# manually created such files. 32 is a safety bound that converts a deadlock
205# into a loud OSError.
206_QUARANTINE_SUFFIX_RETRY_CAP = 32
208# After a successful quarantine, keep at most this many older .corrupt-*
209# files per base path (per side: .faiss.corrupt-* and .pkl.corrupt-* are
210# counted independently). Prevents unbounded disk growth on systems that
211# experience recurring corruption while preserving recent diagnostic
212# artefacts. Keeping 5 means the user has the last 5 corruption events
213# to inspect / submit with a bug report; anything older is dropped.
214_QUARANTINE_KEEP_RECENT = 5
217def _get_faiss_write_lock(username: str, index_path: str) -> _TrackedRLock:
218 """Return the lock for ``(username, index_path)``, creating it on first
219 access. Key is normalised via ``Path.resolve()`` to match
220 ``FileIntegrityManager._normalize_path`` so writers/readers/quarantine
221 all agree on identity.
222 """
223 key = (username, str(Path(index_path).resolve()))
224 with _faiss_write_locks_lock:
225 lock = _faiss_write_locks.get(key)
226 if lock is None:
227 lock = _TrackedRLock(key)
228 _faiss_write_locks[key] = lock
229 return lock
232@contextmanager
233def _hold_faiss_write_lock(username: str, index_path: str):
234 """Context manager that acquires the per-(username, index_path) FAISS
235 write lock, retrying across a pop-during-acquire TOCTOU window.
237 ``_get_faiss_write_lock`` returns a wrapper reference, but
238 ``pop_faiss_locks_for_user`` may run in the gap before our
239 ``acquire()`` adds the key to ``_faiss_active_lock_keys`` and therefore
240 before the key is "active" from ``pop``'s perspective. Without a re-check,
241 the caller would hold a non-canonical wrapper while ``pop`` discarded the
242 canonical entry and the next writer created a fresh lock for the same
243 ``(username, index_path)`` — reintroducing the same-(user, path)
244 concurrent-FAISS-writer race this module exists to eliminate.
246 The window is microseconds and requires a per-user lock cleanup to fire
247 mid-acquire, so this is a defence-in-depth tightening rather than a
248 regression fix — but it is cheap to close. After ``acquire()`` returns,
249 we re-check under ``_faiss_write_locks_lock`` that ``self`` is still the
250 canonical entry; if not, we release and retry with the new canonical
251 lock. Re-entry on the same wrapper is impossible here because each
252 iteration acquires a fresh wrapper reference.
253 """
254 key = (username, str(Path(index_path).resolve()))
255 while True:
256 lock = _get_faiss_write_lock(username, index_path)
257 lock.acquire()
258 with _faiss_write_locks_lock:
259 canonical = _faiss_write_locks.get(key)
260 is_canonical = canonical is lock
261 if is_canonical:
262 try:
263 yield lock
264 finally:
265 lock.release()
266 return
267 # Evicted between _get_faiss_write_lock's lookup and our acquire().
268 # Release our (non-canonical) wrapper — acquire only marks the key
269 # active for the canonical dict entry, so this release cannot erase
270 # a concurrent canonical holder's active-key entry — then retry
271 # with the canonical lock.
272 lock.release()
275def pop_faiss_locks_for_user(username: str) -> None:
276 """Remove inactive FAISS-write locks for a user."""
277 with _faiss_write_locks_lock:
278 stale = [k for k in _faiss_write_locks if k[0] == username]
279 for k in stale:
280 if k not in _faiss_active_lock_keys:
281 _faiss_write_locks.pop(k, None)
284class LibraryRAGService:
285 """Service for managing RAG indexing of library documents."""
287 def __init__(
288 self,
289 username: str,
290 embedding_model: str = DEFAULT_LOCAL_SEARCH_MODEL,
291 embedding_provider: str = DEFAULT_LOCAL_SEARCH_PROVIDER,
292 chunk_size: int = DEFAULT_LOCAL_SEARCH_CHUNK_SIZE,
293 chunk_overlap: int = DEFAULT_LOCAL_SEARCH_CHUNK_OVERLAP,
294 splitter_type: str = DEFAULT_LOCAL_SEARCH_SPLITTER_TYPE,
295 text_separators: Optional[list] = None,
296 distance_metric: str = DEFAULT_LOCAL_SEARCH_DISTANCE_METRIC,
297 normalize_vectors: bool = DEFAULT_LOCAL_SEARCH_NORMALIZE_VECTORS,
298 index_type: str = DEFAULT_LOCAL_SEARCH_INDEX_TYPE,
299 embedding_manager: Optional["LocalEmbeddingManager"] = None,
300 db_password: Optional[str] = None,
301 ):
302 """
303 Initialize library RAG service for a user.
305 Args:
306 username: Username for database access
307 embedding_model: Name of the embedding model to use
308 embedding_provider: Provider type ('sentence_transformers' or 'ollama')
309 chunk_size: Size of text chunks for splitting
310 chunk_overlap: Overlap between consecutive chunks
311 splitter_type: Type of splitter ('recursive', 'token', 'sentence', 'semantic')
312 text_separators: List of text separators for chunking (default: ["\n\n", "\n", ". ", " ", ""])
313 distance_metric: Distance metric ('cosine', 'l2', or 'dot_product')
314 normalize_vectors: Whether to normalize vectors with L2
315 index_type: FAISS index type ('flat', 'hnsw', or 'ivf')
316 embedding_manager: Optional pre-constructed LocalEmbeddingManager for testing/flexibility
317 db_password: Optional database password for background thread access
318 """
319 self.username = username
320 self._db_password = db_password # Can be used for thread access
321 # Initialize optional attributes to None before they're set below
322 # This allows the db_password setter to check them without hasattr
323 self.embedding_manager = None
324 self.integrity_manager = None
325 self.embedding_model = embedding_model
326 self.embedding_provider = embedding_provider
327 self.chunk_size = chunk_size
328 self.chunk_overlap = chunk_overlap
329 self.splitter_type = splitter_type
330 self.text_separators = (
331 text_separators
332 if text_separators is not None
333 else list(DEFAULT_LOCAL_SEARCH_TEXT_SEPARATORS)
334 )
335 self.distance_metric = distance_metric
336 # Ensure normalize_vectors is always a proper boolean
337 self.normalize_vectors = to_bool(normalize_vectors, default=True)
338 self.index_type = index_type
340 # Emit the active configuration so users can confirm their
341 # UI-configured embedding settings are being honored (regression
342 # signal for #3453).
343 logger.info(
344 f"RAG service initialized for user={username}: "
345 f"provider={embedding_provider} model={embedding_model} "
346 f"chunk_size={chunk_size} chunk_overlap={chunk_overlap} "
347 f"splitter={splitter_type} index_type={index_type}"
348 )
350 # Use provided embedding manager or create a new one
351 # (Must be created before text splitter for semantic chunking)
352 # Track ownership so close() only tears down the manager when we
353 # constructed it — a caller-supplied manager stays under caller
354 # control (test fixtures, multi-service callers reusing one manager).
355 self._owns_embedding_manager = embedding_manager is None
356 if embedding_manager is not None:
357 self.embedding_manager = embedding_manager
358 else:
359 # Load the complete user settings snapshot from database using the
360 # proper method.
361 #
362 # STRICT snapshot is mandatory at this seam: a non-strict read
363 # silently falls back to the JSON defaults when the underlying
364 # settings query fails (SQLAlchemyError / stale enum row). Those
365 # defaults can lack an operator-selected cloud embedding
366 # provider's classification inputs and admit a cloud embedder
367 # under a local-only posture. We refuse before any
368 # LocalEmbeddingManager construction instead.
369 from ...settings.manager import SettingsManager
370 from ...security.egress.policy import (
371 Decision,
372 PolicyDeniedError,
373 )
375 # Use proper database session for SettingsManager
376 # Note: using _db_password (backing field) directly here because the
377 # db_password property setter propagates to embedding_manager/integrity_manager,
378 # which are still None at this point in __init__.
379 try:
380 with get_user_db_session(
381 username, self._db_password
382 ) as session:
383 settings_manager = SettingsManager(session)
384 settings_snapshot = settings_manager.get_settings_snapshot(
385 strict=True
386 )
387 except PolicyDeniedError:
388 raise
389 except Exception as exc:
390 raise PolicyDeniedError(
391 Decision(False, "settings_unavailable"),
392 target=embedding_provider,
393 ) from exc
395 # Add the specific settings needed for this RAG service
396 settings_snapshot.update(
397 {
398 "_username": username,
399 "embeddings.provider": embedding_provider,
400 f"embeddings.{embedding_provider}.model": embedding_model,
401 "local_search_chunk_size": chunk_size,
402 "local_search_chunk_overlap": chunk_overlap,
403 }
404 )
406 # Egress policy pre-flight at the constructor boundary so
407 # every direct ``LibraryRAGService(...)`` construction site
408 # is covered, not just the factory. Skipped when an
409 # ``embedding_manager`` is injected (tests / advanced flows)
410 # — those callers vouch for the manager themselves.
411 #
412 # Build the context from the ACTUAL scope (not a hardcoded
413 # BOTH) so PRIVATE_ONLY forces local embeddings even when the
414 # raw embeddings.require_local flag is at its default False —
415 # context_from_snapshot applies that coupling. Resolve the
416 # primary via the shared helper (single source of truth)
417 # instead of the old search.tool + searxng fallback, which
418 # was a fail-OPEN: a missing primary defaulted to the public
419 # searxng so the scope relaxed and a cloud embedder could be
420 # admitted. A missing primary now raises -> fail closed via
421 # the ValueError handler below.
422 from ...security.egress.policy import (
423 context_from_snapshot,
424 evaluate_embeddings,
425 resolve_run_primary_engine,
426 )
428 try:
429 primary = resolve_run_primary_engine(settings_snapshot)
430 policy_ctx = context_from_snapshot(
431 settings_snapshot, primary, username=username
432 )
433 except PolicyDeniedError:
434 raise
435 except ValueError as exc:
436 raise PolicyDeniedError(
437 Decision(False, "invalid_policy_config"),
438 target=embedding_provider,
439 ) from exc
440 if policy_ctx.require_local_embeddings: 440 ↛ 441line 440 didn't jump to line 441 because the condition on line 440 was never true
441 decision = evaluate_embeddings(
442 embedding_provider,
443 policy_ctx,
444 settings_snapshot=settings_snapshot,
445 )
446 if not decision.allowed:
447 logger.bind(policy_audit=True).warning(
448 "LibraryRAGService refused by egress policy",
449 provider=embedding_provider,
450 reason=decision.reason,
451 )
452 raise PolicyDeniedError(decision, target=embedding_provider)
454 self.embedding_manager = LocalEmbeddingManager(
455 embedding_model=embedding_model,
456 embedding_model_type=embedding_provider,
457 settings_snapshot=settings_snapshot,
458 )
460 # Initialize text splitter based on type
461 # (Must be created AFTER embedding_manager for semantic chunking)
462 self.text_splitter = get_text_splitter(
463 splitter_type=self.splitter_type,
464 chunk_size=self.chunk_size,
465 chunk_overlap=self.chunk_overlap,
466 text_separators=self.text_separators,
467 embeddings=self.embedding_manager.embeddings
468 if self.splitter_type == "semantic"
469 else None,
470 )
472 self.rag_index_record = None
474 # Initialize file integrity manager for FAISS indexes
475 self.integrity_manager = FileIntegrityManager(
476 username, password=self._db_password
477 )
478 self.integrity_manager.register_verifier(FAISSIndexVerifier())
480 self._closed = False
482 def close(self):
483 """Release embedding model and index resources."""
484 if self._closed:
485 return
486 self._closed = True
488 # Release embedding manager (which in turn closes the underlying
489 # OllamaEmbeddings httpx clients — see LocalEmbeddingManager.close).
490 # Only when we own it; caller-supplied managers stay under caller
491 # control to avoid double-close / use-after-close.
492 if self.embedding_manager is not None:
493 if self._owns_embedding_manager:
494 self.embedding_manager.close()
495 self.embedding_manager = None
497 # Clear other resources
498 self.rag_index_record = None
499 self.integrity_manager = None
500 self.text_splitter = None
502 def __enter__(self):
503 """Enter context manager."""
504 return self
506 def __exit__(self, exc_type, exc_val, exc_tb):
507 """Exit context manager, ensuring cleanup."""
508 self.close()
509 return False
511 @property
512 def db_password(self):
513 """Get database password."""
514 return self._db_password
516 @db_password.setter
517 def db_password(self, value):
518 """Set database password and propagate to embedding manager and integrity manager."""
519 self._db_password = value
520 if self.embedding_manager:
521 self.embedding_manager.db_password = value
522 if self.integrity_manager:
523 self.integrity_manager.password = value
525 def _get_index_hash(
526 self,
527 collection_name: str,
528 embedding_model: str,
529 embedding_model_type: str,
530 ) -> str:
531 """Generate hash for index identification."""
532 identity = {
533 "chunk_overlap": self.chunk_overlap,
534 "chunk_size": self.chunk_size,
535 "collection_name": collection_name,
536 "distance_metric": self.distance_metric,
537 "embedding_model": embedding_model,
538 "embedding_provider": embedding_model_type,
539 "index_type": self.index_type,
540 "normalize_vectors": self.normalize_vectors,
541 "splitter_type": self.splitter_type,
542 "text_separators": self.text_separators,
543 }
544 hash_input = json.dumps(
545 identity, ensure_ascii=False, separators=(",", ":"), sort_keys=True
546 )
547 return hashlib.sha256(hash_input.encode()).hexdigest()
549 def _matches_current_index_configuration(self, rag_index: RAGIndex) -> bool:
550 stored_provider = rag_index.embedding_model_type
551 if isinstance(stored_provider, EmbeddingProvider): 551 ↛ 553line 551 didn't jump to line 553 because the condition on line 551 was always true
552 stored_provider = stored_provider.value
553 stored_separators = rag_index.text_separators
554 if stored_separators is None:
555 stored_separators = DEFAULT_LOCAL_SEARCH_TEXT_SEPARATORS
556 if not isinstance(stored_separators, list) or not all( 556 ↛ 559line 556 didn't jump to line 559 because the condition on line 556 was never true
557 isinstance(separator, str) for separator in stored_separators
558 ):
559 return False
560 return (
561 rag_index.embedding_model == self.embedding_model
562 and stored_provider == self.embedding_provider
563 and rag_index.chunk_size == self.chunk_size
564 and rag_index.chunk_overlap == self.chunk_overlap
565 and (rag_index.splitter_type or "recursive") == self.splitter_type
566 and stored_separators == self.text_separators
567 and (rag_index.distance_metric or "cosine") == self.distance_metric
568 and (
569 rag_index.normalize_vectors
570 if rag_index.normalize_vectors is not None
571 else True
572 )
573 == self.normalize_vectors
574 and (rag_index.index_type or "flat") == self.index_type
575 )
577 def _find_matching_rag_index(
578 self, db_session: Session, collection_name: str, index_hash: str
579 ) -> RAGIndex | None:
580 rag_index = (
581 db_session.query(RAGIndex).filter_by(index_hash=index_hash).first()
582 )
583 if rag_index is not None:
584 return rag_index
586 candidates = (
587 db_session.query(RAGIndex)
588 .filter_by(
589 collection_name=collection_name,
590 embedding_model=self.embedding_model,
591 embedding_model_type=EmbeddingProvider(self.embedding_provider),
592 )
593 .all()
594 )
595 for candidate in candidates:
596 if self._matches_current_index_configuration(candidate): 596 ↛ 595line 596 didn't jump to line 595 because the condition on line 596 was always true
597 return candidate
598 return None
600 def _resolve_index_for_config(
601 self, db_session: Session, collection_id: str
602 ) -> RAGIndex | None:
603 """The collection's index for this service's configuration, or None.
605 Read-only: unlike ``_get_or_create_rag_index`` it never creates a
606 row or probes the embedding provider.
607 """
608 collection_name = f"collection_{collection_id}"
609 return self._find_matching_rag_index(
610 db_session,
611 collection_name,
612 self._get_index_hash(
613 collection_name, self.embedding_model, self.embedding_provider
614 ),
615 )
617 def _get_index_path(self, index_hash: str) -> Path:
618 """Get path for FAISS index file.
620 Files are scoped to a per-user subdirectory derived from a hash
621 of the username so two users with identical collection name +
622 embedding model don't share the same .faiss/.pkl files on disk.
623 The previous layout used a shared `cache/rag_indices/` for all
624 users, which would let one user's vectors land in another
625 user's load path in any deployment running more than one
626 account against the same data dir.
628 Raises:
629 ValueError: If self.username is empty/None. sha256("") is a
630 known constant, so silently hashing an empty username would
631 let two different "no username" service instances collide
632 into one shared cache directory instead of failing loudly.
633 Defense-in-depth -- see #5481.
634 """
635 if not self.username:
636 raise ValueError(
637 "Cannot compute a RAG index path for an empty username"
638 )
639 # Store in centralized cache directory (respects LDR_DATA_DIR)
640 user_scope = hashlib.sha256(self.username.encode("utf-8")).hexdigest()[
641 :16
642 ]
643 rag_root = get_cache_directory() / "rag_indices"
644 cache_dir = rag_root / user_scope
645 create_directory(
646 cache_dir, context="per-user RAG index cache directory"
647 )
648 # 0o700: filesystem-level defense-in-depth so other local OS accounts
649 # can't browse another LDR user's vector caches. mkdir(parents=True)
650 # silently creates the intermediate rag_indices/ root under the process
651 # umask (typically 0o755, world-traversable), so harden the ROOT too —
652 # chmod'ing only the per-user leaf leaves a 0o755 rag_indices/ that lets
653 # another OS user traverse in and stat/enumerate every account's caches.
654 for d in (rag_root, cache_dir):
655 try:
656 d.chmod(0o700)
657 except OSError:
658 # Best-effort: some filesystems (e.g. Windows shares) don't
659 # honor chmod; the per-user path scope is still applied.
660 pass
661 return cache_dir / f"{index_hash}.faiss"
663 def _migrate_legacy_index_files(
664 self, legacy_path: Path, index_path: Path, rag_index: RAGIndex
665 ) -> None:
666 """Relocate a pre-per-user-scoping index into the per-user directory.
668 Before the per-user path scoping, indexes lived directly in the
669 shared ``cache/rag_indices/`` directory. ``_preflight_index_path``
670 recomputes the path and ignores the stored one, so without this
671 one-time migration an upgrading user's existing index is never
672 found again: semantic search silently returns nothing while
673 ``DocumentCollection.indexed`` still claims everything is indexed,
674 so a non-force re-index no-ops too. Only files directly inside the
675 known legacy shared directory with the expected hash filename are
676 moved — arbitrary stored paths are not followed.
677 """
678 legacy_root = get_cache_directory() / "rag_indices"
679 if (
680 legacy_path == index_path
681 or legacy_path.parent != legacy_root
682 or legacy_path.name != index_path.name
683 or not legacy_path.exists()
684 ):
685 return
686 with _hold_faiss_write_lock(self.username, str(index_path)):
687 if index_path.exists(): 687 ↛ 688line 687 didn't jump to line 688 because the condition on line 687 was never true
688 return
689 try:
690 legacy_pkl = legacy_path.with_suffix(".pkl")
691 # Relocate phase-1's text-free .idmap.json sidecar BEFORE the
692 # .faiss, then move the .faiss LAST. phase-2 rekey looks for the
693 # sidecar next to the .faiss (legacy_rekey._sidecar_for =
694 # <stem>.idmap.json); leaving it at the legacy location orphans
695 # it, and a not-yet-rekeyed index then has no findable sidecar →
696 # rekey sees "no sidecar + old format" and destructively
697 # quarantines a healthy index. Ordering matters for crash
698 # recovery: the re-entry guard keys ONLY on index_path (the
699 # .faiss), so moving the .faiss last makes this idempotent — a
700 # crash after the sidecar move but before the .faiss move leaves
701 # the .faiss at the legacy path (guard still False) and re-entry
702 # simply retries (the sidecar move no-ops, new one exists). The
703 # reverse order would strand index_path.exists()=True with the
704 # sidecar orphaned at the legacy root, which the guard never
705 # re-invokes this to fix.
706 legacy_sidecar = legacy_path.with_name(
707 legacy_path.stem + ".idmap.json"
708 )
709 if legacy_sidecar.exists():
710 new_sidecar = index_path.with_name(
711 index_path.stem + ".idmap.json"
712 )
713 if not new_sidecar.exists(): 713 ↛ 715line 713 didn't jump to line 715 because the condition on line 713 was always true
714 legacy_sidecar.rename(new_sidecar)
715 legacy_path.rename(index_path)
716 # DELETE (don't relocate) any legacy plaintext .pkl companion:
717 # the new store writes NO .pkl, so moving it would just carry the
718 # plaintext chunk text to the new location. Its text lives in the
719 # encrypted DB.
720 if legacy_pkl.exists():
721 try:
722 legacy_pkl.unlink()
723 except OSError:
724 logger.warning(
725 f"Could not remove legacy plaintext .pkl "
726 f"{legacy_pkl} — delete it manually"
727 )
728 except OSError:
729 logger.exception(
730 f"Failed to migrate legacy FAISS index "
731 f"{legacy_path} -> {index_path}"
732 )
733 return
734 # Persist the relocated path so the row stops pointing at the
735 # now-empty legacy location.
736 with get_user_db_session(self.username, self.db_password) as session:
737 idx = session.query(RAGIndex).filter_by(id=rag_index.id).first()
738 if idx: 738 ↛ 741line 738 didn't jump to line 741
739 idx.index_path = str(index_path)
740 session.commit()
741 try:
742 # Refresh the integrity record under the new path (the old
743 # record is keyed by the legacy path and no longer applies).
744 self.integrity_manager.record_file(
745 index_path,
746 related_entity_type="rag_index",
747 related_entity_id=rag_index.id,
748 )
749 except Exception:
750 # Non-fatal: verify_file() creates a record on demand for
751 # unknown paths, so the subsequent load still succeeds.
752 logger.exception(
753 "Failed to refresh integrity record after index migration"
754 )
755 logger.info(
756 f"Migrated legacy shared FAISS index to per-user location: "
757 f"{index_path}"
758 )
760 def _reset_index_state_for_rebuild(
761 self, collection_id: str, rag_index: RAGIndex
762 ) -> None:
763 """Reset per-document indexed state when the index file is gone.
765 Reaching the fresh-build branch with DB rows that still claim
766 indexed content means the on-disk index was lost (missing after an
767 upgrade, quarantined, or deleted on dimension mismatch). Without
768 this reset, ``index_document``/``index_collection`` skip every
769 document (``DocumentCollection.indexed`` is still True), so the
770 only repair is an undocumented force re-index while search
771 silently returns nothing. Clearing the flags makes the regular
772 re-index path rebuild for real and makes the UI show the honest
773 unindexed state. Worst case of the narrow race with an in-flight
774 indexer is re-submitting chunks the merge step dedups by id.
775 """
776 with get_user_db_session(self.username, self.db_password) as session:
777 idx = session.query(RAGIndex).filter_by(id=rag_index.id).first()
778 claims_content = bool(
779 idx and (idx.chunk_count or idx.total_documents)
780 )
781 stale_status = (
782 session.query(RagDocumentStatus)
783 .filter_by(rag_index_id=rag_index.id)
784 .count()
785 )
786 stale_memberships = (
787 session.query(DocumentCollection)
788 .filter_by(collection_id=collection_id, indexed=True)
789 .count()
790 )
791 if not (claims_content or stale_status or stale_memberships):
792 # Genuinely new index — nothing to reset.
793 return
794 if idx: 794 ↛ 797line 794 didn't jump to line 797 because the condition on line 794 was always true
795 idx.chunk_count = 0
796 idx.total_documents = 0
797 session.query(RagDocumentStatus).filter_by(
798 rag_index_id=rag_index.id
799 ).delete()
800 session.query(DocumentCollection).filter_by(
801 collection_id=collection_id, indexed=True
802 ).update({"indexed": False, "chunk_count": 0})
803 session.commit()
804 logger.warning(
805 f"FAISS index file for collection {collection_id} is "
806 f"missing; reset indexed state for {stale_memberships} "
807 f"document(s) so the next index run rebuilds the index."
808 )
810 def _get_or_create_rag_index(
811 self,
812 collection_id: str,
813 *,
814 promote_current: bool = True,
815 db_session: Session | None = None,
816 commit: bool = True,
817 ) -> RAGIndex:
818 """Get or create RAGIndex record for the current configuration.
820 ``promote_current`` (default True) controls whether reusing an existing,
821 not-current row re-promotes it to ``is_current`` (demoting others). Read
822 paths (search) pass False: a mere search under a stale/different config
823 must NOT silently flip which index the collection considers current —
824 only an actual write (index) should. Creating a brand-new index still
825 promotes it regardless (it's the collection's only index for that config).
826 """
827 if db_session is None:
828 with get_user_db_session(
829 self.username, self.db_password
830 ) as session:
831 return self._get_or_create_rag_index_in_session(
832 session,
833 collection_id,
834 promote_current=promote_current,
835 commit=commit,
836 )
837 return self._get_or_create_rag_index_in_session(
838 db_session,
839 collection_id,
840 promote_current=promote_current,
841 commit=commit,
842 )
844 def _get_or_create_rag_index_in_session(
845 self,
846 db_session: Session,
847 collection_id: str,
848 *,
849 promote_current: bool,
850 commit: bool,
851 ) -> RAGIndex:
852 collection_name = f"collection_{collection_id}"
853 index_hash = self._get_index_hash(
854 collection_name, self.embedding_model, self.embedding_provider
855 )
856 rag_index = self._find_matching_rag_index(
857 db_session, collection_name, index_hash
858 )
859 created = False
861 if rag_index is None:
862 index_path = self._get_index_path(index_hash)
863 embedding_manager = self.embedding_manager
864 if embedding_manager is None: 864 ↛ 865line 864 didn't jump to line 865 because the condition on line 864 was never true
865 raise RuntimeError(
866 "Cannot create an index after closing the service"
867 )
868 embedding_dim = len(
869 embedding_manager.embeddings.embed_query("test")
870 )
871 candidate = RAGIndex(
872 collection_name=collection_name,
873 embedding_model=self.embedding_model,
874 embedding_model_type=EmbeddingProvider(self.embedding_provider),
875 embedding_dimension=embedding_dim,
876 index_path=str(index_path),
877 index_hash=index_hash,
878 chunk_size=self.chunk_size,
879 chunk_overlap=self.chunk_overlap,
880 splitter_type=self.splitter_type,
881 text_separators=self.text_separators,
882 distance_metric=self.distance_metric,
883 normalize_vectors=self.normalize_vectors,
884 index_type=self.index_type,
885 chunk_count=0,
886 total_documents=0,
887 status=RAGIndexStatus.ACTIVE,
888 is_current=False,
889 )
890 try:
891 with db_session.begin_nested():
892 db_session.add(candidate)
893 db_session.flush()
894 except IntegrityError:
895 rag_index = self._find_matching_rag_index(
896 db_session, collection_name, index_hash
897 )
898 if rag_index is None: 898 ↛ 899line 898 didn't jump to line 899 because the condition on line 898 was never true
899 raise
900 logger.info(
901 f"Using concurrently created RAG index: {index_hash}"
902 )
903 else:
904 rag_index = candidate
905 created = True
906 logger.info(f"Created new RAG index: {index_hash}")
908 mutated = created or (promote_current and not rag_index.is_current)
909 if mutated:
910 db_session.query(RAGIndex).filter(
911 RAGIndex.collection_name == collection_name,
912 RAGIndex.is_current.is_(True),
913 RAGIndex.id != rag_index.id,
914 ).update({"is_current": False})
915 rag_index.is_current = True
917 # Only commit when this call actually created or promoted an index.
918 # A pure read-through of an existing current index must stay a no-op
919 # so it can't consume a caller's pending write (e.g. the fault the
920 # index-finalize self-heal backstop guards against).
921 if commit and mutated:
922 db_session.commit()
923 db_session.refresh(rag_index)
924 return rag_index
926 def _quarantine_corrupt_index(self, index_path: Path, reason: str) -> None:
927 """Rename a corrupted FAISS index to ``<path>.corrupt-<ns>``
928 instead of deleting it; any legacy plaintext ``.pkl`` companion
929 is deleted outright rather than preserved.
931 Preserves user data for inspection/recovery. The
932 dimension-mismatch branch elsewhere in this method intentionally
933 *deletes* — that case rebuilds from scratch and the old bytes
934 are unreadable with the new model. This helper is for the two
935 "transient or unknown failure" branches (verify_file said no,
936 loading the index into a ``VectorIndex``/faiss raised) where the
937 on-disk bytes may still be usable by a human.
939 Raises ``OSError`` on rename failure (disk full, read-only fs,
940 permission denied). Re-raising prevents silent data loss: if we
941 swallowed the error, the next ``persist()``/``faiss.write_index``
942 would overwrite the corrupt bytes anyway. Caller paths log the
943 exception via the broader try/except around their indexer call.
945 # TODO(#4197-followup): FileIntegrityRecord.consecutive_failures
946 # is not reset by the next record_file call, leaking failure
947 # counts across recovery cycles. Orthogonal to this fix.
948 """
949 ns = time.time_ns()
950 pkl_path = index_path.with_suffix(".pkl")
952 # Same-nanosecond collisions essentially can't happen between
953 # concurrent threads (we hold the per-path lock), but a user
954 # could manually create such files. Loop with a hard cap so a
955 # weird state surfaces as a clean OSError rather than a hang.
956 suffix_n = 0
957 faiss_target = Path(f"{index_path}.corrupt-{ns}")
958 pkl_target = Path(f"{pkl_path}.corrupt-{ns}")
959 while faiss_target.exists() or pkl_target.exists():
960 suffix_n += 1
961 if suffix_n > _QUARANTINE_SUFFIX_RETRY_CAP:
962 raise OSError(
963 f"Quarantine path collisions exceeded "
964 f"{_QUARANTINE_SUFFIX_RETRY_CAP} retries for "
965 f"{index_path}"
966 )
967 faiss_target = Path(f"{index_path}.corrupt-{ns}-{suffix_n}")
968 pkl_target = Path(f"{pkl_path}.corrupt-{ns}-{suffix_n}")
970 try:
971 index_path.rename(faiss_target)
972 logger.warning(
973 f"Quarantined corrupted FAISS index to {faiss_target} "
974 f"(reason: {reason}). Searches against this collection "
975 f"will return empty results until the index is rebuilt. "
976 f"Document chunks are preserved in the database — "
977 f"trigger 'Re-index Collection' (or set "
978 f"research_library.auto_index_enabled=true and re-run "
979 f"indexing) to recover."
980 )
981 if pkl_path.exists():
982 # DELETE (not rename) any .pkl companion. The new store writes
983 # NO .pkl, so a .pkl next to a .faiss is a LEGACY plaintext
984 # docstore — its chunk text lives in the (encrypted) DB, so
985 # renaming it aside to .corrupt-<ns> would needlessly preserve
986 # plaintext chunk text on disk. Plaintext removal is not optional.
987 try:
988 pkl_path.unlink()
989 logger.info(f"Removed legacy plaintext PKL {pkl_path}")
990 except OSError:
991 logger.warning(
992 f"Could not remove legacy plaintext PKL {pkl_path} — "
993 "delete it manually"
994 )
995 else:
996 # Missing .pkl is the normal case for the new store. Not
997 # re-raised: the .faiss is already preserved.
998 logger.debug(f"No PKL companion at {pkl_path}.")
999 except OSError:
1000 if not index_path.exists():
1001 # The source is already gone — most likely another worker
1002 # PROCESS quarantined the same corrupt index concurrently (the
1003 # per-path lock is a threading.Lock, so it serialises only
1004 # threads within ONE process, not across gunicorn workers /
1005 # the scheduler process). The objective — get the corrupt bytes
1006 # off index_path so the caller rebuilds fresh — is already met,
1007 # so don't abort a healthy rebuild with a spurious
1008 # FileNotFoundError. Still purge any lingering plaintext .pkl:
1009 # the no-plaintext-at-rest invariant is not optional even here.
1010 if pkl_path.exists(): 1010 ↛ 1018line 1010 didn't jump to line 1018 because the condition on line 1010 was always true
1011 try:
1012 pkl_path.unlink()
1013 except OSError:
1014 logger.warning(
1015 f"Could not remove legacy plaintext PKL {pkl_path} "
1016 "— delete it manually"
1017 )
1018 logger.warning(
1019 f"Corrupt index {index_path} was already quarantined by "
1020 f"another worker (reason: {reason}); proceeding to rebuild."
1021 )
1022 return
1023 # A genuine failure (disk-full, read-only fs, permission denied)
1024 # with the corrupt bytes STILL in place. Re-raise so the caller
1025 # surfaces a real failure instead of letting the next
1026 # persist()/write_index overwrite the corrupt bytes.
1027 logger.exception(
1028 f"Failed to quarantine corrupted index at {index_path}"
1029 )
1030 raise
1032 # Best-effort retention sweep. The quarantine itself succeeded
1033 # above; failing to prune older files is not a correctness
1034 # issue, just a disk-usage one — log and move on.
1035 self._prune_old_quarantined_files(index_path)
1037 @staticmethod
1038 def _corrupt_sort_key(path: Path) -> Tuple[int, int]:
1039 """Extract ``(ns, suffix_n)`` from a ``.corrupt-<ns>[-<n>]``
1040 filename so retention can sort by the monotonic nanosecond
1041 suffix the quarantine path embeds — *not* by ``st_mtime``.
1043 Filesystem timestamp granularity is sometimes 1s or 2s
1044 (FAT32/ext3/SMB shares), making mtime ordering non-deterministic
1045 when multiple quarantines happen within one tick. The
1046 ``.corrupt-<ns>`` suffix carries the original ``time.time_ns()``
1047 from the quarantine and is reliable across all filesystems.
1049 Returns ``(-1, -1)`` for malformed names (manually-placed
1050 files) so they sort below any real entry under ``reverse=True``
1051 — i.e., they get pruned first.
1052 """
1053 name = path.name
1054 marker = ".corrupt-"
1055 idx = name.rfind(marker)
1056 if idx == -1: 1056 ↛ 1057line 1056 didn't jump to line 1057 because the condition on line 1056 was never true
1057 return (-1, -1)
1058 tail = name[idx + len(marker) :]
1059 parts = tail.split("-")
1060 try:
1061 ns = int(parts[0])
1062 except ValueError:
1063 return (-1, -1)
1064 suffix_n = 0
1065 if len(parts) > 1:
1066 try:
1067 suffix_n = int(parts[-1])
1068 except ValueError:
1069 # Unknown trailing component — keep the file but at the
1070 # base ns ordering.
1071 pass
1072 return (ns, suffix_n)
1074 def _prune_old_quarantined_files(self, index_path: Path) -> None:
1075 """Keep only the ``_QUARANTINE_KEEP_RECENT`` most-recent
1076 ``.corrupt-*`` files for ``index_path`` and its ``.pkl``
1077 companion. Sweeps the two sides independently — pairs share
1078 the same ``-<ns>`` suffix so they're ordered identically.
1080 Ordering uses the embedded ``-<ns>`` from the quarantine
1081 filename, not ``st_mtime``: file systems with 1-2s timestamp
1082 granularity can otherwise produce non-deterministic retention
1083 on bursts.
1085 Best-effort: logs and swallows any error so a sweep failure
1086 never propagates back into the indexing path.
1087 """
1088 parent = index_path.parent
1089 pkl_path = index_path.with_suffix(".pkl")
1091 for base in (index_path, pkl_path):
1092 pattern = f"{base.name}.corrupt-*"
1093 try:
1094 # Sort newest-first by the embedded -<ns>; everything
1095 # past the keep window is stale.
1096 candidates = sorted(
1097 parent.glob(pattern),
1098 key=self._corrupt_sort_key,
1099 reverse=True,
1100 )
1101 except OSError:
1102 logger.warning(
1103 f"Failed to enumerate {pattern} in {parent} for "
1104 f"quarantine retention sweep"
1105 )
1106 continue
1108 for stale in candidates[_QUARANTINE_KEEP_RECENT:]:
1109 try:
1110 stale.unlink()
1111 logger.info(f"Pruned old quarantined file: {stale}")
1112 except OSError:
1113 logger.warning(
1114 f"Failed to prune quarantined file {stale}; continuing"
1115 )
1117 def _preflight_index_path(
1118 self,
1119 collection_id: str,
1120 rag_index: RAGIndex,
1121 *,
1122 reset_stale_state: bool = False,
1123 ) -> Path:
1124 """Verify -> quarantine -> dimension-check pre-flight for ``rag_index``.
1126 Extracted from the pre-cutover ``load_or_create_faiss_index``
1127 (deleted — the old LangChain-FAISS path is gone) so every
1128 index/search/delete call site shares ONE recovery path instead of
1129 duplicating it. Mutates ``rag_index.index_path`` / ``rag_index.embedding_dimension``
1130 in place — matching the pre-refactor behavior of stamping the
1131 recomputed path/dimension back onto the in-memory record so callers
1132 that read ``self.rag_index_record`` afterwards see the current
1133 values — and returns the resolved on-disk path.
1135 Deliberately does NOT load or construct a vector store: that stays
1136 the caller's job (via ``VectorIndex(...)``), so a corrupt/foreign/
1137 dimension-mismatched file surfaces as an exception from THAT
1138 construction, which the caller catches to quarantine + rebuild (see
1139 ``_get_vector_index``).
1140 """
1141 index_path = self._get_index_path(rag_index.index_hash)
1142 stored_path = (
1143 Path(rag_index.index_path) if rag_index.index_path else None
1144 )
1145 rag_index.index_path = str(index_path)
1146 if not index_path.exists() and stored_path is not None:
1147 self._migrate_legacy_index_files(stored_path, index_path, rag_index)
1149 # Hold the per-(username, index_path) write lock across the entire
1150 # verify -> quarantine -> dimension-check sequence — see #4197 (a
1151 # narrower scope would let a concurrent writer's save race the
1152 # verification).
1153 if index_path.exists():
1154 lock = _get_faiss_write_lock(self.username, str(index_path))
1155 with lock:
1156 verified, reason = self.integrity_manager.verify_file(
1157 index_path
1158 )
1159 if not verified:
1160 if reason == NO_INTEGRITY_RECORD: 1160 ↛ 1179line 1160 didn't jump to line 1179 because the condition on line 1160 was never true
1161 # A MISSING record is not proven corruption, and unlike
1162 # checksum_mismatch it cannot be attacker-induced: the
1163 # record lives in the per-user ENCRYPTED DB, so removing
1164 # it needs the DB password (which also decrypts
1165 # everything). It only arises from our own crash between
1166 # the index write and record_file (see legacy_rekey
1167 # Phase B) or a pre-integrity-feature index. ADOPT the
1168 # byte-present file by recording its integrity now — on
1169 # BOTH read and write paths. Adoption is non-destructive
1170 # (it only records a checksum), so a read path may do it:
1171 # refusing here instead would hard-fail EVERY search on a
1172 # byte-healthy, missing-record index forever (a read-only
1173 # collection is never revisited by a write op that would
1174 # otherwise adopt it). If the bytes are actually
1175 # unreadable/foreign/wrong-dim, the VectorIndex load below
1176 # still raises and the caller handles it — corruption is
1177 # caught either way. Only checksum_mismatch (genuine
1178 # tamper) is refused/quarantined below.
1179 logger.warning(
1180 f"No integrity record for {index_path}; adopting the "
1181 "existing index (recording its integrity) rather than "
1182 "quarantining — a missing record is not corruption."
1183 )
1184 self.integrity_manager.record_file(
1185 index_path,
1186 related_entity_type="rag_index",
1187 related_entity_id=rag_index.id,
1188 )
1189 # Adopted = now trusted: fall through to the SAME
1190 # embedding-dimension-drift probe the verified branch
1191 # runs, so an adopted index whose model dimension changed
1192 # is rebuilt rather than silently mismatched at add time.
1193 verified = True
1194 elif not reset_stale_state:
1195 # Read path (search): refuse to serve an index that
1196 # failed verification for a reason OTHER than a missing
1197 # record (a checksum_mismatch = genuine tamper, or a
1198 # transient read failure), but do NOT destroy it —
1199 # quarantine is a WRITE-path recovery decision. A
1200 # subsequent index/write operation quarantines + rebuilds.
1201 raise RuntimeError(
1202 f"Integrity verification failed for {index_path}: "
1203 f"{reason} — refusing on a read path (a subsequent "
1204 "index/write operation will quarantine + rebuild)"
1205 )
1206 elif reason and reason.startswith( 1206 ↛ 1215line 1206 didn't jump to line 1215 because the condition on line 1206 was never true
1207 "checksum_calculation_failed"
1208 ):
1209 # Transient: could not READ the file to compute its
1210 # checksum (disk hiccup, momentary lock). This says
1211 # NOTHING about tampering — unlike checksum_mismatch — so
1212 # quarantining would force a needless full re-embed of a
1213 # healthy collection. Raise so the operation retries (like
1214 # the read path) instead of destroying the index.
1215 raise RuntimeError(
1216 f"Integrity check could not read {index_path} "
1217 f"({reason}) — refusing to quarantine on a transient "
1218 "read failure; retry the operation."
1219 )
1220 else:
1221 logger.error(
1222 f"Integrity verification failed for {index_path}: "
1223 f"{reason}. Quarantining for recovery; creating "
1224 f"new index."
1225 )
1226 self._quarantine_corrupt_index(index_path, reason)
1227 if verified:
1228 # Probe the embedding model OUTSIDE the quarantine
1229 # decision: embed_query fails when the provider itself
1230 # is unreachable (e.g. Ollama down), which says nothing
1231 # about the index file's health, and must propagate
1232 # rather than trigger a destructive rebuild of a
1233 # healthy index.
1234 current_dim = len(
1235 self.embedding_manager.embeddings.embed_query(
1236 "dimension_check"
1237 )
1238 )
1239 stored_dim = rag_index.embedding_dimension
1240 if stored_dim and current_dim != stored_dim:
1241 if not reset_stale_state:
1242 # Read path: surface the mismatch instead of
1243 # deleting the index + wiping DB state (which a mere
1244 # search must never do). A write/index operation
1245 # owns the destructive rebuild.
1246 raise RuntimeError(
1247 f"Embedding dimension mismatch for {index_path}: "
1248 f"index dim={stored_dim}, model dim="
1249 f"{current_dim} — refusing on a read path"
1250 )
1251 logger.warning(
1252 f"Embedding dimension mismatch detected! "
1253 f"Index created with dim={stored_dim}, "
1254 f"current model returns dim={current_dim}. "
1255 f"Deleting old index and rebuilding."
1256 )
1257 try:
1258 index_path.unlink()
1259 pkl_path = index_path.with_suffix(".pkl")
1260 if pkl_path.exists(): 1260 ↛ 1261line 1260 didn't jump to line 1261 because the condition on line 1260 was never true
1261 pkl_path.unlink()
1262 logger.info(
1263 f"Deleted old index files at {index_path}"
1264 )
1265 except Exception:
1266 logger.exception("Failed to delete old index files")
1268 with get_user_db_session(
1269 self.username, self.db_password
1270 ) as session:
1271 idx = (
1272 session.query(RAGIndex)
1273 .filter_by(id=rag_index.id)
1274 .first()
1275 )
1276 if idx: 1276 ↛ 1288line 1276 didn't jump to line 1288
1277 idx.embedding_dimension = current_dim
1278 session.commit()
1279 logger.info(
1280 f"Updated RAGIndex dimension to {current_dim}"
1281 )
1282 # Reset RAGIndex counts, RagDocumentStatus, and
1283 # DocumentCollection.indexed together — a hand-rolled
1284 # partial reset here previously left
1285 # DocumentCollection.indexed=True, which makes
1286 # index_document/index_collection skip every document
1287 # and search silently return nothing after the rebuild.
1288 self._reset_index_state_for_rebuild(
1289 collection_id, rag_index
1290 )
1291 rag_index.embedding_dimension = current_dim
1293 if not index_path.exists():
1294 # The new per-user path is missing — but that alone does NOT mean the
1295 # index is gone. A one-time legacy-path relocation
1296 # (_migrate_legacy_index_files above) can fail transiently (a rename
1297 # hitting disk-full, an AV/backup file lock, a busy network mount),
1298 # leaving the bytes fully intact at the stored legacy path. Treating
1299 # that as "gone" would wipe a healthy collection's indexed bookkeeping
1300 # — even on a mere search — and force a needless full re-embed. Only
1301 # self-heal when the bytes are absent at BOTH the new and the legacy
1302 # location; otherwise leave the state alone so the next write-path
1303 # access retries the relocation.
1304 legacy_bytes_present = bool(
1305 stored_path is not None
1306 and stored_path != index_path
1307 and stored_path.exists()
1308 )
1309 if not legacy_bytes_present:
1310 stale_indexed = bool(
1311 (rag_index.chunk_count or 0) > 0
1312 or (rag_index.total_documents or 0) > 0
1313 )
1314 if stale_indexed:
1315 # DB says indexed, but the file is GONE at both locations
1316 # (cache dir cleared, restored without the cache dir, lost
1317 # volume, crash between quarantine and rebuild). Reset so it
1318 # self-heals — even on a READ path: an absent file has
1319 # nothing to preserve, and without this search silently
1320 # returns [] forever (get_rag_stats still reports 100%
1321 # indexed and a normal reindex short-circuits on
1322 # indexed=True, so nothing ever recovers it).
1323 logger.warning(
1324 f"RAG index file for collection {collection_id} is "
1325 f"MISSING at {index_path} while the collection is marked "
1326 "indexed — resetting indexed state so it rebuilds "
1327 "(search would otherwise return no results indefinitely)."
1328 )
1329 self._reset_index_state_for_rebuild(
1330 collection_id, rag_index
1331 )
1332 elif reset_stale_state: 1332 ↛ 1337line 1332 didn't jump to line 1337 because the condition on line 1332 was always true
1333 self._reset_index_state_for_rebuild(
1334 collection_id, rag_index
1335 )
1337 return index_path
1339 def _get_vector_index(
1340 self,
1341 collection_id: str,
1342 collection_name: str,
1343 *,
1344 reset_stale_state: bool = False,
1345 ) -> VectorIndex:
1346 """Build a :class:`VectorIndex` for ``collection_id``.
1348 Single call site for the index/search/delete paths: runs the
1349 verify/quarantine/dimension-check pre-flight, then constructs the
1350 facade. On a ``ValueError``/``RuntimeError`` from the underlying
1351 store's ``load()`` (corrupt bytes, a foreign/pre-migration format
1352 that isn't an ``IndexIDMap2``, or a dimension mismatch the
1353 pre-flight didn't already catch) the corrupt file is quarantined,
1354 stale per-document indexed state is reset, and construction is
1355 retried once as a fresh, empty store.
1357 NOTE: like ``_get_or_create_rag_index``, this CREATES a RAGIndex
1358 row if none exists for the configured (collection, embedding model)
1359 — callers on a delete-only path that must not create an index
1360 (e.g. ``remove_documents_from_index``) guard with their own
1361 existence check before calling this.
1362 """
1363 # Only WRITE paths (reset_stale_state=True) may re-promote a reused
1364 # index to is_current; a read-path search must not flip the pointer.
1365 rag_index = self._get_or_create_rag_index(
1366 collection_id, promote_current=reset_stale_state
1367 )
1368 self.rag_index_record = rag_index
1369 index_path = self._preflight_index_path(
1370 collection_id, rag_index, reset_stale_state=reset_stale_state
1371 )
1372 lock = _get_faiss_write_lock(self.username, str(index_path))
1374 # nested=True: these run reentrantly inside VectorIndex.apply(), which
1375 # holds just-flushed, uncommitted DocumentChunk rows on this same
1376 # thread-local session. The integrity write must run in a SAVEPOINT and
1377 # let apply()'s caller commit — a plain commit/rollback here would
1378 # settle or DISCARD those pending rows (silently losing indexed text
1379 # while the vectors persist). See FileIntegrityManager._integrity_write.
1380 def _record(path: Path) -> None:
1381 self.integrity_manager.record_file(
1382 path,
1383 related_entity_type="rag_index",
1384 related_entity_id=rag_index.id,
1385 nested=True,
1386 )
1388 def _verify(path: Path):
1389 return self.integrity_manager.verify_file(path, nested=True)
1391 kwargs = {
1392 "username": self.username,
1393 "db_password": self.db_password,
1394 # Lazy: index()/search() resolve the real backend on first embed_*
1395 # call (identical behaviour, one extra attribute hop), while the
1396 # delete-only path (remove_documents_from_index) never materialises
1397 # it at all. See _LazyEmbeddings.
1398 "embeddings": _LazyEmbeddings(self.embedding_manager),
1399 "embedding_model": self.embedding_model,
1400 "embedding_model_type": EmbeddingProvider(self.embedding_provider),
1401 "collection_name": collection_name,
1402 "dimension": rag_index.embedding_dimension,
1403 "path": index_path,
1404 "lock": lock,
1405 "integrity_record": _record,
1406 "integrity_verify": _verify,
1407 # Use the STORED index config, not this service instance's. The
1408 # RAGIndex row is stamped with index_type/metric/normalize at
1409 # creation (see _get_or_create_rag_index); the constructing
1410 # LibraryRAGService may carry different CURRENT defaults (the
1411 # factory falls back to global settings for a collection whose
1412 # metadata was never stored). Building against self.* would load a
1413 # real HNSW file while labelling it "flat" (or vice-versa), and
1414 # FaissVectorStore.supports_delete trusts the label — routing a
1415 # removal through remove_ids on an HNSW index (unsupported → raises)
1416 # or silently rewriting a flat file as HNSW. This mirrors the fix
1417 # already applied to purge_document_vectors below.
1418 "index_type": rag_index.index_type or self.index_type,
1419 "metric": rag_index.distance_metric or self.distance_metric,
1420 "normalize": (
1421 rag_index.normalize_vectors
1422 if rag_index.normalize_vectors is not None
1423 else self.normalize_vectors
1424 ),
1425 }
1426 try:
1427 return VectorIndex(**kwargs)
1428 except (ValueError, RuntimeError):
1429 if not reset_stale_state:
1430 # Read paths (search) must stay side-effect-free. A transient
1431 # load failure (OOM, a disk hiccup, a momentarily-locked file)
1432 # on an existing, byte-healthy index must NOT quarantine the
1433 # file or wipe the collection's indexed state — that would turn
1434 # a retryable blip into a forced full re-embed of the whole
1435 # collection triggered by a mere search. Surface the error;
1436 # quarantine + rebuild is a WRITE-path (reset_stale_state=True)
1437 # decision. (A never-indexed collection never reaches here — no
1438 # file exists, so VectorIndex.create() succeeds above.)
1439 raise
1440 # A load failure here — after _preflight_index_path's verify_file
1441 # just PASSED on the same bytes — is ambiguous: it is either a
1442 # genuine format/dimension mismatch (deterministic — needs a rebuild)
1443 # OR a transient OS hiccup (EMFILE under concurrent multi-user load,
1444 # a momentary AV-scan/backup file lock; the byte-healthy file loads
1445 # fine moments later). faiss surfaces BOTH as the same RuntimeError,
1446 # so quarantining unconditionally destroys a provably-healthy index
1447 # and forces a needless full re-embed on a mere transient blip.
1448 # Distinguish by RETRYING the load once: a transient failure clears,
1449 # a deterministic one repeats. Only quarantine when the retry also
1450 # fails. (Mirrors _preflight_index_path's transient-vs-tamper split.)
1451 logger.warning(
1452 f"Failed to load vector store at {index_path}; retrying once "
1453 "before quarantining (may be a transient OS hiccup)"
1454 )
1455 try:
1456 return VectorIndex(**kwargs)
1457 except (ValueError, RuntimeError):
1458 pass
1459 logger.warning(
1460 f"Vector store at {index_path} failed to load twice; "
1461 "quarantining and rebuilding fresh"
1462 )
1463 with lock:
1464 if index_path.exists():
1465 self._quarantine_corrupt_index(
1466 index_path, "vector_index_load_failed"
1467 )
1468 self._reset_index_state_for_rebuild(collection_id, rag_index)
1469 return VectorIndex(**kwargs)
1471 def search(
1472 self, query: str, collection_id: str, top_k: int
1473 ) -> List[SearchResult]:
1474 """Semantic search within ``collection_id`` via the vector store.
1476 Runs the same verify/quarantine/dimension pre-flight as indexing
1477 (through ``_get_vector_index``), then delegates to
1478 ``VectorIndex.search``. ``reset_stale_state`` is left at its
1479 default False — this is a read path: a transient integrity hiccup
1480 during a mere search must not wipe the collection's indexed state
1481 and force a full re-embed.
1483 Callers are expected to have already confirmed the collection has
1484 indexed documents (e.g. via ``get_rag_stats``) before calling this
1485 — this method will otherwise create an empty RAGIndex/store for a
1486 never-indexed (collection, embedding model) pair and return no
1487 results, rather than raising.
1488 """
1489 collection_name = f"collection_{collection_id}"
1490 vindex = self._get_vector_index(collection_id, collection_name)
1491 return vindex.search(query, top_k)
1493 def get_current_index_info(
1494 self, collection_id: Optional[str] = None
1495 ) -> Optional[Dict[str, Any]]:
1496 """
1497 Get information about the current RAG index for a collection.
1499 Args:
1500 collection_id: UUID of collection (defaults to Library if None)
1501 """
1502 with get_user_db_session(self.username, self.db_password) as session:
1503 # Get collection name in the format stored in RAGIndex (collection_<uuid>)
1504 if collection_id:
1505 collection = (
1506 session.query(Collection)
1507 .filter_by(id=collection_id)
1508 .first()
1509 )
1510 collection_name = (
1511 f"collection_{collection_id}" if collection else "unknown"
1512 )
1513 else:
1514 # Default to Library collection
1515 from ...database.library_init import get_default_library_id
1517 collection_id = get_default_library_id(
1518 self.username, self.db_password
1519 )
1520 collection_name = f"collection_{collection_id}"
1522 rag_index = (
1523 session.query(RAGIndex)
1524 .filter_by(collection_name=collection_name, is_current=True)
1525 .first()
1526 )
1528 if not rag_index:
1529 # Debug: check all RAG indices for this collection
1530 all_indices = session.query(RAGIndex).all()
1531 logger.info(
1532 f"No RAG index found for collection_name='{collection_name}'. All indices: {[(idx.collection_name, idx.is_current) for idx in all_indices]}"
1533 )
1534 return None
1536 # Calculate actual counts from rag_document_status table
1537 from ...database.models.library import RagDocumentStatus
1539 actual_chunk_count = (
1540 session.query(func.sum(RagDocumentStatus.chunk_count))
1541 .filter_by(collection_id=collection_id)
1542 .scalar()
1543 or 0
1544 )
1546 actual_doc_count = (
1547 session.query(RagDocumentStatus)
1548 .filter_by(collection_id=collection_id)
1549 .count()
1550 )
1552 return {
1553 "embedding_model": rag_index.embedding_model,
1554 "embedding_model_type": rag_index.embedding_model_type.value
1555 if rag_index.embedding_model_type
1556 else None,
1557 "embedding_dimension": rag_index.embedding_dimension,
1558 "chunk_size": rag_index.chunk_size,
1559 "chunk_overlap": rag_index.chunk_overlap,
1560 "chunk_count": actual_chunk_count,
1561 "total_documents": actual_doc_count,
1562 "created_at": rag_index.created_at.isoformat(),
1563 "last_updated_at": rag_index.last_updated_at.isoformat(),
1564 }
1566 def _flag_document_for_reindex(
1567 self, session, document_id: str, collection_id: str
1568 ) -> None:
1569 """Mark a document NOT indexed so the reconciler re-embeds it.
1571 Used on index_document failure paths: ``VectorIndex.index()`` runs
1572 ``store.apply()`` — DURABLY mutating the FAISS file — BEFORE the DB
1573 commit, so a failure/rollback after that point can leave the source's
1574 prior chunk rows resurrected with no vectors (silently unsearchable)
1575 while ``DocumentCollection.indexed`` is reverted to its pre-attempt
1576 state. The background reconciler only revisits ``indexed=False`` rows,
1577 so without this flag the document stays permanently, silently
1578 unsearchable. Best-effort + committed in its own statement;
1579 index_document owns per-document commits (see index_all_documents), so
1580 this clobbers no caller.
1581 """
1582 try:
1583 session.query(DocumentCollection).filter_by(
1584 document_id=document_id, collection_id=collection_id
1585 ).update(
1586 {"indexed": False, "chunk_count": 0},
1587 synchronize_session=False,
1588 )
1589 session.query(RagDocumentStatus).filter_by(
1590 document_id=document_id, collection_id=collection_id
1591 ).delete(synchronize_session=False)
1592 session.commit()
1593 except Exception:
1594 safe_rollback(
1595 session, "library_rag_service._flag_document_for_reindex"
1596 )
1597 logger.warning(
1598 f"Could not flag document {document_id} for re-index after a "
1599 f"failed index of collection {collection_id}; a manual reindex "
1600 "may be needed to restore its search."
1601 )
1603 @contextmanager
1604 def _collection_transaction_lock(self, collection_id: str):
1605 """Serialize the DB+FAISS transaction for one collection.
1607 FAISS already uses this same re-entrant lock for reload/apply/persist and
1608 integrity recording. Acquiring it around the complete document operation
1609 also closes the store-persist-before-SQLCipher-commit window to sibling
1610 workers: no worker can verify/quarantine or mutate the file until the
1611 preceding worker's database commit (or rollback/reindex flag) finishes.
1612 """
1613 collection_name = f"collection_{collection_id}"
1614 index_hash = self._get_index_hash(
1615 collection_name, self.embedding_model, self.embedding_provider
1616 )
1617 index_path = self._get_index_path(index_hash)
1618 with _hold_faiss_write_lock(self.username, str(index_path)) as lock:
1619 yield lock
1621 def _prepare_document(
1622 self, document_id: str, collection_id: str, force_reindex: bool
1623 ) -> Dict[str, Any]:
1624 """Load, split, and embed without mutating collection state."""
1625 with get_user_db_session(self.username, self.db_password) as session:
1626 document = session.query(Document).filter_by(id=document_id).first()
1627 if not document: 1627 ↛ 1628line 1627 didn't jump to line 1628 because the condition on line 1627 was never true
1628 return {"status": "error", "error": "Document not found"}
1629 link = (
1630 session.query(DocumentCollection)
1631 .filter_by(document_id=document_id, collection_id=collection_id)
1632 .first()
1633 )
1634 if link is not None and link.indexed and not force_reindex: 1634 ↛ 1635line 1634 didn't jump to line 1635 because the condition on line 1634 was never true
1635 return {
1636 "status": "skipped",
1637 "message": "Document already indexed for this collection",
1638 "chunk_count": link.chunk_count,
1639 }
1640 if not document.text_content: 1640 ↛ 1641line 1640 didn't jump to line 1641 because the condition on line 1640 was never true
1641 return {"status": "needs_serial_write"}
1643 title = document.title or document.filename or "Untitled"
1644 metadata = {
1645 "source": document.original_url,
1646 "document_id": document_id,
1647 "collection_id": collection_id,
1648 "document_title": title,
1649 "title": title,
1650 "authors": document.authors,
1651 "published_date": str(document.published_date)
1652 if document.published_date
1653 else None,
1654 "doi": document.doi,
1655 "arxiv_id": document.arxiv_id,
1656 "pmid": document.pmid,
1657 "pmcid": document.pmcid,
1658 "extraction_method": document.extraction_method,
1659 "word_count": document.word_count,
1660 }
1661 source_doc = LangchainDocument(
1662 page_content=document.text_content, metadata=metadata
1663 )
1665 chunks = self.text_splitter.split_documents([source_doc])
1666 chunk_inputs = [
1667 ChunkInput(
1668 text=chunk.page_content,
1669 metadata={**metadata, "chunk_index": i},
1670 )
1671 for i, chunk in enumerate(chunks)
1672 ]
1673 vectors = np.asarray(
1674 self.embedding_manager.embeddings.embed_documents(
1675 [chunk.text for chunk in chunk_inputs]
1676 ),
1677 dtype="float32",
1678 )
1679 logger.info(
1680 f"Prepared document {document_id}: {len(chunk_inputs)} chunks embedded"
1681 )
1682 return {
1683 "status": "prepared",
1684 "prepared": _PreparedDocument(
1685 document_id=document_id,
1686 collection_id=collection_id,
1687 chunk_inputs=chunk_inputs,
1688 vectors=vectors,
1689 ),
1690 }
1692 def _write_prepared_document(
1693 self, prepared: _PreparedDocument
1694 ) -> Dict[str, Any]:
1695 """Serialize chunk rows, FAISS mutation, status updates, and commit."""
1696 document_id = prepared.document_id
1697 collection_id = prepared.collection_id
1698 chunks = prepared.chunk_inputs
1699 index_hash = self._get_index_hash(
1700 f"collection_{collection_id}",
1701 self.embedding_model,
1702 self.embedding_provider,
1703 )
1704 index_path = self._get_index_path(index_hash)
1705 with _hold_faiss_write_lock(self.username, str(index_path)):
1706 with get_user_db_session(
1707 self.username, self.db_password
1708 ) as session:
1709 try:
1710 ensure_in_collection(session, document_id, collection_id)
1711 collection_name = f"collection_{collection_id}"
1712 vindex = self._get_vector_index(
1713 collection_id, collection_name, reset_stale_state=True
1714 )
1715 stats = vindex.index_prepared(
1716 source_type="document",
1717 source_id=document_id,
1718 chunks=chunks,
1719 vectors=prepared.vectors,
1720 replace=True,
1721 session=session,
1722 )
1724 timestamp = datetime.now(UTC)
1725 existing = (
1726 session.query(RagDocumentStatus)
1727 .filter_by(
1728 document_id=document_id,
1729 collection_id=collection_id,
1730 rag_index_id=self.rag_index_record.id,
1731 )
1732 .first()
1733 )
1734 old_chunks = existing.chunk_count if existing else 0
1735 session.merge(
1736 RagDocumentStatus(
1737 document_id=document_id,
1738 collection_id=collection_id,
1739 rag_index_id=self.rag_index_record.id,
1740 chunk_count=len(chunks),
1741 indexed_at=timestamp,
1742 )
1743 )
1744 session.query(DocumentCollection).filter_by(
1745 document_id=document_id, collection_id=collection_id
1746 ).update(
1747 {
1748 "indexed": True,
1749 "chunk_count": len(chunks),
1750 "last_indexed_at": timestamp,
1751 }
1752 )
1753 rag_index = (
1754 session.query(RAGIndex)
1755 .filter_by(id=self.rag_index_record.id)
1756 .first()
1757 )
1758 if rag_index: 1758 ↛ 1763line 1758 didn't jump to line 1763 because the condition on line 1758 was always true
1759 rag_index.chunk_count += len(chunks) - old_chunks
1760 if existing is None:
1761 rag_index.total_documents += 1
1762 rag_index.last_updated_at = timestamp
1763 session.commit()
1764 logger.info(
1765 f"Persisted prepared document {document_id}: "
1766 f"{stats.added} vectors added"
1767 )
1768 return {"status": "success", "chunk_count": len(chunks)}
1769 except Exception as exc:
1770 safe_rollback(
1771 session,
1772 "library_rag_service._write_prepared_document",
1773 )
1774 self._flag_document_for_reindex(
1775 session, document_id, collection_id
1776 )
1777 logger.exception(
1778 f"Error writing prepared document {document_id}"
1779 )
1780 return {
1781 "status": "error",
1782 "error": f"Operation failed: {type(exc).__name__}",
1783 }
1785 def index_document(
1786 self, document_id: str, collection_id: str, force_reindex: bool = False
1787 ) -> Dict[str, Any]:
1788 collection_name = f"collection_{collection_id}"
1789 index_hash = self._get_index_hash(
1790 collection_name, self.embedding_model, self.embedding_provider
1791 )
1792 index_path = self._get_index_path(index_hash)
1793 with _hold_faiss_write_lock(self.username, str(index_path)):
1794 return self._index_document_locked(
1795 document_id, collection_id, force_reindex
1796 )
1798 def _index_one(
1799 self, document_id: str, collection_id: str, force_reindex: bool
1800 ) -> Dict[str, Any]:
1801 """Single-doc dispatch seam used by :meth:`index_documents_parallel`.
1803 Default implementation delegates to :meth:`index_document` so production
1804 behaviour is unchanged. Tests and subclasses override this method to
1805 short-circuit the full per-doc pipeline (which would require mocking
1806 every internal collaborator) without losing the prepared/serialized
1807 pipeline in production.
1809 The previous parallel runner detected overrides via
1810 ``"index_document" in self.__dict__`` — that only honoured
1811 instance-attribute patches and silently bypassed class-level subclass
1812 overrides. The runner now compares the bound method against the base
1813 class attribute so both instance patches and subclass overrides are
1814 honoured.
1815 """
1816 return self.index_document(document_id, collection_id, force_reindex)
1818 def _index_document_locked(
1819 self, document_id: str, collection_id: str, force_reindex: bool = False
1820 ) -> Dict[str, Any]:
1821 """
1822 Index a single document into RAG for a specific collection.
1824 Args:
1825 document_id: UUID of the Document to index
1826 collection_id: UUID of the Collection to index for
1827 force_reindex: Whether to force reindexing even if already indexed
1829 Returns:
1830 Dict with status, chunk_count, and any errors
1831 """
1832 with get_user_db_session(self.username, self.db_password) as session:
1833 # Get the document
1834 document = session.query(Document).filter_by(id=document_id).first()
1836 if not document:
1837 return {"status": "error", "error": "Document not found"}
1839 # Get or create DocumentCollection entry
1840 doc_collection = ensure_in_collection(
1841 session, document_id, collection_id
1842 )
1844 # Check if already indexed for this collection
1845 if doc_collection.indexed and not force_reindex:
1846 return {
1847 "status": "skipped",
1848 "message": "Document already indexed for this collection",
1849 "chunk_count": doc_collection.chunk_count,
1850 }
1852 # Validate text content
1853 if not document.text_content:
1854 # A previously-indexed document now cleared to empty text (e.g.
1855 # a note edited down to whitespace) must have its stale
1856 # chunks/vectors PURGED, not left searchable. The note-edit path
1857 # already set indexed=False, so detect prior chunks directly and
1858 # remove them via the facade's empty-replace.
1859 collection_name = f"collection_{collection_id}"
1860 has_prior = (
1861 session.query(DocumentChunk.id)
1862 .filter_by(
1863 source_type="document",
1864 source_id=document_id,
1865 collection_name=collection_name,
1866 )
1867 .first()
1868 is not None
1869 )
1870 if has_prior: 1870 ↛ 1879line 1870 didn't jump to line 1879 because the condition on line 1870 was never true
1871 # Same store-durable-before-DB-commit hazard as the main
1872 # reindex path below, but this branch sits OUTSIDE that
1873 # path's try/except: the empty-replace apply() DURABLY
1874 # removes the prior vectors before the commit here. On
1875 # failure, roll back and flag the document for re-index so
1876 # the reconciler re-runs the purge instead of leaving
1877 # resurrected rows whose vectors are gone (silently
1878 # unsearchable) marked indexed.
1879 try:
1880 vindex = self._get_vector_index(
1881 collection_id,
1882 collection_name,
1883 reset_stale_state=True,
1884 )
1885 vindex.index(
1886 source_type="document",
1887 source_id=document_id,
1888 chunks=[],
1889 replace=True,
1890 session=session,
1891 )
1892 doc_collection.indexed = False
1893 doc_collection.chunk_count = 0
1894 # Clear the canonical RagDocumentStatus row too. Row
1895 # existence is the "indexed" marker get_rag_stats / the
1896 # RAG status route read, so it MUST move with
1897 # DocumentCollection.indexed — a leftover row would
1898 # report this just-purged (now empty, unsearchable)
1899 # document as still indexed with a stale chunk_count.
1900 # Mirrors remove_document_from_rag and the three sibling
1901 # not-indexed paths; the cleared branch was the lone
1902 # exception.
1903 session.query(RagDocumentStatus).filter_by(
1904 document_id=document_id, collection_id=collection_id
1905 ).delete(synchronize_session=False)
1906 session.commit()
1907 except Exception as e:
1908 safe_rollback(
1909 session,
1910 "library_rag_service.index_document cleared",
1911 )
1912 self._flag_document_for_reindex(
1913 session, document_id, collection_id
1914 )
1915 logger.exception(
1916 f"Error purging cleared document {document_id} "
1917 f"for collection {collection_id}"
1918 )
1919 return {
1920 "status": "error",
1921 "error": f"Operation failed: {type(e).__name__}",
1922 }
1923 return {
1924 "status": "cleared",
1925 "message": (
1926 "Document has no text content; prior chunks/"
1927 "vectors purged"
1928 ),
1929 "chunk_count": 0,
1930 }
1931 return {
1932 "status": "error",
1933 "error": "Document has no text content",
1934 }
1936 # Replace-on-reindex context: a document being re-indexed (e.g.
1937 # a note whose content changed) has prior chunks with DIFFERENT
1938 # content hashes from the new text. Left alone they accumulate
1939 # in both the DB and the vector store forever — unbounded
1940 # per-edit growth plus stale duplicate search hits.
1941 # ``VectorIndex.index(replace=True)`` below handles this: it
1942 # looks up the source's prior DocumentChunk rows (authoritative,
1943 # queried fresh — not from an ephemeral pre-committed snapshot),
1944 # applies the vector-store removal durably, then deletes those
1945 # rows in the SAME transaction this method already holds open
1946 # (so a failure rolls back together with everything else, never
1947 # orphaning rows or poisoning the shared session).
1949 try:
1950 # Create LangChain Document from text
1951 doc = LangchainDocument(
1952 page_content=document.text_content,
1953 metadata={
1954 "source": document.original_url,
1955 "document_id": document_id, # Add document ID for source linking
1956 "collection_id": collection_id, # Add collection ID
1957 "title": document.title
1958 or document.filename
1959 or "Untitled",
1960 "document_title": document.title
1961 or document.filename
1962 or "Untitled", # Add for compatibility
1963 "authors": document.authors,
1964 "published_date": str(document.published_date)
1965 if document.published_date
1966 else None,
1967 "doi": document.doi,
1968 "arxiv_id": document.arxiv_id,
1969 "pmid": document.pmid,
1970 "pmcid": document.pmcid,
1971 "extraction_method": document.extraction_method,
1972 "word_count": document.word_count,
1973 },
1974 )
1976 # Split into chunks
1977 chunks = self.text_splitter.split_documents([doc])
1978 logger.info(
1979 f"Split document {document_id} into {len(chunks)} chunks"
1980 )
1982 # Get collection name for chunk storage
1983 collection = (
1984 session.query(Collection)
1985 .filter_by(id=collection_id)
1986 .first()
1987 )
1988 # Use collection_<uuid> format for internal storage
1989 collection_name = (
1990 f"collection_{collection_id}" if collection else "unknown"
1991 )
1993 document_title = (
1994 document.title or document.filename or "Untitled"
1995 )
1996 # Mirrors the parent doc's metadata dict built above (source,
1997 # ids, bibliographic fields) so it survives into each
1998 # chunk's ``DocumentChunk.document_metadata`` — this goes
1999 # ONLY to the encrypted DB (never the vector store; see the
2000 # SECURITY INVARIANT in vector_stores/base.py) and future-
2001 # proofs citation/UI features that read per-chunk metadata.
2002 shared_chunk_metadata = {
2003 "source": document.original_url,
2004 "document_id": document_id,
2005 "collection_id": collection_id,
2006 "document_title": document_title,
2007 "title": document_title,
2008 "authors": document.authors,
2009 "published_date": str(document.published_date)
2010 if document.published_date
2011 else None,
2012 "doi": document.doi,
2013 "arxiv_id": document.arxiv_id,
2014 "pmid": document.pmid,
2015 "pmcid": document.pmcid,
2016 "extraction_method": document.extraction_method,
2017 "word_count": document.word_count,
2018 }
2019 chunk_inputs = [
2020 ChunkInput(
2021 text=chunk.page_content,
2022 metadata={
2023 **shared_chunk_metadata,
2024 "chunk_index": i,
2025 },
2026 )
2027 for i, chunk in enumerate(chunks)
2028 ]
2030 # Embed + persist via the new int-id-keyed vector store —
2031 # this writes DocumentChunk rows (text only in the encrypted
2032 # DB) and the vector-store's ids/vectors (never text; see
2033 # the SECURITY INVARIANT in vector_stores/base.py). Passing
2034 # ``session`` means the facade does not commit/rollback —
2035 # this transaction still owns that (below).
2036 #
2037 # replace=True ALWAYS: a full-document (re)index owns all of a
2038 # source's chunks, so it must drop the source's prior
2039 # chunks/vectors before adding the new ones. force_reindex is
2040 # NOT a safe proxy: a note edit clears the indexed STATUS
2041 # (_mark_note_stale_for_reindex_in_session) but leaves the
2042 # pre-edit chunk rows/vectors in place, so the subsequent
2043 # non-forced reindex would reach here with force_reindex=False
2044 # AND prior chunks present — appending would leave stale/deleted
2045 # note text permanently searchable and duplicate every chunk.
2046 # For a genuinely never-indexed source, prior_ids is empty and
2047 # replace is a no-op, so this is always safe.
2048 vindex = self._get_vector_index(
2049 collection_id, collection_name, reset_stale_state=True
2050 )
2051 index_stats = vindex.index(
2052 source_type="document",
2053 source_id=document_id,
2054 chunks=chunk_inputs,
2055 replace=True,
2056 session=session,
2057 )
2058 logger.info(
2059 f"Vector index update for document {document_id}: "
2060 f"{index_stats.added} added, {index_stats.removed} "
2061 f"removed ({index_stats.chunks} chunks total)"
2062 )
2064 # Check if document was already indexed (for stats update)
2065 existing_status = (
2066 session.query(RagDocumentStatus)
2067 .filter_by(
2068 document_id=document_id,
2069 collection_id=collection_id,
2070 rag_index_id=self.rag_index_record.id,
2071 )
2072 .first()
2073 )
2074 was_already_indexed = existing_status is not None
2075 old_chunk_count = (
2076 (existing_status.chunk_count or 0) if existing_status else 0
2077 )
2079 # Mark document as indexed using rag_document_status table
2080 # Row existence = indexed, simple and clean
2081 timestamp = datetime.now(UTC)
2083 # Create or update RagDocumentStatus using ORM merge (atomic upsert)
2084 rag_status = RagDocumentStatus(
2085 document_id=document_id,
2086 collection_id=collection_id,
2087 rag_index_id=self.rag_index_record.id,
2088 chunk_count=len(chunks),
2089 indexed_at=timestamp,
2090 )
2091 session.merge(rag_status)
2093 logger.info(
2094 f"Marked document as indexed in rag_document_status: doc_id={document_id}, coll_id={collection_id}, chunks={len(chunks)}"
2095 )
2097 # Also update DocumentCollection table for backward compatibility
2098 session.query(DocumentCollection).filter_by(
2099 document_id=document_id, collection_id=collection_id
2100 ).update(
2101 {
2102 "indexed": True,
2103 "chunk_count": len(chunks),
2104 "last_indexed_at": timestamp,
2105 }
2106 )
2108 logger.info(
2109 "Also updated DocumentCollection.indexed for backward compatibility"
2110 )
2112 # Keep the aggregate synchronized on both first index and
2113 # replacement reindex. A replacement changes the per-document
2114 # row above, so apply its new-old delta instead of skipping the
2115 # aggregate update entirely.
2116 rag_index_obj = (
2117 session.query(RAGIndex)
2118 .filter_by(id=self.rag_index_record.id)
2119 .first()
2120 )
2121 if rag_index_obj: 2121 ↛ 2140line 2121 didn't jump to line 2140 because the condition on line 2121 was always true
2122 chunk_delta = len(chunks) - old_chunk_count
2123 rag_index_obj.chunk_count += chunk_delta
2124 if not was_already_indexed:
2125 rag_index_obj.total_documents += 1
2126 rag_index_obj.last_updated_at = datetime.now(UTC)
2127 logger.info(
2128 "Updated RAGIndex stats: "
2129 f"chunk_count {chunk_delta:+d}, "
2130 f"total_documents {'+1' if not was_already_indexed else '+0'}"
2131 )
2133 # Replace-on-reindex DB prune: no longer needed here —
2134 # ``VectorIndex.index(replace=True)`` above already deletes
2135 # this source's prior DocumentChunk rows (and their
2136 # vectors) as part of the same transaction/lock when
2137 # ``force_reindex`` is set.
2139 # Flush ORM changes to database before commit
2140 session.flush()
2141 logger.info(f"Flushed ORM changes for document {document_id}")
2143 # Commit the transaction. Durability is provided by
2144 # synchronous=NORMAL (sqlcipher_utils.py); SQLite
2145 # auto-checkpoints WAL at wal_autocheckpoint=250 frames.
2146 # An explicit PRAGMA wal_checkpoint(FULL) here used to
2147 # block other writers long enough to exhaust busy_timeout
2148 # under bulk-download concurrency (#4197).
2149 session.commit()
2151 logger.info(
2152 f"Successfully indexed document {document_id} for collection {collection_id} "
2153 f"with {len(chunks)} chunks"
2154 )
2156 return {
2157 "status": "success",
2158 "chunk_count": len(chunks),
2159 }
2161 except Exception as e:
2162 # The session is shared (thread-local) with the caller.
2163 # If session.flush() or session.commit() raised, the session
2164 # is in PendingRollbackError state until rolled back —
2165 # leaving subsequent operations to cascade. Roll back BEFORE
2166 # returning the error dict so the caller sees a clean
2167 # session. (Same pattern as the #3827 fix.)
2168 safe_rollback(session, "library_rag_service.index_document")
2169 # Backstop for the store-durable-before-DB-commit window: a
2170 # REPLACE reindex ran store.apply() (DURABLY removing this
2171 # source's prior vectors) before the DB commit, so the rollback
2172 # above resurrects rows whose vectors are already gone while
2173 # reverting indexed→its pre-attempt True. Flag for re-index so
2174 # the reconciler re-embeds and store+DB reconverge.
2175 self._flag_document_for_reindex(
2176 session, document_id, collection_id
2177 )
2178 logger.exception(
2179 f"Error indexing document {document_id} for collection {collection_id}"
2180 )
2181 return {
2182 "status": "error",
2183 "error": f"Operation failed: {type(e).__name__}",
2184 }
2186 def index_all_documents(
2187 self,
2188 collection_id: str,
2189 force_reindex: bool = False,
2190 progress_callback=None,
2191 ) -> Dict[str, Any]:
2192 """
2193 Index all documents in a collection into RAG.
2195 Args:
2196 collection_id: UUID of the collection to index
2197 force_reindex: Whether to force reindexing already indexed documents
2198 progress_callback: Optional callback function called after each document with (current, total, doc_title, status)
2200 Returns:
2201 Dict with counts of successful, skipped, and failed documents
2202 """
2203 with get_user_db_session(self.username, self.db_password) as session:
2204 # Get all DocumentCollection entries for this collection
2205 query = session.query(DocumentCollection).filter_by(
2206 collection_id=collection_id
2207 )
2209 if not force_reindex:
2210 # Only index documents that haven't been indexed yet
2211 query = query.filter_by(indexed=False)
2213 doc_collections = query.all()
2215 if not doc_collections:
2216 return {
2217 "status": "info",
2218 "message": "No documents to index",
2219 "successful": 0,
2220 "skipped": 0,
2221 "failed": 0,
2222 }
2224 results = {"successful": 0, "skipped": 0, "failed": 0, "errors": []}
2225 total = len(doc_collections)
2227 for idx, doc_collection in enumerate(doc_collections, 1):
2228 # Get the document for title info
2229 document = (
2230 session.query(Document)
2231 .filter_by(id=doc_collection.document_id)
2232 .first()
2233 )
2234 title = document.title if document else "Unknown"
2236 result = self.index_document(
2237 doc_collection.document_id, collection_id, force_reindex
2238 )
2240 if result["status"] == "success":
2241 results["successful"] += 1
2242 elif result["status"] in ("skipped", "cleared"):
2243 # "cleared" (a previously-indexed doc emptied to no text,
2244 # its stale vectors purged) is a handled non-failure, not an
2245 # error — counting it in "failed" would surface a blank-error
2246 # failure for a correct purge.
2247 results["skipped"] += 1
2248 else:
2249 results["failed"] += 1
2250 results["errors"].append(
2251 {
2252 "doc_id": doc_collection.document_id,
2253 "title": title,
2254 "error": result.get("error"),
2255 }
2256 )
2258 # Call progress callback if provided
2259 if progress_callback:
2260 progress_callback(idx, total, title, result["status"])
2262 logger.info(
2263 f"Indexed collection {collection_id}: "
2264 f"{results['successful']} successful, "
2265 f"{results['skipped']} skipped, "
2266 f"{results['failed']} failed"
2267 )
2269 return results
2271 def remove_document_from_rag(
2272 self, document_id: str, collection_id: str
2273 ) -> Dict[str, Any]:
2274 """
2275 Remove a document's chunks from RAG for a specific collection.
2277 Args:
2278 document_id: UUID of the Document to remove
2279 collection_id: UUID of the Collection to remove from
2281 Returns:
2282 Dict with status and count of removed chunks
2283 """
2284 with get_user_db_session(self.username, self.db_password) as session:
2285 # Get the DocumentCollection entry
2286 doc_collection = (
2287 session.query(DocumentCollection)
2288 .filter_by(document_id=document_id, collection_id=collection_id)
2289 .first()
2290 )
2292 if not doc_collection:
2293 return {
2294 "status": "error",
2295 "error": "Document not found in collection",
2296 }
2298 try:
2299 # Get collection name in the format collection_<uuid>
2300 collection = (
2301 session.query(Collection)
2302 .filter_by(id=collection_id)
2303 .first()
2304 )
2305 # Use collection_<uuid> format for internal storage
2306 collection_name = (
2307 f"collection_{collection_id}" if collection else "unknown"
2308 )
2310 # Remove the document's vectors from the FAISS index FIRST.
2311 # ``remove_documents_from_index`` -> ``VectorIndex.delete``
2312 # looks up the DocumentChunk rows by (source_type, source_id,
2313 # collection_name) to learn which FAISS int ids to remove,
2314 # and deletes those matching rows itself as part of the same
2315 # operation. If the DB rows were deleted first, that lookup
2316 # would find nothing and the vectors (and their persisted
2317 # text) would be stranded in the FAISS store forever — never
2318 # removed and never re-discoverable by a later cleanup.
2319 # Best-effort: a FAISS hiccup must not fail the DB removal
2320 # below.
2321 # ``remove_documents_from_index`` deletes the matching chunk
2322 # rows itself (one row per removed vector) and returns that
2323 # count. Capture it so the reported ``deleted_count`` reflects
2324 # the rows removed here, not just the backstop sweep below: in
2325 # the normal path those rows are already gone by the time the
2326 # backstop runs, so the backstop returns ~0 and — left
2327 # uncounted — would under-report the removal (often as 0) to
2328 # the caller/UI.
2329 removed_vectors = 0
2330 try:
2331 removed_vectors = self.remove_documents_from_index(
2332 [document_id], collection_id
2333 )
2334 except Exception:
2335 logger.exception(
2336 f"Could not remove document {document_id} vectors "
2337 f"from FAISS index for collection {collection_id}"
2338 )
2340 # DB-only cleanup for any rows ``remove_documents_from_index``
2341 # did not match (e.g. no current FAISS index for this
2342 # collection). Disjoint from ``removed_vectors`` above (those
2343 # rows are already committed-deleted), so the two sum to the
2344 # true total removed.
2345 deleted_count = removed_vectors + (
2346 self.embedding_manager._delete_chunks_from_db(
2347 collection_name=collection_name,
2348 source_id=document_id,
2349 )
2350 )
2352 # Update DocumentCollection RAG status
2353 doc_collection.indexed = False
2354 doc_collection.chunk_count = 0
2355 doc_collection.last_indexed_at = None
2356 # Clear the canonical RagDocumentStatus row too. Row-existence is
2357 # the "indexed" marker get_rag_stats / the RAG status route read,
2358 # so it MUST move together with DocumentCollection.indexed — a
2359 # leftover row would report the just-removed document as still
2360 # indexed.
2361 session.query(RagDocumentStatus).filter_by(
2362 document_id=document_id, collection_id=collection_id
2363 ).delete(synchronize_session=False)
2364 session.commit()
2366 logger.info(
2367 f"Removed {deleted_count} chunks for document {document_id} from collection {collection_id}"
2368 )
2370 return {"status": "success", "deleted_count": deleted_count}
2372 except Exception as e:
2373 # session.commit() above can raise; without rollback the
2374 # shared thread-local session stays poisoned for the
2375 # caller's next operation (issue #3827 pattern).
2376 safe_rollback(
2377 session, "library_rag_service.remove_document_from_rag"
2378 )
2379 # Backstop for the store-durable-before-DB-commit window:
2380 # remove_documents_from_index -> VectorIndex.delete ran
2381 # store.apply() (DURABLY removing this document's vectors) BEFORE
2382 # the DB status work above committed, so the rollback can leave
2383 # the row still marked indexed=True with its vectors already gone
2384 # -- indexed-but-unsearchable, and over-reported as indexed by
2385 # get_rag_stats with no auto-recovery. Durably COMPLETE the
2386 # removal's status change so the DB matches the removed vectors:
2387 # mark it not-indexed and drop the canonical RagDocumentStatus
2388 # row. This is a REMOVAL backstop, not a re-index one (contrast
2389 # index_document's _flag_document_for_reindex) -- the vectors are
2390 # intentionally gone, so we finish the removal rather than
2391 # resurrect the document. Best-effort: a retry of the removal
2392 # reconciles it if this too fails.
2393 try:
2394 session.query(DocumentCollection).filter_by(
2395 document_id=document_id, collection_id=collection_id
2396 ).update(
2397 {
2398 "indexed": False,
2399 "chunk_count": 0,
2400 "last_indexed_at": None,
2401 },
2402 synchronize_session=False,
2403 )
2404 session.query(RagDocumentStatus).filter_by(
2405 document_id=document_id, collection_id=collection_id
2406 ).delete(synchronize_session=False)
2407 session.commit()
2408 except Exception:
2409 safe_rollback(
2410 session,
2411 "library_rag_service.remove_document_from_rag backstop",
2412 )
2413 logger.warning(
2414 f"Could not finalize removal state for document "
2415 f"{document_id} in collection {collection_id}; a retry "
2416 "of the removal will reconcile it."
2417 )
2418 logger.exception(
2419 f"Error removing document {document_id} from collection {collection_id}"
2420 )
2421 return {
2422 "status": "error",
2423 "error": f"Operation failed: {type(e).__name__}",
2424 }
2426 def purge_document_chunks(
2427 self, document_id: str, collection_id: str
2428 ) -> int:
2429 """Delete a document's chunk rows from a collection's RAG store
2430 WITHOUT requiring the ``DocumentCollection`` join row to still exist.
2432 ``remove_document_from_rag`` short-circuits (returns "Document not
2433 found in collection") when the join row is gone — which is exactly
2434 the state after a ``Document`` is cascade-deleted. Callers that
2435 delete the Document first (note deletion) must use this to actually
2436 purge the now-orphaned ``DocumentChunk`` rows; otherwise the chunks
2437 (and their embedding ids) linger and semantic search keeps
2438 resolving hits to a Document row that no longer exists.
2440 Mirrors the chunk-delete half of ``remove_document_from_rag``.
2441 FAISS-vector pruning is NOT done here: pair this with
2442 ``purge_document_vectors`` when the document is being deleted
2443 outright — replace-on-reindex can never fire again for a deleted
2444 document id. Returns the number of chunk rows deleted.
2445 """
2446 collection_name = f"collection_{collection_id}"
2447 return self.embedding_manager._delete_chunks_from_db(
2448 collection_name=collection_name,
2449 source_id=document_id,
2450 )
2452 def purge_document_vectors(
2453 self, document_id: str, collection_id: str
2454 ) -> int:
2455 """Remove a deleted document's chunk vectors from the collection's
2456 vector store.
2458 Companion to ``purge_document_chunks`` for callers that delete the
2459 ``Document`` row itself (note deletion). Without this the vectors
2460 lingered until a replace-on-reindex of the SAME document id — which
2461 can never fire again once the document is deleted — so collection
2462 search kept serving the deleted content indefinitely.
2464 Routes to ``VectorIndex.delete(source_type="document",
2465 source_id=document_id)``. Unlike the pre-cutover implementation,
2466 no separate "shared chunk ownership" check is needed here: the new
2467 store is one-row-per-source (see the "One row = one vector" note in
2468 vector_stores/facade.py) — identical text in two documents gets two
2469 independent rows/vectors, so deleting this document's rows can
2470 never remove another document's vector.
2472 Deliberately does NOT go through ``_get_vector_index`` /
2473 ``_get_or_create_rag_index``: a deletion path must not create an
2474 index record (or probe the embedding provider) just to delete from
2475 it, and does not quarantine/rebuild on a failed load — an
2476 unreadable index has nothing purgeable right now; quarantine/
2477 rebuild decisions stay with the indexing paths.
2479 Returns the number of vectors removed.
2480 """
2481 collection_name = f"collection_{collection_id}"
2482 index_hash = self._get_index_hash(
2483 collection_name, self.embedding_model, self.embedding_provider
2484 )
2485 with get_user_db_session(self.username, self.db_password) as session:
2486 rag_index = self._find_matching_rag_index(
2487 session, collection_name, index_hash
2488 )
2489 if rag_index is None:
2490 # Collection was never indexed with this configuration —
2491 # no vectors to purge. Deliberately NOT
2492 # _get_or_create_rag_index: a deletion path must not
2493 # create index records (or probe the embedding provider).
2494 return 0
2495 self.rag_index_record = rag_index
2497 # Resolve where the file ACTUALLY lives, the way the index/search paths
2498 # do — but without _preflight's side effects (this deletion path must
2499 # not relocate/record/create). rag_index.index_path is the authoritative
2500 # DB record (the legacy shared-cache path pre-migration, the per-user
2501 # path after _preflight relocates it); _get_index_path(index_hash) only
2502 # recomputes the POST-migration location. Checking only the recomputed
2503 # path silently no-ops on an un-migrated install — the file is still at
2504 # the legacy index_path, so the deleted document's real vector survives.
2505 recomputed_path = self._get_index_path(rag_index.index_hash)
2506 recorded_path = (
2507 Path(rag_index.index_path) if rag_index.index_path else None
2508 )
2509 if recorded_path is not None and recorded_path.exists(): 2509 ↛ 2510line 2509 didn't jump to line 2510 because the condition on line 2509 was never true
2510 index_path = recorded_path
2511 elif recomputed_path.exists():
2512 index_path = recomputed_path
2513 else:
2514 return 0
2515 verified, reason = self.integrity_manager.verify_file(index_path)
2516 if not verified: 2516 ↛ 2525line 2516 didn't jump to line 2525 because the condition on line 2516 was always true
2517 # Leave quarantine/rebuild decisions to the indexing paths;
2518 # an unreadable index has nothing purgeable right now.
2519 logger.warning(
2520 f"Skipping vector purge for document {document_id}: index "
2521 f"{index_path} fails verification ({reason})"
2522 )
2523 return 0
2525 lock = _get_faiss_write_lock(self.username, str(index_path))
2527 # nested=True: these run reentrantly inside VectorIndex.apply(), which
2528 # holds just-flushed, uncommitted DocumentChunk rows on this same
2529 # thread-local session. The integrity write must run in a SAVEPOINT and
2530 # let apply()'s caller commit — a plain commit/rollback here would
2531 # settle or DISCARD those pending rows (silently losing indexed text
2532 # while the vectors persist). See FileIntegrityManager._integrity_write.
2533 def _record(path: Path) -> None:
2534 self.integrity_manager.record_file(
2535 path,
2536 related_entity_type="rag_index",
2537 related_entity_id=rag_index.id,
2538 nested=True,
2539 )
2541 def _verify(path: Path):
2542 return self.integrity_manager.verify_file(path, nested=True)
2544 try:
2545 vindex = VectorIndex(
2546 username=self.username,
2547 db_password=self.db_password,
2548 # Delete-only path: never materialise the embedding backend.
2549 # VectorIndex.delete() resolves rows by id and never calls
2550 # embed_*; forcing the model here would cost seconds + hundreds
2551 # of MB of RSS per deleted document for nothing.
2552 embeddings=_LazyEmbeddings(self.embedding_manager),
2553 embedding_model=self.embedding_model,
2554 embedding_model_type=EmbeddingProvider(self.embedding_provider),
2555 collection_name=collection_name,
2556 dimension=rag_index.embedding_dimension,
2557 path=index_path,
2558 lock=lock,
2559 integrity_record=_record,
2560 integrity_verify=_verify,
2561 # Use the STORED index config, not the service instance's — a
2562 # service built with defaults (index_type="flat") would evaluate
2563 # supports_delete wrongly for an HNSW collection and silently
2564 # fail the removal. NULL (legacy rows) falls back to self.*.
2565 index_type=rag_index.index_type or self.index_type,
2566 metric=rag_index.distance_metric or self.distance_metric,
2567 normalize=(
2568 rag_index.normalize_vectors
2569 if rag_index.normalize_vectors is not None
2570 else self.normalize_vectors
2571 ),
2572 )
2573 stats = vindex.delete(source_type="document", source_id=document_id)
2574 except (ValueError, RuntimeError, OSError):
2575 # OSError too: this is a BEST-EFFORT purge whose caller goes on to
2576 # delete the DocumentChunk rows. A disk error from the FAISS
2577 # persist() (OSError) must not escape and abort that row deletion —
2578 # it would strand the chunk rows while leaving the vectors, the
2579 # opposite inconsistency. Skip the vector purge and let the caller
2580 # proceed; the vectors reconcile on the next re-index.
2581 logger.warning(
2582 f"Skipping vector purge for document {document_id}: could "
2583 f"not load or update index {index_path}"
2584 )
2585 return 0
2586 return stats.removed
2588 def index_documents_batch(
2589 self,
2590 doc_info: List[tuple],
2591 collection_id: str,
2592 force_reindex: bool = False,
2593 ) -> Dict[str, Dict[str, Any]]:
2594 """
2595 Index multiple documents in a batch for a specific collection.
2597 Args:
2598 doc_info: List of (doc_id, title) tuples
2599 collection_id: UUID of the collection to index for
2600 force_reindex: Whether to force reindexing even if already indexed
2602 Returns:
2603 Dict mapping doc_id to individual result
2604 """
2605 results = {}
2606 doc_ids = [doc_id for doc_id, _ in doc_info]
2608 # Use single database session for querying
2609 with get_user_db_session(self.username, self.db_password) as session:
2610 # Pre-load all documents for this batch
2611 documents = (
2612 session.query(Document).filter(Document.id.in_(doc_ids)).all()
2613 )
2615 # Create lookup for quick access
2616 doc_lookup = {doc.id: doc for doc in documents}
2618 # Pre-load DocumentCollection entries
2619 doc_collections = (
2620 session.query(DocumentCollection)
2621 .filter(
2622 DocumentCollection.document_id.in_(doc_ids),
2623 DocumentCollection.collection_id == collection_id,
2624 )
2625 .all()
2626 )
2627 doc_collection_lookup = {
2628 dc.document_id: dc for dc in doc_collections
2629 }
2631 # Process each document in the batch
2632 for doc_id, title in doc_info:
2633 document = doc_lookup.get(doc_id)
2635 if not document:
2636 results[doc_id] = {
2637 "status": "error",
2638 "error": "Document not found",
2639 }
2640 continue
2642 # Check if already indexed via DocumentCollection
2643 doc_collection = doc_collection_lookup.get(doc_id)
2644 if (
2645 doc_collection
2646 and doc_collection.indexed
2647 and not force_reindex
2648 ):
2649 results[doc_id] = {
2650 "status": "skipped",
2651 "message": "Document already indexed for this collection",
2652 "chunk_count": doc_collection.chunk_count,
2653 }
2654 continue
2656 # Do NOT short-circuit empty text_content here: a document/note
2657 # edited down to empty still has its PRIOR chunks + vectors on
2658 # disk, and only index_document's empty-content branch purges
2659 # them (via an empty-replace). Returning an "error" here left
2660 # that stale (plaintext) content permanently searchable through
2661 # the bulk "Index All" path — and self-perpetuating, since the
2662 # doc's indexed flag never flips. Route it through index_document
2663 # so the batch path gets the same purge-on-clear as the SSE /
2664 # single-document paths; it returns "cleared" (purged) or the
2665 # same "error" when there were no prior chunks.
2666 try:
2667 result = self.index_document(
2668 doc_id, collection_id, force_reindex
2669 )
2670 results[doc_id] = result
2671 except Exception as e:
2672 logger.exception(
2673 f"Error indexing document {doc_id} in batch"
2674 )
2675 results[doc_id] = {
2676 "status": "error",
2677 "error": f"Indexing failed: {type(e).__name__}",
2678 }
2680 return results
2682 def index_documents_parallel(
2683 self,
2684 doc_info: List[Tuple[str, str]],
2685 collection_id: str,
2686 force_reindex: bool = False,
2687 max_workers: int = 4,
2688 progress_callback: Optional[
2689 Callable[[int, int, str, str], None]
2690 ] = None,
2691 is_cancelled: Optional[Callable[[], bool]] = None,
2692 ) -> Dict[str, Any]:
2693 """
2694 Index many documents concurrently with bounded fan-out.
2696 Each per-document task runs the prepared pipeline
2697 (:meth:`_prepare_document` for load/split/embed, then the
2698 coordinator's serialized durable write) unless a subclass overrides
2699 :meth:`_index_one` or :meth:`index_document`, in which case that
2700 per-doc path is invoked directly — the per-doc logic (split, embed,
2701 persist chunks, FAISS-merge-under-lock, mark indexed, update stats)
2702 is unchanged. The bounded fan-out only changes *which* sequential
2703 loop dispatches the work: before this method, six call sites in the
2704 RAG routes (main's ``research_library/routes/rag_routes.py``, now
2705 :mod:`local_deep_research.web.routers.rag`) and the background
2706 reconciler in :mod:`scheduler.background` each did this one doc at
2707 a time. The slow embedding round-trip (outside the FAISS write
2708 lock) is now overlapped across documents.
2710 Concurrency invariants this method relies on (already true):
2712 * :meth:`index_document` opens a fresh per-thread DB session via
2713 :func:`get_user_db_session` — workers do not share sessions.
2714 * All FAISS mutations go through ``VectorIndex.apply()`` in
2715 :mod:`local_deep_research.vector_stores`, which is only
2716 invoked under the per-(username, index_path) write lock
2717 returned by :func:`_get_faiss_write_lock`. The service passes
2718 that lock into the vector-store facade at the call sites
2719 around L422, L879, L1097, and L2025, so concurrent
2720 ``index_document`` workers serialise their FAISS appends on
2721 the same lock without any extra coordination here.
2722 * ``LocalEmbeddingManager.embeddings`` uses double-checked
2723 locking for first-load.
2724 * ``is_cancelled`` is polled between completions AND between
2725 submissions. The bounded-submission loop admits new work into
2726 the pool only when ``is_cancelled()`` returns False, and
2727 each :func:`_run_one` worker re-checks ``is_cancelled()``
2728 immediately before delegating to ``index_document``. This
2729 closes the false-poll-to-submit gap (the cancel signal can
2730 be set between the main thread's poll and ``pool.submit()``)
2731 so post-cancel starts are zero by construction. The final
2732 ``is_cancelled()`` poll in the ``finally`` block also
2733 preserves ``cancelled=True`` in the aggregate when the
2734 signal lands while the last in-flight worker is running.
2735 * In-flight embedding round-trips cannot be interrupted (Python
2736 has no safe way to kill a SentenceTransformer forward pass
2737 mid-array) — they finish naturally before the helper returns
2738 via ``pool.shutdown(wait=True)``.
2740 Args:
2741 doc_info: List of ``(doc_id, title)`` pairs to index.
2742 collection_id: UUID of the collection.
2743 force_reindex: Forwarded to :meth:`index_document`.
2744 max_workers: Bound on the size of the
2745 ``ThreadPoolExecutor`` — i.e. how many documents can be
2746 **preparing** concurrently (load + split + embed). It does
2747 NOT bound the prepared-but-unwritten backlog: completed
2748 futures hold their full vector batches in memory until the
2749 coordinator's serial write stage picks them up, so for very
2750 large batches the resident memory footprint can exceed what
2751 ``max_workers`` alone suggests. Embedding is the slow stage
2752 in practice, so this is acceptable for the default 4-worker
2753 fan-out; a caller with many large documents should either
2754 shrink ``max_workers`` or chunk the submission. ``1`` falls
2755 back to fully sequential behaviour (still useful when a
2756 flaky network makes concurrent httpx calls worse than
2757 serial). Default 4.
2758 progress_callback: Optional hook called once per completed
2759 document with ``(completed, total, title, status)`` —
2760 matches the legacy shape used by
2761 :meth:`index_all_documents` so existing UI progress
2762 paths keep working.
2763 is_cancelled: Optional zero-arg callable polled between
2764 completions AND between submissions. When it returns
2765 ``True``, no further work is admitted into the pool;
2766 in-flight futures finish naturally. The executor is
2767 always shut down with ``wait=True`` so no worker can
2768 outlive the caller's service context.
2770 Returns:
2771 ``{"successful": int, "skipped": int, "failed": int,
2772 "errors": [...], "results": {doc_id: {...}},
2773 "cancelled": bool, "total": int}``.
2774 """
2775 if max_workers < 1:
2776 max_workers = 1
2778 total = len(doc_info)
2779 if total == 0:
2780 return {
2781 "successful": 0,
2782 "skipped": 0,
2783 "failed": 0,
2784 "errors": [],
2785 "results": {},
2786 "cancelled": False,
2787 "total": 0,
2788 }
2790 # Title-lookup keyed by doc_id — futures return in ``as_completed``
2791 # order, not submission order, so we cannot rely on positional
2792 # iteration when emitting progress.
2793 title_lookup = dict(doc_info)
2795 results: Dict[str, Dict[str, Any]] = {}
2796 counters = {"successful": 0, "skipped": 0, "failed": 0}
2797 errors: List[Dict[str, Any]] = []
2798 cancelled = False
2800 def _run_one(doc_id: str) -> Dict[str, Any]:
2801 """Single-doc worker — never raises; returns a status dict."""
2802 try:
2803 # Three seams tell us to bypass the prepared/serialized
2804 # pipeline and call the per-doc method directly:
2805 #
2806 # 1. A class-level subclass override of ``_index_one`` —
2807 # detected via
2808 # ``type(self)._index_one is not LibraryRAGService._index_one``
2809 # so it works regardless of whether the override
2810 # has ``__func__`` (the old
2811 # ``self._index_one.__func__ is not LibraryRAGService._index_one``
2812 # check raised AttributeError on instance-attribute
2813 # patches that don't expose ``__func__``).
2814 # 2. A class-level subclass override of the legacy
2815 # ``index_document`` extension point this seam
2816 # replaced — callers that override ``index_document``
2817 # at class level (not just instance-patch it for
2818 # tests) still get the per-doc bypass, otherwise they
2819 # would silently fall through to the prepared
2820 # pipeline and lose the override.
2821 # 3. An instance-level patch of ``_index_one`` or
2822 # ``index_document`` — the legacy seams used by the
2823 # existing parallel-runner tests that mock the
2824 # per-doc method directly.
2825 #
2826 # If none is present, run the prepared pipeline: split
2827 # and embed concurrently, then hand the prepared document
2828 # off to the coordinator for serialized durable writes.
2829 cls = type(self)
2830 if cls._index_one is not LibraryRAGService._index_one:
2831 return self._index_one(doc_id, collection_id, force_reindex)
2832 if cls.index_document is not LibraryRAGService.index_document:
2833 return self.index_document(
2834 doc_id, collection_id, force_reindex
2835 )
2836 if (
2837 "_index_one" in self.__dict__
2838 or "index_document" in self.__dict__
2839 ):
2840 return self._index_one(doc_id, collection_id, force_reindex)
2841 prepared_result = self._prepare_document(
2842 doc_id, collection_id, force_reindex
2843 )
2844 if prepared_result.get("status") == "needs_serial_write": 2844 ↛ 2845line 2844 didn't jump to line 2845 because the condition on line 2844 was never true
2845 return self.index_document(
2846 doc_id, collection_id, force_reindex
2847 )
2848 return prepared_result
2849 except Exception as exc:
2850 # ``index_document`` already returns an error dict on its
2851 # own internal failures; this catches only bugs / hook
2852 # breaks outside that contract.
2853 logger.exception(f"Parallel index worker for {doc_id} crashed")
2854 return {
2855 "status": "error",
2856 "error": f"Indexing failed: {type(exc).__name__}",
2857 }
2859 # De-duplicated submission queue preserving input order. The
2860 # underlying ``index_document`` is not idempotent at the DB level
2861 # beyond its own checks; without this guard a buggy caller could
2862 # double-insert chunks from the same doc.
2863 pending_doc_ids: List[str] = []
2864 seen: Set[str] = set()
2865 for doc_id, _title in doc_info:
2866 if doc_id in seen:
2867 continue
2868 seen.add(doc_id)
2869 pending_doc_ids.append(doc_id)
2871 pool = ThreadPoolExecutor(
2872 max_workers=max_workers,
2873 thread_name_prefix="ldr-index-doc-",
2874 )
2875 try:
2876 future_to_doc: Dict[Any, str] = {}
2877 submit_idx = 0
2878 completed = 0
2879 # Bounded-submission loop: poll ``is_cancelled`` BEFORE
2880 # admitting any new work, then submit up to ``max_workers``
2881 # futures, then wait for one completion.
2882 #
2883 # This closes the post-cancel submission leak in the
2884 # previous all-upfront submission model: with the old
2885 # loop, all futures were submitted before the first
2886 # ``is_cancelled()`` poll, so queued futures could be
2887 # picked up by idle workers between the caller's cancel
2888 # signal (e.g. ``_sse_cancel.set()`` in
2889 # ``web/routers/rag.py::index_collection``) and the helper's
2890 # ``pool.shutdown(cancel_futures=True)``. Here, new work
2891 # is admitted only when ``is_cancelled()`` returns False,
2892 # so post-cancel starts are zero by construction.
2893 while future_to_doc or submit_idx < len(pending_doc_ids):
2894 # 1. Cancellation fence: stop submitting new work if
2895 # cancelled. Once we observe cancellation we don't
2896 # re-poll — the gate stays closed until the loop
2897 # drains.
2898 if (
2899 not cancelled
2900 and is_cancelled is not None
2901 and is_cancelled()
2902 ):
2903 cancelled = True
2904 logger.info(
2905 f"Parallel indexing cancelled at {completed}/{total}"
2906 )
2907 for f in future_to_doc: 2907 ↛ 2908line 2907 didn't jump to line 2908 because the loop on line 2907 never started
2908 f.cancel()
2910 # 2. Admit new work up to ``max_workers`` in flight
2911 # (skipped once cancelled).
2912 if not cancelled:
2913 while len(future_to_doc) < max_workers and submit_idx < len(
2914 pending_doc_ids
2915 ):
2916 if is_cancelled is not None and is_cancelled():
2917 cancelled = True
2918 logger.info(
2919 f"Parallel indexing cancelled during submission at {completed}/{total}"
2920 )
2921 for f in future_to_doc:
2922 f.cancel()
2923 break
2924 doc_id = pending_doc_ids[submit_idx]
2925 submit_idx += 1
2926 future_to_doc[pool.submit(_run_one, doc_id)] = doc_id
2928 # 3. Nothing in flight and nothing left to submit → done.
2929 if not future_to_doc:
2930 break
2932 # 4. Wait for the next completion, then process it.
2933 try:
2934 done_iter = as_completed(list(future_to_doc.keys()))
2935 fut = next(done_iter, None)
2936 except Exception:
2937 logger.exception(
2938 "as_completed failed unexpectedly during parallel indexing"
2939 )
2940 fut = None
2941 if fut is None: 2941 ↛ 2945line 2941 didn't jump to line 2945 because the condition on line 2941 was never true
2942 # Iterator short-circuited without yielding — loop
2943 # back. The next iteration's cancel poll / submit
2944 # check will either submit more work or exit.
2945 continue
2947 doc_id = future_to_doc.pop(fut)
2948 try:
2949 result = fut.result()
2950 except CancelledError:
2951 result = {
2952 "status": "skipped",
2953 "message": "Cancelled before submission",
2954 }
2955 except Exception as exc:
2956 logger.exception(
2957 f"Future for doc {doc_id} raised unexpectedly"
2958 )
2959 result = {
2960 "status": "error",
2961 "error": f"Worker crashed: {type(exc).__name__}",
2962 }
2963 status = result.get("status", "error")
2964 prepared = result.get("prepared")
2965 if status == "prepared" and isinstance(
2966 prepared, _PreparedDocument
2967 ):
2968 result = self._write_prepared_document(prepared)
2969 status = result.get("status", "error")
2970 results[doc_id] = result
2971 if status == "success":
2972 counters["successful"] += 1
2973 elif status in ("skipped", "cleared"):
2974 counters["skipped"] += 1
2975 else:
2976 counters["failed"] += 1
2977 errors.append(
2978 {
2979 "doc_id": doc_id,
2980 "title": title_lookup.get(doc_id, "Unknown"),
2981 "error": result.get("error"),
2982 }
2983 )
2984 completed += 1
2985 if progress_callback is not None:
2986 try:
2987 progress_callback(
2988 completed,
2989 total,
2990 title_lookup.get(doc_id, "Unknown"),
2991 status,
2992 )
2993 except Exception:
2994 # A buggy callback must not poison the indexing
2995 # run — log and continue.
2996 logger.exception(
2997 "Parallel index progress_callback raised"
2998 )
2999 finally:
3000 if not cancelled and is_cancelled is not None and is_cancelled():
3001 cancelled = True
3002 # Always ``wait=True`` so in-flight workers finish before the
3003 # helper returns. With bounded submission no new work is
3004 # admitted after cancel is observed, so the only in-flight
3005 # work at shutdown is what was already running — these finish
3006 # naturally. ``cancel_futures=cancelled`` is a no-op here
3007 # (no queued futures by construction) but is kept for
3008 # symmetry with other shutdown sites in the codebase and as
3009 # defense-in-depth if a future change relaxes the bounded-
3010 # submission invariant.
3011 pool.shutdown(wait=True, cancel_futures=cancelled)
3013 if cancelled:
3014 # Mark docs that were never submitted (still in the
3015 # pending queue past ``submit_idx``) as skipped. Submitted
3016 # docs all have results — either from processing in the
3017 # loop above or from the ``wait=True`` shutdown drain —
3018 # so we don't need to touch them here.
3019 for doc_id in pending_doc_ids[submit_idx:]:
3020 if doc_id not in results: 3020 ↛ 3019line 3020 didn't jump to line 3019 because the condition on line 3020 was always true
3021 results[doc_id] = {
3022 "status": "skipped",
3023 "message": "Cancelled before submission",
3024 }
3025 counters["skipped"] += 1
3027 logger.info(
3028 f"Parallel indexing complete: "
3029 f"{counters['successful']} successful, "
3030 f"{counters['skipped']} skipped, "
3031 f"{counters['failed']} failed"
3032 f"{' (cancelled)' if cancelled else ''}"
3033 )
3035 return {
3036 "successful": counters["successful"],
3037 "skipped": counters["skipped"],
3038 "failed": counters["failed"],
3039 "errors": errors,
3040 "results": results,
3041 "cancelled": cancelled,
3042 "total": total,
3043 }
3045 def reconcile_collection_index(self, collection_id: str) -> Dict[str, Any]:
3046 """Reconcile indexed flags and stats against durable FAISS membership.
3048 A document is indexed only when it has at least one current-model chunk
3049 row and every such chunk id exists in the live FAISS index. Stale status
3050 rows are removed and compatibility flags/stats are rebuilt atomically.
3052 Safety invariant: if the live FAISS index reports zero ids (transient
3053 load error, missing file, on-disk corruption the verifier just
3054 quarantined) but the DB still has DocumentChunk rows for the collection,
3055 we REFUSE to shrink state. The function returns
3056 ``{"reconciliation_skipped": True, "reason": "..."}`` and the worker
3057 surfaces this as a task failure rather than clearing every indexed flag
3058 and RagDocumentStatus row. A mass-clear from a transient store fault
3059 would orphan every search hit and require a full re-index to recover.
3061 ``RagDocumentStatus.indexed_at`` is preserved for durable docs that
3062 already have a row — only newly durable docs get the current timestamp.
3063 """
3064 collection_name = f"collection_{collection_id}"
3065 with self._collection_transaction_lock(collection_id):
3066 vindex = self._get_vector_index(
3067 collection_id, collection_name, reset_stale_state=False
3068 )
3069 live_ids = set(vindex.live_ids())
3070 rag_index_id = self.rag_index_record.id
3071 with get_user_db_session(
3072 self.username, self.db_password
3073 ) as session:
3074 rows = (
3075 session.query(DocumentChunk.id, DocumentChunk.source_id)
3076 .filter_by(
3077 source_type="document",
3078 collection_name=collection_name,
3079 embedding_model=self.embedding_model,
3080 embedding_model_type=EmbeddingProvider(
3081 self.embedding_provider
3082 ),
3083 )
3084 .all()
3085 )
3086 ids_by_document: Dict[str, Set[int]] = {}
3087 for chunk_id, source_id in rows:
3088 if source_id is not None: 3088 ↛ 3087line 3088 didn't jump to line 3087 because the condition on line 3088 was always true
3089 ids_by_document.setdefault(str(source_id), set()).add(
3090 int(chunk_id)
3091 )
3093 # Refuse-to-shrink guard: a transient empty live_ids while
3094 # chunk rows exist would otherwise mass-clear every indexed
3095 # flag and RagDocumentStatus row, silently killing search
3096 # for every document in the collection. Treat as a
3097 # reconciliation failure; the worker surfaces the skip.
3098 if not live_ids and ids_by_document:
3099 logger.warning(
3100 "reconcile_collection_index: live_ids empty but "
3101 f"{len(ids_by_document)} document(s) have chunk rows; "
3102 "refusing to clear indexed flags"
3103 )
3104 return {
3105 "reconciliation_skipped": True,
3106 "reason": (
3107 "live vector store reported zero ids while "
3108 "DocumentChunk rows exist; refusing to clear "
3109 "indexed flags"
3110 ),
3111 "indexed_documents": 0,
3112 "indexed_chunks": 0,
3113 "live_vectors": 0,
3114 "orphan_vectors": 0,
3115 }
3117 durable = {
3118 doc_id: chunk_ids
3119 for doc_id, chunk_ids in ids_by_document.items()
3120 if chunk_ids and chunk_ids.issubset(live_ids)
3121 }
3123 # Snapshot existing indexed_at BEFORE the bulk delete so
3124 # durable docs that already had a status row keep their
3125 # original timestamp. Only docs newly promoted to durable
3126 # get ``now``.
3127 prior_indexed_at: Dict[str, Any] = {
3128 str(row.document_id): row.indexed_at
3129 for row in session.query(
3130 RagDocumentStatus.document_id,
3131 RagDocumentStatus.indexed_at,
3132 )
3133 .filter(RagDocumentStatus.collection_id == collection_id)
3134 .all()
3135 }
3137 links = (
3138 session.query(DocumentCollection)
3139 .filter_by(collection_id=collection_id)
3140 .all()
3141 )
3142 now = datetime.now(UTC)
3143 for link in links:
3144 chunk_ids = durable.get(str(link.document_id), set())
3145 link.indexed = bool(chunk_ids)
3146 link.chunk_count = len(chunk_ids)
3147 if chunk_ids: 3147 ↛ 3143line 3147 didn't jump to line 3143 because the condition on line 3147 was always true
3148 link.last_indexed_at = link.last_indexed_at or now
3150 session.query(RagDocumentStatus).filter_by(
3151 collection_id=collection_id
3152 ).delete(synchronize_session=False)
3153 for doc_id, chunk_ids in durable.items():
3154 session.add(
3155 RagDocumentStatus(
3156 document_id=doc_id,
3157 collection_id=collection_id,
3158 rag_index_id=rag_index_id,
3159 chunk_count=len(chunk_ids),
3160 indexed_at=prior_indexed_at.get(doc_id, now),
3161 )
3162 )
3164 rag_index = (
3165 session.query(RAGIndex).filter_by(id=rag_index_id).first()
3166 )
3167 if rag_index: 3167 ↛ 3171line 3167 didn't jump to line 3171 because the condition on line 3167 was always true
3168 rag_index.chunk_count = sum(map(len, durable.values()))
3169 rag_index.total_documents = len(durable)
3170 rag_index.last_updated_at = now
3171 session.commit()
3173 durable_ids = {i for ids in durable.values() for i in ids}
3174 return {
3175 "indexed_documents": len(durable),
3176 "indexed_chunks": len(durable_ids),
3177 "live_vectors": len(live_ids),
3178 "orphan_vectors": len(live_ids - durable_ids),
3179 }
3181 def get_rag_stats(
3182 self, collection_id: Optional[str] = None
3183 ) -> Dict[str, Any]:
3184 """
3185 Get RAG statistics for a collection.
3187 Args:
3188 collection_id: UUID of the collection (defaults to Library)
3190 Returns:
3191 Dict with counts and metadata about indexed documents
3192 """
3193 with get_user_db_session(self.username, self.db_password) as session:
3194 # Get collection ID (default to Library)
3195 if not collection_id: 3195 ↛ 3196line 3195 didn't jump to line 3196 because the condition on line 3195 was never true
3196 from ...database.library_init import get_default_library_id
3198 collection_id = get_default_library_id(
3199 self.username, self.db_password
3200 )
3202 # Count total documents in collection
3203 total_docs = (
3204 session.query(DocumentCollection)
3205 .filter_by(collection_id=collection_id)
3206 .count()
3207 )
3209 # Count indexed documents from rag_document_status table
3210 from ...database.models.library import RagDocumentStatus
3212 # One collection can hold status rows from several indexes at
3213 # once, so scope the counts to the index this configuration
3214 # resolves to. No such index means nothing is indexed for it.
3215 rag_index = self._resolve_index_for_config(session, collection_id)
3217 if rag_index is None:
3218 indexed_docs = 0
3219 total_chunks = 0
3220 foreign_rows = (
3221 session.query(RagDocumentStatus)
3222 .filter_by(collection_id=collection_id)
3223 .count()
3224 )
3225 if foreign_rows:
3226 logger.warning(
3227 f"Collection {collection_id} has no index for the "
3228 f"active embedding configuration "
3229 f"({self.embedding_provider}/{self.embedding_model}), "
3230 f"so it reports zero indexed documents. "
3231 f"{foreign_rows} document(s) are indexed under other "
3232 f"configurations and are not searchable from this one."
3233 )
3234 else:
3235 indexed_docs = (
3236 session.query(RagDocumentStatus)
3237 .filter_by(
3238 collection_id=collection_id,
3239 rag_index_id=rag_index.id,
3240 )
3241 .count()
3242 )
3244 # Count total chunks from rag_document_status table
3245 total_chunks = (
3246 session.query(func.sum(RagDocumentStatus.chunk_count))
3247 .filter_by(
3248 collection_id=collection_id,
3249 rag_index_id=rag_index.id,
3250 )
3251 .scalar()
3252 or 0
3253 )
3255 # Describe the index the counts above came from. Sampling any
3256 # chunk of the collection would describe whichever index wrote
3257 # first, which is the same mismatch as the counts.
3258 embedding_info = {}
3259 if rag_index is not None:
3260 stored_provider = rag_index.embedding_model_type
3261 if isinstance(stored_provider, EmbeddingProvider):
3262 stored_provider = stored_provider.value
3263 embedding_info = {
3264 "model": rag_index.embedding_model,
3265 "model_type": stored_provider,
3266 "dimension": rag_index.embedding_dimension,
3267 }
3269 return {
3270 "total_documents": total_docs,
3271 "indexed_documents": indexed_docs,
3272 "unindexed_documents": total_docs - indexed_docs,
3273 "total_chunks": total_chunks,
3274 "embedding_info": embedding_info,
3275 "chunk_size": self.chunk_size,
3276 "chunk_overlap": self.chunk_overlap,
3277 }
3279 def remove_documents_from_index(
3280 self,
3281 document_ids: List[str],
3282 collection_id: str,
3283 *,
3284 chunk_ids_by_document: Optional[Dict[str, List[int]]] = None,
3285 ) -> int:
3286 """Delete given documents' vectors from a collection's vector store.
3288 Vector-store-only: the DB chunk rows are removed separately by the
3289 caller. This closes the gap where deleting a document's DB chunks
3290 left its vectors (and their persisted text — now impossible, see
3291 vector_stores/base.py's SECURITY INVARIANT) in the store, so the
3292 document kept surfacing in semantic search until a full reindex.
3294 For each id, calls ``VectorIndex.delete(source_type="document",
3295 source_id=...)`` — which looks up the matching ``DocumentChunk``
3296 rows by (source_type, source_id, collection_name) — UNLESS
3297 ``chunk_ids_by_document`` supplies that document's chunk ids
3298 directly, in which case ``VectorIndex.delete_ids(ids)`` is used
3299 instead. Callers whose own per-item transactions must delete the DB
3300 rows eagerly (for their own atomicity/rollback semantics) BEFORE
3301 this batched cleanup runs need the latter: the source_id lookup
3302 would find nothing once the rows are already gone (see the Zotero
3303 sync caller). Hardcodes ``source_type="document"``: both current
3304 callers (Zotero sync cleanup, ``remove_document_from_rag``) only
3305 ever pass library-document ids (the pre-cutover implementation
3306 matched a mixed ``document_id``/``source_id`` metadata key
3307 generically across "document" and "user_document" sources, but
3308 nothing ever called it with anything but library documents).
3310 Returns the total number of vectors removed across all ids. No-ops
3311 (returns 0) when the collection has no current vector index — this
3312 guard runs BEFORE building a ``VectorIndex`` so a delete call never
3313 creates one (``_get_vector_index`` would otherwise create the
3314 RAGIndex row via ``_get_or_create_rag_index``).
3315 """
3316 from ...database.models.library import RAGIndex
3317 from ...database.session_context import get_user_db_session
3319 wanted = [d for d in document_ids if d]
3320 if not wanted:
3321 return 0
3323 collection_name = f"collection_{collection_id}"
3324 # Don't create an index just to delete from it — only act if one
3325 # already exists for this collection.
3326 with get_user_db_session(self.username, self.db_password) as session:
3327 has_index = (
3328 session.query(RAGIndex.id)
3329 .filter_by(collection_name=collection_name, is_current=True)
3330 .first()
3331 is not None
3332 )
3333 if not has_index:
3334 return 0
3336 try:
3337 # reset_stale_state=True is REQUIRED on this write path. If the
3338 # pre-flight finds the index corrupt it quarantines the file
3339 # (renames it aside); a fresh EMPTY store is then built and the
3340 # delete loop below persists that empty store over the real path.
3341 # Without resetting the per-document indexed state, the DB would
3342 # still claim every other document is indexed while their vectors
3343 # are gone — silently wiping the whole collection's search with no
3344 # recovery but a manual full reindex. The reset only fires when the
3345 # index was actually quarantined (guarded on the file no longer
3346 # existing), so a healthy remove is unaffected. search() omits this
3347 # flag safely because it never persists an empty store.
3348 vindex = self._get_vector_index(
3349 collection_id, collection_name, reset_stale_state=True
3350 )
3351 except (ValueError, RuntimeError):
3352 logger.warning(
3353 "Could not load vector index for "
3354 f"{collection_name} while removing document vectors"
3355 )
3356 return 0
3358 # Collect ALL removable chunk ids across every document, then remove
3359 # them in ONE apply() (vindex.delete_ids). For an HNSW collection each
3360 # apply() rebuilds the whole index, so a per-document loop is
3361 # O(documents x collection_size); a single batched call rebuilds once.
3362 # Captured ids are used directly; documents without captured ids have
3363 # their chunk ids resolved by (source_type, source_id, collection_name)
3364 # — those rows still exist at this point.
3365 all_ids: List[int] = []
3366 unresolved: List[str] = []
3367 for document_id in wanted:
3368 captured = (
3369 chunk_ids_by_document.get(document_id)
3370 if chunk_ids_by_document
3371 else None
3372 )
3373 if captured: 3373 ↛ 3374line 3373 didn't jump to line 3374 because the condition on line 3373 was never true
3374 all_ids.extend(int(i) for i in captured if i is not None)
3375 else:
3376 unresolved.append(document_id)
3378 if unresolved: 3378 ↛ 3394line 3378 didn't jump to line 3394 because the condition on line 3378 was always true
3379 in_chunk = 500 # keep below SQLITE_MAX_VARIABLE_NUMBER (999)
3380 with get_user_db_session(
3381 self.username, self.db_password
3382 ) as session:
3383 for i in range(0, len(unresolved), in_chunk):
3384 batch = unresolved[i : i + in_chunk]
3385 all_ids.extend(
3386 row.id
3387 for row in session.query(DocumentChunk.id).filter(
3388 DocumentChunk.source_type == "document",
3389 DocumentChunk.collection_name == collection_name,
3390 DocumentChunk.source_id.in_(batch),
3391 )
3392 )
3394 removed = 0
3395 if all_ids: 3395 ↛ 3405line 3395 didn't jump to line 3405 because the condition on line 3395 was always true
3396 try:
3397 stats = vindex.delete_ids(sorted(set(all_ids)))
3398 removed = stats.removed
3399 except Exception:
3400 logger.warning(
3401 f"Could not remove {len(set(all_ids))} vectors for "
3402 f"{len(wanted)} document(s) from {collection_name}"
3403 )
3405 if removed: 3405 ↛ 3410line 3405 didn't jump to line 3410 because the condition on line 3405 was always true
3406 logger.info(
3407 f"Removed {removed} vectors for "
3408 f"{len(wanted)} document(s) from {collection_name}"
3409 )
3410 return removed