Coverage for src/local_deep_research/vector_stores/facade.py: 93%
162 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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.orm import Session
42from ..database.models.library import DocumentChunk
43from ..database.session_context import get_user_db_session
44from .base import IntegrityRecord, IntegrityVerify, WriteLock
45from .config import get_vector_store_class
48class _BindKwargs(TypedDict):
49 """Keyword arguments forwarded to ``store_cls.create()``/``.load()``.
51 A plain ``dict`` would infer a uniform (``object``) value type, erasing
52 each field's real type when unpacked via ``**bind`` — a ``TypedDict``
53 keeps each key's type through the ``**`` unpacking so mypy can check the
54 call against ``BaseVectorStore.create``/``.load``.
55 """
57 dimension: int
58 index_type: str
59 metric: str
60 normalize: bool
61 path: Path
62 lock: WriteLock
63 integrity_record: IntegrityRecord
64 integrity_verify: IntegrityVerify
67@dataclass
68class ChunkInput:
69 """One chunk to index: its text plus optional positional/title metadata.
71 Kept backend- and framework-neutral (not a LangChain ``Document``) so the
72 facade never couples to a specific splitter/vector library.
73 """
75 text: str
76 metadata: dict = field(default_factory=dict)
79@dataclass
80class SearchResult:
81 """A search hit, rehydrated from the encrypted DB by chunk id."""
83 chunk_id: int
84 text: str
85 # Raw backend score; the relevance formula is the caller's choice. ``metric``
86 # says how to read it: for cosine/dot_product (inner product) HIGHER is more
87 # similar; for l2 LOWER is closer. Branch on it or scores invert.
88 distance: float
89 metric: str
90 metadata: dict
91 document_title: Optional[str]
92 source_id: Optional[str]
93 source_type: Optional[str]
96@dataclass
97class IndexStats:
98 """Result of :meth:`VectorIndex.index`."""
100 chunks: int
101 added: int
102 removed: int
105@dataclass
106class DeleteStats:
107 """Result of :meth:`VectorIndex.delete`."""
109 removed: int
110 rows_deleted: int
113class VectorIndex:
114 """Per-collection semantic index. Construct once per (collection, user).
116 Persistence resources (``path`` / ``lock`` / integrity hooks) are injected —
117 the caller (service layer) owns per-user path resolution and the write-lock
118 registry; this facade just wires them into the store.
119 """
121 def __init__(
122 self,
123 *,
124 username: str,
125 db_password: str,
126 embeddings: Embeddings,
127 embedding_model: str,
128 embedding_model_type, # EmbeddingProvider enum (DocumentChunk column)
129 collection_name: str,
130 dimension: int,
131 path: Path,
132 lock: WriteLock,
133 integrity_record: IntegrityRecord,
134 integrity_verify: IntegrityVerify,
135 index_type: str = "flat",
136 metric: str = "cosine",
137 normalize: bool = True,
138 provider: str = "faiss",
139 ) -> None:
140 self.username = username
141 self.db_password = db_password # gitleaks:allow
142 self.embeddings = embeddings
143 self.embedding_model = embedding_model
144 self.embedding_model_type = embedding_model_type
145 self.collection_name = collection_name
146 # Canonicalize the metric so the physical index BUILD (faiss_store:
147 # inner-product iff metric in ("cosine","dot_product")) and the search
148 # engines' relevance formula (which branch on ``metric == "l2"``) agree.
149 # A non-canonical value like "L2"/"Cosine"/" cosine " would otherwise
150 # build one index type but be scored as the other, silently corrupting
151 # relevance. Both the store (via ``bind`` below) and SearchResult.metric
152 # then see the same canonical string.
153 metric = (metric or "cosine").strip().lower()
154 self.metric = metric
156 store_cls = get_vector_store_class(provider)
157 path = Path(path)
158 bind: _BindKwargs = {
159 "dimension": dimension,
160 "index_type": index_type,
161 "metric": metric,
162 "normalize": normalize,
163 "path": path,
164 "lock": lock,
165 "integrity_record": integrity_record,
166 "integrity_verify": integrity_verify,
167 }
168 # ``bind`` already carries ``path``; pass it only via kwargs (load()'s
169 # first param is path) so we don't hand ``path`` twice.
170 self._store = (
171 store_cls.load(**bind)
172 if path.exists()
173 else store_cls.create(**bind)
174 )
176 @contextmanager
177 def _db(self, session: Optional[Session]) -> Iterator[Tuple[Session, bool]]:
178 """Yield ``(session, owns)``.
180 ``get_user_db_session`` returns a THREAD-LOCAL, reused session, so this
181 facade must not blindly ``commit()``/``rollback()`` — it would settle a
182 caller's unrelated pending work. When the caller passes its own
183 ``session``, we use it and yield ``owns=False`` (the caller owns the
184 transaction boundary and the apply()-then-commit durability gap). Only
185 when no session is passed do we open one and own commit/rollback.
186 """
187 if session is not None:
188 yield session, False
189 else:
190 with get_user_db_session(self.username, self.db_password) as s:
191 yield s, True
193 def index_prepared(
194 self,
195 *,
196 source_type: str,
197 source_id: Optional[str],
198 chunks: Sequence[ChunkInput],
199 vectors: Optional[np.ndarray],
200 replace: bool = True,
201 session: Optional[Session] = None,
202 ) -> IndexStats:
203 """Persist pre-embedded chunks without invoking the embedding backend.
205 Bulk indexing uses this after parallel preparation. Database row/id
206 allocation, vector application, prior-row removal, and commit semantics
207 remain identical to :meth:`index`; only embedding is moved outside the
208 collection writer critical section.
209 """
210 chunks = list(chunks)
211 if not chunks and (not replace or source_id is None): 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true
212 return IndexStats(chunks=0, added=0, removed=0)
213 if chunks:
214 if vectors is None:
215 raise ValueError("chunks given without vectors")
216 vectors = np.asarray(vectors, dtype="float32")
217 if len(vectors) != len(chunks):
218 raise ValueError("vector count must match chunk count")
220 sid = str(source_id) if source_id is not None else None
221 with self._db(session) as (sess, owns):
222 prior_ids: List[int] = []
223 if replace and sid is not None: 223 ↛ 235line 223 didn't jump to line 235 because the condition on line 223 was always true
224 prior_ids = [
225 row.id
226 for row in sess.query(DocumentChunk.id).filter_by(
227 source_type=source_type,
228 source_id=sid,
229 collection_name=self.collection_name,
230 embedding_model=self.embedding_model,
231 embedding_model_type=self.embedding_model_type,
232 )
233 ]
235 rows = [self._build_row(c, source_type, sid) for c in chunks]
236 if rows:
237 sess.add_all(rows)
238 sess.flush()
239 new_ids = [cast(int, row.id) for row in rows]
240 try:
241 stats = self._store.apply(
242 add_ids=new_ids,
243 add_vectors=vectors,
244 remove_ids=prior_ids,
245 )
246 except Exception:
247 if owns: 247 ↛ 249line 247 didn't jump to line 249 because the condition on line 247 was always true
248 sess.rollback()
249 raise
250 self._finalize_db(
251 sess, owns, prior_ids, f"index prepared {source_type}/{sid}"
252 )
254 return IndexStats(
255 chunks=len(chunks),
256 added=stats["added"],
257 removed=stats["removed"],
258 )
260 # ------------------------------------------------------------------ #
261 # index
262 # ------------------------------------------------------------------ #
263 def index(
264 self,
265 *,
266 source_type: str,
267 source_id: Optional[str],
268 chunks: Sequence[ChunkInput],
269 replace: bool = True,
270 session: Optional[Session] = None,
271 ) -> IndexStats:
272 """Embed + persist ``chunks`` for one source into this collection.
274 Ordering: write chunk rows and ``flush()`` for their int ids, then
275 ``store.apply()`` (persists vectors atomically under its lock), then
276 delete the source's prior rows, then ``commit()`` (only if we own the
277 session). If ``store.apply()`` fails we roll back — the new rows never
278 existed, so no DB row can point at a missing vector.
280 ``store.apply()`` persists vectors durably BEFORE the DB commit, so if
281 the commit later fails the vectors outlive their (rolled-back) rows as
282 orphans — searchable-by-nothing, cleaned on the next re-index of the
283 source. When ``session`` is passed the caller owns that window.
285 ``replace=True`` removes this source's PRIOR chunks/vectors first;
286 ``False`` appends.
287 """
288 chunks = list(chunks)
289 # An empty chunk list with replace=True still must remove this source's
290 # PRIOR rows/vectors — e.g. a note edited down to whitespace yields zero
291 # chunks, and skipping here would leave the user-cleared content
292 # permanently searchable. Only a genuine no-op (nothing to add AND
293 # nothing to replace) returns early.
294 if not chunks and (not replace or source_id is None): 294 ↛ 295line 294 didn't jump to line 295 because the condition on line 294 was never true
295 return IndexStats(chunks=0, added=0, removed=0)
297 texts = [c.text for c in chunks]
298 vectors = (
299 np.asarray(self.embeddings.embed_documents(texts), dtype="float32")
300 if texts
301 else None
302 )
304 sid = str(source_id) if source_id is not None else None
305 with self._db(session) as (sess, owns):
306 prior_ids: List[int] = []
307 if replace and sid is not None:
308 # Scope the prior-chunk lookup to THIS store's embedding-model
309 # config. collection_name is identical across every model config
310 # for a collection (only RAGIndex.index_hash / the physical
311 # .faiss path differ per model), so a lookup by
312 # (source_type, source_id, collection_name) alone would also
313 # match rows indexed under a DIFFERENT model. Those ids would
314 # then be handed to THIS (physically separate, model-specific)
315 # store as remove_ids — where faiss silently no-ops on unknown
316 # ids — and _finalize_db would still DELETE the rows. Result: a
317 # reindex under model B destroys model A's chunk text (its only
318 # copy) while A's .faiss keeps the now-unresolvable vector,
319 # even though A's index is deliberately kept for model revert.
320 prior_ids = [
321 row.id
322 for row in sess.query(DocumentChunk.id).filter_by(
323 source_type=source_type,
324 source_id=sid,
325 collection_name=self.collection_name,
326 embedding_model=self.embedding_model,
327 embedding_model_type=self.embedding_model_type,
328 )
329 ]
331 rows = [self._build_row(c, source_type, sid) for c in chunks]
332 if rows:
333 sess.add_all(rows)
334 sess.flush() # assigns int PKs to every row in one round-trip
335 # SQLAlchemy's legacy Column() declarative style types instance
336 # attribute access as Column[int] rather than int; cast() is a
337 # no-op at runtime (row.id is already an int post-flush()).
338 new_ids = [cast(int, 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 # cast(): row.id is Column[int] under the legacy Column()
462 # declarative style; a no-op at runtime.
463 cast(int, row.id): row
464 for row in sess.query(DocumentChunk).filter(
465 DocumentChunk.id.in_(ids),
466 DocumentChunk.collection_name == self.collection_name,
467 )
468 }
470 results: List[SearchResult] = []
471 for chunk_id, distance in hits:
472 row = rows.get(chunk_id)
473 if row is None:
474 logger.debug(
475 f"Search hit {chunk_id} has no DocumentChunk row; skipping"
476 )
477 continue
478 # SQLAlchemy's legacy Column() declarative style types instance
479 # attribute access as Column[...] rather than the plain Python
480 # value; cast() to the DB-verified runtime type (no conversion,
481 # so a None stays None rather than becoming e.g. the "None"
482 # string a str()/int() call would produce).
483 results.append(
484 SearchResult(
485 chunk_id=chunk_id,
486 text=cast(str, row.chunk_text),
487 distance=distance,
488 metric=self.metric,
489 metadata=cast(Optional[dict], row.document_metadata) or {},
490 document_title=cast(Optional[str], row.document_title),
491 source_id=cast(Optional[str], row.source_id),
492 source_type=cast(Optional[str], row.source_type),
493 )
494 )
495 return results
497 # ------------------------------------------------------------------ #
498 # delete
499 # ------------------------------------------------------------------ #
500 def delete(
501 self,
502 *,
503 source_type: str,
504 source_id: str,
505 session: Optional[Session] = None,
506 ) -> DeleteStats:
507 """Remove all of one source's chunks/vectors from this collection.
509 ``store.apply()`` (the durable side) runs before the DB delete + commit,
510 matching :meth:`index`'s ordering.
512 Re-queries ``DocumentChunk`` by (source_type, source_id,
513 collection_name) at call time to discover the ids to remove, so the
514 rows must still exist when this is called. If a caller has already
515 deleted those rows elsewhere, this finds nothing to do; use
516 :meth:`delete_ids` instead with the ids captured before deletion.
517 """
518 sid = str(source_id)
519 with self._db(session) as (sess, owns):
520 ids = [
521 row.id
522 for row in sess.query(DocumentChunk.id).filter_by(
523 source_type=source_type,
524 source_id=sid,
525 collection_name=self.collection_name,
526 )
527 ]
528 if not ids:
529 return DeleteStats(removed=0, rows_deleted=0)
531 try:
532 stats = self._store.apply(
533 add_ids=[], add_vectors=None, remove_ids=ids
534 )
535 except Exception:
536 if owns:
537 sess.rollback()
538 raise
539 # Guard the row-delete AND commit together (mirrors index()): a bare
540 # query(...).delete() that raised after the durable apply() above
541 # would leave orphan rows with no rollback and no diagnostic.
542 rows_deleted = self._finalize_db(
543 sess, owns, ids, f"delete {source_type}/{sid}"
544 )
546 return DeleteStats(removed=stats["removed"], rows_deleted=rows_deleted)
548 def delete_ids(
549 self,
550 ids: Sequence[int],
551 *,
552 session: Optional[Session] = None,
553 ) -> DeleteStats:
554 """Remove specific chunk ids' vectors from this collection's store.
556 For callers that already captured the ``DocumentChunk`` int ids
557 BEFORE the matching rows were deleted elsewhere — e.g. a caller
558 whose own per-item transaction must delete the rows eagerly (for
559 its own atomicity/rollback semantics) ahead of a later BATCHED
560 vector-store cleanup pass. Unlike :meth:`delete`, this does NOT
561 re-query ``DocumentChunk`` by (source_type, source_id,
562 collection_name) to discover the ids — by the time this runs the
563 rows the ids came from may already be gone, and that lookup would
564 find nothing (see the "rows must still exist" requirement noted on
565 :meth:`delete`). Any rows matching the given ids that do still
566 exist are cleaned up too (defensive; normally none are left).
567 """
568 wanted = sorted({int(i) for i in ids if i is not None})
569 if not wanted:
570 return DeleteStats(removed=0, rows_deleted=0)
572 with self._db(session) as (sess, owns):
573 try:
574 stats = self._store.apply(
575 add_ids=[], add_vectors=None, remove_ids=wanted
576 )
577 except Exception:
578 if owns:
579 sess.rollback()
580 raise
581 # Guard the row-delete AND commit together (mirrors index()).
582 rows_deleted = self._finalize_db(
583 sess, owns, wanted, f"delete_ids ({len(wanted)} ids)"
584 )
586 return DeleteStats(removed=stats["removed"], rows_deleted=rows_deleted)
588 def live_ids(self) -> List[int]:
589 """Authoritative vector ids currently present in the store."""
590 return self._store.live_ids()
592 def count(self) -> int:
593 """Number of vectors currently in the store."""
594 return self._store.count()