Coverage for src/local_deep_research/research_library/services/library_rag_service.py: 91%
1188 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Library RAG Service
4Handles indexing and searching library documents using RAG:
5- Index text documents into vector database
6- Chunk documents for semantic search
7- Generate embeddings using local models
8- Manage FAISS indices per research
9- Track RAG status in library
10"""
12import threading
13import time
14import json
15import re
16from contextlib import contextmanager
17from dataclasses import dataclass
18from datetime import datetime, UTC
20import numpy as np
21from pathlib import Path
22from typing import (
23 Any,
24 Callable,
25 Dict,
26 List,
27 Literal,
28 Optional,
29 Set,
30 Tuple,
31 Union,
32)
33from concurrent.futures import ThreadPoolExecutor, as_completed, CancelledError
35from langchain_core.documents import Document as LangchainDocument
36from loguru import logger
37from sqlalchemy import exists, func
38from sqlalchemy.exc import IntegrityError
39from sqlalchemy.orm import Session
41from ...config.paths import get_cache_directory
42from ...constants import (
43 DEFAULT_LOCAL_SEARCH_CHUNK_OVERLAP,
44 DEFAULT_LOCAL_SEARCH_CHUNK_SIZE,
45 DEFAULT_LOCAL_SEARCH_DISTANCE_METRIC,
46 DEFAULT_LOCAL_SEARCH_INDEX_TYPE,
47 DEFAULT_LOCAL_SEARCH_MODEL,
48 DEFAULT_LOCAL_SEARCH_NORMALIZE_VECTORS,
49 DEFAULT_LOCAL_SEARCH_PROVIDER,
50 DEFAULT_LOCAL_SEARCH_SPLITTER_TYPE,
51 DEFAULT_LOCAL_SEARCH_TEXT_SEPARATORS,
52)
53from ...database.models.library import (
54 Document,
55 DocumentChunk,
56 DocumentCollection,
57 Collection,
58 RAGIndex,
59 RAGIndexStatus,
60 RagDocumentStatus,
61 EmbeddingProvider,
62)
63from ...database.session_context import get_user_db_session, safe_rollback
64from ...database.thread_local_session import thread_session_manager
65from ...utilities.type_utils import to_bool
66from ..utils import ensure_in_collection
67from ...embeddings.splitters import get_text_splitter
68from ...web_search_engines.engines.local_embedding_manager import (
69 LocalEmbeddingManager,
70)
71from ...security.directory_creation import create_directory
72from ...security.file_integrity import FileIntegrityManager, FAISSIndexVerifier
73from ...security.file_integrity.integrity_manager import NO_INTEGRITY_RECORD
74from ...vector_stores.facade import VectorIndex, ChunkInput, SearchResult
75from .collection_rebuild import (
76 CollectionRebuild,
77 RebuildCancelled,
78 collection_configuration,
79 document_fingerprint,
80)
81from langchain_core.embeddings import Embeddings
82import hashlib
85def _claim_document_snapshot(session, document_id, text_content, document_hash):
86 """Keep the validated SQLite snapshot stable until the index commit.
88 SELECT FOR UPDATE is ignored by SQLite. A conditional no-op write either
89 detects a writer that already won or holds the write lock through chunk
90 and status persistence. Call after vector initialization, which may commit
91 index setup, and never before embedding.
92 """
93 return (
94 session.query(Document)
95 .filter(
96 Document.id == document_id,
97 Document.text_content == text_content,
98 Document.document_hash == document_hash,
99 )
100 .update({Document.id: document_id}, synchronize_session=False)
101 )
104@dataclass
105class _PreparedDocument:
106 document_id: str
107 collection_id: str
108 chunk_inputs: List[ChunkInput]
109 vectors: np.ndarray
110 source_fingerprint: str | None = None
111 collection_configuration: str | None = None
112 content_hash: Optional[str] = None
113 initial_indexed: Optional[bool] = None
114 status_exists: Optional[bool] = None
115 # True iff a DocumentCollection row existed at prepare time. False
116 # means first-time ingest: Phase 3 may create the link. Distinguishes
117 # that from a concurrent unlink of a row that was present then.
118 had_collection_link: Optional[bool] = None
121# Texts per embed_documents() call while a cancellation scope is armed.
122# Embeddings are per-text independent, so splitting changes nothing but
123# adds cancel points between batches. Without a scope the single call
124# below is kept exactly as today (zero behavior change for background /
125# sequential paths that cannot be cancelled).
126_EMBEDDING_CANCEL_BATCH_SIZE = 32
129def _indexing_scope_error() -> Optional[Dict[str, str]]:
130 """Reject borrowed scopes before indexing can commit or roll them back.
132 Indexing commits writes and may recover from failures with a full rollback.
133 A nested scope reuses the caller's Session, so it cannot own either action.
134 Check before entering get_user_db_session, whose exception cleanup also
135 rolls back that Session.
137 Only an open ``get_user_db_session`` block on the calling thread counts
138 (``thread_session_manager.in_scope``). A Session obtained some other way,
139 such as directly from ``get_metrics_session``, is not detected.
140 """
141 if thread_session_manager.in_scope():
142 return {
143 "status": "error",
144 "error": "Indexing must run outside an existing database session",
145 }
146 return None
149# Service settings a Collection stores besides embedding_model, named as
150# LibraryRAGService attributes.
151_COLLECTION_INDEX_SETTINGS = (
152 "embedding_provider",
153 "chunk_size",
154 "chunk_overlap",
155 "splitter_type",
156 "text_separators",
157 "distance_metric",
158 "normalize_vectors",
159 "index_type",
160)
163def resolve_collection_index_settings(
164 collection: Collection, defaults: Dict[str, Any]
165) -> Dict[str, Any]:
166 """Return the settings a service indexing ``collection`` uses.
168 A stored value wins and a NULL column takes ``defaults``. Collections
169 indexed before #1054 never stored the five advanced settings, and
170 migration 0032 adds those columns without a backfill. ``get_rag_service``
171 passes the user's current settings as ``defaults``.
172 """
173 provider = collection.embedding_model_type
174 normalize = collection.normalize_vectors
175 stored = {
176 "embedding_provider": provider.value if provider else None,
177 "chunk_size": collection.chunk_size,
178 "chunk_overlap": collection.chunk_overlap,
179 "splitter_type": collection.splitter_type,
180 "text_separators": collection.text_separators,
181 "distance_metric": collection.distance_metric,
182 "normalize_vectors": None if normalize is None else to_bool(normalize),
183 "index_type": collection.index_type,
184 }
185 return {
186 name: defaults[name] if value is None else value
187 for name, value in stored.items()
188 }
191class _LazyEmbeddings(Embeddings):
192 """Embeddings proxy that defers backend construction until first use.
194 A ``VectorIndex`` built purely to ``delete()`` never touches ``embeddings``
195 (delete resolves rows by id and calls ``store.apply(remove_ids=...)``).
196 ``LocalEmbeddingManager.embeddings`` is a lazy property that, on first
197 access, synchronously loads the real backend (a SentenceTransformer costs
198 seconds and hundreds of MB of RSS, an Ollama client opens live httpx
199 connections). Handing that property straight to the delete path forces that
200 cost once per (deleted-document x collection). This proxy pulls the real
201 embeddings from the manager only if ``embed_*`` is genuinely called, so the
202 delete path pays nothing while index/search paths still work unchanged.
203 """
205 def __init__(self, manager) -> None:
206 self._manager = manager
208 def embed_documents(self, texts):
209 return self._manager.embeddings.embed_documents(texts)
211 def embed_query(self, text):
212 return self._manager.embeddings.embed_query(text)
215_CollectionLockPurpose = Literal["transaction", "rebuild"]
216_IndexLockKey = Tuple[str, Union[str, Tuple[_CollectionLockPurpose, str]]]
219class _TrackedRLock:
220 """RLock proxy that records active keys even for direct acquire() calls.
222 The recursion depth is tracked locally (not via ``threading.RLock._is_owned``,
223 which is a CPython-private API and is not part of the threading contract —
224 absent on PyPy/other implementations). The key set membership is mutated
225 atomically with respect to ``acquire`` and ``pop_faiss_locks_for_user``: the
226 per-(user, index_path) entry is added when ownership transitions 0 → 1 and
227 discarded when it transitions 1 → 0, and both transitions happen under
228 ``_faiss_write_locks_lock``.
230 Without holding ``_faiss_write_locks_lock`` across release, there would be a
231 race window: thread A could release → thread B (blocked in ``acquire``) unblock
232 and add its key → thread A then observe depth 0 and discard B's key. With
233 ``pop_faiss_locks_for_user`` evicting a "free" lock, the next writer would
234 create a fresh lock for the same ``(username, index_path)`` and race B on the
235 FAISS file. Holding ``_faiss_write_locks_lock`` across inner release + depth
236 check + discard ensures no concurrent acquire's key-add can interleave with
237 key-discard.
238 """
240 def __init__(self, key: _IndexLockKey):
241 self._key = key
242 self._lock = threading.RLock()
243 # Per-wrapper recursion depth (independent of the inner RLock so we
244 # never need the CPython-private ``_is_owned`` accessor). ``_owner``
245 # records the outer holder's thread id, or ``None`` when free.
246 # ``_tracks_active_key`` is True only when this wrapper added
247 # ``self._key`` to ``_faiss_active_lock_keys`` on its outer acquire
248 # (canonical wrappers only) so a non-canonical release cannot erase
249 # another holder's active-key entry.
250 self._depth = 0
251 self._owner = None
252 self._tracks_active_key = False
254 def acquire(self, *args, **kwargs):
255 first_acquire = self._lock.acquire(*args, **kwargs)
256 if first_acquire:
257 # ``RLock.acquire`` returns True for the outer acquisition AND for
258 # inner re-acquires. We detect outer-vs-inner via the wrapper's
259 # tracked ``_owner``: if the current thread already holds the lock,
260 # this is a re-entry; otherwise it's a fresh outer acquisition.
261 #
262 # Only the *canonical* wrapper for ``self._key`` (the one currently
263 # stored in ``_faiss_write_locks``) may mark the key active. A
264 # non-canonical wrapper can still be acquired briefly during the
265 # pop-during-acquire TOCTOU window handled by
266 # ``_hold_faiss_write_lock``; if it also added the key, its matching
267 # ``release()`` would discard the key while the *canonical* holder
268 # still owned the file — reopening the concurrent-FAISS-writer race.
269 was_owner = self._owner == threading.get_ident()
270 if not was_owner:
271 with _faiss_write_locks_lock:
272 if _faiss_write_locks.get(self._key) is self:
273 _faiss_active_lock_keys.add(self._key)
274 self._tracks_active_key = True
275 else:
276 self._tracks_active_key = False
277 self._owner = threading.get_ident()
278 self._depth += 1
279 return first_acquire
281 def release(self):
282 # Atomic with respect to acquire's add: hold the key-set lock across
283 # the inner release and the depth check so a concurrent acquire's
284 # key-add cannot interleave with our key-discard. ``_faiss_write_locks_lock``
285 # is taken BEFORE the inner release (the old bug was the reverse order).
286 #
287 # Only discard the active key if *this* wrapper registered it on
288 # acquire. Non-canonical wrappers never register, so their release
289 # cannot erase a concurrent canonical holder's active-key entry.
290 with _faiss_write_locks_lock:
291 self._lock.release()
292 self._depth -= 1
293 if self._depth == 0:
294 if self._tracks_active_key:
295 _faiss_active_lock_keys.discard(self._key)
296 self._tracks_active_key = False
297 self._owner = None
299 def locked(self) -> bool:
300 """True iff the current thread holds the wrapper at depth >= 1."""
301 return self._owner == threading.get_ident() and self._depth > 0
303 def __enter__(self):
304 self.acquire()
305 return self
307 def __exit__(self, *args):
308 self.release()
311# load_or_create verify→quarantine→build sequence per (username, index_path).
312# MUST be module-level: each auto-index / scheduler / search worker constructs
313# its own LibraryRAGService, so an instance-scoped lock would coordinate
314# nothing. Pattern is adapted from web/queue/processor_v2._user_critical_locks
315# — instance-scoped→module-scoped, username-only-key→(username, path)-key.
316# See #4197 for the race this guards (concurrent save_local interleaves bytes,
317# producing checksum_mismatch → destructive unlink, lost data).
318# Logical collection operations share tracked ownership and user cleanup,
319# but use tuple identities distinct from every physical path string.
320_faiss_write_locks: Dict[_IndexLockKey, _TrackedRLock] = {}
321_faiss_write_locks_lock = threading.Lock()
322_faiss_active_lock_keys: Set[_IndexLockKey] = set()
324# Hard cap on suffix-increment retries when generating the .corrupt-<ns> path.
325# Normal case is one attempt — same-ns collisions only happen if the user
326# manually created such files. 32 is a safety bound that converts a deadlock
327# into a loud OSError.
328_QUARANTINE_SUFFIX_RETRY_CAP = 32
330# After a successful quarantine, keep at most this many older .corrupt-*
331# files per base path (per side: .faiss.corrupt-* and .pkl.corrupt-* are
332# counted independently). Prevents unbounded disk growth on systems that
333# experience recurring corruption while preserving recent diagnostic
334# artefacts. Keeping 5 means the user has the last 5 corruption events
335# to inspect / submit with a bug report; anything older is dropped.
336_QUARANTINE_KEEP_RECENT = 5
338# ``FaissVectorStore._prepare`` raises ``ValueError`` with one of these exact,
339# fixed-text messages when it refuses a query vector for a reason OTHER than
340# a dimension mismatch (that one varies by dimension, so it stays a regex in
341# ``search()`` below): non-finite (NaN/Inf), zero/overflow magnitude under L2
342# normalization, and raw magnitude too large for unnormalized float32
343# distance math. ``search()`` matches these verbatim (copied from
344# ``_prepare``, not derived from it) so an unrelated provider ``ValueError``
345# that merely resembles one of them keeps passing through unchanged — same
346# fail-open-by-default posture as the dimension regex.
347_DEGENERATE_QUERY_VECTOR_MESSAGES = frozenset(
348 {
349 "vectors contain non-finite values (NaN/Inf) — bad embedding",
350 (
351 "vectors contain a zero-magnitude or overflow-magnitude row "
352 "that L2 normalization cannot represent — bad embedding"
353 ),
354 (
355 "vectors contain a magnitude too large for float32 distance "
356 "computation (would overflow to a dropped/mis-ranked search "
357 "hit) — normalize the embeddings or use a bounded model"
358 ),
359 }
360)
363def _get_faiss_write_lock(username: str, index_path: str) -> _TrackedRLock:
364 """Return the lock for ``(username, index_path)``, creating it on first
365 access. Key is normalised via ``Path.resolve()`` to match
366 ``FileIntegrityManager._normalize_path`` so writers/readers/quarantine
367 all agree on identity.
368 """
369 key = (username, str(Path(index_path).resolve()))
370 with _faiss_write_locks_lock:
371 lock = _faiss_write_locks.get(key)
372 if lock is None:
373 lock = _TrackedRLock(key)
374 _faiss_write_locks[key] = lock
375 return lock
378@contextmanager
379def _hold_faiss_write_lock(username: str, index_path: str):
380 """Context manager that acquires the per-(username, index_path) FAISS
381 write lock, retrying across a pop-during-acquire TOCTOU window.
383 ``_get_faiss_write_lock`` returns a wrapper reference, but
384 ``pop_faiss_locks_for_user`` may run in the gap before our
385 ``acquire()`` adds the key to ``_faiss_active_lock_keys`` and therefore
386 before the key is "active" from ``pop``'s perspective. Without a re-check,
387 the caller would hold a non-canonical wrapper while ``pop`` discarded the
388 canonical entry and the next writer created a fresh lock for the same
389 ``(username, index_path)`` — reintroducing the same-(user, path)
390 concurrent-FAISS-writer race this module exists to eliminate.
392 The window is microseconds and requires a per-user lock cleanup to fire
393 mid-acquire, so this is a defence-in-depth tightening rather than a
394 regression fix — but it is cheap to close. After ``acquire()`` returns,
395 we re-check under ``_faiss_write_locks_lock`` that ``self`` is still the
396 canonical entry; if not, we release and retry with the new canonical
397 lock. Re-entry on the same wrapper is impossible here because each
398 iteration acquires a fresh wrapper reference.
399 """
400 key = (username, str(Path(index_path).resolve()))
401 while True:
402 lock = _get_faiss_write_lock(username, index_path)
403 lock.acquire()
404 with _faiss_write_locks_lock:
405 canonical = _faiss_write_locks.get(key)
406 is_canonical = canonical is lock
407 if is_canonical:
408 try:
409 yield lock
410 finally:
411 lock.release()
412 return
413 # Evicted between _get_faiss_write_lock's lookup and our acquire().
414 # Release our (non-canonical) wrapper — acquire only marks the key
415 # active for the canonical dict entry, so this release cannot erase
416 # a concurrent canonical holder's active-key entry — then retry
417 # with the canonical lock.
418 lock.release()
421def _get_collection_lock(
422 username: str, collection_id: str, *, purpose: _CollectionLockPurpose
423) -> _TrackedRLock:
424 """Identify a collection operation without interpreting its ID as a path.
426 The tuple namespace cannot collide with physical FAISS path keys. Share
427 tracked-lock cleanup with those locks so closing a user still releases
428 inactive entries while preserving locks held by background workers.
429 """
430 key = (username, (purpose, str(collection_id)))
431 with _faiss_write_locks_lock:
432 lock = _faiss_write_locks.get(key)
433 if lock is None:
434 lock = _TrackedRLock(key)
435 _faiss_write_locks[key] = lock
436 return lock
439@contextmanager
440def _hold_collection_lock(
441 username: str,
442 collection_id: str,
443 *,
444 purpose: _CollectionLockPurpose,
445 blocking: bool = True,
446):
447 """Hold a canonical logical lock, or yield None on nonblocking contention.
449 Retry eviction between lookup and acquisition just as physical FAISS
450 locks do. Collection IDs remain opaque and never reach filesystem APIs.
451 """
452 key = (username, (purpose, str(collection_id)))
453 while True:
454 lock = _get_collection_lock(username, collection_id, purpose=purpose)
455 if not lock.acquire(blocking=blocking):
456 yield None
457 return
458 with _faiss_write_locks_lock:
459 is_canonical = _faiss_write_locks.get(key) is lock
460 if is_canonical:
461 try:
462 yield lock
463 finally:
464 lock.release()
465 return
466 # Evicted before acquire; a stale wrapper cannot protect this key.
467 lock.release()
470def pop_faiss_locks_for_user(username: str) -> None:
471 """Remove inactive physical and logical index locks for a user."""
472 with _faiss_write_locks_lock:
473 stale = [k for k in _faiss_write_locks if k[0] == username]
474 for k in stale:
475 if k not in _faiss_active_lock_keys:
476 _faiss_write_locks.pop(k, None)
479class LibraryRAGService:
480 """Service for managing RAG indexing of library documents."""
482 def __init__(
483 self,
484 username: str,
485 embedding_model: str = DEFAULT_LOCAL_SEARCH_MODEL,
486 embedding_provider: str = DEFAULT_LOCAL_SEARCH_PROVIDER,
487 chunk_size: int = DEFAULT_LOCAL_SEARCH_CHUNK_SIZE,
488 chunk_overlap: int = DEFAULT_LOCAL_SEARCH_CHUNK_OVERLAP,
489 splitter_type: str = DEFAULT_LOCAL_SEARCH_SPLITTER_TYPE,
490 text_separators: Optional[list] = None,
491 distance_metric: str = DEFAULT_LOCAL_SEARCH_DISTANCE_METRIC,
492 normalize_vectors: bool = DEFAULT_LOCAL_SEARCH_NORMALIZE_VECTORS,
493 index_type: str = DEFAULT_LOCAL_SEARCH_INDEX_TYPE,
494 embedding_manager: Optional["LocalEmbeddingManager"] = None,
495 db_password: Optional[str] = None,
496 ):
497 """
498 Initialize library RAG service for a user.
500 Args:
501 username: Username for database access
502 embedding_model: Name of the embedding model to use
503 embedding_provider: Provider type ('sentence_transformers' or 'ollama')
504 chunk_size: Size of text chunks for splitting
505 chunk_overlap: Overlap between consecutive chunks
506 splitter_type: Type of splitter ('recursive', 'token', 'sentence', 'semantic')
507 text_separators: List of text separators for chunking (default: ["\n\n", "\n", ". ", " ", ""])
508 distance_metric: Distance metric ('cosine', 'l2', or 'dot_product')
509 normalize_vectors: Whether to normalize vectors with L2
510 index_type: FAISS index type ('flat', 'hnsw', or 'ivf')
511 embedding_manager: Optional pre-constructed LocalEmbeddingManager for testing/flexibility
512 db_password: Optional database password for background thread access
513 """
514 self.username = username
515 self._db_password = db_password # Can be used for thread access
516 # Initialize optional attributes to None before they're set below
517 # This allows the db_password setter to check them without hasattr
518 self.embedding_manager = None
519 self.integrity_manager = None
520 self.embedding_model = embedding_model
521 self.embedding_provider = embedding_provider
522 self.chunk_size = chunk_size
523 self.chunk_overlap = chunk_overlap
524 self.splitter_type = splitter_type
525 self.text_separators = (
526 text_separators
527 if text_separators is not None
528 else list(DEFAULT_LOCAL_SEARCH_TEXT_SEPARATORS)
529 )
530 self.distance_metric = distance_metric
531 # Ensure normalize_vectors is always a proper boolean
532 self.normalize_vectors = to_bool(normalize_vectors, default=True)
533 self.index_type = index_type
535 # Emit the active configuration so users can confirm their
536 # UI-configured embedding settings are being honored (regression
537 # signal for #3453).
538 logger.info(
539 f"RAG service initialized for user={username}: "
540 f"provider={embedding_provider} model={embedding_model} "
541 f"chunk_size={chunk_size} chunk_overlap={chunk_overlap} "
542 f"splitter={splitter_type} index_type={index_type}"
543 )
545 # Use provided embedding manager or create a new one
546 # (Must be created before text splitter for semantic chunking)
547 # Track ownership so close() only tears down the manager when we
548 # constructed it — a caller-supplied manager stays under caller
549 # control (test fixtures, multi-service callers reusing one manager).
550 self._owns_embedding_manager = embedding_manager is None
551 if embedding_manager is not None:
552 self.embedding_manager = embedding_manager
553 else:
554 # Load the complete user settings snapshot from database using the
555 # proper method.
556 #
557 # STRICT snapshot is mandatory at this seam: a non-strict read
558 # silently falls back to the JSON defaults when the underlying
559 # settings query fails (SQLAlchemyError / stale enum row). Those
560 # defaults can lack an operator-selected cloud embedding
561 # provider's classification inputs and admit a cloud embedder
562 # under a local-only posture. We refuse before any
563 # LocalEmbeddingManager construction instead.
564 from ...settings.manager import SettingsManager
565 from ...security.egress.policy import (
566 Decision,
567 PolicyDeniedError,
568 )
570 # Use proper database session for SettingsManager
571 # Note: using _db_password (backing field) directly here because the
572 # db_password property setter propagates to embedding_manager/integrity_manager,
573 # which are still None at this point in __init__.
574 try:
575 with get_user_db_session(
576 username, self._db_password
577 ) as session:
578 settings_manager = SettingsManager(session)
579 settings_snapshot = settings_manager.get_settings_snapshot(
580 strict=True
581 )
582 except PolicyDeniedError:
583 raise
584 except Exception as exc:
585 raise PolicyDeniedError(
586 Decision(False, "settings_unavailable"),
587 target=embedding_provider,
588 ) from exc
590 # Add the specific settings needed for this RAG service
591 settings_snapshot.update(
592 {
593 "_username": username,
594 "embeddings.provider": embedding_provider,
595 f"embeddings.{embedding_provider}.model": embedding_model,
596 "local_search_chunk_size": chunk_size,
597 "local_search_chunk_overlap": chunk_overlap,
598 }
599 )
601 # Egress policy pre-flight at the constructor boundary so
602 # every direct ``LibraryRAGService(...)`` construction site
603 # is covered, not just the factory. Skipped when an
604 # ``embedding_manager`` is injected (tests / advanced flows)
605 # — those callers vouch for the manager themselves.
606 #
607 # Build the context from the ACTUAL scope (not a hardcoded
608 # BOTH) so PRIVATE_ONLY forces local embeddings even when the
609 # raw embeddings.require_local flag is at its default False —
610 # context_from_snapshot applies that coupling. Resolve the
611 # primary via the shared helper (single source of truth)
612 # instead of the old search.tool + searxng fallback, which
613 # was a fail-OPEN: a missing primary defaulted to the public
614 # searxng so the scope relaxed and a cloud embedder could be
615 # admitted. A missing primary now raises -> fail closed via
616 # the ValueError handler below.
617 from ...security.egress.policy import (
618 context_from_snapshot,
619 evaluate_embeddings,
620 resolve_run_primary_engine,
621 )
623 try:
624 primary = resolve_run_primary_engine(settings_snapshot)
625 policy_ctx = context_from_snapshot(
626 settings_snapshot, primary, username=username
627 )
628 except PolicyDeniedError:
629 raise
630 except ValueError as exc:
631 raise PolicyDeniedError(
632 Decision(False, "invalid_policy_config"),
633 target=embedding_provider,
634 ) from exc
635 if policy_ctx.require_local_embeddings:
636 decision = evaluate_embeddings(
637 embedding_provider,
638 policy_ctx,
639 settings_snapshot=settings_snapshot,
640 )
641 if not decision.allowed: 641 ↛ 642line 641 didn't jump to line 642 because the condition on line 641 was never true
642 logger.bind(policy_audit=True).warning(
643 "LibraryRAGService refused by egress policy",
644 provider=embedding_provider,
645 reason=decision.reason,
646 )
647 raise PolicyDeniedError(decision, target=embedding_provider)
649 self.embedding_manager = LocalEmbeddingManager(
650 embedding_model=embedding_model,
651 embedding_model_type=embedding_provider,
652 settings_snapshot=settings_snapshot,
653 )
655 # Initialize text splitter based on type
656 # (Must be created AFTER embedding_manager for semantic chunking)
657 self.text_splitter = get_text_splitter(
658 splitter_type=self.splitter_type,
659 chunk_size=self.chunk_size,
660 chunk_overlap=self.chunk_overlap,
661 text_separators=self.text_separators,
662 embeddings=self.embedding_manager.embeddings
663 if self.splitter_type == "semantic"
664 else None,
665 )
667 self.rag_index_record = None
669 # Initialize file integrity manager for FAISS indexes
670 self.integrity_manager = FileIntegrityManager(
671 username, password=self._db_password
672 )
673 self.integrity_manager.register_verifier(FAISSIndexVerifier())
675 self._closed = False
677 def close(self):
678 """Release embedding model and index resources."""
679 if self._closed:
680 return
681 self._closed = True
683 # Release embedding manager (which in turn closes the underlying
684 # OllamaEmbeddings httpx clients — see LocalEmbeddingManager.close).
685 # Only when we own it; caller-supplied managers stay under caller
686 # control to avoid double-close / use-after-close.
687 if self.embedding_manager is not None:
688 if self._owns_embedding_manager:
689 self.embedding_manager.close()
690 self.embedding_manager = None
692 # Clear other resources
693 self.rag_index_record = None
694 self.integrity_manager = None
695 self.text_splitter = None
697 def __enter__(self):
698 """Enter context manager."""
699 return self
701 def __exit__(self, exc_type, exc_val, exc_tb):
702 """Exit context manager, ensuring cleanup."""
703 self.close()
704 return False
706 @property
707 def db_password(self):
708 """Get database password."""
709 return self._db_password
711 @db_password.setter
712 def db_password(self, value):
713 """Set database password and propagate to embedding manager and integrity manager."""
714 self._db_password = value
715 if self.embedding_manager:
716 self.embedding_manager.db_password = value
717 if self.integrity_manager:
718 self.integrity_manager.password = value
720 def _get_index_hash(
721 self,
722 collection_name: str,
723 embedding_model: str,
724 embedding_model_type: str,
725 ) -> str:
726 """Generate hash for index identification."""
727 identity = {
728 "chunk_overlap": self.chunk_overlap,
729 "chunk_size": self.chunk_size,
730 "collection_name": collection_name,
731 "distance_metric": self.distance_metric,
732 "embedding_model": embedding_model,
733 "embedding_provider": embedding_model_type,
734 "index_type": self.index_type,
735 "normalize_vectors": self.normalize_vectors,
736 "splitter_type": self.splitter_type,
737 "text_separators": self.text_separators,
738 }
739 hash_input = json.dumps(
740 identity, ensure_ascii=False, separators=(",", ":"), sort_keys=True
741 )
742 return hashlib.sha256(hash_input.encode()).hexdigest()
744 def _matches_collection_configuration(self, collection: Collection) -> bool:
745 """Whether ``collection`` still stores this service's configuration.
747 NULL columns resolve as in ``get_rag_service``, which filled them from
748 the user's settings when it built this service, so they resolve to
749 this service's own values. Only a stored value can differ, such as one
750 a complete rebuild published.
751 """
752 own = {name: getattr(self, name) for name in _COLLECTION_INDEX_SETTINGS}
753 return (
754 collection.embedding_model == self.embedding_model
755 and resolve_collection_index_settings(collection, own) == own
756 )
758 def _matches_current_index_configuration(self, rag_index: RAGIndex) -> bool:
759 stored_provider = rag_index.embedding_model_type
760 if isinstance(stored_provider, EmbeddingProvider): 760 ↛ 762line 760 didn't jump to line 762 because the condition on line 760 was always true
761 stored_provider = stored_provider.value
762 stored_separators = rag_index.text_separators
763 if stored_separators is None:
764 stored_separators = DEFAULT_LOCAL_SEARCH_TEXT_SEPARATORS
765 if not isinstance(stored_separators, list) or not all( 765 ↛ 768line 765 didn't jump to line 768 because the condition on line 765 was never true
766 isinstance(separator, str) for separator in stored_separators
767 ):
768 return False
769 return (
770 rag_index.embedding_model == self.embedding_model
771 and stored_provider == self.embedding_provider
772 and rag_index.chunk_size == self.chunk_size
773 and rag_index.chunk_overlap == self.chunk_overlap
774 and (rag_index.splitter_type or "recursive") == self.splitter_type
775 and stored_separators == self.text_separators
776 and (rag_index.distance_metric or "cosine") == self.distance_metric
777 and (
778 rag_index.normalize_vectors
779 if rag_index.normalize_vectors is not None
780 else True
781 )
782 == self.normalize_vectors
783 and (rag_index.index_type or "flat") == self.index_type
784 )
786 def _find_matching_rag_index(
787 self, db_session: Session, collection_name: str, index_hash: str
788 ) -> RAGIndex | None:
789 rag_index = (
790 db_session.query(RAGIndex).filter_by(index_hash=index_hash).first()
791 )
792 if rag_index is not None:
793 return rag_index
795 candidates = (
796 db_session.query(RAGIndex)
797 .filter_by(
798 collection_name=collection_name,
799 embedding_model=self.embedding_model,
800 embedding_model_type=EmbeddingProvider(self.embedding_provider),
801 )
802 .all()
803 )
804 for candidate in candidates:
805 if self._matches_current_index_configuration(candidate):
806 return candidate
807 return None
809 def _resolve_index_for_config(
810 self, db_session: Session, collection_id: str
811 ) -> RAGIndex | None:
812 """The collection's index for this service's configuration, or None.
814 Read-only: unlike ``_get_or_create_rag_index`` it never creates a
815 row or probes the embedding provider.
816 """
817 collection_name = f"collection_{collection_id}"
818 return self._find_matching_rag_index(
819 db_session,
820 collection_name,
821 self._get_index_hash(
822 collection_name, self.embedding_model, self.embedding_provider
823 ),
824 )
826 def _get_index_path(self, index_hash: str) -> Path:
827 """Get path for FAISS index file.
829 Files are scoped to a per-user subdirectory derived from a hash
830 of the username so two users with identical collection name +
831 embedding model don't share the same .faiss/.pkl files on disk.
832 The previous layout used a shared `cache/rag_indices/` for all
833 users, which would let one user's vectors land in another
834 user's load path in any deployment running more than one
835 account against the same data dir.
837 Raises:
838 ValueError: If self.username is empty/None. sha256("") is a
839 known constant, so silently hashing an empty username would
840 let two different "no username" service instances collide
841 into one shared cache directory instead of failing loudly.
842 Defense-in-depth -- see #5481.
843 """
844 if not self.username:
845 raise ValueError(
846 "Cannot compute a RAG index path for an empty username"
847 )
848 # Store in centralized cache directory (respects LDR_DATA_DIR)
849 user_scope = hashlib.sha256(self.username.encode("utf-8")).hexdigest()[
850 :16
851 ]
852 rag_root = get_cache_directory() / "rag_indices"
853 cache_dir = rag_root / user_scope
854 create_directory(
855 cache_dir, context="per-user RAG index cache directory"
856 )
857 # 0o700: filesystem-level defense-in-depth so other local OS accounts
858 # can't browse another LDR user's vector caches. mkdir(parents=True)
859 # silently creates the intermediate rag_indices/ root under the process
860 # umask (typically 0o755, world-traversable), so harden the ROOT too —
861 # chmod'ing only the per-user leaf leaves a 0o755 rag_indices/ that lets
862 # another OS user traverse in and stat/enumerate every account's caches.
863 for d in (rag_root, cache_dir):
864 try:
865 d.chmod(0o700)
866 except OSError:
867 # Best-effort: some filesystems (e.g. Windows shares) don't
868 # honor chmod; the per-user path scope is still applied.
869 pass
870 return cache_dir / f"{index_hash}.faiss"
872 def _migrate_legacy_index_files(
873 self, legacy_path: Path, index_path: Path, rag_index: RAGIndex
874 ) -> None:
875 """Relocate a pre-per-user-scoping index into the per-user directory.
877 Before the per-user path scoping, indexes lived directly in the
878 shared ``cache/rag_indices/`` directory. ``_preflight_index_path``
879 recomputes the path and ignores the stored one, so without this
880 one-time migration an upgrading user's existing index is never
881 found again: semantic search silently returns nothing while
882 ``DocumentCollection.indexed`` still claims everything is indexed,
883 so a non-force re-index no-ops too. Only files directly inside the
884 known legacy shared directory with the expected hash filename are
885 moved — arbitrary stored paths are not followed.
886 """
887 legacy_root = get_cache_directory() / "rag_indices"
888 if (
889 legacy_path == index_path
890 or legacy_path.parent != legacy_root
891 or legacy_path.name != index_path.name
892 or not legacy_path.exists()
893 ):
894 return
895 with _hold_faiss_write_lock(self.username, str(index_path)):
896 if index_path.exists(): 896 ↛ 897line 896 didn't jump to line 897 because the condition on line 896 was never true
897 return
898 try:
899 legacy_pkl = legacy_path.with_suffix(".pkl")
900 # Relocate phase-1's text-free .idmap.json sidecar BEFORE the
901 # .faiss, then move the .faiss LAST. phase-2 rekey looks for the
902 # sidecar next to the .faiss (legacy_rekey._sidecar_for =
903 # <stem>.idmap.json); leaving it at the legacy location orphans
904 # it, and a not-yet-rekeyed index then has no findable sidecar →
905 # rekey sees "no sidecar + old format" and destructively
906 # quarantines a healthy index. Ordering matters for crash
907 # recovery: the re-entry guard keys ONLY on index_path (the
908 # .faiss), so moving the .faiss last makes this idempotent — a
909 # crash after the sidecar move but before the .faiss move leaves
910 # the .faiss at the legacy path (guard still False) and re-entry
911 # simply retries (the sidecar move no-ops, new one exists). The
912 # reverse order would strand index_path.exists()=True with the
913 # sidecar orphaned at the legacy root, which the guard never
914 # re-invokes this to fix.
915 legacy_sidecar = legacy_path.with_name(
916 legacy_path.stem + ".idmap.json"
917 )
918 if legacy_sidecar.exists():
919 new_sidecar = index_path.with_name(
920 index_path.stem + ".idmap.json"
921 )
922 if not new_sidecar.exists(): 922 ↛ 925line 922 didn't jump to line 925 because the condition on line 922 was always true
923 # nosemgrep: semgrep.rules.path-traversal-risk, reason: The legacy parent and hash filename are checked above; both fixed-suffix sidecars derive from that path and the service-generated per-user index path
924 legacy_sidecar.rename(new_sidecar)
925 legacy_path.rename(index_path)
926 # DELETE (don't relocate) any legacy plaintext .pkl companion:
927 # the new store writes NO .pkl, so moving it would just carry the
928 # plaintext chunk text to the new location. Its text lives in the
929 # encrypted DB.
930 if legacy_pkl.exists():
931 try:
932 legacy_pkl.unlink()
933 except OSError:
934 logger.warning(
935 f"Could not remove legacy plaintext .pkl "
936 f"{legacy_pkl} — delete it manually"
937 )
938 except OSError:
939 logger.exception(
940 f"Failed to migrate legacy FAISS index "
941 f"{legacy_path} -> {index_path}"
942 )
943 return
944 # Persist the relocated path so the row stops pointing at the
945 # now-empty legacy location.
946 with get_user_db_session(self.username, self.db_password) as session:
947 idx = session.query(RAGIndex).filter_by(id=rag_index.id).first()
948 if idx: 948 ↛ 951line 948 didn't jump to line 951
949 idx.index_path = str(index_path)
950 session.commit()
951 try:
952 # Refresh the integrity record under the new path (the old
953 # record is keyed by the legacy path and no longer applies).
954 self.integrity_manager.record_file(
955 index_path,
956 related_entity_type="rag_index",
957 related_entity_id=rag_index.id,
958 )
959 except Exception:
960 # Non-fatal: verify_file() creates a record on demand for
961 # unknown paths, so the subsequent load still succeeds.
962 logger.exception(
963 "Failed to refresh integrity record after index migration"
964 )
965 logger.info(
966 f"Migrated legacy shared FAISS index to per-user location: "
967 f"{index_path}"
968 )
970 def _reset_index_state_for_rebuild(
971 self, collection_id: str, rag_index: RAGIndex
972 ) -> None:
973 """Reset per-document indexed state when the index file is gone.
975 Reaching the fresh-build branch with DB rows that still claim
976 indexed content means the on-disk index was lost (missing after an
977 upgrade, quarantined, or deleted on dimension mismatch). Without
978 this reset, ``index_document``/``index_collection`` skip every
979 document (``DocumentCollection.indexed`` is still True), so the
980 only repair is an undocumented force re-index while search
981 silently returns nothing. Clearing the flags makes the regular
982 re-index path rebuild for real and makes the UI show the honest
983 unindexed state. Worst case of the narrow race with an in-flight
984 indexer is re-submitting chunks the merge step dedups by id.
985 """
986 with get_user_db_session(self.username, self.db_password) as session:
987 idx = session.query(RAGIndex).filter_by(id=rag_index.id).first()
988 claims_content = bool(
989 idx and (idx.chunk_count or idx.total_documents)
990 )
991 stale_status = (
992 session.query(RagDocumentStatus)
993 .filter_by(rag_index_id=rag_index.id)
994 .count()
995 )
996 stale_memberships = (
997 session.query(DocumentCollection)
998 .filter_by(collection_id=collection_id, indexed=True)
999 .count()
1000 )
1001 if not (claims_content or stale_status or stale_memberships):
1002 # Genuinely new index — nothing to reset.
1003 return
1004 if idx: 1004 ↛ 1007line 1004 didn't jump to line 1007 because the condition on line 1004 was always true
1005 idx.chunk_count = 0
1006 idx.total_documents = 0
1007 session.query(RagDocumentStatus).filter_by(
1008 rag_index_id=rag_index.id
1009 ).delete()
1010 session.query(DocumentCollection).filter_by(
1011 collection_id=collection_id, indexed=True
1012 ).update({"indexed": False, "chunk_count": 0})
1013 session.commit()
1014 logger.warning(
1015 f"FAISS index file for collection {collection_id} is "
1016 f"missing; reset indexed state for {stale_memberships} "
1017 f"document(s) so the next index run rebuilds the index."
1018 )
1020 def _get_or_create_rag_index(
1021 self,
1022 collection_id: str,
1023 *,
1024 promote_current: bool = True,
1025 db_session: Session | None = None,
1026 commit: bool = True,
1027 ) -> RAGIndex:
1028 """Get or create RAGIndex record for the current configuration.
1030 ``promote_current`` (default True) controls whether the resolved row,
1031 reused or just created, is promoted to ``is_current`` (demoting
1032 others). Read paths (search) pass False: a mere search under a
1033 stale/different config must NOT silently flip which index the
1034 collection considers current; only an actual write (index) should.
1035 A row created on a read path is still persisted, but stays
1036 non-current until an indexing run promotes it.
1037 """
1038 if db_session is None:
1039 with get_user_db_session(
1040 self.username, self.db_password
1041 ) as session:
1042 return self._get_or_create_rag_index_in_session(
1043 session,
1044 collection_id,
1045 promote_current=promote_current,
1046 commit=commit,
1047 )
1048 return self._get_or_create_rag_index_in_session(
1049 db_session,
1050 collection_id,
1051 promote_current=promote_current,
1052 commit=commit,
1053 )
1055 def _get_or_create_rag_index_in_session(
1056 self,
1057 db_session: Session,
1058 collection_id: str,
1059 *,
1060 promote_current: bool,
1061 commit: bool,
1062 ) -> RAGIndex:
1063 collection_name = f"collection_{collection_id}"
1064 index_hash = self._get_index_hash(
1065 collection_name, self.embedding_model, self.embedding_provider
1066 )
1067 rag_index = self._find_matching_rag_index(
1068 db_session, collection_name, index_hash
1069 )
1070 created = False
1072 if rag_index is None:
1073 index_path = self._get_index_path(index_hash)
1074 embedding_manager = self.embedding_manager
1075 if embedding_manager is None: 1075 ↛ 1076line 1075 didn't jump to line 1076 because the condition on line 1075 was never true
1076 raise RuntimeError(
1077 "Cannot create an index after closing the service"
1078 )
1079 embedding_dim = len(
1080 embedding_manager.embeddings.embed_query("test")
1081 )
1082 candidate = RAGIndex(
1083 collection_name=collection_name,
1084 embedding_model=self.embedding_model,
1085 embedding_model_type=EmbeddingProvider(self.embedding_provider),
1086 embedding_dimension=embedding_dim,
1087 index_path=str(index_path),
1088 index_hash=index_hash,
1089 chunk_size=self.chunk_size,
1090 chunk_overlap=self.chunk_overlap,
1091 splitter_type=self.splitter_type,
1092 text_separators=self.text_separators,
1093 distance_metric=self.distance_metric,
1094 normalize_vectors=self.normalize_vectors,
1095 index_type=self.index_type,
1096 chunk_count=0,
1097 total_documents=0,
1098 status=RAGIndexStatus.ACTIVE,
1099 is_current=False,
1100 )
1101 try:
1102 with db_session.begin_nested():
1103 db_session.add(candidate)
1104 db_session.flush()
1105 except IntegrityError:
1106 rag_index = self._find_matching_rag_index(
1107 db_session, collection_name, index_hash
1108 )
1109 if rag_index is None: 1109 ↛ 1110line 1109 didn't jump to line 1110 because the condition on line 1109 was never true
1110 raise
1111 logger.info(
1112 f"Using concurrently created RAG index: {index_hash}"
1113 )
1114 else:
1115 rag_index = candidate
1116 created = True
1117 logger.info(f"Created new RAG index: {index_hash}")
1119 # Only a write path (promote_current=True) may move the collection's
1120 # is_current pointer. A read path that had to create the row leaves
1121 # it non-current: the collection's real index must keep serving every
1122 # other caller until an actual indexing run switches over.
1123 promoted = promote_current and not rag_index.is_current
1124 if promoted:
1125 db_session.query(RAGIndex).filter(
1126 RAGIndex.collection_name == collection_name,
1127 RAGIndex.is_current.is_(True),
1128 RAGIndex.id != rag_index.id,
1129 ).update({"is_current": False})
1130 rag_index.is_current = True
1132 # Only commit when this call actually created or promoted an index.
1133 # A pure read-through of an existing current index must stay a no-op
1134 # so it can't consume a caller's pending write (e.g. the fault the
1135 # index-finalize self-heal backstop guards against).
1136 if commit and (created or promoted):
1137 db_session.commit()
1138 db_session.refresh(rag_index)
1139 return rag_index
1141 def _quarantine_corrupt_index(self, index_path: Path, reason: str) -> None:
1142 """Rename a corrupted FAISS index to ``<path>.corrupt-<ns>``
1143 instead of deleting it; any legacy plaintext ``.pkl`` companion
1144 is deleted outright rather than preserved.
1146 Preserves user data for inspection/recovery. The
1147 dimension-mismatch branch elsewhere in this method intentionally
1148 *deletes* — that case rebuilds from scratch and the old bytes
1149 are unreadable with the new model. This helper is for the two
1150 "transient or unknown failure" branches (verify_file said no,
1151 loading the index into a ``VectorIndex``/faiss raised) where the
1152 on-disk bytes may still be usable by a human.
1154 The rebuilt index is re-recorded by the caller (``_record``
1155 through ``FileIntegrityManager.record_file``), which resets the
1156 path's verification streaks (``consecutive_failures`` /
1157 ``consecutive_successes``) — the streak accumulated against the
1158 quarantined bytes must not be charged to the new ones.
1160 Raises ``OSError`` on rename failure (disk full, read-only fs,
1161 permission denied). Re-raising prevents silent data loss: if we
1162 swallowed the error, the next ``persist()``/``faiss.write_index``
1163 would overwrite the corrupt bytes anyway. Caller paths log the
1164 exception via the broader try/except around their indexer call.
1165 """
1166 ns = time.time_ns()
1167 pkl_path = index_path.with_suffix(".pkl")
1169 # Same-nanosecond collisions essentially can't happen between
1170 # concurrent threads (we hold the per-path lock), but a user
1171 # could manually create such files. Loop with a hard cap so a
1172 # weird state surfaces as a clean OSError rather than a hang.
1173 suffix_n = 0
1174 faiss_target = Path(f"{index_path}.corrupt-{ns}")
1175 pkl_target = Path(f"{pkl_path}.corrupt-{ns}")
1176 while faiss_target.exists() or pkl_target.exists():
1177 suffix_n += 1
1178 if suffix_n > _QUARANTINE_SUFFIX_RETRY_CAP:
1179 raise OSError(
1180 f"Quarantine path collisions exceeded "
1181 f"{_QUARANTINE_SUFFIX_RETRY_CAP} retries for "
1182 f"{index_path}"
1183 )
1184 faiss_target = Path(f"{index_path}.corrupt-{ns}-{suffix_n}")
1185 pkl_target = Path(f"{pkl_path}.corrupt-{ns}-{suffix_n}")
1187 try:
1188 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Callers supply the index path recomputed by _preflight_index_path; quarantine appends only a generated timestamp and collision counter in the same directory
1189 index_path.rename(faiss_target)
1190 logger.warning(
1191 f"Quarantined corrupted FAISS index to {faiss_target} "
1192 f"(reason: {reason}). Searches against this collection "
1193 f"will return empty results until the index is rebuilt. "
1194 f"Document chunks are preserved in the database — "
1195 f"trigger 'Re-index Collection' (or set "
1196 f"research_library.auto_index_enabled=true and re-run "
1197 f"indexing) to recover."
1198 )
1199 if pkl_path.exists():
1200 # DELETE (not rename) any .pkl companion. The new store writes
1201 # NO .pkl, so a .pkl next to a .faiss is a LEGACY plaintext
1202 # docstore — its chunk text lives in the (encrypted) DB, so
1203 # renaming it aside to .corrupt-<ns> would needlessly preserve
1204 # plaintext chunk text on disk. Plaintext removal is not optional.
1205 try:
1206 pkl_path.unlink()
1207 logger.info(f"Removed legacy plaintext PKL {pkl_path}")
1208 except OSError:
1209 logger.warning(
1210 f"Could not remove legacy plaintext PKL {pkl_path} — "
1211 "delete it manually"
1212 )
1213 else:
1214 # Missing .pkl is the normal case for the new store. Not
1215 # re-raised: the .faiss is already preserved.
1216 logger.debug(f"No PKL companion at {pkl_path}.")
1217 except OSError:
1218 if not index_path.exists():
1219 # The source is already gone — most likely another worker
1220 # PROCESS quarantined the same corrupt index concurrently (the
1221 # per-path lock is a threading.Lock, so it serialises only
1222 # threads within ONE process, not across gunicorn workers /
1223 # the scheduler process). The objective — get the corrupt bytes
1224 # off index_path so the caller rebuilds fresh — is already met,
1225 # so don't abort a healthy rebuild with a spurious
1226 # FileNotFoundError. Still purge any lingering plaintext .pkl:
1227 # the no-plaintext-at-rest invariant is not optional even here.
1228 if pkl_path.exists(): 1228 ↛ 1236line 1228 didn't jump to line 1236 because the condition on line 1228 was always true
1229 try:
1230 pkl_path.unlink()
1231 except OSError:
1232 logger.warning(
1233 f"Could not remove legacy plaintext PKL {pkl_path} "
1234 "— delete it manually"
1235 )
1236 logger.warning(
1237 f"Corrupt index {index_path} was already quarantined by "
1238 f"another worker (reason: {reason}); proceeding to rebuild."
1239 )
1240 return
1241 # A genuine failure (disk-full, read-only fs, permission denied)
1242 # with the corrupt bytes STILL in place. Re-raise so the caller
1243 # surfaces a real failure instead of letting the next
1244 # persist()/write_index overwrite the corrupt bytes.
1245 logger.exception(
1246 f"Failed to quarantine corrupted index at {index_path}"
1247 )
1248 raise
1250 # Best-effort retention sweep. The quarantine itself succeeded
1251 # above; failing to prune older files is not a correctness
1252 # issue, just a disk-usage one — log and move on.
1253 self._prune_old_quarantined_files(index_path)
1255 @staticmethod
1256 def _corrupt_sort_key(path: Path) -> Tuple[int, int]:
1257 """Extract ``(ns, suffix_n)`` from a ``.corrupt-<ns>[-<n>]``
1258 filename so retention can sort by the monotonic nanosecond
1259 suffix the quarantine path embeds — *not* by ``st_mtime``.
1261 Filesystem timestamp granularity is sometimes 1s or 2s
1262 (FAT32/ext3/SMB shares), making mtime ordering non-deterministic
1263 when multiple quarantines happen within one tick. The
1264 ``.corrupt-<ns>`` suffix carries the original ``time.time_ns()``
1265 from the quarantine and is reliable across all filesystems.
1267 Returns ``(-1, -1)`` for malformed names (manually-placed
1268 files) so they sort below any real entry under ``reverse=True``
1269 — i.e., they get pruned first.
1270 """
1271 name = path.name
1272 marker = ".corrupt-"
1273 idx = name.rfind(marker)
1274 if idx == -1: 1274 ↛ 1275line 1274 didn't jump to line 1275 because the condition on line 1274 was never true
1275 return (-1, -1)
1276 tail = name[idx + len(marker) :]
1277 parts = tail.split("-")
1278 try:
1279 ns = int(parts[0])
1280 except ValueError:
1281 return (-1, -1)
1282 suffix_n = 0
1283 if len(parts) > 1:
1284 try:
1285 suffix_n = int(parts[-1])
1286 except ValueError:
1287 # Unknown trailing component — keep the file but at the
1288 # base ns ordering.
1289 pass
1290 return (ns, suffix_n)
1292 def _prune_old_quarantined_files(self, index_path: Path) -> None:
1293 """Keep only the ``_QUARANTINE_KEEP_RECENT`` most-recent
1294 ``.corrupt-*`` files for ``index_path`` and its ``.pkl``
1295 companion. Sweeps the two sides independently — pairs share
1296 the same ``-<ns>`` suffix so they're ordered identically.
1298 Ordering uses the embedded ``-<ns>`` from the quarantine
1299 filename, not ``st_mtime``: file systems with 1-2s timestamp
1300 granularity can otherwise produce non-deterministic retention
1301 on bursts.
1303 Best-effort: logs and swallows any error so a sweep failure
1304 never propagates back into the indexing path.
1305 """
1306 parent = index_path.parent
1307 pkl_path = index_path.with_suffix(".pkl")
1309 for base in (index_path, pkl_path):
1310 pattern = f"{base.name}.corrupt-*"
1311 try:
1312 # Sort newest-first by the embedded -<ns>; everything
1313 # past the keep window is stale.
1314 candidates = sorted(
1315 parent.glob(pattern),
1316 key=self._corrupt_sort_key,
1317 reverse=True,
1318 )
1319 except OSError:
1320 logger.warning(
1321 f"Failed to enumerate {pattern} in {parent} for "
1322 f"quarantine retention sweep"
1323 )
1324 continue
1326 for stale in candidates[_QUARANTINE_KEEP_RECENT:]:
1327 try:
1328 stale.unlink()
1329 logger.info(f"Pruned old quarantined file: {stale}")
1330 except OSError:
1331 logger.warning(
1332 f"Failed to prune quarantined file {stale}; continuing"
1333 )
1335 def _preflight_index_path(
1336 self,
1337 collection_id: str,
1338 rag_index: RAGIndex,
1339 *,
1340 reset_stale_state: bool = False,
1341 probe_dimension: bool | None = None,
1342 ) -> Path:
1343 """Verify -> quarantine -> dimension-check pre-flight for ``rag_index``.
1345 Extracted from the pre-cutover ``load_or_create_faiss_index``
1346 (deleted — the old LangChain-FAISS path is gone) so every
1347 index/search/delete call site shares ONE recovery path instead of
1348 duplicating it. Mutates ``rag_index.index_path`` / ``rag_index.embedding_dimension``
1349 in place — matching the pre-refactor behavior of stamping the
1350 recomputed path/dimension back onto the in-memory record so callers
1351 that read ``self.rag_index_record`` afterwards see the current
1352 values — and returns the resolved on-disk path.
1354 ``probe_dimension`` decides whether the embedding-dimension drift
1355 probe runs (see the ``verified and should_probe_dimension`` block
1356 below). ``None`` (the default) keeps the original rule — probe iff
1357 ``reset_stale_state`` (a WRITE path). Pass an explicit ``True``/
1358 ``False`` to override that rule for a call site that is not itself
1359 a write but still needs (or must skip) the probe — e.g.
1360 ``reconcile_collection_index`` passes ``True``: unlike ``search()``,
1361 it never embeds a query, so it gets no free dimension check from
1362 the vector store's ``_prepare`` and would otherwise miss drift
1363 entirely.
1365 Deliberately does NOT load or construct a vector store: that stays
1366 the caller's job (via ``VectorIndex(...)``), so a corrupt/foreign/
1367 dimension-mismatched file surfaces as an exception from THAT
1368 construction, which the caller catches to quarantine + rebuild (see
1369 ``_get_vector_index``).
1370 """
1371 index_path = self._get_index_path(rag_index.index_hash)
1372 stored_path = (
1373 Path(rag_index.index_path) if rag_index.index_path else None
1374 )
1375 rag_index.index_path = str(index_path)
1376 if not index_path.exists() and stored_path is not None:
1377 self._migrate_legacy_index_files(stored_path, index_path, rag_index)
1379 # Hold the per-(username, index_path) write lock across the entire
1380 # verify -> quarantine -> dimension-check sequence — see #4197 (a
1381 # narrower scope would let a concurrent writer's save race the
1382 # verification).
1383 if index_path.exists():
1384 lock = _get_faiss_write_lock(self.username, str(index_path))
1385 with lock:
1386 verified, reason = self.integrity_manager.verify_file(
1387 index_path
1388 )
1389 if not verified:
1390 if reason == NO_INTEGRITY_RECORD: 1390 ↛ 1409line 1390 didn't jump to line 1409 because the condition on line 1390 was never true
1391 # A MISSING record is not proven corruption, and unlike
1392 # checksum_mismatch it cannot be attacker-induced: the
1393 # record lives in the per-user ENCRYPTED DB, so removing
1394 # it needs the DB password (which also decrypts
1395 # everything). It only arises from our own crash between
1396 # the index write and record_file (see legacy_rekey
1397 # Phase B) or a pre-integrity-feature index. ADOPT the
1398 # byte-present file by recording its integrity now — on
1399 # BOTH read and write paths. Adoption is non-destructive
1400 # (it only records a checksum), so a read path may do it:
1401 # refusing here instead would hard-fail EVERY search on a
1402 # byte-healthy, missing-record index forever (a read-only
1403 # collection is never revisited by a write op that would
1404 # otherwise adopt it). If the bytes are actually
1405 # unreadable/foreign/wrong-dim, the VectorIndex load below
1406 # still raises and the caller handles it — corruption is
1407 # caught either way. Only checksum_mismatch (genuine
1408 # tamper) is refused/quarantined below.
1409 logger.warning(
1410 f"No integrity record for {index_path}; adopting the "
1411 "existing index (recording its integrity) rather than "
1412 "quarantining — a missing record is not corruption."
1413 )
1414 self.integrity_manager.record_file(
1415 index_path,
1416 related_entity_type="rag_index",
1417 related_entity_id=rag_index.id,
1418 )
1419 # Adopted = now trusted: fall through — a WRITE
1420 # path (reset_stale_state=True) then probes the
1421 # embedding dimension and rebuilds on drift; a READ
1422 # path skips the probe (the vector store's
1423 # ``_prepare`` rejects a wrong-dimension query).
1424 verified = True
1425 elif not reset_stale_state:
1426 # Read path (search): refuse to serve an index that
1427 # failed verification for a reason OTHER than a missing
1428 # record (a checksum_mismatch = genuine tamper, or a
1429 # transient read failure), but do NOT destroy it —
1430 # quarantine is a WRITE-path recovery decision. A
1431 # subsequent index/write operation quarantines + rebuilds.
1432 raise RuntimeError(
1433 f"Integrity verification failed for {index_path}: "
1434 f"{reason} — refusing on a read path (a subsequent "
1435 "index/write operation will quarantine + rebuild)"
1436 )
1437 elif reason and reason.startswith( 1437 ↛ 1446line 1437 didn't jump to line 1446 because the condition on line 1437 was never true
1438 "checksum_calculation_failed"
1439 ):
1440 # Transient: could not READ the file to compute its
1441 # checksum (disk hiccup, momentary lock). This says
1442 # NOTHING about tampering — unlike checksum_mismatch — so
1443 # quarantining would force a needless full re-embed of a
1444 # healthy collection. Raise so the operation retries (like
1445 # the read path) instead of destroying the index.
1446 raise RuntimeError(
1447 f"Integrity check could not read {index_path} "
1448 f"({reason}) — refusing to quarantine on a transient "
1449 "read failure; retry the operation."
1450 )
1451 else:
1452 logger.error(
1453 f"Integrity verification failed for {index_path}: "
1454 f"{reason}. Quarantining for recovery; creating "
1455 f"new index."
1456 )
1457 self._quarantine_corrupt_index(index_path, reason)
1458 should_probe_dimension = (
1459 reset_stale_state
1460 if probe_dimension is None
1461 else probe_dimension
1462 )
1463 if verified and should_probe_dimension:
1464 # Probe the embedding model OUTSIDE the quarantine
1465 # decision (WRITE paths only, by default): embed_query
1466 # fails when the provider itself is unreachable (e.g.
1467 # Ollama down), which says nothing about the index
1468 # file's health, and must propagate rather than trigger
1469 # a destructive rebuild of a healthy index. search()
1470 # skips the probe entirely — a wrong-dimension query is
1471 # rejected by the vector store's ``_prepare`` (see
1472 # FaissVectorStore) so reads stay fail-closed without
1473 # the round-trip. ``reconcile_collection_index`` is the
1474 # one exception: it never embeds a query (it only calls
1475 # ``live_ids()``), so it gets none of that free
1476 # protection from ``_prepare`` — it passes
1477 # ``probe_dimension=True`` explicitly to keep detecting
1478 # drift despite ``reset_stale_state=False``.
1479 current_dim = len(
1480 self.embedding_manager.embeddings.embed_query(
1481 "dimension_check"
1482 )
1483 )
1484 stored_dim = rag_index.embedding_dimension
1485 if stored_dim and current_dim != stored_dim:
1486 logger.warning(
1487 f"Embedding dimension mismatch detected! "
1488 f"Index created with dim={stored_dim}, "
1489 f"current model returns dim={current_dim}. "
1490 f"Deleting old index and rebuilding."
1491 )
1492 try:
1493 index_path.unlink()
1494 pkl_path = index_path.with_suffix(".pkl")
1495 if pkl_path.exists(): 1495 ↛ 1496line 1495 didn't jump to line 1496 because the condition on line 1495 was never true
1496 pkl_path.unlink()
1497 logger.info(
1498 f"Deleted old index files at {index_path}"
1499 )
1500 except Exception:
1501 logger.exception("Failed to delete old index files")
1503 with get_user_db_session(
1504 self.username, self.db_password
1505 ) as session:
1506 idx = (
1507 session.query(RAGIndex)
1508 .filter_by(id=rag_index.id)
1509 .first()
1510 )
1511 if idx: 1511 ↛ 1523line 1511 didn't jump to line 1523
1512 idx.embedding_dimension = current_dim
1513 session.commit()
1514 logger.info(
1515 f"Updated RAGIndex dimension to {current_dim}"
1516 )
1517 # Reset RAGIndex counts, RagDocumentStatus, and
1518 # DocumentCollection.indexed together — a hand-rolled
1519 # partial reset here previously left
1520 # DocumentCollection.indexed=True, which makes
1521 # index_document/index_collection skip every document
1522 # and search silently return nothing after the rebuild.
1523 self._reset_index_state_for_rebuild(
1524 collection_id, rag_index
1525 )
1526 rag_index.embedding_dimension = current_dim
1528 if not index_path.exists():
1529 # The new per-user path is missing — but that alone does NOT mean the
1530 # index is gone. A one-time legacy-path relocation
1531 # (_migrate_legacy_index_files above) can fail transiently (a rename
1532 # hitting disk-full, an AV/backup file lock, a busy network mount),
1533 # leaving the bytes fully intact at the stored legacy path. Treating
1534 # that as "gone" would wipe a healthy collection's indexed bookkeeping
1535 # — even on a mere search — and force a needless full re-embed. Only
1536 # self-heal when the bytes are absent at BOTH the new and the legacy
1537 # location; otherwise leave the state alone so the next write-path
1538 # access retries the relocation.
1539 legacy_bytes_present = bool(
1540 stored_path is not None
1541 and stored_path != index_path
1542 and stored_path.exists()
1543 )
1544 if not legacy_bytes_present:
1545 stale_indexed = bool(
1546 (rag_index.chunk_count or 0) > 0
1547 or (rag_index.total_documents or 0) > 0
1548 )
1549 if stale_indexed:
1550 # DB says indexed, but the file is GONE at both locations
1551 # (cache dir cleared, restored without the cache dir, lost
1552 # volume, crash between quarantine and rebuild). Reset so it
1553 # self-heals — even on a READ path: an absent file has
1554 # nothing to preserve, and without this search silently
1555 # returns [] forever (get_rag_stats still reports 100%
1556 # indexed and a normal reindex short-circuits on
1557 # indexed=True, so nothing ever recovers it).
1558 logger.warning(
1559 f"RAG index file for collection {collection_id} is "
1560 f"MISSING at {index_path} while the collection is marked "
1561 "indexed — resetting indexed state so it rebuilds "
1562 "(search would otherwise return no results indefinitely)."
1563 )
1564 self._reset_index_state_for_rebuild(
1565 collection_id, rag_index
1566 )
1567 elif reset_stale_state:
1568 self._reset_index_state_for_rebuild(
1569 collection_id, rag_index
1570 )
1572 return index_path
1574 def _get_vector_index(
1575 self,
1576 collection_id: str,
1577 collection_name: str,
1578 *,
1579 reset_stale_state: bool = False,
1580 probe_dimension: bool | None = None,
1581 ) -> VectorIndex:
1582 """Build a :class:`VectorIndex` for ``collection_id``.
1584 Single call site for the index/search/delete paths: runs the
1585 verify/quarantine/dimension-check pre-flight, then constructs the
1586 facade. On a ``ValueError``/``RuntimeError`` from the underlying
1587 store's ``load()`` (corrupt bytes, a foreign/pre-migration format
1588 that isn't an ``IndexIDMap2``, or a dimension mismatch the
1589 pre-flight didn't already catch) the corrupt file is quarantined,
1590 stale per-document indexed state is reset, and construction is
1591 retried once as a fresh, empty store.
1593 NOTE: like ``_get_or_create_rag_index``, this CREATES a RAGIndex
1594 row if none exists for the configured (collection, embedding model)
1595 — callers on a delete-only path that must not create an index
1596 (e.g. ``remove_documents_from_index``) guard with their own
1597 existence check before calling this.
1599 ``probe_dimension`` is forwarded to ``_preflight_index_path`` (see
1600 its docstring) unchanged. Every caller but
1601 ``reconcile_collection_index`` leaves it at the default ``None``.
1602 """
1603 # Only WRITE paths (reset_stale_state=True) may re-promote a reused
1604 # index to is_current; a read-path search must not flip the pointer.
1605 rag_index = self._get_or_create_rag_index(
1606 collection_id, promote_current=reset_stale_state
1607 )
1608 self.rag_index_record = rag_index
1609 index_path = self._preflight_index_path(
1610 collection_id,
1611 rag_index,
1612 reset_stale_state=reset_stale_state,
1613 probe_dimension=probe_dimension,
1614 )
1615 lock = _get_faiss_write_lock(self.username, str(index_path))
1617 # nested=True: these run reentrantly inside VectorIndex.apply(), which
1618 # holds just-flushed, uncommitted DocumentChunk rows on this same
1619 # thread-local session. The integrity write must run in a SAVEPOINT and
1620 # let apply()'s caller commit — a plain commit/rollback here would
1621 # settle or DISCARD those pending rows (silently losing indexed text
1622 # while the vectors persist). See FileIntegrityManager._integrity_write.
1623 def _record(path: Path) -> None:
1624 self.integrity_manager.record_file(
1625 path,
1626 related_entity_type="rag_index",
1627 related_entity_id=rag_index.id,
1628 nested=True,
1629 )
1631 def _verify(path: Path):
1632 return self.integrity_manager.verify_file(path, nested=True)
1634 kwargs = {
1635 "username": self.username,
1636 "db_password": self.db_password,
1637 # Lazy: index()/search() resolve the real backend on first embed_*
1638 # call (identical behaviour, one extra attribute hop), while the
1639 # delete-only path (remove_documents_from_index) never materialises
1640 # it at all. See _LazyEmbeddings.
1641 "embeddings": _LazyEmbeddings(self.embedding_manager),
1642 "embedding_model": self.embedding_model,
1643 "embedding_model_type": EmbeddingProvider(self.embedding_provider),
1644 "collection_name": collection_name,
1645 "dimension": rag_index.embedding_dimension,
1646 "path": index_path,
1647 "lock": lock,
1648 "integrity_record": _record,
1649 "integrity_verify": _verify,
1650 # Use the STORED index config, not this service instance's. The
1651 # RAGIndex row is stamped with index_type/metric/normalize at
1652 # creation (see _get_or_create_rag_index); the constructing
1653 # LibraryRAGService may carry different CURRENT defaults (the
1654 # factory falls back to global settings for a collection whose
1655 # metadata was never stored). Building against self.* would load a
1656 # real HNSW file while labelling it "flat" (or vice-versa), and
1657 # FaissVectorStore.supports_delete trusts the label — routing a
1658 # removal through remove_ids on an HNSW index (unsupported → raises)
1659 # or silently rewriting a flat file as HNSW. This mirrors the fix
1660 # already applied to purge_document_vectors below.
1661 "index_type": rag_index.index_type or self.index_type,
1662 "metric": rag_index.distance_metric or self.distance_metric,
1663 "normalize": (
1664 rag_index.normalize_vectors
1665 if rag_index.normalize_vectors is not None
1666 else self.normalize_vectors
1667 ),
1668 }
1669 try:
1670 return VectorIndex(**kwargs)
1671 except (ValueError, RuntimeError):
1672 if not reset_stale_state: 1672 ↛ 1694line 1672 didn't jump to line 1694 because the condition on line 1672 was always true
1673 # Read paths (search) must stay side-effect-free. A transient
1674 # load failure (OOM, a disk hiccup, a momentarily-locked file)
1675 # on an existing, byte-healthy index must NOT quarantine the
1676 # file or wipe the collection's indexed state — that would turn
1677 # a retryable blip into a forced full re-embed of the whole
1678 # collection triggered by a mere search. Surface the error;
1679 # quarantine + rebuild is a WRITE-path (reset_stale_state=True)
1680 # decision. (A never-indexed collection never reaches here — no
1681 # file exists, so VectorIndex.create() succeeds above.)
1682 raise
1683 # A load failure here — after _preflight_index_path's verify_file
1684 # just PASSED on the same bytes — is ambiguous: it is either a
1685 # genuine format/dimension mismatch (deterministic — needs a rebuild)
1686 # OR a transient OS hiccup (EMFILE under concurrent multi-user load,
1687 # a momentary AV-scan/backup file lock; the byte-healthy file loads
1688 # fine moments later). faiss surfaces BOTH as the same RuntimeError,
1689 # so quarantining unconditionally destroys a provably-healthy index
1690 # and forces a needless full re-embed on a mere transient blip.
1691 # Distinguish by RETRYING the load once: a transient failure clears,
1692 # a deterministic one repeats. Only quarantine when the retry also
1693 # fails. (Mirrors _preflight_index_path's transient-vs-tamper split.)
1694 logger.warning(
1695 f"Failed to load vector store at {index_path}; retrying once "
1696 "before quarantining (may be a transient OS hiccup)"
1697 )
1698 try:
1699 return VectorIndex(**kwargs)
1700 except (ValueError, RuntimeError):
1701 pass
1702 logger.warning(
1703 f"Vector store at {index_path} failed to load twice; "
1704 "quarantining and rebuilding fresh"
1705 )
1706 with lock:
1707 if index_path.exists():
1708 self._quarantine_corrupt_index(
1709 index_path, "vector_index_load_failed"
1710 )
1711 self._reset_index_state_for_rebuild(collection_id, rag_index)
1712 return VectorIndex(**kwargs)
1714 def search(
1715 self, query: str, collection_id: str, top_k: int
1716 ) -> List[SearchResult]:
1717 """Search one published generation through the metadata switch."""
1718 with _hold_collection_lock(
1719 self.username, collection_id, purpose="transaction"
1720 ):
1721 return self._search_published(query, collection_id, top_k)
1723 def _search_published(
1724 self, query: str, collection_id: str, top_k: int
1725 ) -> List[SearchResult]:
1726 """Semantic search within ``collection_id`` via the vector store.
1728 Runs the same verify/quarantine pre-flight as indexing (through
1729 ``_get_vector_index``), then delegates to ``VectorIndex.search``.
1730 The embedding-dimension drift probe is SKIPPED on this read path:
1731 the vector store's ``_prepare`` (see FaissVectorStore) rejects a
1732 wrong-dimension, non-finite, zero-magnitude, or overflow-magnitude
1733 query vector before it ever reaches the FAISS index — even on an
1734 EMPTY store, since ``_prepare`` runs before the ``ntotal == 0``
1735 check. Reads stay fail-closed without the probe's round-trip.
1736 Every such rejection is translated below into a service-level
1737 ``RuntimeError`` carrying the collection id: a degenerate or
1738 wrong-dimension query embedding is a PROVIDER fault worth surfacing
1739 loudly, not one to silently swallow. ``reset_stale_state`` is left
1740 at its default False: a transient integrity hiccup during a mere
1741 search must not wipe the collection's indexed state and force a
1742 full re-embed.
1744 Callers are expected to have already confirmed the collection has
1745 indexed documents (e.g. via ``has_indexed_documents``) before
1746 calling this — for an ORDINARY query, this method will otherwise
1747 create an empty RAGIndex/store for a never-indexed (collection,
1748 embedding model) pair and return no results, rather than raising.
1749 A DEGENERATE query embedding (see above) is refused even on that
1750 never-indexed store — it is never swallowed into an empty result
1751 either.
1752 """
1753 collection_name = f"collection_{collection_id}"
1754 vindex = self._get_vector_index(collection_id, collection_name)
1755 try:
1756 return vindex.search(query, top_k)
1757 except ValueError as exc:
1758 # _prepare() in FaissVectorStore raises ValueError when it
1759 # refuses the query vector outright: either a dimension
1760 # mismatch (both dimensions in the message, so it needs a
1761 # regex — see below) or one of the fixed-text "degenerate
1762 # embedding" messages in _DEGENERATE_QUERY_VECTOR_MESSAGES
1763 # above (non-finite, zero/overflow magnitude). Matching the
1764 # exact formats (not a substring/keyword check) means an
1765 # unrelated provider ValueError that merely mentions
1766 # "dimension" or "magnitude" still passes through unchanged.
1767 # Translate every match to a service-level RuntimeError with
1768 # collection context so callers that catch the legacy
1769 # exception type, and logs that report the affected
1770 # collection, still work.
1771 msg = str(exc)
1772 if re.fullmatch(
1773 r"vector dimension [0-9]+ != index dimension [0-9]+", msg
1774 ):
1775 raise RuntimeError(
1776 f"Embedding dimension mismatch — the index for "
1777 f"collection {collection_id} was built with a different "
1778 f"embedding model. Re-index the collection to fix this. "
1779 f"(underlying error: {exc})"
1780 ) from exc
1781 if msg in _DEGENERATE_QUERY_VECTOR_MESSAGES:
1782 raise RuntimeError(
1783 f"Query embedding rejected for collection "
1784 f"{collection_id}: the embedding provider returned a "
1785 f"degenerate vector for this query (non-finite, "
1786 f"zero-magnitude, or overflow-magnitude) that cannot be "
1787 f"searched. (underlying error: {exc})"
1788 ) from exc
1789 raise
1791 def get_current_index_info(
1792 self, collection_id: Optional[str] = None
1793 ) -> Optional[Dict[str, Any]]:
1794 """
1795 Get information about the current RAG index for a collection.
1797 Args:
1798 collection_id: UUID of collection (defaults to Library if None)
1799 """
1800 with get_user_db_session(self.username, self.db_password) as session:
1801 # Get collection name in the format stored in RAGIndex (collection_<uuid>)
1802 if collection_id:
1803 collection = (
1804 session.query(Collection)
1805 .filter_by(id=collection_id)
1806 .first()
1807 )
1808 collection_name = (
1809 f"collection_{collection_id}" if collection else "unknown"
1810 )
1811 else:
1812 # Default to Library collection
1813 from ...database.library_init import get_default_library_id
1815 collection_id = get_default_library_id(
1816 self.username, self.db_password
1817 )
1818 collection_name = f"collection_{collection_id}"
1820 rag_index = (
1821 session.query(RAGIndex)
1822 .filter_by(collection_name=collection_name, is_current=True)
1823 .first()
1824 )
1826 if not rag_index:
1827 # Debug: check all RAG indices for this collection
1828 all_indices = session.query(RAGIndex).all()
1829 logger.info(
1830 f"No RAG index found for collection_name='{collection_name}'. All indices: {[(idx.collection_name, idx.is_current) for idx in all_indices]}"
1831 )
1832 return None
1834 # Calculate actual counts from rag_document_status table
1835 from ...database.models.library import RagDocumentStatus
1837 actual_chunk_count = (
1838 session.query(func.sum(RagDocumentStatus.chunk_count))
1839 .filter_by(collection_id=collection_id)
1840 .scalar()
1841 or 0
1842 )
1844 actual_doc_count = (
1845 session.query(RagDocumentStatus)
1846 .filter_by(collection_id=collection_id)
1847 .count()
1848 )
1850 return {
1851 "embedding_model": rag_index.embedding_model,
1852 "embedding_model_type": rag_index.embedding_model_type.value
1853 if rag_index.embedding_model_type
1854 else None,
1855 "embedding_dimension": rag_index.embedding_dimension,
1856 "chunk_size": rag_index.chunk_size,
1857 "chunk_overlap": rag_index.chunk_overlap,
1858 "chunk_count": actual_chunk_count,
1859 "total_documents": actual_doc_count,
1860 "created_at": rag_index.created_at.isoformat(),
1861 "last_updated_at": rag_index.last_updated_at.isoformat(),
1862 }
1864 def _flag_document_for_reindex(
1865 self, session, document_id: str, collection_id: str
1866 ) -> None:
1867 """Mark a document NOT indexed so the reconciler re-embeds it.
1869 Used on index_document failure paths: ``VectorIndex.index()`` runs
1870 ``store.apply()`` — DURABLY mutating the FAISS file — BEFORE the DB
1871 commit, so a failure/rollback after that point can leave the source's
1872 prior chunk rows resurrected with no vectors (silently unsearchable)
1873 while ``DocumentCollection.indexed`` is reverted to its pre-attempt
1874 state. The background reconciler only revisits ``indexed=False`` rows,
1875 so without this flag the document stays permanently, silently
1876 unsearchable. Best-effort + committed in its own statement;
1877 index_document owns per-document commits (see index_all_documents), so
1878 this clobbers no caller.
1879 """
1880 try:
1881 session.query(DocumentCollection).filter_by(
1882 document_id=document_id, collection_id=collection_id
1883 ).update(
1884 {"indexed": False, "chunk_count": 0},
1885 synchronize_session=False,
1886 )
1887 session.query(RagDocumentStatus).filter_by(
1888 document_id=document_id, collection_id=collection_id
1889 ).delete(synchronize_session=False)
1890 session.commit()
1891 except Exception:
1892 safe_rollback(
1893 session, "library_rag_service._flag_document_for_reindex"
1894 )
1895 logger.warning(
1896 f"Could not flag document {document_id} for re-index after a "
1897 f"failed index of collection {collection_id}; a manual reindex "
1898 "may be needed to restore its search."
1899 )
1901 @contextmanager
1902 def _collection_transaction_lock(self, collection_id: str):
1903 """Serialize the DB+FAISS transaction for one collection.
1905 FAISS already uses this same re-entrant lock for reload/apply/persist and
1906 integrity recording. Acquiring it around the complete document operation
1907 also closes the store-persist-before-SQLCipher-commit window to sibling
1908 workers: no worker can verify/quarantine or mutate the file until the
1909 preceding worker's database commit (or rollback/reindex flag) finishes.
1910 """
1911 # This identity survives model/configuration changes and publication
1912 # of a new generation. Acquire it before resolving the physical path.
1913 with _hold_collection_lock(
1914 self.username, collection_id, purpose="transaction"
1915 ):
1916 with get_user_db_session(
1917 self.username, self.db_password
1918 ) as session:
1919 index = self._resolve_index_for_config(session, collection_id)
1920 index_hash = (
1921 index.index_hash
1922 if index is not None
1923 else self._get_index_hash(
1924 f"collection_{collection_id}",
1925 self.embedding_model,
1926 self.embedding_provider,
1927 )
1928 )
1929 index_path = self._get_index_path(index_hash)
1930 with _hold_faiss_write_lock(self.username, str(index_path)) as lock:
1931 yield lock
1933 def _embed_texts_cancellable(
1934 self,
1935 texts: List[str],
1936 is_cancelled: Optional[Callable[[], bool]] = None,
1937 ) -> Optional[List[List[float]]]:
1938 """Embed texts, honouring cancellation between batches.
1940 Returns the full vector list, or ``None`` when ``is_cancelled()``
1941 fired before or during embedding — the caller must abandon the
1942 document (nothing was persisted yet at the prepare stage). The
1943 in-flight batch finishes naturally: an HTTP round-trip or a local
1944 forward pass cannot be safely interrupted mid-array, so batching
1945 bounds the uncancellable tail to one batch instead of the whole
1946 document.
1947 """
1948 embeddings = self.embedding_manager.embeddings
1949 if is_cancelled is None or len(texts) <= _EMBEDDING_CANCEL_BATCH_SIZE:
1950 if is_cancelled is not None and is_cancelled(): 1950 ↛ 1951line 1950 didn't jump to line 1951 because the condition on line 1950 was never true
1951 return None
1952 return embeddings.embed_documents(texts)
1953 vectors: List[List[float]] = []
1954 for start in range(0, len(texts), _EMBEDDING_CANCEL_BATCH_SIZE):
1955 if is_cancelled():
1956 return None
1957 vectors.extend(
1958 embeddings.embed_documents(
1959 texts[start : start + _EMBEDDING_CANCEL_BATCH_SIZE]
1960 )
1961 )
1962 if is_cancelled():
1963 return None
1964 return vectors
1966 def _prepare_document(
1967 self,
1968 document_id: str,
1969 collection_id: str,
1970 force_reindex: bool,
1971 is_cancelled: Optional[Callable[[], bool]] = None,
1972 ) -> Dict[str, Any]:
1973 """Load, split, and embed without mutating collection state."""
1974 scope_error = _indexing_scope_error()
1975 if scope_error is not None:
1976 return scope_error
1977 with get_user_db_session(self.username, self.db_password) as session:
1978 document = session.query(Document).filter_by(id=document_id).first()
1979 if not document: 1979 ↛ 1980line 1979 didn't jump to line 1980 because the condition on line 1979 was never true
1980 return {"status": "error", "error": "Document not found"}
1981 link = (
1982 session.query(DocumentCollection)
1983 .filter_by(document_id=document_id, collection_id=collection_id)
1984 .first()
1985 )
1986 if link is not None and link.indexed and not force_reindex:
1987 return {
1988 "status": "skipped",
1989 "message": "Document already indexed for this collection",
1990 "chunk_count": link.chunk_count,
1991 }
1992 source_fingerprint = document_fingerprint(document)
1993 collection = session.get(Collection, collection_id)
1994 configuration = (
1995 collection_configuration(collection) if collection else None
1996 )
1997 if not document.text_content:
1998 return {
1999 "status": "needs_serial_write",
2000 "source_fingerprint": source_fingerprint,
2001 }
2003 title = document.title or document.filename or "Untitled"
2004 metadata = {
2005 "source": document.original_url,
2006 "document_id": document_id,
2007 "collection_id": collection_id,
2008 "document_title": title,
2009 "title": title,
2010 "authors": document.authors,
2011 "published_date": str(document.published_date)
2012 if document.published_date
2013 else None,
2014 "doi": document.doi,
2015 "arxiv_id": document.arxiv_id,
2016 "pmid": document.pmid,
2017 "pmcid": document.pmcid,
2018 "extraction_method": document.extraction_method,
2019 "word_count": document.word_count,
2020 }
2021 source_doc = LangchainDocument(
2022 page_content=document.text_content, metadata=metadata
2023 )
2024 prep_content_hash = hashlib.sha256(
2025 (document.text_content or "").encode("utf-8")
2026 ).hexdigest()
2027 prep_initial_indexed = link.indexed if link is not None else False
2028 prep_had_collection_link = link is not None
2029 prep_status_exists = (
2030 session.query(RagDocumentStatus)
2031 .filter_by(
2032 document_id=document_id,
2033 collection_id=collection_id,
2034 rag_index_id=self.rag_index_record.id,
2035 )
2036 .first()
2037 is not None
2038 if self.rag_index_record
2039 else False
2040 )
2042 chunks = self.text_splitter.split_documents([source_doc])
2043 chunk_inputs = [
2044 ChunkInput(
2045 text=chunk.page_content,
2046 metadata={**metadata, "chunk_index": i},
2047 )
2048 for i, chunk in enumerate(chunks)
2049 ]
2050 vectors_list = self._embed_texts_cancellable(
2051 [chunk.text for chunk in chunk_inputs], is_cancelled
2052 )
2053 if vectors_list is None:
2054 logger.info(
2055 f"Embedding cancelled for document {document_id} "
2056 f"({len(chunk_inputs)} chunks)"
2057 )
2058 return {
2059 "status": "cancelled",
2060 "message": "Cancelled during embedding",
2061 }
2062 vectors = np.asarray(
2063 vectors_list,
2064 dtype="float32",
2065 )
2066 logger.info(
2067 f"Prepared document {document_id}: {len(chunk_inputs)} chunks embedded"
2068 )
2069 return {
2070 "status": "prepared",
2071 "prepared": _PreparedDocument(
2072 document_id=document_id,
2073 collection_id=collection_id,
2074 chunk_inputs=chunk_inputs,
2075 vectors=vectors,
2076 source_fingerprint=source_fingerprint,
2077 collection_configuration=configuration,
2078 content_hash=prep_content_hash,
2079 initial_indexed=prep_initial_indexed,
2080 status_exists=prep_status_exists,
2081 had_collection_link=prep_had_collection_link,
2082 ),
2083 }
2085 def _write_prepared_document(
2086 self, prepared: _PreparedDocument
2087 ) -> Dict[str, Any]:
2088 """Serialize chunk rows, FAISS mutation, status updates, and commit."""
2089 scope_error = _indexing_scope_error()
2090 if scope_error is not None:
2091 return scope_error
2092 document_id = prepared.document_id
2093 collection_id = prepared.collection_id
2094 chunks = prepared.chunk_inputs
2095 with self._collection_transaction_lock(collection_id):
2096 with get_user_db_session(
2097 self.username, self.db_password
2098 ) as session:
2099 # Embedding runs outside the collection lock. A complete
2100 # rebuild may have published a different configuration while
2101 # this document was being prepared. Reject that stale work
2102 # before any write or recovery flag can touch the new index.
2103 if prepared.source_fingerprint is not None:
2104 document = session.get(Document, document_id)
2105 if ( 2105 ↛ 2110line 2105 didn't jump to line 2110 because the condition on line 2105 was never true
2106 document is None
2107 or document_fingerprint(document)
2108 != prepared.source_fingerprint
2109 ):
2110 return {
2111 "status": "error",
2112 "error": "Document changed; retry indexing",
2113 }
2114 if prepared.collection_configuration is not None:
2115 collection = session.get(Collection, collection_id)
2116 if (
2117 collection is None
2118 or collection_configuration(collection)
2119 != prepared.collection_configuration
2120 or (
2121 collection.embedding_model is not None
2122 and not self._matches_collection_configuration(
2123 collection
2124 )
2125 )
2126 ):
2127 return {
2128 "status": "error",
2129 "error": "Collection configuration changed; retry indexing",
2130 }
2131 try:
2132 # Phase 3 Revalidation Gate: verify collection link, content, and indexed status.
2133 # The membership check runs unconditionally so legacy
2134 # prepared objects without a content_hash snapshot still
2135 # get first-time vs concurrent-unlink handling instead of
2136 # skipping the gate entirely.
2137 link = (
2138 session.query(DocumentCollection)
2139 .filter_by(
2140 document_id=document_id,
2141 collection_id=collection_id,
2142 )
2143 .first()
2144 )
2145 if link is None:
2146 if not getattr(prepared, "had_collection_link", None):
2147 # Reconciler case (b): research orphans have
2148 # no membership at prepare. Create it now
2149 # rather than treating absence as a concurrent
2150 # unlink. A missing snapshot (None) preserves
2151 # main's unconditional-create behavior.
2152 link = ensure_in_collection(
2153 session, document_id, collection_id
2154 )
2155 session.flush()
2156 else:
2157 logger.info(
2158 f"Document {document_id} was removed from collection {collection_id} "
2159 "during embedding; aborting Phase 3 persist."
2160 )
2161 self._flag_document_for_reindex(
2162 session, document_id, collection_id
2163 )
2164 return {
2165 "status": "skipped",
2166 "message": "Document removed from collection during embedding",
2167 }
2169 current_doc = (
2170 session.query(Document)
2171 .filter_by(id=document_id)
2172 .first()
2173 )
2174 if not current_doc: 2174 ↛ 2175line 2174 didn't jump to line 2175 because the condition on line 2174 was never true
2175 logger.info(
2176 f"Document {document_id} was deleted during embedding; "
2177 "aborting Phase 3 persist."
2178 )
2179 self._flag_document_for_reindex(
2180 session, document_id, collection_id
2181 )
2182 return {
2183 "status": "error",
2184 "error": "Document deleted during indexing",
2185 }
2187 if isinstance(getattr(prepared, "content_hash", None), str):
2188 current_hash = hashlib.sha256(
2189 (current_doc.text_content or "").encode("utf-8")
2190 ).hexdigest()
2191 if current_hash != prepared.content_hash:
2192 logger.info(
2193 f"Document {document_id} content changed during embedding; "
2194 "aborting persist and flagging for reindex."
2195 )
2196 self._flag_document_for_reindex(
2197 session, document_id, collection_id
2198 )
2199 return {
2200 "status": "stale",
2201 "message": "Document modified during indexing, flagged for reindex",
2202 }
2204 current_status_exists = (
2205 session.query(RagDocumentStatus)
2206 .filter_by(
2207 document_id=document_id,
2208 collection_id=collection_id,
2209 rag_index_id=self.rag_index_record.id,
2210 )
2211 .first()
2212 is not None
2213 if self.rag_index_record
2214 else False
2215 )
2216 if ( 2216 ↛ 2224line 2216 didn't jump to line 2224 because the condition on line 2216 was never true
2217 isinstance(prepared.initial_indexed, bool)
2218 and isinstance(link.indexed, bool)
2219 and link.indexed != prepared.initial_indexed
2220 ) or (
2221 isinstance(prepared.status_exists, bool)
2222 and current_status_exists != prepared.status_exists
2223 ):
2224 logger.info(
2225 f"Document {document_id} indexing status modified during embedding; "
2226 "aborting persist and flagging for reindex."
2227 )
2228 self._flag_document_for_reindex(
2229 session, document_id, collection_id
2230 )
2231 return {
2232 "status": "stale",
2233 "message": "Document indexed state changed during indexing, flagged for reindex",
2234 }
2236 collection_name = f"collection_{collection_id}"
2237 observed_text = current_doc.text_content
2238 observed_hash = current_doc.document_hash
2239 vindex = self._get_vector_index(
2240 collection_id, collection_name, reset_stale_state=True
2241 )
2242 if (
2243 _claim_document_snapshot(
2244 session, document_id, observed_text, observed_hash
2245 )
2246 == 0
2247 ):
2248 self._flag_document_for_reindex(
2249 session, document_id, collection_id
2250 )
2251 return {
2252 "status": "stale",
2253 "message": "Document modified during indexing, flagged for reindex",
2254 }
2255 stats = vindex.index_prepared(
2256 source_type="document",
2257 source_id=document_id,
2258 chunks=chunks,
2259 vectors=prepared.vectors,
2260 replace=True,
2261 session=session,
2262 )
2264 timestamp = datetime.now(UTC)
2265 existing = (
2266 session.query(RagDocumentStatus)
2267 .filter_by(
2268 document_id=document_id,
2269 collection_id=collection_id,
2270 rag_index_id=self.rag_index_record.id,
2271 )
2272 .first()
2273 )
2274 old_chunks = existing.chunk_count if existing else 0
2275 session.merge(
2276 RagDocumentStatus(
2277 document_id=document_id,
2278 collection_id=collection_id,
2279 rag_index_id=self.rag_index_record.id,
2280 chunk_count=len(chunks),
2281 indexed_at=timestamp,
2282 )
2283 )
2284 updated_rows = (
2285 session.query(DocumentCollection)
2286 .filter_by(
2287 document_id=document_id,
2288 collection_id=collection_id,
2289 )
2290 .update(
2291 {
2292 "indexed": True,
2293 "chunk_count": len(chunks),
2294 "last_indexed_at": timestamp,
2295 }
2296 )
2297 )
2298 if isinstance(updated_rows, int) and updated_rows == 0:
2299 # An unlink committed between the Phase-3 gate query
2300 # and these persist writes (the unlink DB phase does
2301 # not hold the FAISS lock). Rolling back avoids
2302 # re-creating status/aggregate rows for a non-member
2303 # and leaves phantom vectors to the unlink's purge.
2304 logger.info(
2305 f"Document {document_id} was removed from collection {collection_id} "
2306 "during persist; aborting Phase 3 commit."
2307 )
2308 safe_rollback(
2309 session,
2310 "library_rag_service._write_prepared_document link vanished",
2311 )
2312 self._flag_document_for_reindex(
2313 session, document_id, collection_id
2314 )
2315 return {
2316 "status": "skipped",
2317 "message": "Document removed from collection during embedding",
2318 }
2319 rag_index = (
2320 session.query(RAGIndex)
2321 .filter_by(id=self.rag_index_record.id)
2322 .first()
2323 )
2324 if rag_index: 2324 ↛ 2329line 2324 didn't jump to line 2329 because the condition on line 2324 was always true
2325 rag_index.chunk_count += len(chunks) - old_chunks
2326 if existing is None:
2327 rag_index.total_documents += 1
2328 rag_index.last_updated_at = timestamp
2329 session.commit()
2330 logger.info(
2331 f"Persisted prepared document {document_id}: "
2332 f"{stats.added} vectors added"
2333 )
2334 return {"status": "success", "chunk_count": len(chunks)}
2335 except Exception as exc:
2336 safe_rollback(
2337 session,
2338 "library_rag_service._write_prepared_document",
2339 )
2340 self._flag_document_for_reindex(
2341 session, document_id, collection_id
2342 )
2343 logger.exception(
2344 f"Error writing prepared document {document_id}"
2345 )
2346 return {
2347 "status": "error",
2348 "error": f"Operation failed: {type(exc).__name__}",
2349 }
2351 def index_document(
2352 self, document_id: str, collection_id: str, force_reindex: bool = False
2353 ) -> Dict[str, Any]:
2354 # Reject before acquiring the collection lock: a caller holding a DB
2355 # write transaction could otherwise deadlock a writer holding this lock.
2356 scope_error = _indexing_scope_error()
2357 if scope_error is not None:
2358 return scope_error
2359 with self._collection_transaction_lock(collection_id):
2360 return self._index_document_locked(
2361 document_id, collection_id, force_reindex
2362 )
2364 def _index_one(
2365 self, document_id: str, collection_id: str, force_reindex: bool
2366 ) -> Dict[str, Any]:
2367 """Single-doc dispatch seam used by :meth:`index_documents_parallel`.
2369 Default implementation delegates to :meth:`index_document` so production
2370 behaviour is unchanged. Tests and subclasses override this method to
2371 short-circuit the full per-doc pipeline (which would require mocking
2372 every internal collaborator) without losing the prepared/serialized
2373 pipeline in production.
2375 The previous parallel runner detected overrides via
2376 ``"index_document" in self.__dict__`` — that only honoured
2377 instance-attribute patches and silently bypassed class-level subclass
2378 overrides. The runner now compares the bound method against the base
2379 class attribute so both instance patches and subclass overrides are
2380 honoured.
2381 """
2382 return self.index_document(document_id, collection_id, force_reindex)
2384 def _index_document_locked(
2385 self, document_id: str, collection_id: str, force_reindex: bool = False
2386 ) -> Dict[str, Any]:
2387 """
2388 Index a single document into RAG for a specific collection.
2390 Args:
2391 document_id: UUID of the Document to index
2392 collection_id: UUID of the Collection to index for
2393 force_reindex: Whether to force reindexing even if already indexed
2395 Returns:
2396 Dict with status, chunk_count, and any errors
2397 """
2398 scope_error = _indexing_scope_error()
2399 if scope_error is not None:
2400 return scope_error
2402 # Phase 1: Read & chunk doc
2403 with get_user_db_session(self.username, self.db_password) as session:
2404 # Get the document
2405 document = session.query(Document).filter_by(id=document_id).first()
2407 if not document:
2408 return {"status": "error", "error": "Document not found"}
2410 # A service may have been constructed before a complete rebuild
2411 # published a different model. Check under the collection lock,
2412 # before membership changes, empty-document purges or embedding.
2413 collection = session.get(Collection, collection_id)
2414 if (
2415 collection is not None
2416 and collection.embedding_model is not None
2417 and not self._matches_collection_configuration(collection)
2418 ):
2419 return {
2420 "status": "error",
2421 "error": "Collection configuration changed; retry indexing",
2422 }
2424 # Get or create DocumentCollection entry
2425 doc_collection = ensure_in_collection(
2426 session, document_id, collection_id
2427 )
2428 # If doc_collection was newly created in this session, commit it so the link
2429 # exists durably before Phase 2 without holding a write lock during embedding.
2430 # No test-session import is needed here: `in session.new` is
2431 # False for unittest.mock sessions (MagicMock __contains__
2432 # defaults to False), so mocks safely skip the pre-commit.
2433 try:
2434 _is_new_link = (
2435 getattr(session, "new", None) is not None
2436 and doc_collection in session.new
2437 )
2438 except Exception:
2439 _is_new_link = False
2440 if _is_new_link:
2441 try:
2442 session.commit()
2443 except Exception: # allow: silent-exception
2444 logger.debug(
2445 f"Could not pre-commit newly created collection link for {document_id} in {collection_id}"
2446 )
2448 # Check if already indexed for this collection
2449 if doc_collection.indexed and not force_reindex:
2450 return {
2451 "status": "skipped",
2452 "message": "Document already indexed for this collection",
2453 "chunk_count": doc_collection.chunk_count,
2454 }
2456 # Validate text content
2457 if not document.text_content:
2458 # A previously-indexed document now cleared to empty text (e.g.
2459 # a note edited down to whitespace) must have its stale
2460 # chunks/vectors PURGED, not left searchable. The note-edit path
2461 # already set indexed=False, so detect prior chunks directly and
2462 # remove them via the facade's empty-replace.
2463 collection_name = f"collection_{collection_id}"
2464 has_prior = (
2465 session.query(DocumentChunk.id)
2466 .filter_by(
2467 source_type="document",
2468 source_id=document_id,
2469 collection_name=collection_name,
2470 )
2471 .first()
2472 is not None
2473 )
2474 if has_prior:
2475 # Same store-durable-before-DB-commit hazard as the main
2476 # reindex path below, but this branch sits OUTSIDE that
2477 # path's try/except: the empty-replace apply() DURABLY
2478 # removes the prior vectors before the commit here. On
2479 # failure, roll back and flag the document for re-index so
2480 # the reconciler re-runs the purge instead of leaving
2481 # resurrected rows whose vectors are gone (silently
2482 # unsearchable) marked indexed.
2483 try:
2484 vindex = self._get_vector_index(
2485 collection_id,
2486 collection_name,
2487 reset_stale_state=True,
2488 )
2489 vindex.index(
2490 source_type="document",
2491 source_id=document_id,
2492 chunks=[],
2493 replace=True,
2494 session=session,
2495 )
2496 doc_collection.indexed = False
2497 doc_collection.chunk_count = 0
2498 # Clear the canonical RagDocumentStatus row too. Row
2499 # existence is the "indexed" marker get_rag_stats / the
2500 # RAG status route read, so it MUST move with
2501 # DocumentCollection.indexed — a leftover row would
2502 # report this just-purged (now empty, unsearchable)
2503 # document as still indexed with a stale chunk_count.
2504 # Mirrors remove_document_from_rag and the three sibling
2505 # not-indexed paths; the cleared branch was the lone
2506 # exception.
2507 session.query(RagDocumentStatus).filter_by(
2508 document_id=document_id, collection_id=collection_id
2509 ).delete(synchronize_session=False)
2510 session.commit()
2511 except Exception as e:
2512 safe_rollback(
2513 session,
2514 "library_rag_service.index_document cleared",
2515 )
2516 self._flag_document_for_reindex(
2517 session, document_id, collection_id
2518 )
2519 logger.exception(
2520 f"Error purging cleared document {document_id} "
2521 f"for collection {collection_id}"
2522 )
2523 return {
2524 "status": "error",
2525 "error": f"Operation failed: {type(e).__name__}",
2526 }
2527 return {
2528 "status": "cleared",
2529 "message": (
2530 "Document has no text content; prior chunks/"
2531 "vectors purged"
2532 ),
2533 "chunk_count": 0,
2534 }
2535 return {
2536 "status": "error",
2537 "error": "Document has no text content",
2538 }
2540 # Replace-on-reindex context: a document being re-indexed (e.g.
2541 # a note whose content changed) has prior chunks with DIFFERENT
2542 # content hashes from the new text. Left alone they accumulate
2543 # in both the DB and the vector store forever — unbounded
2544 # per-edit growth plus stale duplicate search hits.
2545 # ``VectorIndex.index(replace=True)`` below handles this: it
2546 # looks up the source's prior DocumentChunk rows (authoritative,
2547 # queried fresh — not from an ephemeral pre-committed snapshot),
2548 # applies the vector-store removal durably, then deletes those
2549 # rows in the SAME transaction this method already holds open
2550 # (so a failure rolls back together with everything else, never
2551 # orphaning rows or poisoning the shared session).
2553 try:
2554 phase1_content_hash = hashlib.sha256(
2555 (document.text_content or "").encode("utf-8")
2556 ).hexdigest()
2557 phase1_doc_hash = document.document_hash
2558 phase1_indexed = doc_collection.indexed
2559 phase1_status_exists = (
2560 session.query(RagDocumentStatus)
2561 .filter_by(
2562 document_id=document_id,
2563 collection_id=collection_id,
2564 rag_index_id=self.rag_index_record.id,
2565 )
2566 .first()
2567 is not None
2568 if self.rag_index_record
2569 else False
2570 )
2572 # Create LangChain Document from text
2573 doc = LangchainDocument(
2574 page_content=document.text_content,
2575 metadata={
2576 "source": document.original_url,
2577 "document_id": document_id, # Add document ID for source linking
2578 "collection_id": collection_id, # Add collection ID
2579 "title": document.title
2580 or document.filename
2581 or "Untitled",
2582 "document_title": document.title
2583 or document.filename
2584 or "Untitled", # Add for compatibility
2585 "authors": document.authors,
2586 "published_date": str(document.published_date)
2587 if document.published_date
2588 else None,
2589 "doi": document.doi,
2590 "arxiv_id": document.arxiv_id,
2591 "pmid": document.pmid,
2592 "pmcid": document.pmcid,
2593 "extraction_method": document.extraction_method,
2594 "word_count": document.word_count,
2595 },
2596 )
2598 # Split into chunks
2599 chunks = self.text_splitter.split_documents([doc])
2600 logger.info(
2601 f"Split document {document_id} into {len(chunks)} chunks"
2602 )
2604 # Get collection name for chunk storage
2605 collection = (
2606 session.query(Collection)
2607 .filter_by(id=collection_id)
2608 .first()
2609 )
2610 # Use collection_<uuid> format for internal storage
2611 collection_name = (
2612 f"collection_{collection_id}" if collection else "unknown"
2613 )
2615 document_title = (
2616 document.title or document.filename or "Untitled"
2617 )
2618 # Mirrors the parent doc's metadata dict built above (source,
2619 # ids, bibliographic fields) so it survives into each
2620 # chunk's ``DocumentChunk.document_metadata`` — this goes
2621 # ONLY to the encrypted DB (never the vector store; see the
2622 # SECURITY INVARIANT in vector_stores/base.py) and future-
2623 # proofs citation/UI features that read per-chunk metadata.
2624 shared_chunk_metadata = {
2625 "source": document.original_url,
2626 "document_id": document_id,
2627 "collection_id": collection_id,
2628 "document_title": document_title,
2629 "title": document_title,
2630 "authors": document.authors,
2631 "published_date": str(document.published_date)
2632 if document.published_date
2633 else None,
2634 "doi": document.doi,
2635 "arxiv_id": document.arxiv_id,
2636 "pmid": document.pmid,
2637 "pmcid": document.pmcid,
2638 "extraction_method": document.extraction_method,
2639 "word_count": document.word_count,
2640 }
2641 chunk_inputs = [
2642 ChunkInput(
2643 text=chunk.page_content,
2644 metadata={
2645 **shared_chunk_metadata,
2646 "chunk_index": i,
2647 },
2648 )
2649 for i, chunk in enumerate(chunks)
2650 ]
2652 # Roll back uncommitted changes (such as ensure_in_collection staging
2653 # and autoflush from query(Collection)) so no SQLite transaction
2654 # or write lock is held during Phase 2 embedding.
2655 safe_rollback(
2656 session, "library_rag_service.index_document phase 1"
2657 )
2659 except Exception as e:
2660 # The session is shared (thread-local) with the caller.
2661 # If session.flush() or session.commit() raised, the session
2662 # is in PendingRollbackError state until rolled back —
2663 # leaving subsequent operations to cascade. Roll back BEFORE
2664 # returning the error dict so the caller sees a clean
2665 # session. (Same pattern as the #3827 fix.)
2666 safe_rollback(session, "library_rag_service.index_document")
2667 # Backstop for the store-durable-before-DB-commit window: a
2668 # REPLACE reindex ran store.apply() (DURABLY removing this
2669 # source's prior vectors) before the DB commit, so the rollback
2670 # above resurrects rows whose vectors are already gone while
2671 # reverting indexed→its pre-attempt True. Flag for re-index so
2672 # the reconciler re-embeds and store+DB reconverge.
2673 self._flag_document_for_reindex(
2674 session, document_id, collection_id
2675 )
2676 logger.exception(
2677 f"Error indexing document {document_id} for collection {collection_id}"
2678 )
2679 return {
2680 "status": "error",
2681 "error": f"Operation failed: {type(e).__name__}",
2682 }
2684 # Phase 2: Embed vectors outside any DB session
2685 try:
2686 texts = [c.text for c in chunk_inputs]
2687 vectors = (
2688 np.asarray(
2689 self.embedding_manager.embeddings.embed_documents(texts),
2690 dtype="float32",
2691 )
2692 if texts
2693 else None
2694 )
2695 except Exception as e:
2696 logger.exception(
2697 f"Error embedding document {document_id} for collection {collection_id}"
2698 )
2699 with get_user_db_session(
2700 self.username, self.db_password
2701 ) as session:
2702 self._flag_document_for_reindex(
2703 session, document_id, collection_id
2704 )
2705 return {
2706 "status": "error",
2707 "error": f"Operation failed: {type(e).__name__}",
2708 }
2710 # Phase 3: Open short session to persist chunks, apply vectors, and commit
2711 with get_user_db_session(self.username, self.db_password) as session:
2712 try:
2713 # Phase 3 Revalidation Gate: verify collection link, content, and indexed status
2714 link = (
2715 session.query(DocumentCollection)
2716 .filter_by(
2717 document_id=document_id,
2718 collection_id=collection_id,
2719 )
2720 .first()
2721 )
2722 if link is None:
2723 logger.info(
2724 f"Document {document_id} was removed from collection {collection_id} "
2725 "during embedding; aborting Phase 3 persist."
2726 )
2727 self._flag_document_for_reindex(
2728 session, document_id, collection_id
2729 )
2730 return {
2731 "status": "skipped",
2732 "message": "Document removed from collection during embedding",
2733 }
2735 current_doc = (
2736 session.query(Document).filter_by(id=document_id).first()
2737 )
2738 if not current_doc:
2739 logger.info(
2740 f"Document {document_id} was deleted during embedding; "
2741 "aborting Phase 3 persist."
2742 )
2743 self._flag_document_for_reindex(
2744 session, document_id, collection_id
2745 )
2746 return {
2747 "status": "error",
2748 "error": "Document deleted during indexing",
2749 }
2751 current_content_hash = hashlib.sha256(
2752 (current_doc.text_content or "").encode("utf-8")
2753 ).hexdigest()
2754 if (
2755 current_content_hash != phase1_content_hash
2756 or current_doc.document_hash != phase1_doc_hash
2757 ):
2758 logger.info(
2759 f"Document {document_id} content changed during embedding; "
2760 "aborting persist and flagging for reindex."
2761 )
2762 self._flag_document_for_reindex(
2763 session, document_id, collection_id
2764 )
2765 return {
2766 "status": "stale",
2767 "message": "Document modified during indexing, flagged for reindex",
2768 }
2770 current_status_exists = (
2771 session.query(RagDocumentStatus)
2772 .filter_by(
2773 document_id=document_id,
2774 collection_id=collection_id,
2775 rag_index_id=self.rag_index_record.id,
2776 )
2777 .first()
2778 is not None
2779 if self.rag_index_record
2780 else False
2781 )
2782 if ( 2782 ↛ 2791line 2782 didn't jump to line 2791 because the condition on line 2782 was never true
2783 isinstance(link.indexed, bool)
2784 and isinstance(phase1_indexed, bool)
2785 and link.indexed != phase1_indexed
2786 ) or (
2787 isinstance(current_status_exists, bool)
2788 and isinstance(phase1_status_exists, bool)
2789 and current_status_exists != phase1_status_exists
2790 ):
2791 logger.info(
2792 f"Document {document_id} indexing status modified during embedding; "
2793 "aborting persist and flagging for reindex."
2794 )
2795 self._flag_document_for_reindex(
2796 session, document_id, collection_id
2797 )
2798 return {
2799 "status": "stale",
2800 "message": "Document indexed state changed during indexing, flagged for reindex",
2801 }
2803 observed_text = current_doc.text_content
2804 observed_hash = current_doc.document_hash
2805 vindex = self._get_vector_index(
2806 collection_id, collection_name, reset_stale_state=True
2807 )
2809 if (
2810 _claim_document_snapshot(
2811 session, document_id, observed_text, observed_hash
2812 )
2813 == 0
2814 ):
2815 self._flag_document_for_reindex(
2816 session, document_id, collection_id
2817 )
2818 return {
2819 "status": "stale",
2820 "message": "Document modified during indexing, flagged for reindex",
2821 }
2823 if hasattr(vindex, "index_prepared"):
2824 index_stats = vindex.index_prepared(
2825 source_type="document",
2826 source_id=document_id,
2827 chunks=chunk_inputs,
2828 vectors=vectors,
2829 replace=True,
2830 session=session,
2831 )
2832 else:
2833 index_stats = vindex.index(
2834 source_type="document",
2835 source_id=document_id,
2836 chunks=chunk_inputs,
2837 replace=True,
2838 session=session,
2839 )
2841 logger.info(
2842 f"Vector index update for document {document_id}: "
2843 f"{index_stats.added} added, {index_stats.removed} "
2844 f"removed ({index_stats.chunks} chunks total)"
2845 )
2847 # Check if document was already indexed (for stats update)
2848 existing_status = (
2849 session.query(RagDocumentStatus)
2850 .filter_by(
2851 document_id=document_id,
2852 collection_id=collection_id,
2853 rag_index_id=self.rag_index_record.id,
2854 )
2855 .first()
2856 )
2857 was_already_indexed = existing_status is not None
2858 old_chunk_count = (
2859 (existing_status.chunk_count or 0) if existing_status else 0
2860 )
2862 # Mark document as indexed using rag_document_status table
2863 # Row existence = indexed, simple and clean
2864 timestamp = datetime.now(UTC)
2866 # Create or update RagDocumentStatus using ORM merge (atomic upsert)
2867 rag_status = RagDocumentStatus(
2868 document_id=document_id,
2869 collection_id=collection_id,
2870 rag_index_id=self.rag_index_record.id,
2871 chunk_count=len(chunks),
2872 indexed_at=timestamp,
2873 )
2874 session.merge(rag_status)
2876 logger.info(
2877 f"Marked document as indexed in rag_document_status: doc_id={document_id}, coll_id={collection_id}, chunks={len(chunks)}"
2878 )
2880 # Also update DocumentCollection table for backward compatibility
2881 _serial_updated = (
2882 session.query(DocumentCollection)
2883 .filter_by(
2884 document_id=document_id,
2885 collection_id=collection_id,
2886 )
2887 .update(
2888 {
2889 "indexed": True,
2890 "chunk_count": len(chunks),
2891 "last_indexed_at": timestamp,
2892 }
2893 )
2894 )
2895 if isinstance(_serial_updated, int) and _serial_updated == 0:
2896 # An unlink committed between the Phase-3 gate query and
2897 # these persist writes. Roll back the staged status merge
2898 # above so no non-member row/aggregate is resurrected.
2899 logger.info(
2900 f"Document {document_id} was removed from collection {collection_id} "
2901 "during persist; aborting Phase 3 commit."
2902 )
2903 safe_rollback(
2904 session,
2905 "library_rag_service.index_document link vanished",
2906 )
2907 self._flag_document_for_reindex(
2908 session, document_id, collection_id
2909 )
2910 return {
2911 "status": "skipped",
2912 "message": "Document removed from collection during embedding",
2913 }
2915 logger.info(
2916 "Also updated DocumentCollection.indexed for backward compatibility"
2917 )
2919 # Keep the aggregate synchronized on both first index and
2920 # replacement reindex. A replacement changes the per-document
2921 # row above, so apply its new-old delta instead of skipping the
2922 # aggregate update entirely.
2923 rag_index_obj = (
2924 session.query(RAGIndex)
2925 .filter_by(id=self.rag_index_record.id)
2926 .first()
2927 )
2928 if rag_index_obj: 2928 ↛ 2941line 2928 didn't jump to line 2941 because the condition on line 2928 was always true
2929 chunk_delta = len(chunks) - old_chunk_count
2930 rag_index_obj.chunk_count += chunk_delta
2931 if not was_already_indexed:
2932 rag_index_obj.total_documents += 1
2933 rag_index_obj.last_updated_at = datetime.now(UTC)
2934 logger.info(
2935 "Updated RAGIndex stats: "
2936 f"chunk_count {chunk_delta:+d}, "
2937 f"total_documents {'+1' if not was_already_indexed else '+0'}"
2938 )
2940 # Flush ORM changes to database before commit
2941 session.flush()
2942 logger.info(f"Flushed ORM changes for document {document_id}")
2944 # Commit the transaction. Durability is provided by
2945 # synchronous=NORMAL (sqlcipher_utils.py); SQLite
2946 # auto-checkpoints WAL at wal_autocheckpoint=250 frames.
2947 # An explicit PRAGMA wal_checkpoint(FULL) here used to
2948 # block other writers long enough to exhaust busy_timeout
2949 # under bulk-download concurrency (#4197).
2950 session.commit()
2952 logger.info(
2953 f"Successfully indexed document {document_id} for collection {collection_id} "
2954 f"with {len(chunks)} chunks"
2955 )
2957 return {
2958 "status": "success",
2959 "chunk_count": len(chunks),
2960 }
2962 except Exception as e:
2963 # The session is shared (thread-local) with the caller.
2964 # If session.flush() or session.commit() raised, the session
2965 # is in PendingRollbackError state until rolled back —
2966 # leaving subsequent operations to cascade. Roll back BEFORE
2967 # returning the error dict so the caller sees a clean
2968 # session. (Same pattern as the #3827 fix.)
2969 safe_rollback(session, "library_rag_service.index_document")
2970 # Backstop for the store-durable-before-DB-commit window: a
2971 # REPLACE reindex ran store.apply() (DURABLY removing this
2972 # source's prior vectors) before the DB commit, so the rollback
2973 # above resurrects rows whose vectors are already gone while
2974 # reverting indexed→its pre-attempt True. Flag for re-index so
2975 # the reconciler re-embeds and store+DB reconverge.
2976 self._flag_document_for_reindex(
2977 session, document_id, collection_id
2978 )
2979 logger.exception(
2980 f"Error indexing document {document_id} for collection {collection_id}"
2981 )
2982 return {
2983 "status": "error",
2984 "error": f"Operation failed: {type(e).__name__}",
2985 }
2987 def index_all_documents(
2988 self,
2989 collection_id: str,
2990 force_reindex: bool = False,
2991 progress_callback=None,
2992 ) -> Dict[str, Any]:
2993 """
2994 Index all documents in a collection into RAG.
2996 Args:
2997 collection_id: UUID of the collection to index
2998 force_reindex: Whether to force reindexing already indexed documents
2999 progress_callback: Optional callback function called after each document with (current, total, doc_title, status)
3001 Returns:
3002 Dict with counts of successful, skipped, and failed documents.
3003 Inside an existing ``get_user_db_session`` block on this thread
3004 it indexes nothing and returns ``status`` ``"error"`` with zero
3005 counts.
3006 """
3007 scope_error = _indexing_scope_error()
3008 if scope_error is not None:
3009 # Reject before the read scope below: nested in the caller's
3010 # scope, its exception cleanup would roll back the caller's
3011 # Session.
3012 return {
3013 **scope_error,
3014 "successful": 0,
3015 "skipped": 0,
3016 "failed": 0,
3017 "errors": [],
3018 }
3019 with get_user_db_session(self.username, self.db_password) as session:
3020 # Get all DocumentCollection entries for this collection
3021 query = session.query(DocumentCollection).filter_by(
3022 collection_id=collection_id
3023 )
3025 if not force_reindex:
3026 # Only index documents that haven't been indexed yet
3027 query = query.filter_by(indexed=False)
3029 # Materialize the values needed by the loop before leaving the
3030 # read scope. A thread-local session is reused by nested calls;
3031 # keeping this scope open during index_document would retain its
3032 # connection throughout splitting, embedding, and callbacks.
3033 doc_info = []
3034 for doc_collection in query.all():
3035 document_id = doc_collection.document_id
3036 document = (
3037 session.query(Document).filter_by(id=document_id).first()
3038 )
3039 doc_info.append(
3040 (document_id, document.title if document else "Unknown")
3041 )
3043 if not doc_info:
3044 return {
3045 "status": "info",
3046 "message": "No documents to index",
3047 "successful": 0,
3048 "skipped": 0,
3049 "failed": 0,
3050 }
3052 results = {"successful": 0, "skipped": 0, "failed": 0, "errors": []}
3053 total = len(doc_info)
3055 for idx, (document_id, title) in enumerate(doc_info, 1):
3056 result = self.index_document(
3057 document_id, collection_id, force_reindex
3058 )
3060 if result["status"] == "success":
3061 results["successful"] += 1
3062 elif result["status"] in ("skipped", "cleared"):
3063 # "cleared" (a previously-indexed doc emptied to no text,
3064 # its stale vectors purged) is a handled non-failure, not an
3065 # error — counting it in "failed" would surface a blank-error
3066 # failure for a correct purge.
3067 results["skipped"] += 1
3068 else:
3069 results["failed"] += 1
3070 results["errors"].append(
3071 {
3072 "doc_id": document_id,
3073 "title": title,
3074 "error": result.get("error"),
3075 }
3076 )
3078 # Call progress callback if provided
3079 if progress_callback:
3080 progress_callback(idx, total, title, result["status"])
3082 logger.info(
3083 f"Indexed collection {collection_id}: "
3084 f"{results['successful']} successful, "
3085 f"{results['skipped']} skipped, "
3086 f"{results['failed']} failed"
3087 )
3089 return results
3091 def remove_document_from_rag(
3092 self, document_id: str, collection_id: str
3093 ) -> Dict[str, Any]:
3094 with _hold_collection_lock(
3095 self.username, collection_id, purpose="transaction"
3096 ):
3097 return self._remove_document_from_rag_locked(
3098 document_id, collection_id
3099 )
3101 def _remove_document_from_rag_locked(
3102 self, document_id: str, collection_id: str
3103 ) -> Dict[str, Any]:
3104 """
3105 Remove a document's chunks from RAG for a specific collection.
3107 Args:
3108 document_id: UUID of the Document to remove
3109 collection_id: UUID of the Collection to remove from
3111 Returns:
3112 Dict with status and count of removed chunks
3113 """
3114 with get_user_db_session(self.username, self.db_password) as session:
3115 # Get the DocumentCollection entry
3116 doc_collection = (
3117 session.query(DocumentCollection)
3118 .filter_by(document_id=document_id, collection_id=collection_id)
3119 .first()
3120 )
3122 if not doc_collection:
3123 return {
3124 "status": "error",
3125 "error": "Document not found in collection",
3126 }
3128 try:
3129 # Get collection name in the format collection_<uuid>
3130 collection = (
3131 session.query(Collection)
3132 .filter_by(id=collection_id)
3133 .first()
3134 )
3135 # Use collection_<uuid> format for internal storage
3136 collection_name = (
3137 f"collection_{collection_id}" if collection else "unknown"
3138 )
3140 # Remove the document's vectors from the FAISS index FIRST.
3141 # ``remove_documents_from_index`` -> ``VectorIndex.delete``
3142 # looks up the DocumentChunk rows by (source_type, source_id,
3143 # collection_name) to learn which FAISS int ids to remove,
3144 # and deletes those matching rows itself as part of the same
3145 # operation. If the DB rows were deleted first, that lookup
3146 # would find nothing and the vectors (and their persisted
3147 # text) would be stranded in the FAISS store forever — never
3148 # removed and never re-discoverable by a later cleanup.
3149 # Best-effort: a FAISS hiccup must not fail the DB removal
3150 # below.
3151 # ``remove_documents_from_index`` deletes the matching chunk
3152 # rows itself (one row per removed vector) and returns that
3153 # count. Capture it so the reported ``deleted_count`` reflects
3154 # the rows removed here, not just the backstop sweep below: in
3155 # the normal path those rows are already gone by the time the
3156 # backstop runs, so the backstop returns ~0 and — left
3157 # uncounted — would under-report the removal (often as 0) to
3158 # the caller/UI.
3159 removed_vectors = 0
3160 try:
3161 removed_vectors = self.remove_documents_from_index(
3162 [document_id], collection_id
3163 )
3164 except Exception:
3165 logger.exception(
3166 f"Could not remove document {document_id} vectors "
3167 f"from FAISS index for collection {collection_id}"
3168 )
3170 # DB-only cleanup for any rows ``remove_documents_from_index``
3171 # did not match (e.g. no current FAISS index for this
3172 # collection). Disjoint from ``removed_vectors`` above (those
3173 # rows are already committed-deleted), so the two sum to the
3174 # true total removed.
3175 deleted_count = removed_vectors + (
3176 self.embedding_manager._delete_chunks_from_db(
3177 collection_name=collection_name,
3178 source_id=document_id,
3179 )
3180 )
3182 # Update DocumentCollection RAG status
3183 doc_collection.indexed = False
3184 doc_collection.chunk_count = 0
3185 doc_collection.last_indexed_at = None
3186 # Clear the canonical RagDocumentStatus row too. Row-existence is
3187 # the "indexed" marker get_rag_stats / the RAG status route read,
3188 # so it MUST move together with DocumentCollection.indexed — a
3189 # leftover row would report the just-removed document as still
3190 # indexed.
3191 session.query(RagDocumentStatus).filter_by(
3192 document_id=document_id, collection_id=collection_id
3193 ).delete(synchronize_session=False)
3194 session.commit()
3196 logger.info(
3197 f"Removed {deleted_count} chunks for document {document_id} from collection {collection_id}"
3198 )
3200 return {"status": "success", "deleted_count": deleted_count}
3202 except Exception as e:
3203 # session.commit() above can raise; without rollback the
3204 # shared thread-local session stays poisoned for the
3205 # caller's next operation (issue #3827 pattern).
3206 safe_rollback(
3207 session, "library_rag_service.remove_document_from_rag"
3208 )
3209 # Backstop for the store-durable-before-DB-commit window:
3210 # remove_documents_from_index -> VectorIndex.delete ran
3211 # store.apply() (DURABLY removing this document's vectors) BEFORE
3212 # the DB status work above committed, so the rollback can leave
3213 # the row still marked indexed=True with its vectors already gone
3214 # -- indexed-but-unsearchable, and over-reported as indexed by
3215 # get_rag_stats with no auto-recovery. Durably COMPLETE the
3216 # removal's status change so the DB matches the removed vectors:
3217 # mark it not-indexed and drop the canonical RagDocumentStatus
3218 # row. This is a REMOVAL backstop, not a re-index one (contrast
3219 # index_document's _flag_document_for_reindex) -- the vectors are
3220 # intentionally gone, so we finish the removal rather than
3221 # resurrect the document. Best-effort: a retry of the removal
3222 # reconciles it if this too fails.
3223 try:
3224 session.query(DocumentCollection).filter_by(
3225 document_id=document_id, collection_id=collection_id
3226 ).update(
3227 {
3228 "indexed": False,
3229 "chunk_count": 0,
3230 "last_indexed_at": None,
3231 },
3232 synchronize_session=False,
3233 )
3234 session.query(RagDocumentStatus).filter_by(
3235 document_id=document_id, collection_id=collection_id
3236 ).delete(synchronize_session=False)
3237 session.commit()
3238 except Exception:
3239 safe_rollback(
3240 session,
3241 "library_rag_service.remove_document_from_rag backstop",
3242 )
3243 logger.warning(
3244 f"Could not finalize removal state for document "
3245 f"{document_id} in collection {collection_id}; a retry "
3246 "of the removal will reconcile it."
3247 )
3248 logger.exception(
3249 f"Error removing document {document_id} from collection {collection_id}"
3250 )
3251 return {
3252 "status": "error",
3253 "error": f"Operation failed: {type(e).__name__}",
3254 }
3256 def purge_document_chunks(
3257 self, document_id: str, collection_id: str
3258 ) -> int:
3259 with _hold_collection_lock(
3260 self.username, collection_id, purpose="transaction"
3261 ):
3262 return self._purge_document_chunks_locked(
3263 document_id, collection_id
3264 )
3266 def _purge_document_chunks_locked(
3267 self, document_id: str, collection_id: str
3268 ) -> int:
3269 """Delete a document's chunk rows from a collection's RAG store
3270 WITHOUT requiring the ``DocumentCollection`` join row to still exist.
3272 ``remove_document_from_rag`` short-circuits (returns "Document not
3273 found in collection") when the join row is gone — which is exactly
3274 the state after a ``Document`` is cascade-deleted. Callers that
3275 delete the Document first (note deletion) must use this to actually
3276 purge the now-orphaned ``DocumentChunk`` rows; otherwise the chunks
3277 (and their embedding ids) linger and semantic search keeps
3278 resolving hits to a Document row that no longer exists.
3280 Mirrors the chunk-delete half of ``remove_document_from_rag``.
3281 FAISS-vector pruning is NOT done here: pair this with
3282 ``purge_document_vectors`` when the document is being deleted
3283 outright — replace-on-reindex can never fire again for a deleted
3284 document id. Returns the number of chunk rows deleted.
3285 """
3286 collection_name = f"collection_{collection_id}"
3287 return self.embedding_manager._delete_chunks_from_db(
3288 collection_name=collection_name,
3289 source_id=document_id,
3290 )
3292 def purge_document_vectors(
3293 self, document_id: str, collection_id: str
3294 ) -> int:
3295 with _hold_collection_lock(
3296 self.username, collection_id, purpose="transaction"
3297 ):
3298 return self._purge_document_vectors_locked(
3299 document_id, collection_id
3300 )
3302 def _purge_document_vectors_locked(
3303 self, document_id: str, collection_id: str
3304 ) -> int:
3305 """Remove a deleted document's chunk vectors from the collection's
3306 vector store.
3308 Companion to ``purge_document_chunks`` for callers that delete the
3309 ``Document`` row itself (note deletion). Without this the vectors
3310 lingered until a replace-on-reindex of the SAME document id — which
3311 can never fire again once the document is deleted — so collection
3312 search kept serving the deleted content indefinitely.
3314 Routes to ``VectorIndex.delete(source_type="document",
3315 source_id=document_id)``. Unlike the pre-cutover implementation,
3316 no separate "shared chunk ownership" check is needed here: the new
3317 store is one-row-per-source (see the "One row = one vector" note in
3318 vector_stores/facade.py) — identical text in two documents gets two
3319 independent rows/vectors, so deleting this document's rows can
3320 never remove another document's vector.
3322 Deliberately does NOT go through ``_get_vector_index`` /
3323 ``_get_or_create_rag_index``: a deletion path must not create an
3324 index record (or probe the embedding provider) just to delete from
3325 it, and does not quarantine/rebuild on a failed load — an
3326 unreadable index has nothing purgeable right now; quarantine/
3327 rebuild decisions stay with the indexing paths.
3329 Returns the number of vectors removed.
3330 """
3331 collection_name = f"collection_{collection_id}"
3332 index_hash = self._get_index_hash(
3333 collection_name, self.embedding_model, self.embedding_provider
3334 )
3335 with get_user_db_session(self.username, self.db_password) as session:
3336 rag_index = self._find_matching_rag_index(
3337 session, collection_name, index_hash
3338 )
3339 if rag_index is None:
3340 # Collection was never indexed with this configuration —
3341 # no vectors to purge. Deliberately NOT
3342 # _get_or_create_rag_index: a deletion path must not
3343 # create index records (or probe the embedding provider).
3344 return 0
3345 self.rag_index_record = rag_index
3347 # Resolve where the file ACTUALLY lives, the way the index/search paths
3348 # do — but without _preflight's side effects (this deletion path must
3349 # not relocate/record/create). rag_index.index_path is the authoritative
3350 # DB record (the legacy shared-cache path pre-migration, the per-user
3351 # path after _preflight relocates it); _get_index_path(index_hash) only
3352 # recomputes the POST-migration location. Checking only the recomputed
3353 # path silently no-ops on an un-migrated install — the file is still at
3354 # the legacy index_path, so the deleted document's real vector survives.
3355 recomputed_path = self._get_index_path(rag_index.index_hash)
3356 recorded_path = (
3357 Path(rag_index.index_path) if rag_index.index_path else None
3358 )
3359 if recorded_path is not None and recorded_path.exists(): 3359 ↛ 3360line 3359 didn't jump to line 3360 because the condition on line 3359 was never true
3360 index_path = recorded_path
3361 elif recomputed_path.exists():
3362 index_path = recomputed_path
3363 else:
3364 return 0
3365 verified, reason = self.integrity_manager.verify_file(index_path)
3366 if not verified: 3366 ↛ 3375line 3366 didn't jump to line 3375 because the condition on line 3366 was always true
3367 # Leave quarantine/rebuild decisions to the indexing paths;
3368 # an unreadable index has nothing purgeable right now.
3369 logger.warning(
3370 f"Skipping vector purge for document {document_id}: index "
3371 f"{index_path} fails verification ({reason})"
3372 )
3373 return 0
3375 lock = _get_faiss_write_lock(self.username, str(index_path))
3377 # nested=True: these run reentrantly inside VectorIndex.apply(), which
3378 # holds just-flushed, uncommitted DocumentChunk rows on this same
3379 # thread-local session. The integrity write must run in a SAVEPOINT and
3380 # let apply()'s caller commit — a plain commit/rollback here would
3381 # settle or DISCARD those pending rows (silently losing indexed text
3382 # while the vectors persist). See FileIntegrityManager._integrity_write.
3383 def _record(path: Path) -> None:
3384 self.integrity_manager.record_file(
3385 path,
3386 related_entity_type="rag_index",
3387 related_entity_id=rag_index.id,
3388 nested=True,
3389 )
3391 def _verify(path: Path):
3392 return self.integrity_manager.verify_file(path, nested=True)
3394 try:
3395 vindex = VectorIndex(
3396 username=self.username,
3397 db_password=self.db_password,
3398 # Delete-only path: never materialise the embedding backend.
3399 # VectorIndex.delete() resolves rows by id and never calls
3400 # embed_*; forcing the model here would cost seconds + hundreds
3401 # of MB of RSS per deleted document for nothing.
3402 embeddings=_LazyEmbeddings(self.embedding_manager),
3403 embedding_model=self.embedding_model,
3404 embedding_model_type=EmbeddingProvider(self.embedding_provider),
3405 collection_name=collection_name,
3406 dimension=rag_index.embedding_dimension,
3407 path=index_path,
3408 lock=lock,
3409 integrity_record=_record,
3410 integrity_verify=_verify,
3411 # Use the STORED index config, not the service instance's — a
3412 # service built with defaults (index_type="flat") would evaluate
3413 # supports_delete wrongly for an HNSW collection and silently
3414 # fail the removal. NULL (legacy rows) falls back to self.*.
3415 index_type=rag_index.index_type or self.index_type,
3416 metric=rag_index.distance_metric or self.distance_metric,
3417 normalize=(
3418 rag_index.normalize_vectors
3419 if rag_index.normalize_vectors is not None
3420 else self.normalize_vectors
3421 ),
3422 )
3423 stats = vindex.delete(source_type="document", source_id=document_id)
3424 except (ValueError, RuntimeError, OSError):
3425 # OSError too: this is a BEST-EFFORT purge whose caller goes on to
3426 # delete the DocumentChunk rows. A disk error from the FAISS
3427 # persist() (OSError) must not escape and abort that row deletion —
3428 # it would strand the chunk rows while leaving the vectors, the
3429 # opposite inconsistency. Skip the vector purge and let the caller
3430 # proceed; the vectors reconcile on the next re-index.
3431 logger.warning(
3432 f"Skipping vector purge for document {document_id}: could "
3433 f"not load or update index {index_path}"
3434 )
3435 return 0
3436 return stats.removed
3438 def index_documents_batch(
3439 self,
3440 doc_info: List[tuple],
3441 collection_id: str,
3442 force_reindex: bool = False,
3443 ) -> Dict[str, Dict[str, Any]]:
3444 """
3445 Index multiple documents in a batch for a specific collection.
3447 Args:
3448 doc_info: List of (doc_id, title) tuples
3449 collection_id: UUID of the collection to index for
3450 force_reindex: Whether to force reindexing even if already indexed
3452 Returns:
3453 Dict mapping doc_id to individual result. Inside an existing
3454 ``get_user_db_session`` block on this thread every result is the
3455 scope error and nothing is read or indexed.
3456 """
3457 scope_error = _indexing_scope_error()
3458 if scope_error is not None:
3459 # Reject before the read scope below, for the same reason as
3460 # index_all_documents.
3461 return {doc_id: dict(scope_error) for doc_id, _title in doc_info}
3462 results = {}
3463 doc_ids = [doc_id for doc_id, _ in doc_info]
3465 # Use one short read scope to collect only scalar state. Keep the
3466 # potentially slow per-document operations outside this scope.
3467 with get_user_db_session(self.username, self.db_password) as session:
3468 # Pre-load all documents for this batch
3469 documents = (
3470 session.query(Document).filter(Document.id.in_(doc_ids)).all()
3471 )
3473 # Create lookup for quick access
3474 existing_doc_ids = {doc.id for doc in documents}
3476 # Pre-load DocumentCollection entries
3477 doc_collections = (
3478 session.query(DocumentCollection)
3479 .filter(
3480 DocumentCollection.document_id.in_(doc_ids),
3481 DocumentCollection.collection_id == collection_id,
3482 )
3483 .all()
3484 )
3485 doc_collection_lookup = {
3486 dc.document_id: (dc.indexed, dc.chunk_count)
3487 for dc in doc_collections
3488 }
3490 # Process each document after the read scope releases its connection.
3491 for doc_id, title in doc_info:
3492 if doc_id not in existing_doc_ids:
3493 results[doc_id] = {
3494 "status": "error",
3495 "error": "Document not found",
3496 }
3497 continue
3499 # Check if already indexed via DocumentCollection
3500 doc_collection = doc_collection_lookup.get(doc_id)
3501 if doc_collection and doc_collection[0] and not force_reindex:
3502 results[doc_id] = {
3503 "status": "skipped",
3504 "message": "Document already indexed for this collection",
3505 "chunk_count": doc_collection[1],
3506 }
3507 continue
3509 # Do NOT short-circuit empty text_content here: a document/note
3510 # edited down to empty still has its PRIOR chunks + vectors on
3511 # disk, and only index_document's empty-content branch purges
3512 # them (via an empty-replace). Returning an "error" here left
3513 # that stale (plaintext) content permanently searchable through
3514 # the bulk "Index All" path — and self-perpetuating, since the
3515 # doc's indexed flag never flips. Route it through index_document
3516 # so the batch path gets the same purge-on-clear as the SSE /
3517 # single-document paths; it returns "cleared" (purged) or the
3518 # same "error" when there were no prior chunks.
3519 try:
3520 result = self.index_document(
3521 doc_id, collection_id, force_reindex
3522 )
3523 results[doc_id] = result
3524 except Exception as e:
3525 logger.exception(f"Error indexing document {doc_id} in batch")
3526 results[doc_id] = {
3527 "status": "error",
3528 "error": f"Indexing failed: {type(e).__name__}",
3529 }
3531 return results
3533 def rebuild_collection(
3534 self,
3535 doc_info,
3536 collection_id,
3537 *,
3538 max_workers=4,
3539 progress_callback=None,
3540 is_cancelled=None,
3541 ):
3542 """Replace the complete collection only after every document is ready."""
3543 scope_error = _indexing_scope_error()
3544 if scope_error is not None:
3545 raise RuntimeError(scope_error["error"])
3547 with _hold_collection_lock(
3548 self.username, collection_id, purpose="rebuild"
3549 ):
3550 with CollectionRebuild(
3551 self,
3552 collection_id,
3553 doc_info,
3554 session_factory=lambda: get_user_db_session(
3555 self.username, self.db_password
3556 ),
3557 ) as rebuild:
3558 aggregate = self.index_documents_parallel(
3559 doc_info,
3560 collection_id,
3561 force_reindex=True,
3562 max_workers=max_workers,
3563 progress_callback=progress_callback,
3564 is_cancelled=is_cancelled,
3565 _rebuild=rebuild,
3566 )
3567 if aggregate["cancelled"]:
3568 return aggregate
3569 if aggregate["failed"]:
3570 raise RuntimeError(
3571 "Collection rebuild failed; the previous index was preserved"
3572 )
3573 try:
3574 rebuild.publish(is_cancelled)
3575 except RebuildCancelled:
3576 aggregate["cancelled"] = True
3577 return aggregate
3579 def index_documents_parallel(
3580 self,
3581 doc_info: List[Tuple[str, str]],
3582 collection_id: str,
3583 force_reindex: bool = False,
3584 max_workers: int = 4,
3585 progress_callback: Optional[
3586 Callable[[int, int, str, str], None]
3587 ] = None,
3588 is_cancelled: Optional[Callable[[], bool]] = None,
3589 *,
3590 _rebuild: CollectionRebuild | None = None,
3591 ) -> Dict[str, Any]:
3592 """
3593 Index many documents concurrently with bounded fan-out.
3595 Each per-document task runs the prepared pipeline
3596 (:meth:`_prepare_document` for load/split/embed, then the
3597 coordinator's serialized durable write) unless a subclass overrides
3598 :meth:`_index_one` or :meth:`index_document`, in which case that
3599 per-doc path is invoked directly — the per-doc logic (split, embed,
3600 persist chunks, FAISS-merge-under-lock, mark indexed, update stats)
3601 is unchanged. The bounded fan-out only changes *which* sequential
3602 loop dispatches the work: before this method, six call sites in the
3603 RAG routes (main's ``research_library/routes/rag_routes.py``, now
3604 :mod:`local_deep_research.web.routers.rag`) and the background
3605 reconciler in :mod:`scheduler.background` each did this one doc at
3606 a time. The slow embedding round-trip (outside the FAISS write
3607 lock) is now overlapped across documents.
3609 Concurrency invariants this method relies on (already true):
3611 * :meth:`index_document` opens a fresh per-thread DB session via
3612 :func:`get_user_db_session` — workers do not share sessions.
3613 * All FAISS mutations go through ``VectorIndex.apply()`` in
3614 :mod:`local_deep_research.vector_stores`, which is only
3615 invoked under the per-(username, index_path) write lock
3616 returned by :func:`_get_faiss_write_lock`. The service passes
3617 that lock into the vector-store facade at the call sites
3618 around L422, L879, L1097, and L2025, so concurrent
3619 ``index_document`` workers serialise their FAISS appends on
3620 the same lock without any extra coordination here.
3621 * ``LocalEmbeddingManager.embeddings`` uses double-checked
3622 locking for first-load.
3623 * ``is_cancelled`` is polled between completions AND between
3624 submissions. The bounded-submission loop admits new work into
3625 the pool only when ``is_cancelled()`` returns False, and
3626 each :func:`_run_one` worker re-checks ``is_cancelled()``
3627 immediately before delegating to ``index_document``. This
3628 closes the false-poll-to-submit gap (the cancel signal can
3629 be set between the main thread's poll and ``pool.submit()``)
3630 so post-cancel starts are zero by construction. The final
3631 ``is_cancelled()`` poll in the ``finally`` block also
3632 preserves ``cancelled=True`` in the aggregate when the
3633 signal lands while the last in-flight worker is running.
3634 * In-flight embedding batches cannot be interrupted (Python
3635 has no safe way to kill an HTTP round-trip or a
3636 SentenceTransformer forward pass mid-array) — but
3637 ``_prepare_document`` embeds in bounded batches
3638 (``_EMBEDDING_CANCEL_BATCH_SIZE``) with an ``is_cancelled()``
3639 check between them, so the uncancellable tail is one batch,
3640 not the whole document. The helper returns via
3641 ``pool.shutdown(wait=True)`` after in-flight batches finish
3642 naturally.
3644 Args:
3645 doc_info: List of ``(doc_id, title)`` pairs to index.
3646 collection_id: UUID of the collection.
3647 force_reindex: Forwarded to :meth:`index_document`.
3648 max_workers: Bound on the size of the
3649 ``ThreadPoolExecutor`` — i.e. how many documents can be
3650 **preparing** concurrently (load + split + embed). It does
3651 NOT bound the prepared-but-unwritten backlog: completed
3652 futures hold their full vector batches in memory until the
3653 coordinator's serial write stage picks them up, so for very
3654 large batches the resident memory footprint can exceed what
3655 ``max_workers`` alone suggests. Embedding is the slow stage
3656 in practice, so this is acceptable for the default 4-worker
3657 fan-out; a caller with many large documents should either
3658 shrink ``max_workers`` or chunk the submission. ``1`` falls
3659 back to fully sequential behaviour (still useful when a
3660 flaky network makes concurrent httpx calls worse than
3661 serial). Default 4.
3662 progress_callback: Optional hook called once per completed
3663 document with ``(completed, total, title, status)`` —
3664 matches the legacy shape used by
3665 :meth:`index_all_documents` so existing UI progress
3666 paths keep working.
3667 is_cancelled: Optional zero-arg callable polled between
3668 completions AND between submissions. When it returns
3669 ``True``, no further work is admitted into the pool;
3670 in-flight futures finish naturally. The executor is
3671 always shut down with ``wait=True`` so no worker can
3672 outlive the caller's service context.
3674 Returns:
3675 ``{"successful": int, "skipped": int, "failed": int,
3676 "errors": [...], "results": {doc_id: {...}},
3677 "cancelled": bool, "total": int}``. Inside an existing
3678 ``get_user_db_session`` block on this thread nothing is loaded,
3679 embedded or written and no callback runs: every document fails
3680 with the scope error.
3681 """
3682 scope_error = _indexing_scope_error()
3683 if scope_error is not None:
3684 # Pool workers run on their own threads, where this caller's
3685 # scope is invisible to _prepare_document's guard. Without this
3686 # check every document would be loaded and embedded before the
3687 # coordinator rejected its write.
3688 title_lookup = dict(doc_info)
3689 return {
3690 "successful": 0,
3691 "skipped": 0,
3692 "failed": len(title_lookup),
3693 "errors": [
3694 {
3695 "doc_id": doc_id,
3696 "title": title,
3697 "error": scope_error["error"],
3698 }
3699 for doc_id, title in title_lookup.items()
3700 ],
3701 "results": {
3702 doc_id: dict(scope_error) for doc_id in title_lookup
3703 },
3704 "cancelled": False,
3705 "total": len(doc_info),
3706 }
3707 if max_workers < 1:
3708 max_workers = 1
3710 total = len(doc_info)
3711 if total == 0:
3712 return {
3713 "successful": 0,
3714 "skipped": 0,
3715 "failed": 0,
3716 "errors": [],
3717 "results": {},
3718 "cancelled": False,
3719 "total": 0,
3720 }
3722 # Title-lookup keyed by doc_id — futures return in ``as_completed``
3723 # order, not submission order, so we cannot rely on positional
3724 # iteration when emitting progress.
3725 title_lookup = dict(doc_info)
3727 results: Dict[str, Dict[str, Any]] = {}
3728 counters = {"successful": 0, "skipped": 0, "failed": 0}
3729 errors: List[Dict[str, Any]] = []
3730 cancelled = False
3732 def _run_one(doc_id: str) -> Dict[str, Any]:
3733 """Single-doc worker — never raises; returns a status dict."""
3734 try:
3735 if _rebuild is not None:
3736 return _rebuild.prepare(doc_id, is_cancelled)
3737 # Three seams tell us to bypass the prepared/serialized
3738 # pipeline and call the per-doc method directly:
3739 #
3740 # 1. A class-level subclass override of ``_index_one`` —
3741 # detected via
3742 # ``type(self)._index_one is not LibraryRAGService._index_one``
3743 # so it works regardless of whether the override
3744 # has ``__func__`` (the old
3745 # ``self._index_one.__func__ is not LibraryRAGService._index_one``
3746 # check raised AttributeError on instance-attribute
3747 # patches that don't expose ``__func__``).
3748 # 2. A class-level subclass override of the legacy
3749 # ``index_document`` extension point this seam
3750 # replaced — callers that override ``index_document``
3751 # at class level (not just instance-patch it for
3752 # tests) still get the per-doc bypass, otherwise they
3753 # would silently fall through to the prepared
3754 # pipeline and lose the override.
3755 # 3. An instance-level patch of ``_index_one`` or
3756 # ``index_document`` — the legacy seams used by the
3757 # existing parallel-runner tests that mock the
3758 # per-doc method directly.
3759 #
3760 # If none is present, run the prepared pipeline: split
3761 # and embed concurrently, then hand the prepared document
3762 # off to the coordinator for serialized durable writes.
3763 cls = type(self)
3764 if cls._index_one is not LibraryRAGService._index_one:
3765 return self._index_one(doc_id, collection_id, force_reindex)
3766 if cls.index_document is not LibraryRAGService.index_document:
3767 return self.index_document(
3768 doc_id, collection_id, force_reindex
3769 )
3770 if (
3771 "_index_one" in self.__dict__
3772 or "index_document" in self.__dict__
3773 ):
3774 return self._index_one(doc_id, collection_id, force_reindex)
3775 prepared_result = self._prepare_document(
3776 doc_id, collection_id, force_reindex, is_cancelled
3777 )
3778 if prepared_result.get("status") == "needs_serial_write": 3778 ↛ 3779line 3778 didn't jump to line 3779 because the condition on line 3778 was never true
3779 return self.index_document(
3780 doc_id, collection_id, force_reindex
3781 )
3782 return prepared_result
3783 except Exception as exc:
3784 # ``index_document`` already returns an error dict on its
3785 # own internal failures; this catches only bugs / hook
3786 # breaks outside that contract.
3787 logger.exception(f"Parallel index worker for {doc_id} crashed")
3788 return {
3789 "status": "error",
3790 "error": f"Indexing failed: {type(exc).__name__}",
3791 }
3793 # De-duplicated submission queue preserving input order. The
3794 # underlying ``index_document`` is not idempotent at the DB level
3795 # beyond its own checks; without this guard a buggy caller could
3796 # double-insert chunks from the same doc.
3797 pending_doc_ids: List[str] = []
3798 seen: Set[str] = set()
3799 for doc_id, _title in doc_info:
3800 if doc_id in seen:
3801 continue
3802 seen.add(doc_id)
3803 pending_doc_ids.append(doc_id)
3805 pool = ThreadPoolExecutor(
3806 max_workers=max_workers,
3807 thread_name_prefix="ldr-index-doc-",
3808 )
3809 try:
3810 future_to_doc: Dict[Any, str] = {}
3811 submit_idx = 0
3812 completed = 0
3813 # Bounded-submission loop: poll ``is_cancelled`` BEFORE
3814 # admitting any new work, then submit up to ``max_workers``
3815 # futures, then wait for one completion.
3816 #
3817 # This closes the post-cancel submission leak in the
3818 # previous all-upfront submission model: with the old
3819 # loop, all futures were submitted before the first
3820 # ``is_cancelled()`` poll, so queued futures could be
3821 # picked up by idle workers between the caller's cancel
3822 # signal (e.g. ``_sse_cancel.set()`` in
3823 # ``web/routers/rag.py::index_collection``) and the helper's
3824 # ``pool.shutdown(cancel_futures=True)``. Here, new work
3825 # is admitted only when ``is_cancelled()`` returns False,
3826 # so post-cancel starts are zero by construction.
3827 while future_to_doc or submit_idx < len(pending_doc_ids):
3828 # 1. Cancellation fence: stop submitting new work if
3829 # cancelled. Once we observe cancellation we don't
3830 # re-poll — the gate stays closed until the loop
3831 # drains.
3832 if (
3833 not cancelled
3834 and is_cancelled is not None
3835 and is_cancelled()
3836 ):
3837 cancelled = True
3838 logger.info(
3839 f"Parallel indexing cancelled at {completed}/{total}"
3840 )
3841 for f in future_to_doc: 3841 ↛ 3842line 3841 didn't jump to line 3842 because the loop on line 3841 never started
3842 f.cancel()
3844 # 2. Admit new work up to ``max_workers`` in flight
3845 # (skipped once cancelled).
3846 if not cancelled:
3847 while len(future_to_doc) < max_workers and submit_idx < len(
3848 pending_doc_ids
3849 ):
3850 if is_cancelled is not None and is_cancelled():
3851 cancelled = True
3852 logger.info(
3853 f"Parallel indexing cancelled during submission at {completed}/{total}"
3854 )
3855 for f in future_to_doc:
3856 f.cancel()
3857 break
3858 doc_id = pending_doc_ids[submit_idx]
3859 submit_idx += 1
3860 future_to_doc[pool.submit(_run_one, doc_id)] = doc_id
3862 # 3. Nothing in flight and nothing left to submit → done.
3863 if not future_to_doc:
3864 break
3866 # 4. Wait for the next completion, then process it.
3867 try:
3868 done_iter = as_completed(list(future_to_doc.keys()))
3869 fut = next(done_iter, None)
3870 except Exception:
3871 logger.exception(
3872 "as_completed failed unexpectedly during parallel indexing"
3873 )
3874 fut = None
3875 if fut is None: 3875 ↛ 3879line 3875 didn't jump to line 3879 because the condition on line 3875 was never true
3876 # Iterator short-circuited without yielding — loop
3877 # back. The next iteration's cancel poll / submit
3878 # check will either submit more work or exit.
3879 continue
3881 doc_id = future_to_doc.pop(fut)
3882 try:
3883 result = fut.result()
3884 except CancelledError:
3885 result = {
3886 "status": "skipped",
3887 "message": "Cancelled before submission",
3888 }
3889 except Exception as exc:
3890 logger.exception(
3891 f"Future for doc {doc_id} raised unexpectedly"
3892 )
3893 result = {
3894 "status": "error",
3895 "error": f"Worker crashed: {type(exc).__name__}",
3896 }
3897 status = result.get("status", "error")
3898 prepared = result.get("prepared")
3899 if status == "prepared" and isinstance(
3900 prepared, _PreparedDocument
3901 ):
3902 result = (
3903 _rebuild.write(prepared)
3904 if _rebuild is not None
3905 else self._write_prepared_document(prepared)
3906 )
3907 status = result.get("status", "error")
3908 results[doc_id] = result
3909 if status == "success":
3910 counters["successful"] += 1
3911 elif status in ("skipped", "cleared"):
3912 counters["skipped"] += 1
3913 elif status == "cancelled":
3914 # A worker observed the cancel signal mid-document
3915 # (see _prepare_document). Count it as skipped — same
3916 # as "Cancelled before submission" below — and flip
3917 # the aggregate even if the main-thread polls have not
3918 # observed the signal yet.
3919 counters["skipped"] += 1
3920 cancelled = True
3921 else:
3922 counters["failed"] += 1
3923 errors.append(
3924 {
3925 "doc_id": doc_id,
3926 "title": title_lookup.get(doc_id, "Unknown"),
3927 "error": result.get("error"),
3928 }
3929 )
3930 completed += 1
3931 if progress_callback is not None:
3932 try:
3933 progress_callback(
3934 completed,
3935 total,
3936 title_lookup.get(doc_id, "Unknown"),
3937 status,
3938 )
3939 except Exception:
3940 # A buggy callback must not poison the indexing
3941 # run — log and continue.
3942 logger.exception(
3943 "Parallel index progress_callback raised"
3944 )
3945 finally:
3946 if not cancelled and is_cancelled is not None and is_cancelled():
3947 cancelled = True
3948 # Always ``wait=True`` so in-flight workers finish before the
3949 # helper returns. With bounded submission no new work is
3950 # admitted after cancel is observed, so the only in-flight
3951 # work at shutdown is what was already running — these finish
3952 # naturally. ``cancel_futures=cancelled`` is a no-op here
3953 # (no queued futures by construction) but is kept for
3954 # symmetry with other shutdown sites in the codebase and as
3955 # defense-in-depth if a future change relaxes the bounded-
3956 # submission invariant.
3957 pool.shutdown(wait=True, cancel_futures=cancelled)
3959 if cancelled:
3960 # Mark docs that were never submitted (still in the
3961 # pending queue past ``submit_idx``) as skipped. Submitted
3962 # docs all have results — either from processing in the
3963 # loop above or from the ``wait=True`` shutdown drain —
3964 # so we don't need to touch them here.
3965 for doc_id in pending_doc_ids[submit_idx:]:
3966 if doc_id not in results: 3966 ↛ 3965line 3966 didn't jump to line 3965 because the condition on line 3966 was always true
3967 results[doc_id] = {
3968 "status": "skipped",
3969 "message": "Cancelled before submission",
3970 }
3971 counters["skipped"] += 1
3973 logger.info(
3974 f"Parallel indexing complete: "
3975 f"{counters['successful']} successful, "
3976 f"{counters['skipped']} skipped, "
3977 f"{counters['failed']} failed"
3978 f"{' (cancelled)' if cancelled else ''}"
3979 )
3981 return {
3982 "successful": counters["successful"],
3983 "skipped": counters["skipped"],
3984 "failed": counters["failed"],
3985 "errors": errors,
3986 "results": results,
3987 "cancelled": cancelled,
3988 "total": total,
3989 }
3991 def reconcile_collection_index(self, collection_id: str) -> Dict[str, Any]:
3992 """Reconcile indexed flags and stats against durable FAISS membership.
3994 A document is indexed only when it has at least one current-model chunk
3995 row, none were invalidated by text replacement, and every such chunk
3996 id exists in the live FAISS index. Stale status rows are removed and
3997 compatibility flags/stats are rebuilt atomically.
3999 Safety invariant: if the live FAISS index reports zero ids (transient
4000 load error, missing file, on-disk corruption the verifier just
4001 quarantined) but the DB still has DocumentChunk rows for the collection,
4002 we REFUSE to shrink state. The function returns
4003 ``{"reconciliation_skipped": True, "reason": "..."}`` and the worker
4004 surfaces this as a task failure rather than clearing every indexed flag
4005 and RagDocumentStatus row. A mass-clear from a transient store fault
4006 would orphan every search hit and require a full re-index to recover.
4008 ``RagDocumentStatus.indexed_at`` is preserved for durable docs that
4009 already have a row — only newly durable docs get the current timestamp.
4010 """
4011 collection_name = f"collection_{collection_id}"
4012 with self._collection_transaction_lock(collection_id):
4013 vindex = self._get_vector_index(
4014 collection_id,
4015 collection_name,
4016 reset_stale_state=False,
4017 # This call site never embeds a query — it only calls
4018 # ``live_ids()`` below — so it gets none of the fail-closed
4019 # dimension check a real search gets for free from
4020 # ``FaissVectorStore._prepare``. Without explicitly opting
4021 # back in here, an embedding-model dimension drift would go
4022 # completely undetected: the file still loads (its stored
4023 # ``dimension`` came from the DB row, which always matches
4024 # ``index.d``), and a non-empty live index means the
4025 # refuse-to-shrink guard below never fires either.
4026 probe_dimension=True,
4027 )
4028 live_ids = set(vindex.live_ids())
4029 rag_index_id = self.rag_index_record.id
4030 with get_user_db_session(
4031 self.username, self.db_password
4032 ) as session:
4033 # Serialize this SQLite repair with text invalidation. The
4034 # vector store was initialized above; no nested setup commit
4035 # may release this write lock before flags/status are saved.
4036 session.query(RAGIndex).filter_by(id=rag_index_id).update(
4037 {RAGIndex.id: rag_index_id}, synchronize_session=False
4038 )
4039 rows = (
4040 session.query(
4041 DocumentChunk.id,
4042 DocumentChunk.source_id,
4043 DocumentChunk.document_metadata,
4044 )
4045 .filter_by(
4046 source_type="document",
4047 collection_name=collection_name,
4048 embedding_model=self.embedding_model,
4049 embedding_model_type=EmbeddingProvider(
4050 self.embedding_provider
4051 ),
4052 )
4053 # Chunk rows can outlive their document: source_id has
4054 # no foreign key, and the purge after a delete is best
4055 # effort. A status row for such a document would break
4056 # the commit below on RagDocumentStatus's foreign key to
4057 # documents (#6748), so only live documents count.
4058 .filter(
4059 exists().where(Document.id == DocumentChunk.source_id)
4060 )
4061 .all()
4062 )
4063 ids_by_document: Dict[str, Set[int]] = {}
4064 invalidated_documents: Set[str] = set()
4065 for chunk_id, source_id, metadata in rows:
4066 if source_id is not None: 4066 ↛ 4065line 4066 didn't jump to line 4065 because the condition on line 4066 was always true
4067 ids_by_document.setdefault(str(source_id), set()).add(
4068 int(chunk_id)
4069 )
4070 if (metadata or {}).get("_ldr_text_invalidated"):
4071 invalidated_documents.add(str(source_id))
4073 # Refuse-to-shrink guard: a transient empty live_ids while
4074 # chunk rows exist would otherwise mass-clear every indexed
4075 # flag and RagDocumentStatus row, silently killing search
4076 # for every document in the collection. Treat as a
4077 # reconciliation failure; the worker surfaces the skip.
4078 if not live_ids and ids_by_document:
4079 logger.warning(
4080 "reconcile_collection_index: live_ids empty but "
4081 f"{len(ids_by_document)} document(s) have chunk rows; "
4082 "refusing to clear indexed flags"
4083 )
4084 return {
4085 "reconciliation_skipped": True,
4086 "reason": (
4087 "live vector store reported zero ids while "
4088 "DocumentChunk rows exist; refusing to clear "
4089 "indexed flags"
4090 ),
4091 "indexed_documents": 0,
4092 "indexed_chunks": 0,
4093 "live_vectors": 0,
4094 "orphan_vectors": 0,
4095 }
4097 durable = {
4098 doc_id: chunk_ids
4099 for doc_id, chunk_ids in ids_by_document.items()
4100 if chunk_ids
4101 and chunk_ids.issubset(live_ids)
4102 and doc_id not in invalidated_documents
4103 }
4105 # Snapshot existing indexed_at BEFORE the bulk delete so
4106 # durable docs that already had a status row keep their
4107 # original timestamp. Only docs newly promoted to durable
4108 # get ``now``.
4109 prior_indexed_at: Dict[str, Any] = {
4110 str(row.document_id): row.indexed_at
4111 for row in session.query(
4112 RagDocumentStatus.document_id,
4113 RagDocumentStatus.indexed_at,
4114 )
4115 .filter(RagDocumentStatus.collection_id == collection_id)
4116 .all()
4117 }
4119 links = (
4120 session.query(DocumentCollection)
4121 .filter_by(collection_id=collection_id)
4122 .all()
4123 )
4124 now = datetime.now(UTC)
4125 for link in links:
4126 chunk_ids = durable.get(str(link.document_id), set())
4127 link.indexed = bool(chunk_ids)
4128 link.chunk_count = len(chunk_ids)
4129 if chunk_ids:
4130 link.last_indexed_at = link.last_indexed_at or now
4132 session.query(RagDocumentStatus).filter_by(
4133 collection_id=collection_id
4134 ).delete(synchronize_session=False)
4135 for doc_id, chunk_ids in durable.items():
4136 session.add(
4137 RagDocumentStatus(
4138 document_id=doc_id,
4139 collection_id=collection_id,
4140 rag_index_id=rag_index_id,
4141 chunk_count=len(chunk_ids),
4142 indexed_at=prior_indexed_at.get(doc_id, now),
4143 )
4144 )
4146 rag_index = (
4147 session.query(RAGIndex).filter_by(id=rag_index_id).first()
4148 )
4149 if rag_index: 4149 ↛ 4153line 4149 didn't jump to line 4153 because the condition on line 4149 was always true
4150 rag_index.chunk_count = sum(map(len, durable.values()))
4151 rag_index.total_documents = len(durable)
4152 rag_index.last_updated_at = now
4153 session.commit()
4155 durable_ids = {i for ids in durable.values() for i in ids}
4156 return {
4157 "indexed_documents": len(durable),
4158 "indexed_chunks": len(durable_ids),
4159 "live_vectors": len(live_ids),
4160 "orphan_vectors": len(live_ids - durable_ids),
4161 }
4163 def has_indexed_documents(self, collection_id: str) -> bool:
4164 """Check search eligibility without collecting unrelated RAG statistics.
4166 Use the same exact-hash and legacy-configuration resolution as search
4167 and stats. Rows belonging to other indexes cannot make this one eligible.
4168 Resolving an index is read-only and never probes the embedding provider.
4169 """
4170 with get_user_db_session(self.username, self.db_password) as session:
4171 rag_index = self._resolve_index_for_config(session, collection_id)
4172 if rag_index is None:
4173 # Ineligible path only, so the extra query here doesn't
4174 # touch the hot (eligible) path. Mirrors the same
4175 # foreign-rows diagnostic get_rag_stats logs, so "I indexed
4176 # this and search returns nothing" is still explained.
4177 foreign_rows = (
4178 session.query(RagDocumentStatus)
4179 .filter_by(collection_id=collection_id)
4180 .count()
4181 )
4182 if foreign_rows:
4183 logger.warning(
4184 f"Collection {collection_id} has no index for the "
4185 f"active embedding configuration "
4186 f"({self.embedding_provider}/{self.embedding_model}), "
4187 f"so it reports zero indexed documents. "
4188 f"{foreign_rows} document(s) are indexed under other "
4189 f"configurations and are not searchable from this one."
4190 )
4191 return False
4192 return bool(
4193 session.query(
4194 exists().where(
4195 RagDocumentStatus.collection_id == collection_id,
4196 RagDocumentStatus.rag_index_id == rag_index.id,
4197 )
4198 ).scalar()
4199 )
4201 def get_rag_stats(
4202 self, collection_id: Optional[str] = None
4203 ) -> Dict[str, Any]:
4204 """
4205 Get RAG statistics for a collection.
4207 Args:
4208 collection_id: UUID of the collection (defaults to Library)
4210 Returns:
4211 Dict with counts and metadata about indexed documents
4212 """
4213 with get_user_db_session(self.username, self.db_password) as session:
4214 # Get collection ID (default to Library)
4215 if not collection_id: 4215 ↛ 4216line 4215 didn't jump to line 4216 because the condition on line 4215 was never true
4216 from ...database.library_init import get_default_library_id
4218 collection_id = get_default_library_id(
4219 self.username, self.db_password
4220 )
4222 # Count total documents in collection
4223 total_docs = (
4224 session.query(DocumentCollection)
4225 .filter_by(collection_id=collection_id)
4226 .count()
4227 )
4229 # Count indexed documents from rag_document_status table
4230 from ...database.models.library import RagDocumentStatus
4232 # One collection can hold status rows from several indexes at
4233 # once, so scope the counts to the index this configuration
4234 # resolves to. No such index means nothing is indexed for it.
4235 rag_index = self._resolve_index_for_config(session, collection_id)
4237 if rag_index is None:
4238 indexed_docs = 0
4239 total_chunks = 0
4240 foreign_rows = (
4241 session.query(RagDocumentStatus)
4242 .filter_by(collection_id=collection_id)
4243 .count()
4244 )
4245 if foreign_rows:
4246 logger.warning(
4247 f"Collection {collection_id} has no index for the "
4248 f"active embedding configuration "
4249 f"({self.embedding_provider}/{self.embedding_model}), "
4250 f"so it reports zero indexed documents. "
4251 f"{foreign_rows} document(s) are indexed under other "
4252 f"configurations and are not searchable from this one."
4253 )
4254 else:
4255 indexed_docs = (
4256 session.query(RagDocumentStatus)
4257 .filter_by(
4258 collection_id=collection_id,
4259 rag_index_id=rag_index.id,
4260 )
4261 .count()
4262 )
4264 # Count total chunks from rag_document_status table
4265 total_chunks = (
4266 session.query(func.sum(RagDocumentStatus.chunk_count))
4267 .filter_by(
4268 collection_id=collection_id,
4269 rag_index_id=rag_index.id,
4270 )
4271 .scalar()
4272 or 0
4273 )
4275 # Describe the index the counts above came from. Sampling any
4276 # chunk of the collection would describe whichever index wrote
4277 # first, which is the same mismatch as the counts.
4278 embedding_info = {}
4279 if rag_index is not None:
4280 stored_provider = rag_index.embedding_model_type
4281 if isinstance(stored_provider, EmbeddingProvider):
4282 stored_provider = stored_provider.value
4283 embedding_info = {
4284 "model": rag_index.embedding_model,
4285 "model_type": stored_provider,
4286 "dimension": rag_index.embedding_dimension,
4287 }
4289 return {
4290 "total_documents": total_docs,
4291 "indexed_documents": indexed_docs,
4292 "unindexed_documents": total_docs - indexed_docs,
4293 "total_chunks": total_chunks,
4294 "embedding_info": embedding_info,
4295 "chunk_size": self.chunk_size,
4296 "chunk_overlap": self.chunk_overlap,
4297 }
4299 def remove_documents_from_index(
4300 self,
4301 document_ids: List[str],
4302 collection_id: str,
4303 *,
4304 chunk_ids_by_document: Optional[Dict[str, List[int]]] = None,
4305 ) -> int:
4306 with _hold_collection_lock(
4307 self.username, collection_id, purpose="transaction"
4308 ):
4309 return self._remove_documents_from_index_locked(
4310 document_ids,
4311 collection_id,
4312 chunk_ids_by_document=chunk_ids_by_document,
4313 )
4315 def _remove_documents_from_index_locked(
4316 self,
4317 document_ids: List[str],
4318 collection_id: str,
4319 *,
4320 chunk_ids_by_document: Optional[Dict[str, List[int]]] = None,
4321 ) -> int:
4322 """Delete given documents' vectors from a collection's vector store.
4324 Vector-store-only: the DB chunk rows are removed separately by the
4325 caller. This closes the gap where deleting a document's DB chunks
4326 left its vectors (and their persisted text — now impossible, see
4327 vector_stores/base.py's SECURITY INVARIANT) in the store, so the
4328 document kept surfacing in semantic search until a full reindex.
4330 For each id, calls ``VectorIndex.delete(source_type="document",
4331 source_id=...)`` — which looks up the matching ``DocumentChunk``
4332 rows by (source_type, source_id, collection_name) — UNLESS
4333 ``chunk_ids_by_document`` supplies that document's chunk ids
4334 directly, in which case ``VectorIndex.delete_ids(ids)`` is used
4335 instead. Callers whose own per-item transactions must delete the DB
4336 rows eagerly (for their own atomicity/rollback semantics) BEFORE
4337 this batched cleanup runs need the latter: the source_id lookup
4338 would find nothing once the rows are already gone (see the Zotero
4339 sync caller). Hardcodes ``source_type="document"``: both current
4340 callers (Zotero sync cleanup, ``remove_document_from_rag``) only
4341 ever pass library-document ids (the pre-cutover implementation
4342 matched a mixed ``document_id``/``source_id`` metadata key
4343 generically across "document" and "user_document" sources, but
4344 nothing ever called it with anything but library documents).
4346 Returns the total number of vectors removed across all ids. No-ops
4347 (returns 0) when the collection has no current vector index — this
4348 guard runs BEFORE building a ``VectorIndex`` so a delete call never
4349 creates one (``_get_vector_index`` would otherwise create the
4350 RAGIndex row via ``_get_or_create_rag_index``).
4351 """
4352 from ...database.models.library import RAGIndex
4353 from ...database.session_context import get_user_db_session
4355 wanted = [d for d in document_ids if d]
4356 if not wanted:
4357 return 0
4359 collection_name = f"collection_{collection_id}"
4360 # Don't create an index just to delete from it — only act if one
4361 # already exists for this collection.
4362 with get_user_db_session(self.username, self.db_password) as session:
4363 has_index = (
4364 session.query(RAGIndex.id)
4365 .filter_by(collection_name=collection_name, is_current=True)
4366 .first()
4367 is not None
4368 )
4369 if not has_index:
4370 return 0
4372 try:
4373 # reset_stale_state=True is REQUIRED on this write path. If the
4374 # pre-flight finds the index corrupt it quarantines the file
4375 # (renames it aside); a fresh EMPTY store is then built and the
4376 # delete loop below persists that empty store over the real path.
4377 # Without resetting the per-document indexed state, the DB would
4378 # still claim every other document is indexed while their vectors
4379 # are gone — silently wiping the whole collection's search with no
4380 # recovery but a manual full reindex. The reset only fires when the
4381 # index was actually quarantined (guarded on the file no longer
4382 # existing), so a healthy remove is unaffected. search() omits this
4383 # flag safely because it never persists an empty store.
4384 vindex = self._get_vector_index(
4385 collection_id, collection_name, reset_stale_state=True
4386 )
4387 except (ValueError, RuntimeError):
4388 logger.warning(
4389 "Could not load vector index for "
4390 f"{collection_name} while removing document vectors"
4391 )
4392 return 0
4394 # Collect ALL removable chunk ids across every document, then remove
4395 # them in ONE apply() (vindex.delete_ids). For an HNSW collection each
4396 # apply() rebuilds the whole index, so a per-document loop is
4397 # O(documents x collection_size); a single batched call rebuilds once.
4398 # Captured ids are used directly; documents without captured ids have
4399 # their chunk ids resolved by (source_type, source_id, collection_name)
4400 # — those rows still exist at this point.
4401 all_ids: List[int] = []
4402 unresolved: List[str] = []
4403 for document_id in wanted:
4404 captured = (
4405 chunk_ids_by_document.get(document_id)
4406 if chunk_ids_by_document
4407 else None
4408 )
4409 if captured: 4409 ↛ 4410line 4409 didn't jump to line 4410 because the condition on line 4409 was never true
4410 all_ids.extend(int(i) for i in captured if i is not None)
4411 else:
4412 unresolved.append(document_id)
4414 if unresolved: 4414 ↛ 4430line 4414 didn't jump to line 4430 because the condition on line 4414 was always true
4415 in_chunk = 500 # keep below SQLITE_MAX_VARIABLE_NUMBER (999)
4416 with get_user_db_session(
4417 self.username, self.db_password
4418 ) as session:
4419 for i in range(0, len(unresolved), in_chunk):
4420 batch = unresolved[i : i + in_chunk]
4421 all_ids.extend(
4422 row.id
4423 for row in session.query(DocumentChunk.id).filter(
4424 DocumentChunk.source_type == "document",
4425 DocumentChunk.collection_name == collection_name,
4426 DocumentChunk.source_id.in_(batch),
4427 )
4428 )
4430 removed = 0
4431 if all_ids: 4431 ↛ 4441line 4431 didn't jump to line 4441 because the condition on line 4431 was always true
4432 try:
4433 stats = vindex.delete_ids(sorted(set(all_ids)))
4434 removed = stats.removed
4435 except Exception:
4436 logger.warning(
4437 f"Could not remove {len(set(all_ids))} vectors for "
4438 f"{len(wanted)} document(s) from {collection_name}"
4439 )
4441 if removed: 4441 ↛ 4446line 4441 didn't jump to line 4446 because the condition on line 4441 was always true
4442 logger.info(
4443 f"Removed {removed} vectors for "
4444 f"{len(wanted)} document(s) from {collection_name}"
4445 )
4446 return removed