Coverage for src/local_deep_research/vector_stores/facade.py: 94%
163 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""High-level semantic-index facade: ``index`` / ``search`` / ``delete``.
3This is the single simple entry point to the vector subsystem. It is the ONE
4place in the ``vector_stores`` package that bridges three things:
6 * the embedding model (a LangChain ``Embeddings``),
7 * the encrypted DB (``DocumentChunk`` — the authoritative home of chunk text
8 and metadata), and
9 * a swappable :class:`BaseVectorStore` (ids + vectors only).
11``base.py`` and ``implementations/`` stay DB-agnostic; only this module imports
12``DocumentChunk``. The vector store holds an int64 id per chunk
13(``DocumentChunk.id``); text/metadata are rehydrated from the DB by that id on
14search — the store never receives text (see the "SECURITY INVARIANT" in
15:mod:`.base`).
17One row = one vector. Identical text in two documents produces two rows/vectors
18(no cross-document sharing); this keeps provenance correct and deletes trivial.
19Exact-duplicate *files* are an ingestion-time concern, handled upstream.
20"""
22import hashlib
23import uuid
24from contextlib import contextmanager
25from dataclasses import dataclass, field
26from pathlib import Path
27from typing import (
28 Iterator,
29 List,
30 Optional,
31 Sequence,
32 Tuple,
33 TypedDict,
34 cast,
35)
37import numpy as np
38from langchain_core.embeddings import Embeddings
39from loguru import logger
40from sqlalchemy import select
41from sqlalchemy.orm import Session
43from ..database.models.library import DocumentChunk
44from ..database.session_context import get_user_db_session
45from .base import IntegrityRecord, IntegrityVerify, WriteLock
46from .config import get_vector_store_class
49class _BindKwargs(TypedDict):
50 """Keyword arguments forwarded to ``store_cls.create()``/``.load()``.
52 A plain ``dict`` would infer a uniform (``object``) value type, erasing
53 each field's real type when unpacked via ``**bind`` — a ``TypedDict``
54 keeps each key's type through the ``**`` unpacking so mypy can check the
55 call against ``BaseVectorStore.create``/``.load``.
56 """
58 dimension: int
59 index_type: str
60 metric: str
61 normalize: bool
62 path: Path
63 lock: WriteLock
64 integrity_record: IntegrityRecord
65 integrity_verify: IntegrityVerify
68@dataclass
69class ChunkInput:
70 """One chunk to index: its text plus optional positional/title metadata.
72 Kept backend- and framework-neutral (not a LangChain ``Document``) so the
73 facade never couples to a specific splitter/vector library.
74 """
76 text: str
77 metadata: dict = field(default_factory=dict)
80@dataclass
81class SearchResult:
82 """A search hit, rehydrated from the encrypted DB by chunk id."""
84 chunk_id: int
85 text: str
86 # Raw backend score; the relevance formula is the caller's choice. ``metric``
87 # says how to read it: for cosine/dot_product (inner product) HIGHER is more
88 # similar; for l2 LOWER is closer. Branch on it or scores invert.
89 distance: float
90 metric: str
91 metadata: dict
92 document_title: Optional[str]
93 source_id: Optional[str]
94 source_type: Optional[str]
97@dataclass
98class IndexStats:
99 """Result of :meth:`VectorIndex.index`."""
101 chunks: int
102 added: int
103 removed: int
106@dataclass
107class DeleteStats:
108 """Result of :meth:`VectorIndex.delete`."""
110 removed: int
111 rows_deleted: int
114class VectorIndex:
115 """Per-collection semantic index. Construct once per (collection, user).
117 Persistence resources (``path`` / ``lock`` / integrity hooks) are injected —
118 the caller (service layer) owns per-user path resolution and the write-lock
119 registry; this facade just wires them into the store.
120 """
122 def __init__(
123 self,
124 *,
125 username: str,
126 db_password: str,
127 embeddings: Embeddings,
128 embedding_model: str,
129 embedding_model_type, # EmbeddingProvider enum (DocumentChunk column)
130 collection_name: str,
131 dimension: int,
132 path: Path,
133 lock: WriteLock,
134 integrity_record: IntegrityRecord,
135 integrity_verify: IntegrityVerify,
136 index_type: str = "flat",
137 metric: str = "cosine",
138 normalize: bool = True,
139 provider: str = "faiss",
140 ) -> None:
141 self.username = username
142 self.db_password = db_password # gitleaks:allow
143 self.embeddings = embeddings
144 self.embedding_model = embedding_model
145 self.embedding_model_type = embedding_model_type
146 self.collection_name = collection_name
147 # Canonicalize the metric so the physical index BUILD (faiss_store:
148 # inner-product iff metric in ("cosine","dot_product")) and the search
149 # engines' relevance formula (which branch on ``metric == "l2"``) agree.
150 # A non-canonical value like "L2"/"Cosine"/" cosine " would otherwise
151 # build one index type but be scored as the other, silently corrupting
152 # relevance. Both the store (via ``bind`` below) and SearchResult.metric
153 # then see the same canonical string.
154 metric = (metric or "cosine").strip().lower()
155 self.metric = metric
157 store_cls = get_vector_store_class(provider)
158 path = Path(path)
159 bind: _BindKwargs = {
160 "dimension": dimension,
161 "index_type": index_type,
162 "metric": metric,
163 "normalize": normalize,
164 "path": path,
165 "lock": lock,
166 "integrity_record": integrity_record,
167 "integrity_verify": integrity_verify,
168 }
169 # ``bind`` already carries ``path``; pass it only via kwargs (load()'s
170 # first param is path) so we don't hand ``path`` twice.
171 self._store = (
172 store_cls.load(**bind)
173 if path.exists()
174 else store_cls.create(**bind)
175 )
177 @contextmanager
178 def _db(self, session: Optional[Session]) -> Iterator[Tuple[Session, bool]]:
179 """Yield ``(session, owns)``.
181 ``get_user_db_session`` returns a THREAD-LOCAL, reused session, so this
182 facade must not blindly ``commit()``/``rollback()`` — it would settle a
183 caller's unrelated pending work. When the caller passes its own
184 ``session``, we use it and yield ``owns=False`` (the caller owns the
185 transaction boundary and the apply()-then-commit durability gap). Only
186 when no session is passed do we open one and own commit/rollback.
187 """
188 if session is not None:
189 yield session, False
190 else:
191 with get_user_db_session(self.username, self.db_password) as s:
192 yield s, True
194 def index_prepared(
195 self,
196 *,
197 source_type: str,
198 source_id: Optional[str],
199 chunks: Sequence[ChunkInput],
200 vectors: Optional[np.ndarray],
201 replace: bool = True,
202 session: Optional[Session] = None,
203 ) -> IndexStats:
204 """Persist pre-embedded chunks without invoking the embedding backend.
206 Bulk indexing uses this after parallel preparation. Database row/id
207 allocation, vector application, prior-row removal, and commit semantics
208 remain identical to :meth:`index`; only embedding is moved outside the
209 collection writer critical section.
210 """
211 chunks = list(chunks)
212 if not chunks and (not replace or source_id is None): 212 ↛ 213line 212 didn't jump to line 213 because the condition on line 212 was never true
213 return IndexStats(chunks=0, added=0, removed=0)
214 if chunks:
215 if vectors is None:
216 raise ValueError("chunks given without vectors")
217 vectors = np.asarray(vectors, dtype="float32")
218 if len(vectors) != len(chunks):
219 raise ValueError("vector count must match chunk count")
221 sid = str(source_id) if source_id is not None else None
222 with self._db(session) as (sess, owns):
223 prior_ids: List[int] = []
224 if replace and sid is not None: 224 ↛ 237line 224 didn't jump to line 237 because the condition on line 224 was always true
225 prior_ids = list(
226 sess.scalars(
227 select(DocumentChunk.id).filter_by(
228 source_type=source_type,
229 source_id=sid,
230 collection_name=self.collection_name,
231 embedding_model=self.embedding_model,
232 embedding_model_type=self.embedding_model_type,
233 )
234 )
235 )
237 rows = [self._build_row(c, source_type, sid) for c in chunks]
238 if rows:
239 sess.add_all(rows)
240 sess.flush()
241 new_ids = [row.id for row in rows]
242 try:
243 stats = self._store.apply(
244 add_ids=new_ids,
245 add_vectors=vectors,
246 remove_ids=prior_ids,
247 )
248 except Exception:
249 if owns:
250 sess.rollback()
251 raise
252 self._finalize_db(
253 sess, owns, prior_ids, f"index prepared {source_type}/{sid}"
254 )
256 return IndexStats(
257 chunks=len(chunks),
258 added=stats["added"],
259 removed=stats["removed"],
260 )
262 # ------------------------------------------------------------------ #
263 # index
264 # ------------------------------------------------------------------ #
265 def index(
266 self,
267 *,
268 source_type: str,
269 source_id: Optional[str],
270 chunks: Sequence[ChunkInput],
271 replace: bool = True,
272 session: Optional[Session] = None,
273 ) -> IndexStats:
274 """Embed + persist ``chunks`` for one source into this collection.
276 Ordering: write chunk rows and ``flush()`` for their int ids, then
277 ``store.apply()`` (persists vectors atomically under its lock), then
278 delete the source's prior rows, then ``commit()`` (only if we own the
279 session). If ``store.apply()`` fails we roll back — the new rows never
280 existed, so no DB row can point at a missing vector.
282 ``store.apply()`` persists vectors durably BEFORE the DB commit, so if
283 the commit later fails the vectors outlive their (rolled-back) rows as
284 orphans — searchable-by-nothing, cleaned on the next re-index of the
285 source. When ``session`` is passed the caller owns that window.
287 ``replace=True`` removes this source's PRIOR chunks/vectors first;
288 ``False`` appends.
289 """
290 chunks = list(chunks)
291 # An empty chunk list with replace=True still must remove this source's
292 # PRIOR rows/vectors — e.g. a note edited down to whitespace yields zero
293 # chunks, and skipping here would leave the user-cleared content
294 # permanently searchable. Only a genuine no-op (nothing to add AND
295 # nothing to replace) returns early.
296 if not chunks and (not replace or source_id is None): 296 ↛ 297line 296 didn't jump to line 297 because the condition on line 296 was never true
297 return IndexStats(chunks=0, added=0, removed=0)
299 texts = [c.text for c in chunks]
300 vectors = (
301 np.asarray(self.embeddings.embed_documents(texts), dtype="float32")
302 if texts
303 else None
304 )
306 sid = str(source_id) if source_id is not None else None
307 with self._db(session) as (sess, owns):
308 prior_ids: List[int] = []
309 if replace and sid is not None:
310 # Scope the prior-chunk lookup to THIS store's embedding-model
311 # config. collection_name is identical across every model config
312 # for a collection (only RAGIndex.index_hash / the physical
313 # .faiss path differ per model), so a lookup by
314 # (source_type, source_id, collection_name) alone would also
315 # match rows indexed under a DIFFERENT model. Those ids would
316 # then be handed to THIS (physically separate, model-specific)
317 # store as remove_ids — where faiss silently no-ops on unknown
318 # ids — and _finalize_db would still DELETE the rows. Result: a
319 # reindex under model B destroys model A's chunk text (its only
320 # copy) while A's .faiss keeps the now-unresolvable vector,
321 # even though A's index is deliberately kept for model revert.
322 prior_ids = list(
323 sess.scalars(
324 select(DocumentChunk.id).filter_by(
325 source_type=source_type,
326 source_id=sid,
327 collection_name=self.collection_name,
328 embedding_model=self.embedding_model,
329 embedding_model_type=self.embedding_model_type,
330 )
331 )
332 )
334 rows = [self._build_row(c, source_type, sid) for c in chunks]
335 if rows:
336 sess.add_all(rows)
337 sess.flush() # assigns int PKs to every row in one round-trip
338 new_ids = [row.id for row in rows]
340 try:
341 stats = self._store.apply(
342 add_ids=new_ids,
343 add_vectors=vectors,
344 remove_ids=prior_ids,
345 )
346 except Exception:
347 # apply() failed without a durable change (fail-closed reload /
348 # add error / atomic-save) — a clean DB rollback keeps them
349 # consistent.
350 if owns: 350 ↛ 352line 350 didn't jump to line 352 because the condition on line 350 was always true
351 sess.rollback()
352 raise
353 # apply() is DURABLE from here; a DB-side failure below leaves the
354 # store and DB possibly disagreeing (orphan vectors/rows).
355 self._finalize_db(
356 sess, owns, prior_ids, f"index {source_type}/{sid}"
357 )
359 return IndexStats(
360 chunks=len(chunks),
361 added=stats["added"],
362 removed=stats["removed"],
363 )
365 @staticmethod
366 def _finalize_db(
367 session: Session,
368 owns: bool,
369 delete_ids: List[int],
370 what: str,
371 ) -> int:
372 """Delete rows + commit (if owned), AFTER a durable ``store.apply()``.
374 Covers the whole post-apply DB window in one place: if either the
375 prior-row delete OR the commit fails, the store has already persisted
376 vectors that the DB won't reflect (orphan vectors / rows) — a
377 known-inconsistent state that reconciles on the next re-index of the
378 source. Log it loudly, roll back if we own the session, then re-raise.
379 ``synchronize_session="fetch"`` keeps the reused thread-local session's
380 identity map from serving stale, already-deleted objects. Returns the
381 number of rows deleted.
382 """
383 deleted = 0
384 try:
385 if delete_ids:
386 deleted = (
387 session.query(DocumentChunk)
388 .filter(DocumentChunk.id.in_(delete_ids))
389 .delete(synchronize_session="fetch")
390 )
391 if owns:
392 session.commit()
393 except Exception:
394 logger.exception(
395 f"{what}: vectors were persisted but the DB update failed; "
396 "store and DB will reconcile on the next re-index of the source"
397 )
398 if owns: 398 ↛ 400line 398 didn't jump to line 400 because the condition on line 398 was always true
399 session.rollback()
400 raise
401 return deleted
403 def _build_row(
404 self, chunk: ChunkInput, source_type: str, source_id: Optional[str]
405 ) -> DocumentChunk:
406 meta = chunk.metadata or {}
407 return DocumentChunk(
408 chunk_hash=hashlib.sha256(chunk.text.encode()).hexdigest(),
409 source_type=source_type,
410 source_id=source_id,
411 collection_name=self.collection_name,
412 chunk_text=chunk.text,
413 chunk_index=meta.get("chunk_index", 0),
414 start_char=meta.get("start_char", 0),
415 end_char=meta.get("end_char", len(chunk.text)),
416 word_count=len(chunk.text.split()),
417 # embedding_id stays populated: it's the old-uuid -> new-int join key
418 # the one-time index-format migration needs. (Legacy NOT NULL UNIQUE
419 # column; drop it only in a later, separate migration.)
420 embedding_id=uuid.uuid4().hex,
421 embedding_model=self.embedding_model,
422 embedding_model_type=self.embedding_model_type,
423 embedding_dimension=self._store.dimension,
424 document_title=meta.get("document_title") or meta.get("title"),
425 document_metadata=meta or None,
426 )
428 # ------------------------------------------------------------------ #
429 # search
430 # ------------------------------------------------------------------ #
431 def search(
432 self,
433 query: str,
434 top_k: int,
435 session: Optional[Session] = None,
436 ) -> List[SearchResult]:
437 """Embed ``query``, search the store, rehydrate hits from the DB by id.
439 Orphan hits (a vector whose DB row is gone) are skipped — there is no
440 in-store text to fall back to, and skipping is the correct signal.
441 """
442 if top_k <= 0:
443 return []
444 query_vec = np.asarray(
445 self.embeddings.embed_query(query), dtype="float32"
446 )
447 hits = self._store.search(query_vec, top_k) # [(chunk_id, distance)]
448 if not hits:
449 return []
451 ids = [cid for cid, _ in hits]
452 with self._db(session) as (sess, _owns):
453 # Scope rehydration to THIS collection. A hit id is a DocumentChunk.id
454 # from this collection's index, but an orphan/stale vector (e.g. a
455 # rolled-back insert left a vector whose id was later reused by a
456 # chunk in a DIFFERENT collection) could otherwise match a row that
457 # belongs elsewhere and leak that other document's text into this
458 # collection's results. Filtering by collection_name makes such a
459 # hit resolve to no row here and be skipped as an orphan.
460 rows = {
461 row.id: row
462 for row in sess.query(DocumentChunk).filter(
463 DocumentChunk.id.in_(ids),
464 DocumentChunk.collection_name == self.collection_name,
465 )
466 }
468 results: List[SearchResult] = []
469 for chunk_id, distance in hits:
470 row = rows.get(chunk_id)
471 if row is None:
472 logger.debug(
473 f"Search hit {chunk_id} has no DocumentChunk row; skipping"
474 )
475 continue
476 # SQLAlchemy's legacy Column() declarative style types instance
477 # attribute access as Column[...] rather than the plain Python
478 # value; cast() to the DB-verified runtime type (no conversion,
479 # so a None stays None rather than becoming e.g. the "None"
480 # string a str()/int() call would produce).
481 results.append(
482 SearchResult(
483 chunk_id=chunk_id,
484 text=cast(str, row.chunk_text),
485 distance=distance,
486 metric=self.metric,
487 metadata=cast(Optional[dict], row.document_metadata) or {},
488 document_title=cast(Optional[str], row.document_title),
489 source_id=cast(Optional[str], row.source_id),
490 source_type=cast(Optional[str], row.source_type),
491 )
492 )
493 return results
495 # ------------------------------------------------------------------ #
496 # delete
497 # ------------------------------------------------------------------ #
498 def delete(
499 self,
500 *,
501 source_type: str,
502 source_id: str,
503 session: Optional[Session] = None,
504 ) -> DeleteStats:
505 """Remove all of one source's chunks/vectors from this collection.
507 ``store.apply()`` (the durable side) runs before the DB delete + commit,
508 matching :meth:`index`'s ordering.
510 Re-queries ``DocumentChunk`` by (source_type, source_id,
511 collection_name) at call time to discover the ids to remove, so the
512 rows must still exist when this is called. If a caller has already
513 deleted those rows elsewhere, this finds nothing to do; use
514 :meth:`delete_ids` instead with the ids captured before deletion.
515 """
516 sid = str(source_id)
517 with self._db(session) as (sess, owns):
518 ids = list(
519 sess.scalars(
520 select(DocumentChunk.id).filter_by(
521 source_type=source_type,
522 source_id=sid,
523 collection_name=self.collection_name,
524 )
525 )
526 )
527 if not ids:
528 return DeleteStats(removed=0, rows_deleted=0)
530 try:
531 stats = self._store.apply(
532 add_ids=[], add_vectors=None, remove_ids=ids
533 )
534 except Exception:
535 if owns:
536 sess.rollback()
537 raise
538 # Guard the row-delete AND commit together (mirrors index()): a bare
539 # query(...).delete() that raised after the durable apply() above
540 # would leave orphan rows with no rollback and no diagnostic.
541 rows_deleted = self._finalize_db(
542 sess, owns, ids, f"delete {source_type}/{sid}"
543 )
545 return DeleteStats(removed=stats["removed"], rows_deleted=rows_deleted)
547 def delete_ids(
548 self,
549 ids: Sequence[int],
550 *,
551 session: Optional[Session] = None,
552 ) -> DeleteStats:
553 """Remove specific chunk ids' vectors from this collection's store.
555 For callers that already captured the ``DocumentChunk`` int ids
556 BEFORE the matching rows were deleted elsewhere — e.g. a caller
557 whose own per-item transaction must delete the rows eagerly (for
558 its own atomicity/rollback semantics) ahead of a later BATCHED
559 vector-store cleanup pass. Unlike :meth:`delete`, this does NOT
560 re-query ``DocumentChunk`` by (source_type, source_id,
561 collection_name) to discover the ids — by the time this runs the
562 rows the ids came from may already be gone, and that lookup would
563 find nothing (see the "rows must still exist" requirement noted on
564 :meth:`delete`). Any rows matching the given ids that do still
565 exist are cleaned up too (defensive; normally none are left).
566 """
567 wanted = sorted({int(i) for i in ids if i is not None})
568 if not wanted:
569 return DeleteStats(removed=0, rows_deleted=0)
571 with self._db(session) as (sess, owns):
572 try:
573 stats = self._store.apply(
574 add_ids=[], add_vectors=None, remove_ids=wanted
575 )
576 except Exception:
577 if owns:
578 sess.rollback()
579 raise
580 # Guard the row-delete AND commit together (mirrors index()).
581 rows_deleted = self._finalize_db(
582 sess, owns, wanted, f"delete_ids ({len(wanted)} ids)"
583 )
585 return DeleteStats(removed=stats["removed"], rows_deleted=rows_deleted)
587 def live_ids(self) -> List[int]:
588 """Authoritative vector ids currently present in the store."""
589 return self._store.live_ids()
591 def count(self) -> int:
592 """Number of vectors currently in the store."""
593 return self._store.count()