Coverage for src/local_deep_research/vector_stores/facade.py: 94%
132 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 07:11 +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 threading
24import uuid
25from contextlib import contextmanager
26from dataclasses import dataclass, field
27from pathlib import Path
28from typing import (
29 Callable,
30 Iterator,
31 List,
32 Optional,
33 Sequence,
34 Tuple,
35 TypedDict,
36 cast,
37)
39import numpy as np
40from langchain_core.embeddings import Embeddings
41from loguru import logger
42from sqlalchemy.orm import Session
44from ..database.models.library import DocumentChunk
45from ..database.session_context import get_user_db_session
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: threading.Lock
64 integrity_record: Callable[[Path], None]
65 integrity_verify: Callable[[Path], Tuple[bool, Optional[str]]]
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: threading.Lock,
134 integrity_record: Callable[[Path], None],
135 integrity_verify: Callable[[Path], Tuple[bool, Optional[str]]],
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 # ------------------------------------------------------------------ #
195 # index
196 # ------------------------------------------------------------------ #
197 def index(
198 self,
199 *,
200 source_type: str,
201 source_id: Optional[str],
202 chunks: Sequence[ChunkInput],
203 replace: bool = True,
204 session: Optional[Session] = None,
205 ) -> IndexStats:
206 """Embed + persist ``chunks`` for one source into this collection.
208 Ordering: write chunk rows and ``flush()`` for their int ids, then
209 ``store.apply()`` (persists vectors atomically under its lock), then
210 delete the source's prior rows, then ``commit()`` (only if we own the
211 session). If ``store.apply()`` fails we roll back — the new rows never
212 existed, so no DB row can point at a missing vector.
214 ``store.apply()`` persists vectors durably BEFORE the DB commit, so if
215 the commit later fails the vectors outlive their (rolled-back) rows as
216 orphans — searchable-by-nothing, cleaned on the next re-index of the
217 source. When ``session`` is passed the caller owns that window.
219 ``replace=True`` removes this source's PRIOR chunks/vectors first;
220 ``False`` appends.
221 """
222 chunks = list(chunks)
223 # An empty chunk list with replace=True still must remove this source's
224 # PRIOR rows/vectors — e.g. a note edited down to whitespace yields zero
225 # chunks, and skipping here would leave the user-cleared content
226 # permanently searchable. Only a genuine no-op (nothing to add AND
227 # nothing to replace) returns early.
228 if not chunks and (not replace or source_id is None): 228 ↛ 229line 228 didn't jump to line 229 because the condition on line 228 was never true
229 return IndexStats(chunks=0, added=0, removed=0)
231 texts = [c.text for c in chunks]
232 vectors = (
233 np.asarray(self.embeddings.embed_documents(texts), dtype="float32")
234 if texts
235 else None
236 )
238 sid = str(source_id) if source_id is not None else None
239 with self._db(session) as (sess, owns):
240 prior_ids: List[int] = []
241 if replace and sid is not None:
242 # Scope the prior-chunk lookup to THIS store's embedding-model
243 # config. collection_name is identical across every model config
244 # for a collection (only RAGIndex.index_hash / the physical
245 # .faiss path differ per model), so a lookup by
246 # (source_type, source_id, collection_name) alone would also
247 # match rows indexed under a DIFFERENT model. Those ids would
248 # then be handed to THIS (physically separate, model-specific)
249 # store as remove_ids — where faiss silently no-ops on unknown
250 # ids — and _finalize_db would still DELETE the rows. Result: a
251 # reindex under model B destroys model A's chunk text (its only
252 # copy) while A's .faiss keeps the now-unresolvable vector,
253 # even though A's index is deliberately kept for model revert.
254 prior_ids = [
255 row.id
256 for row in sess.query(DocumentChunk.id).filter_by(
257 source_type=source_type,
258 source_id=sid,
259 collection_name=self.collection_name,
260 embedding_model=self.embedding_model,
261 embedding_model_type=self.embedding_model_type,
262 )
263 ]
265 rows = [self._build_row(c, source_type, sid) for c in chunks]
266 if rows:
267 sess.add_all(rows)
268 sess.flush() # assigns int PKs to every row in one round-trip
269 # SQLAlchemy's legacy Column() declarative style types instance
270 # attribute access as Column[int] rather than int; cast() is a
271 # no-op at runtime (row.id is already an int post-flush()).
272 new_ids = [cast(int, row.id) for row in rows]
274 try:
275 stats = self._store.apply(
276 add_ids=new_ids,
277 add_vectors=vectors,
278 remove_ids=prior_ids,
279 )
280 except Exception:
281 # apply() failed without a durable change (fail-closed reload /
282 # add error / atomic-save) — a clean DB rollback keeps them
283 # consistent.
284 if owns: 284 ↛ 286line 284 didn't jump to line 286 because the condition on line 284 was always true
285 sess.rollback()
286 raise
287 # apply() is DURABLE from here; a DB-side failure below leaves the
288 # store and DB possibly disagreeing (orphan vectors/rows).
289 self._finalize_db(
290 sess, owns, prior_ids, f"index {source_type}/{sid}"
291 )
293 return IndexStats(
294 chunks=len(chunks),
295 added=stats["added"],
296 removed=stats["removed"],
297 )
299 @staticmethod
300 def _finalize_db(
301 session: Session,
302 owns: bool,
303 delete_ids: List[int],
304 what: str,
305 ) -> int:
306 """Delete rows + commit (if owned), AFTER a durable ``store.apply()``.
308 Covers the whole post-apply DB window in one place: if either the
309 prior-row delete OR the commit fails, the store has already persisted
310 vectors that the DB won't reflect (orphan vectors / rows) — a
311 known-inconsistent state that reconciles on the next re-index of the
312 source. Log it loudly, roll back if we own the session, then re-raise.
313 ``synchronize_session="fetch"`` keeps the reused thread-local session's
314 identity map from serving stale, already-deleted objects. Returns the
315 number of rows deleted.
316 """
317 deleted = 0
318 try:
319 if delete_ids:
320 deleted = (
321 session.query(DocumentChunk)
322 .filter(DocumentChunk.id.in_(delete_ids))
323 .delete(synchronize_session="fetch")
324 )
325 if owns:
326 session.commit()
327 except Exception:
328 logger.exception(
329 f"{what}: vectors were persisted but the DB update failed; "
330 "store and DB will reconcile on the next re-index of the source"
331 )
332 if owns: 332 ↛ 334line 332 didn't jump to line 334 because the condition on line 332 was always true
333 session.rollback()
334 raise
335 return deleted
337 def _build_row(
338 self, chunk: ChunkInput, source_type: str, source_id: Optional[str]
339 ) -> DocumentChunk:
340 meta = chunk.metadata or {}
341 return DocumentChunk(
342 chunk_hash=hashlib.sha256(chunk.text.encode()).hexdigest(),
343 source_type=source_type,
344 source_id=source_id,
345 collection_name=self.collection_name,
346 chunk_text=chunk.text,
347 chunk_index=meta.get("chunk_index", 0),
348 start_char=meta.get("start_char", 0),
349 end_char=meta.get("end_char", len(chunk.text)),
350 word_count=len(chunk.text.split()),
351 # embedding_id stays populated: it's the old-uuid -> new-int join key
352 # the one-time index-format migration needs. (Legacy NOT NULL UNIQUE
353 # column; drop it only in a later, separate migration.)
354 embedding_id=uuid.uuid4().hex,
355 embedding_model=self.embedding_model,
356 embedding_model_type=self.embedding_model_type,
357 embedding_dimension=self._store.dimension,
358 document_title=meta.get("document_title") or meta.get("title"),
359 document_metadata=meta or None,
360 )
362 # ------------------------------------------------------------------ #
363 # search
364 # ------------------------------------------------------------------ #
365 def search(
366 self,
367 query: str,
368 top_k: int,
369 session: Optional[Session] = None,
370 ) -> List[SearchResult]:
371 """Embed ``query``, search the store, rehydrate hits from the DB by id.
373 Orphan hits (a vector whose DB row is gone) are skipped — there is no
374 in-store text to fall back to, and skipping is the correct signal.
375 """
376 if top_k <= 0:
377 return []
378 query_vec = np.asarray(
379 self.embeddings.embed_query(query), dtype="float32"
380 )
381 hits = self._store.search(query_vec, top_k) # [(chunk_id, distance)]
382 if not hits:
383 return []
385 ids = [cid for cid, _ in hits]
386 with self._db(session) as (sess, _owns):
387 # Scope rehydration to THIS collection. A hit id is a DocumentChunk.id
388 # from this collection's index, but an orphan/stale vector (e.g. a
389 # rolled-back insert left a vector whose id was later reused by a
390 # chunk in a DIFFERENT collection) could otherwise match a row that
391 # belongs elsewhere and leak that other document's text into this
392 # collection's results. Filtering by collection_name makes such a
393 # hit resolve to no row here and be skipped as an orphan.
394 rows = {
395 # cast(): row.id is Column[int] under the legacy Column()
396 # declarative style; a no-op at runtime.
397 cast(int, row.id): row
398 for row in sess.query(DocumentChunk).filter(
399 DocumentChunk.id.in_(ids),
400 DocumentChunk.collection_name == self.collection_name,
401 )
402 }
404 results: List[SearchResult] = []
405 for chunk_id, distance in hits:
406 row = rows.get(chunk_id)
407 if row is None:
408 logger.debug(
409 f"Search hit {chunk_id} has no DocumentChunk row; skipping"
410 )
411 continue
412 # SQLAlchemy's legacy Column() declarative style types instance
413 # attribute access as Column[...] rather than the plain Python
414 # value; cast() to the DB-verified runtime type (no conversion,
415 # so a None stays None rather than becoming e.g. the "None"
416 # string a str()/int() call would produce).
417 results.append(
418 SearchResult(
419 chunk_id=chunk_id,
420 text=cast(str, row.chunk_text),
421 distance=distance,
422 metric=self.metric,
423 metadata=cast(Optional[dict], row.document_metadata) or {},
424 document_title=cast(Optional[str], row.document_title),
425 source_id=cast(Optional[str], row.source_id),
426 source_type=cast(Optional[str], row.source_type),
427 )
428 )
429 return results
431 # ------------------------------------------------------------------ #
432 # delete
433 # ------------------------------------------------------------------ #
434 def delete(
435 self,
436 *,
437 source_type: str,
438 source_id: str,
439 session: Optional[Session] = None,
440 ) -> DeleteStats:
441 """Remove all of one source's chunks/vectors from this collection.
443 ``store.apply()`` (the durable side) runs before the DB delete + commit,
444 matching :meth:`index`'s ordering.
446 Re-queries ``DocumentChunk`` by (source_type, source_id,
447 collection_name) at call time to discover the ids to remove, so the
448 rows must still exist when this is called. If a caller has already
449 deleted those rows elsewhere, this finds nothing to do; use
450 :meth:`delete_ids` instead with the ids captured before deletion.
451 """
452 sid = str(source_id)
453 with self._db(session) as (sess, owns):
454 ids = [
455 row.id
456 for row in sess.query(DocumentChunk.id).filter_by(
457 source_type=source_type,
458 source_id=sid,
459 collection_name=self.collection_name,
460 )
461 ]
462 if not ids:
463 return DeleteStats(removed=0, rows_deleted=0)
465 try:
466 stats = self._store.apply(
467 add_ids=[], add_vectors=None, remove_ids=ids
468 )
469 except Exception:
470 if owns:
471 sess.rollback()
472 raise
473 # Guard the row-delete AND commit together (mirrors index()): a bare
474 # query(...).delete() that raised after the durable apply() above
475 # would leave orphan rows with no rollback and no diagnostic.
476 rows_deleted = self._finalize_db(
477 sess, owns, ids, f"delete {source_type}/{sid}"
478 )
480 return DeleteStats(removed=stats["removed"], rows_deleted=rows_deleted)
482 def delete_ids(
483 self,
484 ids: Sequence[int],
485 *,
486 session: Optional[Session] = None,
487 ) -> DeleteStats:
488 """Remove specific chunk ids' vectors from this collection's store.
490 For callers that already captured the ``DocumentChunk`` int ids
491 BEFORE the matching rows were deleted elsewhere — e.g. a caller
492 whose own per-item transaction must delete the rows eagerly (for
493 its own atomicity/rollback semantics) ahead of a later BATCHED
494 vector-store cleanup pass. Unlike :meth:`delete`, this does NOT
495 re-query ``DocumentChunk`` by (source_type, source_id,
496 collection_name) to discover the ids — by the time this runs the
497 rows the ids came from may already be gone, and that lookup would
498 find nothing (see the "rows must still exist" requirement noted on
499 :meth:`delete`). Any rows matching the given ids that do still
500 exist are cleaned up too (defensive; normally none are left).
501 """
502 wanted = sorted({int(i) for i in ids if i is not None})
503 if not wanted:
504 return DeleteStats(removed=0, rows_deleted=0)
506 with self._db(session) as (sess, owns):
507 try:
508 stats = self._store.apply(
509 add_ids=[], add_vectors=None, remove_ids=wanted
510 )
511 except Exception:
512 if owns:
513 sess.rollback()
514 raise
515 # Guard the row-delete AND commit together (mirrors index()).
516 rows_deleted = self._finalize_db(
517 sess, owns, wanted, f"delete_ids ({len(wanted)} ids)"
518 )
520 return DeleteStats(removed=stats["removed"], rows_deleted=rows_deleted)
522 def count(self) -> int:
523 """Number of vectors currently in the store."""
524 return self._store.count()