Coverage for src/local_deep_research/research_library/services/library_rag_service.py: 85%
611 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +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
14from pathlib import Path
15from typing import Any, Dict, List, Optional, Tuple
17from langchain_core.documents import Document as LangchainDocument
18from loguru import logger
19from sqlalchemy import func
21from ...config.paths import get_cache_directory
22from ...database.models.library import (
23 Document,
24 DocumentChunk,
25 DocumentCollection,
26 Collection,
27 RAGIndex,
28 RagDocumentStatus,
29 EmbeddingProvider,
30)
31from ...database.session_context import get_user_db_session, safe_rollback
32from ...utilities.type_utils import to_bool
33from ..utils import ensure_in_collection
34from ...embeddings.splitters import get_text_splitter
35from ...web_search_engines.engines.local_embedding_manager import (
36 LocalEmbeddingManager,
37)
38from ...security.file_integrity import FileIntegrityManager, FAISSIndexVerifier
39from ...security.file_integrity.integrity_manager import NO_INTEGRITY_RECORD
40from ...vector_stores.facade import VectorIndex, ChunkInput, SearchResult
41from langchain_core.embeddings import Embeddings
42import hashlib
45class _LazyEmbeddings(Embeddings):
46 """Embeddings proxy that defers backend construction until first use.
48 A ``VectorIndex`` built purely to ``delete()`` never touches ``embeddings``
49 (delete resolves rows by id and calls ``store.apply(remove_ids=...)``).
50 ``LocalEmbeddingManager.embeddings`` is a lazy property that, on first
51 access, synchronously loads the real backend (a SentenceTransformer costs
52 seconds and hundreds of MB of RSS, an Ollama client opens live httpx
53 connections). Handing that property straight to the delete path forces that
54 cost once per (deleted-document x collection). This proxy pulls the real
55 embeddings from the manager only if ``embed_*`` is genuinely called, so the
56 delete path pays nothing while index/search paths still work unchanged.
57 """
59 def __init__(self, manager) -> None:
60 self._manager = manager
62 def embed_documents(self, texts):
63 return self._manager.embeddings.embed_documents(texts)
65 def embed_query(self, text):
66 return self._manager.embeddings.embed_query(text)
69# Module-level locks serialise the FAISS save+record critical section and the
70# load_or_create verify→quarantine→build sequence per (username, index_path).
71# MUST be module-level: each auto-index / scheduler / search worker constructs
72# its own LibraryRAGService, so an instance-scoped lock would coordinate
73# nothing. Pattern is adapted from web/queue/processor_v2._user_critical_locks
74# — instance-scoped→module-scoped, username-only-key→(username, path)-key.
75# See #4197 for the race this guards (concurrent save_local interleaves bytes,
76# producing checksum_mismatch → destructive unlink, lost data).
77_faiss_write_locks: Dict[Tuple[str, str], threading.Lock] = {}
78_faiss_write_locks_lock = threading.Lock()
80# Hard cap on suffix-increment retries when generating the .corrupt-<ns> path.
81# Normal case is one attempt — same-ns collisions only happen if the user
82# manually created such files. 32 is a safety bound that converts a deadlock
83# into a loud OSError.
84_QUARANTINE_SUFFIX_RETRY_CAP = 32
86# After a successful quarantine, keep at most this many older .corrupt-*
87# files per base path (per side: .faiss.corrupt-* and .pkl.corrupt-* are
88# counted independently). Prevents unbounded disk growth on systems that
89# experience recurring corruption while preserving recent diagnostic
90# artefacts. Keeping 5 means the user has the last 5 corruption events
91# to inspect / submit with a bug report; anything older is dropped.
92_QUARANTINE_KEEP_RECENT = 5
95def _get_faiss_write_lock(username: str, index_path: str) -> threading.Lock:
96 """Return the lock for ``(username, index_path)``, creating it on first
97 access. Key is normalised via ``Path.resolve()`` to match
98 ``FileIntegrityManager._normalize_path`` so writers/readers/quarantine
99 all agree on identity.
100 """
101 key = (username, str(Path(index_path).resolve()))
102 with _faiss_write_locks_lock:
103 lock = _faiss_write_locks.get(key)
104 if lock is None:
105 lock = threading.Lock()
106 _faiss_write_locks[key] = lock
107 return lock
110def pop_faiss_locks_for_user(username: str) -> None:
111 """Remove FAISS-write locks belonging to ``username`` that are NOT in use.
113 Called from the user-close paths (connection_cleanup) so the dict doesn't
114 grow one entry per (user × collection) across the process lifetime.
116 A lock that is currently HELD must NOT be evicted: the next writer would
117 then create a FRESH lock and run concurrently with the current holder on
118 the same index file (lost update / torn write). So each candidate is probed
119 with a non-blocking acquire — only locks we can immediately acquire (hence
120 not held) are removed; a held lock is left in place and cleaned up by a
121 later call once its write finishes. (No deadlock: a writer never holds a
122 faiss lock while waiting on ``_faiss_write_locks_lock``.)
123 """
124 with _faiss_write_locks_lock:
125 stale = [k for k in _faiss_write_locks if k[0] == username]
126 for k in stale:
127 lock = _faiss_write_locks.get(k)
128 if lock is not None and lock.acquire(blocking=False):
129 try:
130 _faiss_write_locks.pop(k, None)
131 finally:
132 lock.release()
135class LibraryRAGService:
136 """Service for managing RAG indexing of library documents."""
138 def __init__(
139 self,
140 username: str,
141 embedding_model: str = "all-MiniLM-L6-v2",
142 embedding_provider: str = "sentence_transformers",
143 chunk_size: int = 1000,
144 chunk_overlap: int = 200,
145 splitter_type: str = "recursive",
146 text_separators: Optional[list] = None,
147 distance_metric: str = "cosine",
148 normalize_vectors: bool = True,
149 index_type: str = "flat",
150 embedding_manager: Optional["LocalEmbeddingManager"] = None,
151 db_password: Optional[str] = None,
152 ):
153 """
154 Initialize library RAG service for a user.
156 Args:
157 username: Username for database access
158 embedding_model: Name of the embedding model to use
159 embedding_provider: Provider type ('sentence_transformers' or 'ollama')
160 chunk_size: Size of text chunks for splitting
161 chunk_overlap: Overlap between consecutive chunks
162 splitter_type: Type of splitter ('recursive', 'token', 'sentence', 'semantic')
163 text_separators: List of text separators for chunking (default: ["\n\n", "\n", ". ", " ", ""])
164 distance_metric: Distance metric ('cosine', 'l2', or 'dot_product')
165 normalize_vectors: Whether to normalize vectors with L2
166 index_type: FAISS index type ('flat', 'hnsw', or 'ivf')
167 embedding_manager: Optional pre-constructed LocalEmbeddingManager for testing/flexibility
168 db_password: Optional database password for background thread access
169 """
170 self.username = username
171 self._db_password = db_password # Can be used for thread access
172 # Initialize optional attributes to None before they're set below
173 # This allows the db_password setter to check them without hasattr
174 self.embedding_manager = None
175 self.integrity_manager = None
176 self.embedding_model = embedding_model
177 self.embedding_provider = embedding_provider
178 self.chunk_size = chunk_size
179 self.chunk_overlap = chunk_overlap
180 self.splitter_type = splitter_type
181 self.text_separators = (
182 text_separators
183 if text_separators is not None
184 else ["\n\n", "\n", ". ", " ", ""]
185 )
186 self.distance_metric = distance_metric
187 # Ensure normalize_vectors is always a proper boolean
188 self.normalize_vectors = to_bool(normalize_vectors, default=True)
189 self.index_type = index_type
191 # Emit the active configuration so users can confirm their
192 # UI-configured embedding settings are being honored (regression
193 # signal for #3453).
194 logger.info(
195 f"RAG service initialized for user={username}: "
196 f"provider={embedding_provider} model={embedding_model} "
197 f"chunk_size={chunk_size} chunk_overlap={chunk_overlap} "
198 f"splitter={splitter_type} index_type={index_type}"
199 )
201 # Use provided embedding manager or create a new one
202 # (Must be created before text splitter for semantic chunking)
203 # Track ownership so close() only tears down the manager when we
204 # constructed it — a caller-supplied manager stays under caller
205 # control (test fixtures, multi-service callers reusing one manager).
206 self._owns_embedding_manager = embedding_manager is None
207 if embedding_manager is not None:
208 self.embedding_manager = embedding_manager
209 else:
210 # Initialize embedding manager with library collection
211 # Load the complete user settings snapshot from database using the proper method
212 from ...settings.manager import SettingsManager
214 # Use proper database session for SettingsManager
215 # Note: using _db_password (backing field) directly here because the
216 # db_password property setter propagates to embedding_manager/integrity_manager,
217 # which are still None at this point in __init__.
218 with get_user_db_session(username, self._db_password) as session:
219 settings_manager = SettingsManager(session)
220 settings_snapshot = settings_manager.get_settings_snapshot()
222 # Add the specific settings needed for this RAG service
223 settings_snapshot.update(
224 {
225 "_username": username,
226 "embeddings.provider": embedding_provider,
227 f"embeddings.{embedding_provider}.model": embedding_model,
228 "local_search_chunk_size": chunk_size,
229 "local_search_chunk_overlap": chunk_overlap,
230 }
231 )
233 # Egress policy pre-flight at the constructor boundary so
234 # every direct ``LibraryRAGService(...)`` construction site
235 # is covered, not just the factory. Skipped when an
236 # ``embedding_manager`` is injected (tests / advanced flows)
237 # — those callers vouch for the manager themselves.
238 from ...security.egress.policy import (
239 Decision,
240 PolicyDeniedError,
241 context_from_snapshot,
242 evaluate_embeddings,
243 resolve_run_primary_engine,
244 )
246 # Build the context from the ACTUAL scope (not a hardcoded
247 # BOTH) so PRIVATE_ONLY forces local embeddings even when the
248 # raw embeddings.require_local flag is at its default False —
249 # context_from_snapshot applies that coupling. Resolve the primary
250 # via the shared helper (single source of truth) instead of the old
251 # search.tool + searxng fallback, which was a fail-OPEN: a missing
252 # primary defaulted to the public searxng so the scope relaxed and a
253 # cloud embedder could be admitted. A missing primary now raises ->
254 # fail closed via the ValueError handler below.
255 try:
256 primary = resolve_run_primary_engine(settings_snapshot)
257 policy_ctx = context_from_snapshot(
258 settings_snapshot, primary, username=username
259 )
260 except PolicyDeniedError:
261 raise
262 except ValueError as exc:
263 raise PolicyDeniedError(
264 Decision(False, "invalid_policy_config"),
265 target=embedding_provider,
266 ) from exc
267 if policy_ctx.require_local_embeddings: 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true
268 decision = evaluate_embeddings(
269 embedding_provider,
270 policy_ctx,
271 settings_snapshot=settings_snapshot,
272 )
273 if not decision.allowed:
274 logger.bind(policy_audit=True).warning(
275 "LibraryRAGService refused by egress policy",
276 provider=embedding_provider,
277 reason=decision.reason,
278 )
279 raise PolicyDeniedError(decision, target=embedding_provider)
281 self.embedding_manager = LocalEmbeddingManager(
282 embedding_model=embedding_model,
283 embedding_model_type=embedding_provider,
284 settings_snapshot=settings_snapshot,
285 )
287 # Initialize text splitter based on type
288 # (Must be created AFTER embedding_manager for semantic chunking)
289 self.text_splitter = get_text_splitter(
290 splitter_type=self.splitter_type,
291 chunk_size=self.chunk_size,
292 chunk_overlap=self.chunk_overlap,
293 text_separators=self.text_separators,
294 embeddings=self.embedding_manager.embeddings
295 if self.splitter_type == "semantic"
296 else None,
297 )
299 self.rag_index_record = None
301 # Initialize file integrity manager for FAISS indexes
302 self.integrity_manager = FileIntegrityManager(
303 username, password=self._db_password
304 )
305 self.integrity_manager.register_verifier(FAISSIndexVerifier())
307 self._closed = False
309 def close(self):
310 """Release embedding model and index resources."""
311 if self._closed:
312 return
313 self._closed = True
315 # Release embedding manager (which in turn closes the underlying
316 # OllamaEmbeddings httpx clients — see LocalEmbeddingManager.close).
317 # Only when we own it; caller-supplied managers stay under caller
318 # control to avoid double-close / use-after-close.
319 if self.embedding_manager is not None:
320 if self._owns_embedding_manager:
321 self.embedding_manager.close()
322 self.embedding_manager = None
324 # Clear other resources
325 self.rag_index_record = None
326 self.integrity_manager = None
327 self.text_splitter = None
329 def __enter__(self):
330 """Enter context manager."""
331 return self
333 def __exit__(self, exc_type, exc_val, exc_tb):
334 """Exit context manager, ensuring cleanup."""
335 self.close()
336 return False
338 @property
339 def db_password(self):
340 """Get database password."""
341 return self._db_password
343 @db_password.setter
344 def db_password(self, value):
345 """Set database password and propagate to embedding manager and integrity manager."""
346 self._db_password = value
347 if self.embedding_manager:
348 self.embedding_manager.db_password = value
349 if self.integrity_manager:
350 self.integrity_manager.password = value
352 def _get_index_hash(
353 self,
354 collection_name: str,
355 embedding_model: str,
356 embedding_model_type: str,
357 ) -> str:
358 """Generate hash for index identification."""
359 hash_input = (
360 f"{collection_name}:{embedding_model}:{embedding_model_type}"
361 )
362 return hashlib.sha256(hash_input.encode()).hexdigest()
364 def _get_index_path(self, index_hash: str) -> Path:
365 """Get path for FAISS index file.
367 Files are scoped to a per-user subdirectory derived from a hash
368 of the username so two users with identical collection name +
369 embedding model don't share the same .faiss/.pkl files on disk.
370 The previous layout used a shared `cache/rag_indices/` for all
371 users, which would let one user's vectors land in another
372 user's load path in any deployment running more than one
373 account against the same data dir.
374 """
375 # Store in centralized cache directory (respects LDR_DATA_DIR)
376 user_scope = hashlib.sha256(self.username.encode("utf-8")).hexdigest()[
377 :16
378 ]
379 rag_root = get_cache_directory() / "rag_indices"
380 cache_dir = rag_root / user_scope
381 cache_dir.mkdir(parents=True, exist_ok=True)
382 # 0o700: filesystem-level defense-in-depth so other local OS accounts
383 # can't browse another LDR user's vector caches. mkdir(parents=True)
384 # silently creates the intermediate rag_indices/ root under the process
385 # umask (typically 0o755, world-traversable), so harden the ROOT too —
386 # chmod'ing only the per-user leaf leaves a 0o755 rag_indices/ that lets
387 # another OS user traverse in and stat/enumerate every account's caches.
388 for d in (rag_root, cache_dir):
389 try:
390 d.chmod(0o700)
391 except OSError:
392 # Best-effort: some filesystems (e.g. Windows shares) don't
393 # honor chmod; the per-user path scope is still applied.
394 pass
395 return cache_dir / f"{index_hash}.faiss"
397 def _migrate_legacy_index_files(
398 self, legacy_path: Path, index_path: Path, rag_index: RAGIndex
399 ) -> None:
400 """Relocate a pre-per-user-scoping index into the per-user directory.
402 Before the per-user path scoping, indexes lived directly in the
403 shared ``cache/rag_indices/`` directory. ``_preflight_index_path``
404 recomputes the path and ignores the stored one, so without this
405 one-time migration an upgrading user's existing index is never
406 found again: semantic search silently returns nothing while
407 ``DocumentCollection.indexed`` still claims everything is indexed,
408 so a non-force re-index no-ops too. Only files directly inside the
409 known legacy shared directory with the expected hash filename are
410 moved — arbitrary stored paths are not followed.
411 """
412 legacy_root = get_cache_directory() / "rag_indices"
413 if (
414 legacy_path == index_path
415 or legacy_path.parent != legacy_root
416 or legacy_path.name != index_path.name
417 or not legacy_path.exists()
418 ):
419 return
420 lock = _get_faiss_write_lock(self.username, str(index_path))
421 with lock:
422 if index_path.exists(): 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true
423 return
424 try:
425 legacy_pkl = legacy_path.with_suffix(".pkl")
426 # Relocate phase-1's text-free .idmap.json sidecar BEFORE the
427 # .faiss, then move the .faiss LAST. phase-2 rekey looks for the
428 # sidecar next to the .faiss (legacy_rekey._sidecar_for =
429 # <stem>.idmap.json); leaving it at the legacy location orphans
430 # it, and a not-yet-rekeyed index then has no findable sidecar →
431 # rekey sees "no sidecar + old format" and destructively
432 # quarantines a healthy index. Ordering matters for crash
433 # recovery: the re-entry guard keys ONLY on index_path (the
434 # .faiss), so moving the .faiss last makes this idempotent — a
435 # crash after the sidecar move but before the .faiss move leaves
436 # the .faiss at the legacy path (guard still False) and re-entry
437 # simply retries (the sidecar move no-ops, new one exists). The
438 # reverse order would strand index_path.exists()=True with the
439 # sidecar orphaned at the legacy root, which the guard never
440 # re-invokes this to fix.
441 legacy_sidecar = legacy_path.with_name(
442 legacy_path.stem + ".idmap.json"
443 )
444 if legacy_sidecar.exists():
445 new_sidecar = index_path.with_name(
446 index_path.stem + ".idmap.json"
447 )
448 if not new_sidecar.exists(): 448 ↛ 450line 448 didn't jump to line 450 because the condition on line 448 was always true
449 legacy_sidecar.rename(new_sidecar)
450 legacy_path.rename(index_path)
451 # DELETE (don't relocate) any legacy plaintext .pkl companion:
452 # the new store writes NO .pkl, so moving it would just carry the
453 # plaintext chunk text to the new location. Its text lives in the
454 # encrypted DB.
455 if legacy_pkl.exists():
456 try:
457 legacy_pkl.unlink()
458 except OSError:
459 logger.warning(
460 f"Could not remove legacy plaintext .pkl "
461 f"{legacy_pkl} — delete it manually"
462 )
463 except OSError:
464 logger.exception(
465 f"Failed to migrate legacy FAISS index "
466 f"{legacy_path} -> {index_path}"
467 )
468 return
469 # Persist the relocated path so the row stops pointing at the
470 # now-empty legacy location.
471 with get_user_db_session(self.username, self.db_password) as session:
472 idx = session.query(RAGIndex).filter_by(id=rag_index.id).first()
473 if idx: 473 ↛ 476line 473 didn't jump to line 476
474 idx.index_path = str(index_path)
475 session.commit()
476 try:
477 # Refresh the integrity record under the new path (the old
478 # record is keyed by the legacy path and no longer applies).
479 self.integrity_manager.record_file(
480 index_path,
481 related_entity_type="rag_index",
482 related_entity_id=rag_index.id,
483 )
484 except Exception:
485 # Non-fatal: verify_file() creates a record on demand for
486 # unknown paths, so the subsequent load still succeeds.
487 logger.exception(
488 "Failed to refresh integrity record after index migration"
489 )
490 logger.info(
491 f"Migrated legacy shared FAISS index to per-user location: "
492 f"{index_path}"
493 )
495 def _reset_index_state_for_rebuild(
496 self, collection_id: str, rag_index: RAGIndex
497 ) -> None:
498 """Reset per-document indexed state when the index file is gone.
500 Reaching the fresh-build branch with DB rows that still claim
501 indexed content means the on-disk index was lost (missing after an
502 upgrade, quarantined, or deleted on dimension mismatch). Without
503 this reset, ``index_document``/``index_collection`` skip every
504 document (``DocumentCollection.indexed`` is still True), so the
505 only repair is an undocumented force re-index while search
506 silently returns nothing. Clearing the flags makes the regular
507 re-index path rebuild for real and makes the UI show the honest
508 unindexed state. Worst case of the narrow race with an in-flight
509 indexer is re-submitting chunks the merge step dedups by id.
510 """
511 with get_user_db_session(self.username, self.db_password) as session:
512 idx = session.query(RAGIndex).filter_by(id=rag_index.id).first()
513 claims_content = bool(
514 idx and (idx.chunk_count or idx.total_documents)
515 )
516 stale_status = (
517 session.query(RagDocumentStatus)
518 .filter_by(rag_index_id=rag_index.id)
519 .count()
520 )
521 stale_memberships = (
522 session.query(DocumentCollection)
523 .filter_by(collection_id=collection_id, indexed=True)
524 .count()
525 )
526 if not (claims_content or stale_status or stale_memberships):
527 # Genuinely new index — nothing to reset.
528 return
529 if idx: 529 ↛ 532line 529 didn't jump to line 532 because the condition on line 529 was always true
530 idx.chunk_count = 0
531 idx.total_documents = 0
532 session.query(RagDocumentStatus).filter_by(
533 rag_index_id=rag_index.id
534 ).delete()
535 session.query(DocumentCollection).filter_by(
536 collection_id=collection_id, indexed=True
537 ).update({"indexed": False, "chunk_count": 0})
538 session.commit()
539 logger.warning(
540 f"FAISS index file for collection {collection_id} is "
541 f"missing; reset indexed state for {stale_memberships} "
542 f"document(s) so the next index run rebuilds the index."
543 )
545 def _get_or_create_rag_index(
546 self, collection_id: str, *, promote_current: bool = True
547 ) -> RAGIndex:
548 """Get or create RAGIndex record for the current configuration.
550 ``promote_current`` (default True) controls whether reusing an existing,
551 not-current row re-promotes it to ``is_current`` (demoting others). Read
552 paths (search) pass False: a mere search under a stale/different config
553 must NOT silently flip which index the collection considers current —
554 only an actual write (index) should. Creating a brand-new index still
555 promotes it regardless (it's the collection's only index for that config).
556 """
557 with get_user_db_session(self.username, self.db_password) as session:
558 # Use collection_<uuid> format
559 collection_name = f"collection_{collection_id}"
560 index_hash = self._get_index_hash(
561 collection_name, self.embedding_model, self.embedding_provider
562 )
564 # Try to get existing index
565 rag_index = (
566 session.query(RAGIndex).filter_by(index_hash=index_hash).first()
567 )
569 if not rag_index:
570 # A new index (e.g. an embedding-model change yields a new
571 # index_hash) supersedes any prior current index for this
572 # collection. Demote the old current row(s) FIRST so exactly one
573 # row is ever is_current=True — every read path resolves "the"
574 # current index via filter_by(is_current=True).first(), which
575 # would otherwise return an arbitrary one of several.
576 session.query(RAGIndex).filter_by(
577 collection_name=collection_name, is_current=True
578 ).update({"is_current": False})
580 # Create new index record
581 index_path = self._get_index_path(index_hash)
583 # Get embedding dimension by embedding a test string
584 test_embedding = self.embedding_manager.embeddings.embed_query(
585 "test"
586 )
587 embedding_dim = len(test_embedding)
589 rag_index = RAGIndex(
590 collection_name=collection_name,
591 embedding_model=self.embedding_model,
592 embedding_model_type=EmbeddingProvider(
593 self.embedding_provider
594 ),
595 embedding_dimension=embedding_dim,
596 index_path=str(index_path),
597 index_hash=index_hash,
598 chunk_size=self.chunk_size,
599 chunk_overlap=self.chunk_overlap,
600 splitter_type=self.splitter_type,
601 text_separators=self.text_separators,
602 distance_metric=self.distance_metric,
603 normalize_vectors=self.normalize_vectors,
604 index_type=self.index_type,
605 chunk_count=0,
606 total_documents=0,
607 status="active",
608 is_current=True,
609 )
610 session.add(rag_index)
611 session.commit()
612 session.refresh(rag_index)
613 logger.info(f"Created new RAG index: {index_hash}")
614 elif promote_current and not rag_index.is_current:
615 # An existing index for THIS (collection, model) config is being
616 # reused (on a WRITE path) but is not marked current. After an
617 # embedding-model switch-and-revert (A -> B -> back to A), row A
618 # was demoted when B was created and never re-promoted, so "the
619 # current index" bookkeeping (filter_by(is_current=True)) would
620 # keep pointing at the now-inactive B. Demote any other current
621 # row for this collection and promote this one, mirroring the
622 # create branch. Gated on promote_current so a read-path search
623 # under a stale config can't steal the current pointer.
624 session.query(RAGIndex).filter(
625 RAGIndex.collection_name == collection_name,
626 RAGIndex.is_current.is_(True),
627 RAGIndex.id != rag_index.id,
628 ).update({"is_current": False})
629 rag_index.is_current = True
630 session.commit()
631 session.refresh(rag_index)
633 return rag_index
635 def _quarantine_corrupt_index(self, index_path: Path, reason: str) -> None:
636 """Rename a corrupted FAISS index to ``<path>.corrupt-<ns>``
637 instead of deleting it; any legacy plaintext ``.pkl`` companion
638 is deleted outright rather than preserved.
640 Preserves user data for inspection/recovery. The
641 dimension-mismatch branch elsewhere in this method intentionally
642 *deletes* — that case rebuilds from scratch and the old bytes
643 are unreadable with the new model. This helper is for the two
644 "transient or unknown failure" branches (verify_file said no,
645 loading the index into a ``VectorIndex``/faiss raised) where the
646 on-disk bytes may still be usable by a human.
648 Raises ``OSError`` on rename failure (disk full, read-only fs,
649 permission denied). Re-raising prevents silent data loss: if we
650 swallowed the error, the next ``persist()``/``faiss.write_index``
651 would overwrite the corrupt bytes anyway. Caller paths log the
652 exception via the broader try/except around their indexer call.
654 # TODO(#4197-followup): FileIntegrityRecord.consecutive_failures
655 # is not reset by the next record_file call, leaking failure
656 # counts across recovery cycles. Orthogonal to this fix.
657 """
658 ns = time.time_ns()
659 pkl_path = index_path.with_suffix(".pkl")
661 # Same-nanosecond collisions essentially can't happen between
662 # concurrent threads (we hold the per-path lock), but a user
663 # could manually create such files. Loop with a hard cap so a
664 # weird state surfaces as a clean OSError rather than a hang.
665 suffix_n = 0
666 faiss_target = Path(f"{index_path}.corrupt-{ns}")
667 pkl_target = Path(f"{pkl_path}.corrupt-{ns}")
668 while faiss_target.exists() or pkl_target.exists():
669 suffix_n += 1
670 if suffix_n > _QUARANTINE_SUFFIX_RETRY_CAP:
671 raise OSError(
672 f"Quarantine path collisions exceeded "
673 f"{_QUARANTINE_SUFFIX_RETRY_CAP} retries for "
674 f"{index_path}"
675 )
676 faiss_target = Path(f"{index_path}.corrupt-{ns}-{suffix_n}")
677 pkl_target = Path(f"{pkl_path}.corrupt-{ns}-{suffix_n}")
679 try:
680 index_path.rename(faiss_target)
681 logger.warning(
682 f"Quarantined corrupted FAISS index to {faiss_target} "
683 f"(reason: {reason}). Searches against this collection "
684 f"will return empty results until the index is rebuilt. "
685 f"Document chunks are preserved in the database — "
686 f"trigger 'Re-index Collection' (or set "
687 f"research_library.auto_index_enabled=true and re-run "
688 f"indexing) to recover."
689 )
690 if pkl_path.exists():
691 # DELETE (not rename) any .pkl companion. The new store writes
692 # NO .pkl, so a .pkl next to a .faiss is a LEGACY plaintext
693 # docstore — its chunk text lives in the (encrypted) DB, so
694 # renaming it aside to .corrupt-<ns> would needlessly preserve
695 # plaintext chunk text on disk. Plaintext removal is not optional.
696 try:
697 pkl_path.unlink()
698 logger.info(f"Removed legacy plaintext PKL {pkl_path}")
699 except OSError:
700 logger.warning(
701 f"Could not remove legacy plaintext PKL {pkl_path} — "
702 "delete it manually"
703 )
704 else:
705 # Missing .pkl is the normal case for the new store. Not
706 # re-raised: the .faiss is already preserved.
707 logger.debug(f"No PKL companion at {pkl_path}.")
708 except OSError:
709 if not index_path.exists():
710 # The source is already gone — most likely another worker
711 # PROCESS quarantined the same corrupt index concurrently (the
712 # per-path lock is a threading.Lock, so it serialises only
713 # threads within ONE process, not across gunicorn workers /
714 # the scheduler process). The objective — get the corrupt bytes
715 # off index_path so the caller rebuilds fresh — is already met,
716 # so don't abort a healthy rebuild with a spurious
717 # FileNotFoundError. Still purge any lingering plaintext .pkl:
718 # the no-plaintext-at-rest invariant is not optional even here.
719 if pkl_path.exists(): 719 ↛ 727line 719 didn't jump to line 727 because the condition on line 719 was always true
720 try:
721 pkl_path.unlink()
722 except OSError:
723 logger.warning(
724 f"Could not remove legacy plaintext PKL {pkl_path} "
725 "— delete it manually"
726 )
727 logger.warning(
728 f"Corrupt index {index_path} was already quarantined by "
729 f"another worker (reason: {reason}); proceeding to rebuild."
730 )
731 return
732 # A genuine failure (disk-full, read-only fs, permission denied)
733 # with the corrupt bytes STILL in place. Re-raise so the caller
734 # surfaces a real failure instead of letting the next
735 # persist()/write_index overwrite the corrupt bytes.
736 logger.exception(
737 f"Failed to quarantine corrupted index at {index_path}"
738 )
739 raise
741 # Best-effort retention sweep. The quarantine itself succeeded
742 # above; failing to prune older files is not a correctness
743 # issue, just a disk-usage one — log and move on.
744 self._prune_old_quarantined_files(index_path)
746 @staticmethod
747 def _corrupt_sort_key(path: Path) -> Tuple[int, int]:
748 """Extract ``(ns, suffix_n)`` from a ``.corrupt-<ns>[-<n>]``
749 filename so retention can sort by the monotonic nanosecond
750 suffix the quarantine path embeds — *not* by ``st_mtime``.
752 Filesystem timestamp granularity is sometimes 1s or 2s
753 (FAT32/ext3/SMB shares), making mtime ordering non-deterministic
754 when multiple quarantines happen within one tick. The
755 ``.corrupt-<ns>`` suffix carries the original ``time.time_ns()``
756 from the quarantine and is reliable across all filesystems.
758 Returns ``(-1, -1)`` for malformed names (manually-placed
759 files) so they sort below any real entry under ``reverse=True``
760 — i.e., they get pruned first.
761 """
762 name = path.name
763 marker = ".corrupt-"
764 idx = name.rfind(marker)
765 if idx == -1: 765 ↛ 766line 765 didn't jump to line 766 because the condition on line 765 was never true
766 return (-1, -1)
767 tail = name[idx + len(marker) :]
768 parts = tail.split("-")
769 try:
770 ns = int(parts[0])
771 except ValueError:
772 return (-1, -1)
773 suffix_n = 0
774 if len(parts) > 1:
775 try:
776 suffix_n = int(parts[-1])
777 except ValueError:
778 # Unknown trailing component — keep the file but at the
779 # base ns ordering.
780 pass
781 return (ns, suffix_n)
783 def _prune_old_quarantined_files(self, index_path: Path) -> None:
784 """Keep only the ``_QUARANTINE_KEEP_RECENT`` most-recent
785 ``.corrupt-*`` files for ``index_path`` and its ``.pkl``
786 companion. Sweeps the two sides independently — pairs share
787 the same ``-<ns>`` suffix so they're ordered identically.
789 Ordering uses the embedded ``-<ns>`` from the quarantine
790 filename, not ``st_mtime``: file systems with 1-2s timestamp
791 granularity can otherwise produce non-deterministic retention
792 on bursts.
794 Best-effort: logs and swallows any error so a sweep failure
795 never propagates back into the indexing path.
796 """
797 parent = index_path.parent
798 pkl_path = index_path.with_suffix(".pkl")
800 for base in (index_path, pkl_path):
801 pattern = f"{base.name}.corrupt-*"
802 try:
803 # Sort newest-first by the embedded -<ns>; everything
804 # past the keep window is stale.
805 candidates = sorted(
806 parent.glob(pattern),
807 key=self._corrupt_sort_key,
808 reverse=True,
809 )
810 except OSError:
811 logger.warning(
812 f"Failed to enumerate {pattern} in {parent} for "
813 f"quarantine retention sweep"
814 )
815 continue
817 for stale in candidates[_QUARANTINE_KEEP_RECENT:]:
818 try:
819 stale.unlink()
820 logger.info(f"Pruned old quarantined file: {stale}")
821 except OSError:
822 logger.warning(
823 f"Failed to prune quarantined file {stale}; continuing"
824 )
826 def _preflight_index_path(
827 self,
828 collection_id: str,
829 rag_index: RAGIndex,
830 *,
831 reset_stale_state: bool = False,
832 ) -> Path:
833 """Verify -> quarantine -> dimension-check pre-flight for ``rag_index``.
835 Extracted from the pre-cutover ``load_or_create_faiss_index``
836 (deleted — the old LangChain-FAISS path is gone) so every
837 index/search/delete call site shares ONE recovery path instead of
838 duplicating it. Mutates ``rag_index.index_path`` / ``rag_index.embedding_dimension``
839 in place — matching the pre-refactor behavior of stamping the
840 recomputed path/dimension back onto the in-memory record so callers
841 that read ``self.rag_index_record`` afterwards see the current
842 values — and returns the resolved on-disk path.
844 Deliberately does NOT load or construct a vector store: that stays
845 the caller's job (via ``VectorIndex(...)``), so a corrupt/foreign/
846 dimension-mismatched file surfaces as an exception from THAT
847 construction, which the caller catches to quarantine + rebuild (see
848 ``_get_vector_index``).
849 """
850 index_path = self._get_index_path(rag_index.index_hash)
851 stored_path = (
852 Path(rag_index.index_path) if rag_index.index_path else None
853 )
854 rag_index.index_path = str(index_path)
855 if not index_path.exists() and stored_path is not None:
856 self._migrate_legacy_index_files(stored_path, index_path, rag_index)
858 # Hold the per-(username, index_path) write lock across the entire
859 # verify -> quarantine -> dimension-check sequence — see #4197 (a
860 # narrower scope would let a concurrent writer's save race the
861 # verification).
862 if index_path.exists():
863 lock = _get_faiss_write_lock(self.username, str(index_path))
864 with lock:
865 verified, reason = self.integrity_manager.verify_file(
866 index_path
867 )
868 if not verified:
869 if reason == NO_INTEGRITY_RECORD: 869 ↛ 888line 869 didn't jump to line 888 because the condition on line 869 was never true
870 # A MISSING record is not proven corruption, and unlike
871 # checksum_mismatch it cannot be attacker-induced: the
872 # record lives in the per-user ENCRYPTED DB, so removing
873 # it needs the DB password (which also decrypts
874 # everything). It only arises from our own crash between
875 # the index write and record_file (see legacy_rekey
876 # Phase B) or a pre-integrity-feature index. ADOPT the
877 # byte-present file by recording its integrity now — on
878 # BOTH read and write paths. Adoption is non-destructive
879 # (it only records a checksum), so a read path may do it:
880 # refusing here instead would hard-fail EVERY search on a
881 # byte-healthy, missing-record index forever (a read-only
882 # collection is never revisited by a write op that would
883 # otherwise adopt it). If the bytes are actually
884 # unreadable/foreign/wrong-dim, the VectorIndex load below
885 # still raises and the caller handles it — corruption is
886 # caught either way. Only checksum_mismatch (genuine
887 # tamper) is refused/quarantined below.
888 logger.warning(
889 f"No integrity record for {index_path}; adopting the "
890 "existing index (recording its integrity) rather than "
891 "quarantining — a missing record is not corruption."
892 )
893 self.integrity_manager.record_file(
894 index_path,
895 related_entity_type="rag_index",
896 related_entity_id=rag_index.id,
897 )
898 # Adopted = now trusted: fall through to the SAME
899 # embedding-dimension-drift probe the verified branch
900 # runs, so an adopted index whose model dimension changed
901 # is rebuilt rather than silently mismatched at add time.
902 verified = True
903 elif not reset_stale_state:
904 # Read path (search): refuse to serve an index that
905 # failed verification for a reason OTHER than a missing
906 # record (a checksum_mismatch = genuine tamper, or a
907 # transient read failure), but do NOT destroy it —
908 # quarantine is a WRITE-path recovery decision. A
909 # subsequent index/write operation quarantines + rebuilds.
910 raise RuntimeError(
911 f"Integrity verification failed for {index_path}: "
912 f"{reason} — refusing on a read path (a subsequent "
913 "index/write operation will quarantine + rebuild)"
914 )
915 elif reason and reason.startswith( 915 ↛ 924line 915 didn't jump to line 924 because the condition on line 915 was never true
916 "checksum_calculation_failed"
917 ):
918 # Transient: could not READ the file to compute its
919 # checksum (disk hiccup, momentary lock). This says
920 # NOTHING about tampering — unlike checksum_mismatch — so
921 # quarantining would force a needless full re-embed of a
922 # healthy collection. Raise so the operation retries (like
923 # the read path) instead of destroying the index.
924 raise RuntimeError(
925 f"Integrity check could not read {index_path} "
926 f"({reason}) — refusing to quarantine on a transient "
927 "read failure; retry the operation."
928 )
929 else:
930 logger.error(
931 f"Integrity verification failed for {index_path}: "
932 f"{reason}. Quarantining for recovery; creating "
933 f"new index."
934 )
935 self._quarantine_corrupt_index(index_path, reason)
936 if verified:
937 # Probe the embedding model OUTSIDE the quarantine
938 # decision: embed_query fails when the provider itself
939 # is unreachable (e.g. Ollama down), which says nothing
940 # about the index file's health, and must propagate
941 # rather than trigger a destructive rebuild of a
942 # healthy index.
943 current_dim = len(
944 self.embedding_manager.embeddings.embed_query(
945 "dimension_check"
946 )
947 )
948 stored_dim = rag_index.embedding_dimension
949 if stored_dim and current_dim != stored_dim:
950 if not reset_stale_state:
951 # Read path: surface the mismatch instead of
952 # deleting the index + wiping DB state (which a mere
953 # search must never do). A write/index operation
954 # owns the destructive rebuild.
955 raise RuntimeError(
956 f"Embedding dimension mismatch for {index_path}: "
957 f"index dim={stored_dim}, model dim="
958 f"{current_dim} — refusing on a read path"
959 )
960 logger.warning(
961 f"Embedding dimension mismatch detected! "
962 f"Index created with dim={stored_dim}, "
963 f"current model returns dim={current_dim}. "
964 f"Deleting old index and rebuilding."
965 )
966 try:
967 index_path.unlink()
968 pkl_path = index_path.with_suffix(".pkl")
969 if pkl_path.exists(): 969 ↛ 970line 969 didn't jump to line 970 because the condition on line 969 was never true
970 pkl_path.unlink()
971 logger.info(
972 f"Deleted old index files at {index_path}"
973 )
974 except Exception:
975 logger.exception("Failed to delete old index files")
977 with get_user_db_session(
978 self.username, self.db_password
979 ) as session:
980 idx = (
981 session.query(RAGIndex)
982 .filter_by(id=rag_index.id)
983 .first()
984 )
985 if idx: 985 ↛ 997line 985 didn't jump to line 997
986 idx.embedding_dimension = current_dim
987 session.commit()
988 logger.info(
989 f"Updated RAGIndex dimension to {current_dim}"
990 )
991 # Reset RAGIndex counts, RagDocumentStatus, and
992 # DocumentCollection.indexed together — a hand-rolled
993 # partial reset here previously left
994 # DocumentCollection.indexed=True, which makes
995 # index_document/index_collection skip every document
996 # and search silently return nothing after the rebuild.
997 self._reset_index_state_for_rebuild(
998 collection_id, rag_index
999 )
1000 rag_index.embedding_dimension = current_dim
1002 if not index_path.exists():
1003 # The new per-user path is missing — but that alone does NOT mean the
1004 # index is gone. A one-time legacy-path relocation
1005 # (_migrate_legacy_index_files above) can fail transiently (a rename
1006 # hitting disk-full, an AV/backup file lock, a busy network mount),
1007 # leaving the bytes fully intact at the stored legacy path. Treating
1008 # that as "gone" would wipe a healthy collection's indexed bookkeeping
1009 # — even on a mere search — and force a needless full re-embed. Only
1010 # self-heal when the bytes are absent at BOTH the new and the legacy
1011 # location; otherwise leave the state alone so the next write-path
1012 # access retries the relocation.
1013 legacy_bytes_present = bool(
1014 stored_path is not None
1015 and stored_path != index_path
1016 and stored_path.exists()
1017 )
1018 if not legacy_bytes_present:
1019 stale_indexed = bool(
1020 (rag_index.chunk_count or 0) > 0
1021 or (rag_index.total_documents or 0) > 0
1022 )
1023 if stale_indexed:
1024 # DB says indexed, but the file is GONE at both locations
1025 # (cache dir cleared, restored without the cache dir, lost
1026 # volume, crash between quarantine and rebuild). Reset so it
1027 # self-heals — even on a READ path: an absent file has
1028 # nothing to preserve, and without this search silently
1029 # returns [] forever (get_rag_stats still reports 100%
1030 # indexed and a normal reindex short-circuits on
1031 # indexed=True, so nothing ever recovers it).
1032 logger.warning(
1033 f"RAG index file for collection {collection_id} is "
1034 f"MISSING at {index_path} while the collection is marked "
1035 "indexed — resetting indexed state so it rebuilds "
1036 "(search would otherwise return no results indefinitely)."
1037 )
1038 self._reset_index_state_for_rebuild(
1039 collection_id, rag_index
1040 )
1041 elif reset_stale_state: 1041 ↛ 1046line 1041 didn't jump to line 1046 because the condition on line 1041 was always true
1042 self._reset_index_state_for_rebuild(
1043 collection_id, rag_index
1044 )
1046 return index_path
1048 def _get_vector_index(
1049 self,
1050 collection_id: str,
1051 collection_name: str,
1052 *,
1053 reset_stale_state: bool = False,
1054 ) -> VectorIndex:
1055 """Build a :class:`VectorIndex` for ``collection_id``.
1057 Single call site for the index/search/delete paths: runs the
1058 verify/quarantine/dimension-check pre-flight, then constructs the
1059 facade. On a ``ValueError``/``RuntimeError`` from the underlying
1060 store's ``load()`` (corrupt bytes, a foreign/pre-migration format
1061 that isn't an ``IndexIDMap2``, or a dimension mismatch the
1062 pre-flight didn't already catch) the corrupt file is quarantined,
1063 stale per-document indexed state is reset, and construction is
1064 retried once as a fresh, empty store.
1066 NOTE: like ``_get_or_create_rag_index``, this CREATES a RAGIndex
1067 row if none exists for the configured (collection, embedding model)
1068 — callers on a delete-only path that must not create an index
1069 (e.g. ``remove_documents_from_index``) guard with their own
1070 existence check before calling this.
1071 """
1072 # Only WRITE paths (reset_stale_state=True) may re-promote a reused
1073 # index to is_current; a read-path search must not flip the pointer.
1074 rag_index = self._get_or_create_rag_index(
1075 collection_id, promote_current=reset_stale_state
1076 )
1077 self.rag_index_record = rag_index
1078 index_path = self._preflight_index_path(
1079 collection_id, rag_index, reset_stale_state=reset_stale_state
1080 )
1081 lock = _get_faiss_write_lock(self.username, str(index_path))
1083 # nested=True: these run reentrantly inside VectorIndex.apply(), which
1084 # holds just-flushed, uncommitted DocumentChunk rows on this same
1085 # thread-local session. The integrity write must run in a SAVEPOINT and
1086 # let apply()'s caller commit — a plain commit/rollback here would
1087 # settle or DISCARD those pending rows (silently losing indexed text
1088 # while the vectors persist). See FileIntegrityManager._integrity_write.
1089 def _record(path: Path) -> None:
1090 self.integrity_manager.record_file(
1091 path,
1092 related_entity_type="rag_index",
1093 related_entity_id=rag_index.id,
1094 nested=True,
1095 )
1097 def _verify(path: Path):
1098 return self.integrity_manager.verify_file(path, nested=True)
1100 kwargs = {
1101 "username": self.username,
1102 "db_password": self.db_password,
1103 # Lazy: index()/search() resolve the real backend on first embed_*
1104 # call (identical behaviour, one extra attribute hop), while the
1105 # delete-only path (remove_documents_from_index) never materialises
1106 # it at all. See _LazyEmbeddings.
1107 "embeddings": _LazyEmbeddings(self.embedding_manager),
1108 "embedding_model": self.embedding_model,
1109 "embedding_model_type": EmbeddingProvider(self.embedding_provider),
1110 "collection_name": collection_name,
1111 "dimension": rag_index.embedding_dimension,
1112 "path": index_path,
1113 "lock": lock,
1114 "integrity_record": _record,
1115 "integrity_verify": _verify,
1116 # Use the STORED index config, not this service instance's. The
1117 # RAGIndex row is stamped with index_type/metric/normalize at
1118 # creation (see _get_or_create_rag_index); the constructing
1119 # LibraryRAGService may carry different CURRENT defaults (the
1120 # factory falls back to global settings for a collection whose
1121 # metadata was never stored). Building against self.* would load a
1122 # real HNSW file while labelling it "flat" (or vice-versa), and
1123 # FaissVectorStore.supports_delete trusts the label — routing a
1124 # removal through remove_ids on an HNSW index (unsupported → raises)
1125 # or silently rewriting a flat file as HNSW. This mirrors the fix
1126 # already applied to purge_document_vectors below.
1127 "index_type": rag_index.index_type or self.index_type,
1128 "metric": rag_index.distance_metric or self.distance_metric,
1129 "normalize": (
1130 rag_index.normalize_vectors
1131 if rag_index.normalize_vectors is not None
1132 else self.normalize_vectors
1133 ),
1134 }
1135 try:
1136 return VectorIndex(**kwargs)
1137 except (ValueError, RuntimeError):
1138 if not reset_stale_state:
1139 # Read paths (search) must stay side-effect-free. A transient
1140 # load failure (OOM, a disk hiccup, a momentarily-locked file)
1141 # on an existing, byte-healthy index must NOT quarantine the
1142 # file or wipe the collection's indexed state — that would turn
1143 # a retryable blip into a forced full re-embed of the whole
1144 # collection triggered by a mere search. Surface the error;
1145 # quarantine + rebuild is a WRITE-path (reset_stale_state=True)
1146 # decision. (A never-indexed collection never reaches here — no
1147 # file exists, so VectorIndex.create() succeeds above.)
1148 raise
1149 # A load failure here — after _preflight_index_path's verify_file
1150 # just PASSED on the same bytes — is ambiguous: it is either a
1151 # genuine format/dimension mismatch (deterministic — needs a rebuild)
1152 # OR a transient OS hiccup (EMFILE under concurrent multi-user load,
1153 # a momentary AV-scan/backup file lock; the byte-healthy file loads
1154 # fine moments later). faiss surfaces BOTH as the same RuntimeError,
1155 # so quarantining unconditionally destroys a provably-healthy index
1156 # and forces a needless full re-embed on a mere transient blip.
1157 # Distinguish by RETRYING the load once: a transient failure clears,
1158 # a deterministic one repeats. Only quarantine when the retry also
1159 # fails. (Mirrors _preflight_index_path's transient-vs-tamper split.)
1160 logger.warning(
1161 f"Failed to load vector store at {index_path}; retrying once "
1162 "before quarantining (may be a transient OS hiccup)"
1163 )
1164 try:
1165 return VectorIndex(**kwargs)
1166 except (ValueError, RuntimeError):
1167 pass
1168 logger.warning(
1169 f"Vector store at {index_path} failed to load twice; "
1170 "quarantining and rebuilding fresh"
1171 )
1172 with lock:
1173 if index_path.exists():
1174 self._quarantine_corrupt_index(
1175 index_path, "vector_index_load_failed"
1176 )
1177 self._reset_index_state_for_rebuild(collection_id, rag_index)
1178 return VectorIndex(**kwargs)
1180 def search(
1181 self, query: str, collection_id: str, top_k: int
1182 ) -> List[SearchResult]:
1183 """Semantic search within ``collection_id`` via the vector store.
1185 Runs the same verify/quarantine/dimension pre-flight as indexing
1186 (through ``_get_vector_index``), then delegates to
1187 ``VectorIndex.search``. ``reset_stale_state`` is left at its
1188 default False — this is a read path: a transient integrity hiccup
1189 during a mere search must not wipe the collection's indexed state
1190 and force a full re-embed.
1192 Callers are expected to have already confirmed the collection has
1193 indexed documents (e.g. via ``get_rag_stats``) before calling this
1194 — this method will otherwise create an empty RAGIndex/store for a
1195 never-indexed (collection, embedding model) pair and return no
1196 results, rather than raising.
1197 """
1198 collection_name = f"collection_{collection_id}"
1199 vindex = self._get_vector_index(collection_id, collection_name)
1200 return vindex.search(query, top_k)
1202 def get_current_index_info(
1203 self, collection_id: Optional[str] = None
1204 ) -> Optional[Dict[str, Any]]:
1205 """
1206 Get information about the current RAG index for a collection.
1208 Args:
1209 collection_id: UUID of collection (defaults to Library if None)
1210 """
1211 with get_user_db_session(self.username, self.db_password) as session:
1212 # Get collection name in the format stored in RAGIndex (collection_<uuid>)
1213 if collection_id:
1214 collection = (
1215 session.query(Collection)
1216 .filter_by(id=collection_id)
1217 .first()
1218 )
1219 collection_name = (
1220 f"collection_{collection_id}" if collection else "unknown"
1221 )
1222 else:
1223 # Default to Library collection
1224 from ...database.library_init import get_default_library_id
1226 collection_id = get_default_library_id(
1227 self.username, self.db_password
1228 )
1229 collection_name = f"collection_{collection_id}"
1231 rag_index = (
1232 session.query(RAGIndex)
1233 .filter_by(collection_name=collection_name, is_current=True)
1234 .first()
1235 )
1237 if not rag_index:
1238 # Debug: check all RAG indices for this collection
1239 all_indices = session.query(RAGIndex).all()
1240 logger.info(
1241 f"No RAG index found for collection_name='{collection_name}'. All indices: {[(idx.collection_name, idx.is_current) for idx in all_indices]}"
1242 )
1243 return None
1245 # Calculate actual counts from rag_document_status table
1246 from ...database.models.library import RagDocumentStatus
1248 actual_chunk_count = (
1249 session.query(func.sum(RagDocumentStatus.chunk_count))
1250 .filter_by(collection_id=collection_id)
1251 .scalar()
1252 or 0
1253 )
1255 actual_doc_count = (
1256 session.query(RagDocumentStatus)
1257 .filter_by(collection_id=collection_id)
1258 .count()
1259 )
1261 return {
1262 "embedding_model": rag_index.embedding_model,
1263 "embedding_model_type": rag_index.embedding_model_type.value
1264 if rag_index.embedding_model_type
1265 else None,
1266 "embedding_dimension": rag_index.embedding_dimension,
1267 "chunk_size": rag_index.chunk_size,
1268 "chunk_overlap": rag_index.chunk_overlap,
1269 "chunk_count": actual_chunk_count,
1270 "total_documents": actual_doc_count,
1271 "created_at": rag_index.created_at.isoformat(),
1272 "last_updated_at": rag_index.last_updated_at.isoformat(),
1273 }
1275 def _flag_document_for_reindex(
1276 self, session, document_id: str, collection_id: str
1277 ) -> None:
1278 """Mark a document NOT indexed so the reconciler re-embeds it.
1280 Used on index_document failure paths: ``VectorIndex.index()`` runs
1281 ``store.apply()`` — DURABLY mutating the FAISS file — BEFORE the DB
1282 commit, so a failure/rollback after that point can leave the source's
1283 prior chunk rows resurrected with no vectors (silently unsearchable)
1284 while ``DocumentCollection.indexed`` is reverted to its pre-attempt
1285 state. The background reconciler only revisits ``indexed=False`` rows,
1286 so without this flag the document stays permanently, silently
1287 unsearchable. Best-effort + committed in its own statement;
1288 index_document owns per-document commits (see index_all_documents), so
1289 this clobbers no caller.
1290 """
1291 try:
1292 session.query(DocumentCollection).filter_by(
1293 document_id=document_id, collection_id=collection_id
1294 ).update(
1295 {"indexed": False, "chunk_count": 0},
1296 synchronize_session=False,
1297 )
1298 session.query(RagDocumentStatus).filter_by(
1299 document_id=document_id, collection_id=collection_id
1300 ).delete(synchronize_session=False)
1301 session.commit()
1302 except Exception:
1303 safe_rollback(
1304 session, "library_rag_service._flag_document_for_reindex"
1305 )
1306 logger.warning(
1307 f"Could not flag document {document_id} for re-index after a "
1308 f"failed index of collection {collection_id}; a manual reindex "
1309 "may be needed to restore its search."
1310 )
1312 def index_document(
1313 self, document_id: str, collection_id: str, force_reindex: bool = False
1314 ) -> Dict[str, Any]:
1315 """
1316 Index a single document into RAG for a specific collection.
1318 Args:
1319 document_id: UUID of the Document to index
1320 collection_id: UUID of the Collection to index for
1321 force_reindex: Whether to force reindexing even if already indexed
1323 Returns:
1324 Dict with status, chunk_count, and any errors
1325 """
1326 with get_user_db_session(self.username, self.db_password) as session:
1327 # Get the document
1328 document = session.query(Document).filter_by(id=document_id).first()
1330 if not document:
1331 return {"status": "error", "error": "Document not found"}
1333 # Get or create DocumentCollection entry
1334 doc_collection = ensure_in_collection(
1335 session, document_id, collection_id
1336 )
1338 # Check if already indexed for this collection
1339 if doc_collection.indexed and not force_reindex:
1340 return {
1341 "status": "skipped",
1342 "message": "Document already indexed for this collection",
1343 "chunk_count": doc_collection.chunk_count,
1344 }
1346 # Validate text content
1347 if not document.text_content:
1348 # A previously-indexed document now cleared to empty text (e.g.
1349 # a note edited down to whitespace) must have its stale
1350 # chunks/vectors PURGED, not left searchable. The note-edit path
1351 # already set indexed=False, so detect prior chunks directly and
1352 # remove them via the facade's empty-replace.
1353 collection_name = f"collection_{collection_id}"
1354 has_prior = (
1355 session.query(DocumentChunk.id)
1356 .filter_by(
1357 source_type="document",
1358 source_id=document_id,
1359 collection_name=collection_name,
1360 )
1361 .first()
1362 is not None
1363 )
1364 if has_prior: 1364 ↛ 1373line 1364 didn't jump to line 1373 because the condition on line 1364 was never true
1365 # Same store-durable-before-DB-commit hazard as the main
1366 # reindex path below, but this branch sits OUTSIDE that
1367 # path's try/except: the empty-replace apply() DURABLY
1368 # removes the prior vectors before the commit here. On
1369 # failure, roll back and flag the document for re-index so
1370 # the reconciler re-runs the purge instead of leaving
1371 # resurrected rows whose vectors are gone (silently
1372 # unsearchable) marked indexed.
1373 try:
1374 vindex = self._get_vector_index(
1375 collection_id,
1376 collection_name,
1377 reset_stale_state=True,
1378 )
1379 vindex.index(
1380 source_type="document",
1381 source_id=document_id,
1382 chunks=[],
1383 replace=True,
1384 session=session,
1385 )
1386 doc_collection.indexed = False
1387 doc_collection.chunk_count = 0
1388 # Clear the canonical RagDocumentStatus row too. Row
1389 # existence is the "indexed" marker get_rag_stats / the
1390 # RAG status route read, so it MUST move with
1391 # DocumentCollection.indexed — a leftover row would
1392 # report this just-purged (now empty, unsearchable)
1393 # document as still indexed with a stale chunk_count.
1394 # Mirrors remove_document_from_rag and the three sibling
1395 # not-indexed paths; the cleared branch was the lone
1396 # exception.
1397 session.query(RagDocumentStatus).filter_by(
1398 document_id=document_id, collection_id=collection_id
1399 ).delete(synchronize_session=False)
1400 session.commit()
1401 except Exception as e:
1402 safe_rollback(
1403 session,
1404 "library_rag_service.index_document cleared",
1405 )
1406 self._flag_document_for_reindex(
1407 session, document_id, collection_id
1408 )
1409 logger.exception(
1410 f"Error purging cleared document {document_id} "
1411 f"for collection {collection_id}"
1412 )
1413 return {
1414 "status": "error",
1415 "error": f"Operation failed: {type(e).__name__}",
1416 }
1417 return {
1418 "status": "cleared",
1419 "message": (
1420 "Document has no text content; prior chunks/"
1421 "vectors purged"
1422 ),
1423 "chunk_count": 0,
1424 }
1425 return {
1426 "status": "error",
1427 "error": "Document has no text content",
1428 }
1430 # Replace-on-reindex context: a document being re-indexed (e.g.
1431 # a note whose content changed) has prior chunks with DIFFERENT
1432 # content hashes from the new text. Left alone they accumulate
1433 # in both the DB and the vector store forever — unbounded
1434 # per-edit growth plus stale duplicate search hits.
1435 # ``VectorIndex.index(replace=True)`` below handles this: it
1436 # looks up the source's prior DocumentChunk rows (authoritative,
1437 # queried fresh — not from an ephemeral pre-committed snapshot),
1438 # applies the vector-store removal durably, then deletes those
1439 # rows in the SAME transaction this method already holds open
1440 # (so a failure rolls back together with everything else, never
1441 # orphaning rows or poisoning the shared session).
1443 try:
1444 # Create LangChain Document from text
1445 doc = LangchainDocument(
1446 page_content=document.text_content,
1447 metadata={
1448 "source": document.original_url,
1449 "document_id": document_id, # Add document ID for source linking
1450 "collection_id": collection_id, # Add collection ID
1451 "title": document.title
1452 or document.filename
1453 or "Untitled",
1454 "document_title": document.title
1455 or document.filename
1456 or "Untitled", # Add for compatibility
1457 "authors": document.authors,
1458 "published_date": str(document.published_date)
1459 if document.published_date
1460 else None,
1461 "doi": document.doi,
1462 "arxiv_id": document.arxiv_id,
1463 "pmid": document.pmid,
1464 "pmcid": document.pmcid,
1465 "extraction_method": document.extraction_method,
1466 "word_count": document.word_count,
1467 },
1468 )
1470 # Split into chunks
1471 chunks = self.text_splitter.split_documents([doc])
1472 logger.info(
1473 f"Split document {document_id} into {len(chunks)} chunks"
1474 )
1476 # Get collection name for chunk storage
1477 collection = (
1478 session.query(Collection)
1479 .filter_by(id=collection_id)
1480 .first()
1481 )
1482 # Use collection_<uuid> format for internal storage
1483 collection_name = (
1484 f"collection_{collection_id}" if collection else "unknown"
1485 )
1487 document_title = (
1488 document.title or document.filename or "Untitled"
1489 )
1490 # Mirrors the parent doc's metadata dict built above (source,
1491 # ids, bibliographic fields) so it survives into each
1492 # chunk's ``DocumentChunk.document_metadata`` — this goes
1493 # ONLY to the encrypted DB (never the vector store; see the
1494 # SECURITY INVARIANT in vector_stores/base.py) and future-
1495 # proofs citation/UI features that read per-chunk metadata.
1496 shared_chunk_metadata = {
1497 "source": document.original_url,
1498 "document_id": document_id,
1499 "collection_id": collection_id,
1500 "document_title": document_title,
1501 "title": document_title,
1502 "authors": document.authors,
1503 "published_date": str(document.published_date)
1504 if document.published_date
1505 else None,
1506 "doi": document.doi,
1507 "arxiv_id": document.arxiv_id,
1508 "pmid": document.pmid,
1509 "pmcid": document.pmcid,
1510 "extraction_method": document.extraction_method,
1511 "word_count": document.word_count,
1512 }
1513 chunk_inputs = [
1514 ChunkInput(
1515 text=chunk.page_content,
1516 metadata={
1517 **shared_chunk_metadata,
1518 "chunk_index": i,
1519 },
1520 )
1521 for i, chunk in enumerate(chunks)
1522 ]
1524 # Embed + persist via the new int-id-keyed vector store —
1525 # this writes DocumentChunk rows (text only in the encrypted
1526 # DB) and the vector-store's ids/vectors (never text; see
1527 # the SECURITY INVARIANT in vector_stores/base.py). Passing
1528 # ``session`` means the facade does not commit/rollback —
1529 # this transaction still owns that (below).
1530 #
1531 # replace=True ALWAYS: a full-document (re)index owns all of a
1532 # source's chunks, so it must drop the source's prior
1533 # chunks/vectors before adding the new ones. force_reindex is
1534 # NOT a safe proxy: a note edit clears the indexed STATUS
1535 # (_mark_note_stale_for_reindex_in_session) but leaves the
1536 # pre-edit chunk rows/vectors in place, so the subsequent
1537 # non-forced reindex would reach here with force_reindex=False
1538 # AND prior chunks present — appending would leave stale/deleted
1539 # note text permanently searchable and duplicate every chunk.
1540 # For a genuinely never-indexed source, prior_ids is empty and
1541 # replace is a no-op, so this is always safe.
1542 vindex = self._get_vector_index(
1543 collection_id, collection_name, reset_stale_state=True
1544 )
1545 index_stats = vindex.index(
1546 source_type="document",
1547 source_id=document_id,
1548 chunks=chunk_inputs,
1549 replace=True,
1550 session=session,
1551 )
1552 logger.info(
1553 f"Vector index update for document {document_id}: "
1554 f"{index_stats.added} added, {index_stats.removed} "
1555 f"removed ({index_stats.chunks} chunks total)"
1556 )
1558 from datetime import datetime, UTC
1560 # Check if document was already indexed (for stats update)
1561 existing_status = (
1562 session.query(RagDocumentStatus)
1563 .filter_by(
1564 document_id=document_id, collection_id=collection_id
1565 )
1566 .first()
1567 )
1568 was_already_indexed = existing_status is not None
1570 # Mark document as indexed using rag_document_status table
1571 # Row existence = indexed, simple and clean
1572 timestamp = datetime.now(UTC)
1574 # Create or update RagDocumentStatus using ORM merge (atomic upsert)
1575 rag_status = RagDocumentStatus(
1576 document_id=document_id,
1577 collection_id=collection_id,
1578 rag_index_id=self.rag_index_record.id,
1579 chunk_count=len(chunks),
1580 indexed_at=timestamp,
1581 )
1582 session.merge(rag_status)
1584 logger.info(
1585 f"Marked document as indexed in rag_document_status: doc_id={document_id}, coll_id={collection_id}, chunks={len(chunks)}"
1586 )
1588 # Also update DocumentCollection table for backward compatibility
1589 session.query(DocumentCollection).filter_by(
1590 document_id=document_id, collection_id=collection_id
1591 ).update(
1592 {
1593 "indexed": True,
1594 "chunk_count": len(chunks),
1595 "last_indexed_at": timestamp,
1596 }
1597 )
1599 logger.info(
1600 "Also updated DocumentCollection.indexed for backward compatibility"
1601 )
1603 # Update RAGIndex statistics (only if not already indexed)
1604 rag_index_obj = (
1605 session.query(RAGIndex)
1606 .filter_by(id=self.rag_index_record.id)
1607 .first()
1608 )
1609 if rag_index_obj and not was_already_indexed:
1610 rag_index_obj.chunk_count += len(chunks)
1611 rag_index_obj.total_documents += 1
1612 rag_index_obj.last_updated_at = datetime.now(UTC)
1613 logger.info(
1614 f"Updated RAGIndex stats: chunk_count +{len(chunks)}, total_documents +1"
1615 )
1617 # Replace-on-reindex DB prune: no longer needed here —
1618 # ``VectorIndex.index(replace=True)`` above already deletes
1619 # this source's prior DocumentChunk rows (and their
1620 # vectors) as part of the same transaction/lock when
1621 # ``force_reindex`` is set.
1623 # Flush ORM changes to database before commit
1624 session.flush()
1625 logger.info(f"Flushed ORM changes for document {document_id}")
1627 # Commit the transaction. Durability is provided by
1628 # synchronous=NORMAL (sqlcipher_utils.py); SQLite
1629 # auto-checkpoints WAL at wal_autocheckpoint=250 frames.
1630 # An explicit PRAGMA wal_checkpoint(FULL) here used to
1631 # block other writers long enough to exhaust busy_timeout
1632 # under bulk-download concurrency (#4197).
1633 session.commit()
1635 logger.info(
1636 f"Successfully indexed document {document_id} for collection {collection_id} "
1637 f"with {len(chunks)} chunks"
1638 )
1640 return {
1641 "status": "success",
1642 "chunk_count": len(chunks),
1643 }
1645 except Exception as e:
1646 # The session is shared (thread-local) with the caller.
1647 # If session.flush() or session.commit() raised, the session
1648 # is in PendingRollbackError state until rolled back —
1649 # leaving subsequent operations to cascade. Roll back BEFORE
1650 # returning the error dict so the caller sees a clean
1651 # session. (Same pattern as the #3827 fix.)
1652 safe_rollback(session, "library_rag_service.index_document")
1653 # Backstop for the store-durable-before-DB-commit window: a
1654 # REPLACE reindex ran store.apply() (DURABLY removing this
1655 # source's prior vectors) before the DB commit, so the rollback
1656 # above resurrects rows whose vectors are already gone while
1657 # reverting indexed→its pre-attempt True. Flag for re-index so
1658 # the reconciler re-embeds and store+DB reconverge.
1659 self._flag_document_for_reindex(
1660 session, document_id, collection_id
1661 )
1662 logger.exception(
1663 f"Error indexing document {document_id} for collection {collection_id}"
1664 )
1665 return {
1666 "status": "error",
1667 "error": f"Operation failed: {type(e).__name__}",
1668 }
1670 def index_all_documents(
1671 self,
1672 collection_id: str,
1673 force_reindex: bool = False,
1674 progress_callback=None,
1675 ) -> Dict[str, Any]:
1676 """
1677 Index all documents in a collection into RAG.
1679 Args:
1680 collection_id: UUID of the collection to index
1681 force_reindex: Whether to force reindexing already indexed documents
1682 progress_callback: Optional callback function called after each document with (current, total, doc_title, status)
1684 Returns:
1685 Dict with counts of successful, skipped, and failed documents
1686 """
1687 with get_user_db_session(self.username, self.db_password) as session:
1688 # Get all DocumentCollection entries for this collection
1689 query = session.query(DocumentCollection).filter_by(
1690 collection_id=collection_id
1691 )
1693 if not force_reindex:
1694 # Only index documents that haven't been indexed yet
1695 query = query.filter_by(indexed=False)
1697 doc_collections = query.all()
1699 if not doc_collections:
1700 return {
1701 "status": "info",
1702 "message": "No documents to index",
1703 "successful": 0,
1704 "skipped": 0,
1705 "failed": 0,
1706 }
1708 results = {"successful": 0, "skipped": 0, "failed": 0, "errors": []}
1709 total = len(doc_collections)
1711 for idx, doc_collection in enumerate(doc_collections, 1):
1712 # Get the document for title info
1713 document = (
1714 session.query(Document)
1715 .filter_by(id=doc_collection.document_id)
1716 .first()
1717 )
1718 title = document.title if document else "Unknown"
1720 result = self.index_document(
1721 doc_collection.document_id, collection_id, force_reindex
1722 )
1724 if result["status"] == "success":
1725 results["successful"] += 1
1726 elif result["status"] in ("skipped", "cleared"):
1727 # "cleared" (a previously-indexed doc emptied to no text,
1728 # its stale vectors purged) is a handled non-failure, not an
1729 # error — counting it in "failed" would surface a blank-error
1730 # failure for a correct purge.
1731 results["skipped"] += 1
1732 else:
1733 results["failed"] += 1
1734 results["errors"].append(
1735 {
1736 "doc_id": doc_collection.document_id,
1737 "title": title,
1738 "error": result.get("error"),
1739 }
1740 )
1742 # Call progress callback if provided
1743 if progress_callback:
1744 progress_callback(idx, total, title, result["status"])
1746 logger.info(
1747 f"Indexed collection {collection_id}: "
1748 f"{results['successful']} successful, "
1749 f"{results['skipped']} skipped, "
1750 f"{results['failed']} failed"
1751 )
1753 return results
1755 def remove_document_from_rag(
1756 self, document_id: str, collection_id: str
1757 ) -> Dict[str, Any]:
1758 """
1759 Remove a document's chunks from RAG for a specific collection.
1761 Args:
1762 document_id: UUID of the Document to remove
1763 collection_id: UUID of the Collection to remove from
1765 Returns:
1766 Dict with status and count of removed chunks
1767 """
1768 with get_user_db_session(self.username, self.db_password) as session:
1769 # Get the DocumentCollection entry
1770 doc_collection = (
1771 session.query(DocumentCollection)
1772 .filter_by(document_id=document_id, collection_id=collection_id)
1773 .first()
1774 )
1776 if not doc_collection:
1777 return {
1778 "status": "error",
1779 "error": "Document not found in collection",
1780 }
1782 try:
1783 # Get collection name in the format collection_<uuid>
1784 collection = (
1785 session.query(Collection)
1786 .filter_by(id=collection_id)
1787 .first()
1788 )
1789 # Use collection_<uuid> format for internal storage
1790 collection_name = (
1791 f"collection_{collection_id}" if collection else "unknown"
1792 )
1794 # Remove the document's vectors from the FAISS index FIRST.
1795 # ``remove_documents_from_index`` -> ``VectorIndex.delete``
1796 # looks up the DocumentChunk rows by (source_type, source_id,
1797 # collection_name) to learn which FAISS int ids to remove,
1798 # and deletes those matching rows itself as part of the same
1799 # operation. If the DB rows were deleted first, that lookup
1800 # would find nothing and the vectors (and their persisted
1801 # text) would be stranded in the FAISS store forever — never
1802 # removed and never re-discoverable by a later cleanup.
1803 # Best-effort: a FAISS hiccup must not fail the DB removal
1804 # below.
1805 # ``remove_documents_from_index`` deletes the matching chunk
1806 # rows itself (one row per removed vector) and returns that
1807 # count. Capture it so the reported ``deleted_count`` reflects
1808 # the rows removed here, not just the backstop sweep below: in
1809 # the normal path those rows are already gone by the time the
1810 # backstop runs, so the backstop returns ~0 and — left
1811 # uncounted — would under-report the removal (often as 0) to
1812 # the caller/UI.
1813 removed_vectors = 0
1814 try:
1815 removed_vectors = self.remove_documents_from_index(
1816 [document_id], collection_id
1817 )
1818 except Exception:
1819 logger.exception(
1820 f"Could not remove document {document_id} vectors "
1821 f"from FAISS index for collection {collection_id}"
1822 )
1824 # DB-only cleanup for any rows ``remove_documents_from_index``
1825 # did not match (e.g. no current FAISS index for this
1826 # collection). Disjoint from ``removed_vectors`` above (those
1827 # rows are already committed-deleted), so the two sum to the
1828 # true total removed.
1829 deleted_count = removed_vectors + (
1830 self.embedding_manager._delete_chunks_from_db(
1831 collection_name=collection_name,
1832 source_id=document_id,
1833 )
1834 )
1836 # Update DocumentCollection RAG status
1837 doc_collection.indexed = False
1838 doc_collection.chunk_count = 0
1839 doc_collection.last_indexed_at = None
1840 # Clear the canonical RagDocumentStatus row too. Row-existence is
1841 # the "indexed" marker get_rag_stats / the RAG status route read,
1842 # so it MUST move together with DocumentCollection.indexed — a
1843 # leftover row would report the just-removed document as still
1844 # indexed.
1845 session.query(RagDocumentStatus).filter_by(
1846 document_id=document_id, collection_id=collection_id
1847 ).delete(synchronize_session=False)
1848 session.commit()
1850 logger.info(
1851 f"Removed {deleted_count} chunks for document {document_id} from collection {collection_id}"
1852 )
1854 return {"status": "success", "deleted_count": deleted_count}
1856 except Exception as e:
1857 # session.commit() above can raise; without rollback the
1858 # shared thread-local session stays poisoned for the
1859 # caller's next operation (issue #3827 pattern).
1860 safe_rollback(
1861 session, "library_rag_service.remove_document_from_rag"
1862 )
1863 # Backstop for the store-durable-before-DB-commit window:
1864 # remove_documents_from_index -> VectorIndex.delete ran
1865 # store.apply() (DURABLY removing this document's vectors) BEFORE
1866 # the DB status work above committed, so the rollback can leave
1867 # the row still marked indexed=True with its vectors already gone
1868 # -- indexed-but-unsearchable, and over-reported as indexed by
1869 # get_rag_stats with no auto-recovery. Durably COMPLETE the
1870 # removal's status change so the DB matches the removed vectors:
1871 # mark it not-indexed and drop the canonical RagDocumentStatus
1872 # row. This is a REMOVAL backstop, not a re-index one (contrast
1873 # index_document's _flag_document_for_reindex) -- the vectors are
1874 # intentionally gone, so we finish the removal rather than
1875 # resurrect the document. Best-effort: a retry of the removal
1876 # reconciles it if this too fails.
1877 try:
1878 session.query(DocumentCollection).filter_by(
1879 document_id=document_id, collection_id=collection_id
1880 ).update(
1881 {
1882 "indexed": False,
1883 "chunk_count": 0,
1884 "last_indexed_at": None,
1885 },
1886 synchronize_session=False,
1887 )
1888 session.query(RagDocumentStatus).filter_by(
1889 document_id=document_id, collection_id=collection_id
1890 ).delete(synchronize_session=False)
1891 session.commit()
1892 except Exception:
1893 safe_rollback(
1894 session,
1895 "library_rag_service.remove_document_from_rag backstop",
1896 )
1897 logger.warning(
1898 f"Could not finalize removal state for document "
1899 f"{document_id} in collection {collection_id}; a retry "
1900 "of the removal will reconcile it."
1901 )
1902 logger.exception(
1903 f"Error removing document {document_id} from collection {collection_id}"
1904 )
1905 return {
1906 "status": "error",
1907 "error": f"Operation failed: {type(e).__name__}",
1908 }
1910 def purge_document_chunks(
1911 self, document_id: str, collection_id: str
1912 ) -> int:
1913 """Delete a document's chunk rows from a collection's RAG store
1914 WITHOUT requiring the ``DocumentCollection`` join row to still exist.
1916 ``remove_document_from_rag`` short-circuits (returns "Document not
1917 found in collection") when the join row is gone — which is exactly
1918 the state after a ``Document`` is cascade-deleted. Callers that
1919 delete the Document first (note deletion) must use this to actually
1920 purge the now-orphaned ``DocumentChunk`` rows; otherwise the chunks
1921 (and their embedding ids) linger and semantic search keeps
1922 resolving hits to a Document row that no longer exists.
1924 Mirrors the chunk-delete half of ``remove_document_from_rag``.
1925 FAISS-vector pruning is NOT done here: pair this with
1926 ``purge_document_vectors`` when the document is being deleted
1927 outright — replace-on-reindex can never fire again for a deleted
1928 document id. Returns the number of chunk rows deleted.
1929 """
1930 collection_name = f"collection_{collection_id}"
1931 return self.embedding_manager._delete_chunks_from_db(
1932 collection_name=collection_name,
1933 source_id=document_id,
1934 )
1936 def purge_document_vectors(
1937 self, document_id: str, collection_id: str
1938 ) -> int:
1939 """Remove a deleted document's chunk vectors from the collection's
1940 vector store.
1942 Companion to ``purge_document_chunks`` for callers that delete the
1943 ``Document`` row itself (note deletion). Without this the vectors
1944 lingered until a replace-on-reindex of the SAME document id — which
1945 can never fire again once the document is deleted — so collection
1946 search kept serving the deleted content indefinitely.
1948 Routes to ``VectorIndex.delete(source_type="document",
1949 source_id=document_id)``. Unlike the pre-cutover implementation,
1950 no separate "shared chunk ownership" check is needed here: the new
1951 store is one-row-per-source (see the "One row = one vector" note in
1952 vector_stores/facade.py) — identical text in two documents gets two
1953 independent rows/vectors, so deleting this document's rows can
1954 never remove another document's vector.
1956 Deliberately does NOT go through ``_get_vector_index`` /
1957 ``_get_or_create_rag_index``: a deletion path must not create an
1958 index record (or probe the embedding provider) just to delete from
1959 it, and does not quarantine/rebuild on a failed load — an
1960 unreadable index has nothing purgeable right now; quarantine/
1961 rebuild decisions stay with the indexing paths.
1963 Returns the number of vectors removed.
1964 """
1965 collection_name = f"collection_{collection_id}"
1966 index_hash = self._get_index_hash(
1967 collection_name, self.embedding_model, self.embedding_provider
1968 )
1969 with get_user_db_session(self.username, self.db_password) as session:
1970 rag_index = (
1971 session.query(RAGIndex).filter_by(index_hash=index_hash).first()
1972 )
1973 if rag_index is None:
1974 # Collection was never indexed with this configuration —
1975 # no vectors to purge. Deliberately NOT
1976 # _get_or_create_rag_index: a deletion path must not
1977 # create index records (or probe the embedding provider).
1978 return 0
1979 self.rag_index_record = rag_index
1981 # Resolve where the file ACTUALLY lives, the way the index/search paths
1982 # do — but without _preflight's side effects (this deletion path must
1983 # not relocate/record/create). rag_index.index_path is the authoritative
1984 # DB record (the legacy shared-cache path pre-migration, the per-user
1985 # path after _preflight relocates it); _get_index_path(index_hash) only
1986 # recomputes the POST-migration location. Checking only the recomputed
1987 # path silently no-ops on an un-migrated install — the file is still at
1988 # the legacy index_path, so the deleted document's real vector survives.
1989 recomputed_path = self._get_index_path(rag_index.index_hash)
1990 recorded_path = (
1991 Path(rag_index.index_path) if rag_index.index_path else None
1992 )
1993 if recorded_path is not None and recorded_path.exists(): 1993 ↛ 1994line 1993 didn't jump to line 1994 because the condition on line 1993 was never true
1994 index_path = recorded_path
1995 elif recomputed_path.exists():
1996 index_path = recomputed_path
1997 else:
1998 return 0
1999 verified, reason = self.integrity_manager.verify_file(index_path)
2000 if not verified: 2000 ↛ 2009line 2000 didn't jump to line 2009 because the condition on line 2000 was always true
2001 # Leave quarantine/rebuild decisions to the indexing paths;
2002 # an unreadable index has nothing purgeable right now.
2003 logger.warning(
2004 f"Skipping vector purge for document {document_id}: index "
2005 f"{index_path} fails verification ({reason})"
2006 )
2007 return 0
2009 lock = _get_faiss_write_lock(self.username, str(index_path))
2011 # nested=True: these run reentrantly inside VectorIndex.apply(), which
2012 # holds just-flushed, uncommitted DocumentChunk rows on this same
2013 # thread-local session. The integrity write must run in a SAVEPOINT and
2014 # let apply()'s caller commit — a plain commit/rollback here would
2015 # settle or DISCARD those pending rows (silently losing indexed text
2016 # while the vectors persist). See FileIntegrityManager._integrity_write.
2017 def _record(path: Path) -> None:
2018 self.integrity_manager.record_file(
2019 path,
2020 related_entity_type="rag_index",
2021 related_entity_id=rag_index.id,
2022 nested=True,
2023 )
2025 def _verify(path: Path):
2026 return self.integrity_manager.verify_file(path, nested=True)
2028 try:
2029 vindex = VectorIndex(
2030 username=self.username,
2031 db_password=self.db_password,
2032 # Delete-only path: never materialise the embedding backend.
2033 # VectorIndex.delete() resolves rows by id and never calls
2034 # embed_*; forcing the model here would cost seconds + hundreds
2035 # of MB of RSS per deleted document for nothing.
2036 embeddings=_LazyEmbeddings(self.embedding_manager),
2037 embedding_model=self.embedding_model,
2038 embedding_model_type=EmbeddingProvider(self.embedding_provider),
2039 collection_name=collection_name,
2040 dimension=rag_index.embedding_dimension,
2041 path=index_path,
2042 lock=lock,
2043 integrity_record=_record,
2044 integrity_verify=_verify,
2045 # Use the STORED index config, not the service instance's — a
2046 # service built with defaults (index_type="flat") would evaluate
2047 # supports_delete wrongly for an HNSW collection and silently
2048 # fail the removal. NULL (legacy rows) falls back to self.*.
2049 index_type=rag_index.index_type or self.index_type,
2050 metric=rag_index.distance_metric or self.distance_metric,
2051 normalize=(
2052 rag_index.normalize_vectors
2053 if rag_index.normalize_vectors is not None
2054 else self.normalize_vectors
2055 ),
2056 )
2057 stats = vindex.delete(source_type="document", source_id=document_id)
2058 except (ValueError, RuntimeError, OSError):
2059 # OSError too: this is a BEST-EFFORT purge whose caller goes on to
2060 # delete the DocumentChunk rows. A disk error from the FAISS
2061 # persist() (OSError) must not escape and abort that row deletion —
2062 # it would strand the chunk rows while leaving the vectors, the
2063 # opposite inconsistency. Skip the vector purge and let the caller
2064 # proceed; the vectors reconcile on the next re-index.
2065 logger.warning(
2066 f"Skipping vector purge for document {document_id}: could "
2067 f"not load or update index {index_path}"
2068 )
2069 return 0
2070 return stats.removed
2072 def index_documents_batch(
2073 self,
2074 doc_info: List[tuple],
2075 collection_id: str,
2076 force_reindex: bool = False,
2077 ) -> Dict[str, Dict[str, Any]]:
2078 """
2079 Index multiple documents in a batch for a specific collection.
2081 Args:
2082 doc_info: List of (doc_id, title) tuples
2083 collection_id: UUID of the collection to index for
2084 force_reindex: Whether to force reindexing even if already indexed
2086 Returns:
2087 Dict mapping doc_id to individual result
2088 """
2089 results = {}
2090 doc_ids = [doc_id for doc_id, _ in doc_info]
2092 # Use single database session for querying
2093 with get_user_db_session(self.username, self.db_password) as session:
2094 # Pre-load all documents for this batch
2095 documents = (
2096 session.query(Document).filter(Document.id.in_(doc_ids)).all()
2097 )
2099 # Create lookup for quick access
2100 doc_lookup = {doc.id: doc for doc in documents}
2102 # Pre-load DocumentCollection entries
2103 doc_collections = (
2104 session.query(DocumentCollection)
2105 .filter(
2106 DocumentCollection.document_id.in_(doc_ids),
2107 DocumentCollection.collection_id == collection_id,
2108 )
2109 .all()
2110 )
2111 doc_collection_lookup = {
2112 dc.document_id: dc for dc in doc_collections
2113 }
2115 # Process each document in the batch
2116 for doc_id, title in doc_info:
2117 document = doc_lookup.get(doc_id)
2119 if not document:
2120 results[doc_id] = {
2121 "status": "error",
2122 "error": "Document not found",
2123 }
2124 continue
2126 # Check if already indexed via DocumentCollection
2127 doc_collection = doc_collection_lookup.get(doc_id)
2128 if (
2129 doc_collection
2130 and doc_collection.indexed
2131 and not force_reindex
2132 ):
2133 results[doc_id] = {
2134 "status": "skipped",
2135 "message": "Document already indexed for this collection",
2136 "chunk_count": doc_collection.chunk_count,
2137 }
2138 continue
2140 # Do NOT short-circuit empty text_content here: a document/note
2141 # edited down to empty still has its PRIOR chunks + vectors on
2142 # disk, and only index_document's empty-content branch purges
2143 # them (via an empty-replace). Returning an "error" here left
2144 # that stale (plaintext) content permanently searchable through
2145 # the bulk "Index All" path — and self-perpetuating, since the
2146 # doc's indexed flag never flips. Route it through index_document
2147 # so the batch path gets the same purge-on-clear as the SSE /
2148 # single-document paths; it returns "cleared" (purged) or the
2149 # same "error" when there were no prior chunks.
2150 try:
2151 result = self.index_document(
2152 doc_id, collection_id, force_reindex
2153 )
2154 results[doc_id] = result
2155 except Exception as e:
2156 logger.exception(
2157 f"Error indexing document {doc_id} in batch"
2158 )
2159 results[doc_id] = {
2160 "status": "error",
2161 "error": f"Indexing failed: {type(e).__name__}",
2162 }
2164 return results
2166 def get_rag_stats(
2167 self, collection_id: Optional[str] = None
2168 ) -> Dict[str, Any]:
2169 """
2170 Get RAG statistics for a collection.
2172 Args:
2173 collection_id: UUID of the collection (defaults to Library)
2175 Returns:
2176 Dict with counts and metadata about indexed documents
2177 """
2178 with get_user_db_session(self.username, self.db_password) as session:
2179 # Get collection ID (default to Library)
2180 if not collection_id: 2180 ↛ 2181line 2180 didn't jump to line 2181 because the condition on line 2180 was never true
2181 from ...database.library_init import get_default_library_id
2183 collection_id = get_default_library_id(
2184 self.username, self.db_password
2185 )
2187 # Count total documents in collection
2188 total_docs = (
2189 session.query(DocumentCollection)
2190 .filter_by(collection_id=collection_id)
2191 .count()
2192 )
2194 # Count indexed documents from rag_document_status table
2195 from ...database.models.library import RagDocumentStatus
2197 indexed_docs = (
2198 session.query(RagDocumentStatus)
2199 .filter_by(collection_id=collection_id)
2200 .count()
2201 )
2203 # Count total chunks from rag_document_status table
2204 total_chunks = (
2205 session.query(func.sum(RagDocumentStatus.chunk_count))
2206 .filter_by(collection_id=collection_id)
2207 .scalar()
2208 or 0
2209 )
2211 # Get collection name in the format stored in DocumentChunk (collection_<uuid>)
2212 collection = (
2213 session.query(Collection).filter_by(id=collection_id).first()
2214 )
2215 collection_name = (
2216 f"collection_{collection_id}" if collection else "library"
2217 )
2219 # Get embedding model info from chunks
2220 chunk_sample = (
2221 session.query(DocumentChunk)
2222 .filter_by(collection_name=collection_name)
2223 .first()
2224 )
2226 embedding_info = {}
2227 if chunk_sample:
2228 embedding_info = {
2229 "model": chunk_sample.embedding_model,
2230 "model_type": chunk_sample.embedding_model_type.value
2231 if chunk_sample.embedding_model_type
2232 else None,
2233 "dimension": chunk_sample.embedding_dimension,
2234 }
2236 return {
2237 "total_documents": total_docs,
2238 "indexed_documents": indexed_docs,
2239 "unindexed_documents": total_docs - indexed_docs,
2240 "total_chunks": total_chunks,
2241 "embedding_info": embedding_info,
2242 "chunk_size": self.chunk_size,
2243 "chunk_overlap": self.chunk_overlap,
2244 }
2246 def remove_documents_from_index(
2247 self,
2248 document_ids: List[str],
2249 collection_id: str,
2250 *,
2251 chunk_ids_by_document: Optional[Dict[str, List[int]]] = None,
2252 ) -> int:
2253 """Delete given documents' vectors from a collection's vector store.
2255 Vector-store-only: the DB chunk rows are removed separately by the
2256 caller. This closes the gap where deleting a document's DB chunks
2257 left its vectors (and their persisted text — now impossible, see
2258 vector_stores/base.py's SECURITY INVARIANT) in the store, so the
2259 document kept surfacing in semantic search until a full reindex.
2261 For each id, calls ``VectorIndex.delete(source_type="document",
2262 source_id=...)`` — which looks up the matching ``DocumentChunk``
2263 rows by (source_type, source_id, collection_name) — UNLESS
2264 ``chunk_ids_by_document`` supplies that document's chunk ids
2265 directly, in which case ``VectorIndex.delete_ids(ids)`` is used
2266 instead. Callers whose own per-item transactions must delete the DB
2267 rows eagerly (for their own atomicity/rollback semantics) BEFORE
2268 this batched cleanup runs need the latter: the source_id lookup
2269 would find nothing once the rows are already gone (see the Zotero
2270 sync caller). Hardcodes ``source_type="document"``: both current
2271 callers (Zotero sync cleanup, ``remove_document_from_rag``) only
2272 ever pass library-document ids (the pre-cutover implementation
2273 matched a mixed ``document_id``/``source_id`` metadata key
2274 generically across "document" and "user_document" sources, but
2275 nothing ever called it with anything but library documents).
2277 Returns the total number of vectors removed across all ids. No-ops
2278 (returns 0) when the collection has no current vector index — this
2279 guard runs BEFORE building a ``VectorIndex`` so a delete call never
2280 creates one (``_get_vector_index`` would otherwise create the
2281 RAGIndex row via ``_get_or_create_rag_index``).
2282 """
2283 from ...database.models.library import RAGIndex
2284 from ...database.session_context import get_user_db_session
2286 wanted = [d for d in document_ids if d]
2287 if not wanted:
2288 return 0
2290 collection_name = f"collection_{collection_id}"
2291 # Don't create an index just to delete from it — only act if one
2292 # already exists for this collection.
2293 with get_user_db_session(self.username, self.db_password) as session:
2294 has_index = (
2295 session.query(RAGIndex.id)
2296 .filter_by(collection_name=collection_name, is_current=True)
2297 .first()
2298 is not None
2299 )
2300 if not has_index:
2301 return 0
2303 try:
2304 # reset_stale_state=True is REQUIRED on this write path. If the
2305 # pre-flight finds the index corrupt it quarantines the file
2306 # (renames it aside); a fresh EMPTY store is then built and the
2307 # delete loop below persists that empty store over the real path.
2308 # Without resetting the per-document indexed state, the DB would
2309 # still claim every other document is indexed while their vectors
2310 # are gone — silently wiping the whole collection's search with no
2311 # recovery but a manual full reindex. The reset only fires when the
2312 # index was actually quarantined (guarded on the file no longer
2313 # existing), so a healthy remove is unaffected. search() omits this
2314 # flag safely because it never persists an empty store.
2315 vindex = self._get_vector_index(
2316 collection_id, collection_name, reset_stale_state=True
2317 )
2318 except (ValueError, RuntimeError):
2319 logger.warning(
2320 "Could not load vector index for "
2321 f"{collection_name} while removing document vectors"
2322 )
2323 return 0
2325 # Collect ALL removable chunk ids across every document, then remove
2326 # them in ONE apply() (vindex.delete_ids). For an HNSW collection each
2327 # apply() rebuilds the whole index, so a per-document loop is
2328 # O(documents x collection_size); a single batched call rebuilds once.
2329 # Captured ids are used directly; documents without captured ids have
2330 # their chunk ids resolved by (source_type, source_id, collection_name)
2331 # — those rows still exist at this point.
2332 all_ids: List[int] = []
2333 unresolved: List[str] = []
2334 for document_id in wanted:
2335 captured = (
2336 chunk_ids_by_document.get(document_id)
2337 if chunk_ids_by_document
2338 else None
2339 )
2340 if captured: 2340 ↛ 2341line 2340 didn't jump to line 2341 because the condition on line 2340 was never true
2341 all_ids.extend(int(i) for i in captured if i is not None)
2342 else:
2343 unresolved.append(document_id)
2345 if unresolved: 2345 ↛ 2361line 2345 didn't jump to line 2361 because the condition on line 2345 was always true
2346 in_chunk = 500 # keep below SQLITE_MAX_VARIABLE_NUMBER (999)
2347 with get_user_db_session(
2348 self.username, self.db_password
2349 ) as session:
2350 for i in range(0, len(unresolved), in_chunk):
2351 batch = unresolved[i : i + in_chunk]
2352 all_ids.extend(
2353 row.id
2354 for row in session.query(DocumentChunk.id).filter(
2355 DocumentChunk.source_type == "document",
2356 DocumentChunk.collection_name == collection_name,
2357 DocumentChunk.source_id.in_(batch),
2358 )
2359 )
2361 removed = 0
2362 if all_ids: 2362 ↛ 2372line 2362 didn't jump to line 2372 because the condition on line 2362 was always true
2363 try:
2364 stats = vindex.delete_ids(sorted(set(all_ids)))
2365 removed = stats.removed
2366 except Exception:
2367 logger.warning(
2368 f"Could not remove {len(set(all_ids))} vectors for "
2369 f"{len(wanted)} document(s) from {collection_name}"
2370 )
2372 if removed: 2372 ↛ 2377line 2372 didn't jump to line 2377 because the condition on line 2372 was always true
2373 logger.info(
2374 f"Removed {removed} vectors for "
2375 f"{len(wanted)} document(s) from {collection_name}"
2376 )
2377 return removed