cbrkit.indexable.sqlite_vec

  1"""SQLite ``sqlite-vec`` tabular storage backend (async-first, sync facade).
  2
  3Extends the dialect-agnostic :class:`cbrkit.indexable.sqlalchemy_async`
  4base with vector search powered by the `sqlite-vec
  5<https://github.com/asg017/sqlite-vec>`_ loadable extension.  The layout is
  6shaped to SQLite's own primitives rather than mirroring pgvector:
  7
  8- **Dense** embeddings live in a ``vec0`` *virtual table* (``<table>_vec``)
  9  keyed by the casebase key, queried with ``sqlite-vec``'s KNN
 10  ``... MATCH ... AND k = N``.  Building on ``vec0`` (rather than a plain
 11  BLOB column) means the backend inherits future ``vec0`` capabilities —
 12  approximate nearest-neighbor indexing as it lands upstream — and supports
 13  quantized element types today via :paramref:`vector_type`.
 14- **Sparse** full-text search uses SQLite's built-in **FTS5** in a second
 15  shadow table (``<table>_fts``).
 16
 17Both shadow tables are kept in sync with the main table: the FTS shadow
 18entirely by triggers, and the ``vec0`` shadow by an ``AFTER DELETE`` trigger
 19plus a Python-side insert on write (embeddings are computed in Python, so a
 20SQL trigger cannot populate them).  The ``sqlite-vec`` extension is loaded on
 21every connection via a SQLAlchemy ``checkout`` event reaching the
 22``aiosqlite`` worker thread.
 23
 24The shadow tables are cbrkit-owned auxiliary indexes, distinct from the
 25main data table: unlike pgvector (where the vector/FTS targets are *columns*
 26a host model can declare), a ``vec0`` virtual table cannot live inside a
 27SQLAlchemy model.  cbrkit therefore creates and maintains the shadows
 28*regardless* of :paramref:`manage_schema` — which governs only the main
 29data table.  A custom-schema app brings its own ``model`` / ``table`` and
 30calls :meth:`reindex` once to backfill the shadows from pre-existing rows;
 31data written through cbrkit afterwards stays in sync automatically.
 32
 33Filtering composes with KNN by joining the ``vec0`` matches back to the main
 34table; because ``vec0`` returns a fixed ``k`` *before* the join, a ``where``
 35filter oversamples candidates (see the retriever's ``hybrid_oversample``)
 36and may under-fill the limit on highly selective filters.
 37
 38The simplest setup stores plain strings (``V = str``) in a cbrkit-built
 39table::
 40
 41    storage = sqlite_vec(
 42        url="sqlite+aiosqlite:///cases.db",
 43        value_column="text",
 44        vector_dim=384,
 45        index_type="hybrid",
 46        conversion_func=embed,
 47    )
 48    storage.put_index({"a": "red sedan"})  # index -> {"a": "red sedan"}
 49"""
 50
 51from collections.abc import Mapping
 52from dataclasses import dataclass
 53from typing import Any, Literal, cast
 54from weakref import WeakSet
 55
 56import numpy as np
 57import sqlalchemy as sa
 58import sqlite_vec as sqlite_vec_ext
 59from aiosqlite import Connection as AiosqliteConnection
 60from sqlalchemy import Engine, event
 61from sqlalchemy.engine import AdaptedConnection
 62from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 63from sqlalchemy.pool import ConnectionPoolEntry, PoolProxiedConnection
 64
 65from ..helpers import chunkify, forward_fields, run_coroutine
 66from ..typing import BatchConversionFunc, NumpyArray
 67from ._common import SQLITE_VEC_METRICS, SQLITE_VEC_TYPES
 68from .sqlalchemy import build_indexable_table, sqlalchemy, sqlalchemy_async
 69
 70# Row-dict key carrying the serialized vector from _populate_system_columns to
 71# _insert_vectors.  Deliberately not a column name: the vec0 shadow's columns
 72# are an independent namespace from the data table's, so keying by
 73# vector_column would make an unrelated host column of the same name collide.
 74_VECTOR_KEY = "__cbrkit_vector__"
 75
 76# Engines whose checkout hook is already attached, tracked weakly so a
 77# discarded engine does not keep its entry alive.
 78_hooked_engines: WeakSet[Engine] = WeakSet()
 79
 80
 81def serialize_vector(vec: NumpyArray) -> bytes:
 82    """Pack an embedding into the ``sqlite-vec`` ``float32`` BLOB format."""
 83    return sqlite_vec_ext.serialize_float32(np.asarray(vec, dtype=np.float32).tolist())
 84
 85
 86def _attach_sqlite_vec_loader(engine: AsyncEngine) -> None:
 87    """Load ``sqlite-vec`` on every connection of a SQLite async engine.
 88
 89    The real ``sqlite3`` connection lives in ``aiosqlite``'s worker thread,
 90    so the extension must be loaded in that thread through the adapter's
 91    ``run_async`` method.
 92
 93    Hooked on ``checkout`` rather than ``connect`` so connections a host engine
 94    opened before handing it over get loaded too, instead of failing with ``no
 95    such module: vec0``.  The ``record.info`` latch keeps the load to once per
 96    real connection, and attaching is idempotent per engine.
 97    """
 98    sync_engine = engine.sync_engine
 99
100    if sync_engine.dialect.name != "sqlite" or sync_engine in _hooked_engines:
101        return
102
103    _hooked_engines.add(sync_engine)
104    ext_path = sqlite_vec_ext.loadable_path()
105
106    async def _load_extension(driver: AiosqliteConnection) -> None:
107        await driver.enable_load_extension(True)
108
109        try:
110            await driver.load_extension(ext_path)
111        finally:
112            await driver.enable_load_extension(False)
113
114    @event.listens_for(sync_engine, "checkout")
115    def _load(
116        dbapi_conn: AdaptedConnection,
117        record: ConnectionPoolEntry,
118        _: PoolProxiedConnection,
119    ) -> None:
120        if record.info.get("cbrkit_sqlite_vec_loaded"):
121            return
122
123        dbapi_conn.run_async(_load_extension)
124        record.info["cbrkit_sqlite_vec_loaded"] = True
125
126
127@dataclass(slots=True)
128class sqlite_vec_async[K: int | str, V = Mapping[str, Any]](sqlalchemy_async[K, V]):
129    """Async SQLite/``sqlite-vec`` tabular storage.
130
131    Extends :class:`sqlalchemy_async` with a ``vec0`` shadow table for dense
132    KNN and an FTS5 shadow table for sparse search, both derived from the
133    inherited :paramref:`value_column` (the embeddable text column).
134
135    Args:
136        vector_column: Name of the vector column inside the ``vec0`` table.
137        vector_dim: Embedding dimension.  Required for *index_type* ∈
138            {``"dense"``, ``"hybrid"``} (``vec0`` declares the column as
139            ``<type>[dim]``).
140        vector_type: ``vec0`` element type — ``"float32"`` (exact) or
141            ``"int8"`` (quantized, ~4x smaller; embeddings are quantized by
142            ``sqlite-vec`` assuming unit-normalized vectors).
143        index_type: ``"dense"`` (vector KNN only), ``"sparse"`` (FTS5 only),
144            or ``"hybrid"`` (both).
145        metric_type: Distance metric for the ``vec0`` table (and mirrored at
146            search time by the retriever wrapper).
147        fts_tokenizer: Optional FTS5 ``tokenize=`` directive (e.g.
148            ``"porter unicode61"``); ``None`` uses the FTS5 default.
149        conversion_func: Embedding function.  Required for ``"dense"`` /
150            ``"hybrid"`` index types.
151    """
152
153    vector_column: str = "embedding"
154    vector_dim: int | None = None
155    vector_type: Literal["float32", "int8"] = "float32"
156    index_type: Literal["dense", "sparse", "hybrid"] = "dense"
157    metric_type: Literal["cosine", "l2", "l1"] = "cosine"
158    fts_tokenizer: str | None = None
159    conversion_func: BatchConversionFunc[str, NumpyArray] | None = None
160
161    @property
162    def has_dense(self) -> bool:
163        """Whether this storage maintains a dense ``vec0`` index."""
164        return self.index_type in ("dense", "hybrid")
165
166    @property
167    def has_sparse(self) -> bool:
168        """Whether this storage maintains a sparse FTS5 index."""
169        return self.index_type in ("sparse", "hybrid")
170
171    @property
172    def vec_table_name(self) -> str:
173        """Name of the ``vec0`` shadow table (``<table>_vec``)."""
174        return f"{self.sa_table.name}_vec"
175
176    @property
177    def fts_table_name(self) -> str:
178        """Name of the FTS5 shadow table (``<table>_fts``)."""
179        return f"{self.sa_table.name}_fts"
180
181    @property
182    def vector_value_sql(self) -> str:
183        """SQL template wrapping a bound float32 BLOB for the element type."""
184        return SQLITE_VEC_TYPES[self.vector_type].value_template
185
186    @property
187    def fts_table(self) -> sa.Table:
188        """A lightweight :class:`sa.Table` over the FTS5 shadow for queries.
189
190        Lives on its own :class:`sa.MetaData` so it never participates in
191        the main table's DDL; the retriever joins it back to the main table
192        on the key column.
193        """
194        assert self.value_column is not None
195        return build_indexable_table(
196            self.fts_table_name,
197            metadata=sa.MetaData(),
198            key_column=self.key_column,
199            key_type=self.key_type,
200            columns={self.value_column: sa.Text()},
201        )
202
203    def _validate_init(self) -> None:
204        super(sqlite_vec_async, self)._validate_init()
205        if self.value_column is None:
206            raise ValueError("value_column is required for sqlite_vec")
207        if self.table is not None and self.value_column not in self.table.columns:
208            raise ValueError(
209                f"value_column={self.value_column!r} must be a column of the "
210                "model / table"
211            )
212        if self.has_dense and self.conversion_func is None:
213            raise ValueError(
214                f"conversion_func is required for index_type={self.index_type!r}"
215            )
216        if self.has_dense and self.vector_dim is None:
217            raise ValueError(
218                f"vector_dim is required for index_type={self.index_type!r} "
219                "(the vec0 shadow declares the column as <type>[dim])"
220            )
221        if self.vector_dim is not None and self.vector_dim <= 0:
222            raise ValueError(
223                f"vector_dim must be a positive int (got {self.vector_dim!r})"
224            )
225
226    def _init_engine(self) -> None:
227        super(sqlite_vec_async, self)._init_engine()
228        _attach_sqlite_vec_loader(self._engine)
229
230    def _create_schema(self, sync_conn: sa.Connection) -> None:
231        # The base creates the main data table only when manage_schema=True;
232        # the shadow indexes are cbrkit-owned and created either way.
233        super(sqlite_vec_async, self)._create_schema(sync_conn)
234        self._create_shadows(sync_conn)
235
236    def _create_shadows(self, sync_conn: sa.Connection) -> None:
237        assert self.value_column is not None
238        key, val, main = self.key_column, self.value_column, self._table.name
239
240        if self.has_dense:
241            assert self.vector_dim is not None
242            vec = self.vec_table_name
243            pk_type = "integer" if self.key_type == "int" else "text"
244            col_type = SQLITE_VEC_TYPES[self.vector_type].column_type
245            metric = SQLITE_VEC_METRICS[self.metric_type].distance_metric
246            # vec0's own DDL parser does not accept quoted column identifiers.
247            sync_conn.execute(
248                sa.text(
249                    f'CREATE VIRTUAL TABLE IF NOT EXISTS "{vec}" USING vec0('
250                    f"{key} {pk_type} primary key, "
251                    f"{self.vector_column} {col_type}[{self.vector_dim}] "
252                    f"distance_metric={metric})"
253                )
254            )
255            # The embedding is computed in Python, so only deletes are
256            # trigger-maintained; inserts happen in _do_upsert.
257            sync_conn.execute(
258                sa.text(
259                    f'CREATE TRIGGER IF NOT EXISTS "{vec}_ad" AFTER DELETE ON "{main}" '
260                    f'BEGIN DELETE FROM "{vec}" WHERE "{key}" = old."{key}"; END'
261                )
262            )
263
264        if self.has_sparse:
265            fts = self.fts_table_name
266            tok = f", tokenize='{self.fts_tokenizer}'" if self.fts_tokenizer else ""
267            sync_conn.execute(
268                sa.text(
269                    f'CREATE VIRTUAL TABLE IF NOT EXISTS "{fts}" '
270                    f'USING fts5("{key}" UNINDEXED, "{val}"{tok})'
271                )
272            )
273            # Triggers keep the FTS shadow in sync with every write path
274            # (the base upserts via delete-then-insert; UPDATE is covered too).
275            sync_conn.execute(
276                sa.text(
277                    f'CREATE TRIGGER IF NOT EXISTS "{fts}_ai" AFTER INSERT ON "{main}" '
278                    f'BEGIN INSERT INTO "{fts}"("{key}", "{val}") '
279                    f'VALUES (new."{key}", new."{val}"); END'
280                )
281            )
282            sync_conn.execute(
283                sa.text(
284                    f'CREATE TRIGGER IF NOT EXISTS "{fts}_ad" AFTER DELETE ON "{main}" '
285                    f'BEGIN DELETE FROM "{fts}" WHERE "{key}" = old."{key}"; END'
286                )
287            )
288            sync_conn.execute(
289                sa.text(
290                    f'CREATE TRIGGER IF NOT EXISTS "{fts}_au" AFTER UPDATE ON "{main}" '
291                    f'BEGIN DELETE FROM "{fts}" WHERE "{key}" = old."{key}"; '
292                    f'INSERT INTO "{fts}"("{key}", "{val}") '
293                    f'VALUES (new."{key}", new."{val}"); END'
294                )
295            )
296
297    def _populate_system_columns(self, rows: list[dict[str, Any]]) -> None:
298        """Embed the text column into a serialized vector on each row.
299
300        The vector belongs to the ``vec0`` shadow rather than the main table, so
301        it rides along under :data:`_VECTOR_KEY` and the base narrows it back
302        out (see :meth:`_table_rows`).
303        """
304        if not self.has_dense:
305            return
306        assert self.conversion_func is not None
307        assert self.value_column is not None
308        texts = [row[self.value_column] for row in rows]
309        for row, vec in zip(rows, self.conversion_func(texts), strict=True):
310            row[_VECTOR_KEY] = serialize_vector(vec)
311
312    async def _do_upsert(
313        self, conn: AsyncConnection, rows: list[dict[str, Any]]
314    ) -> None:
315        # Write the main table (triggers drop stale vec0/fts rows and refresh
316        # fts), then repopulate the vec0 shadow from the same rows.
317        await super(sqlite_vec_async, self)._do_upsert(conn, rows)
318        if self.has_dense:
319            await self._insert_vectors(conn, rows)
320
321    async def _insert_vectors(
322        self, conn: AsyncConnection, rows: list[dict[str, Any]]
323    ) -> None:
324        """Insert the vectors prepared on *rows* into the ``vec0`` shadow."""
325        stmt = sa.text(
326            f'INSERT INTO "{self.vec_table_name}"'
327            f'("{self.key_column}", "{self.vector_column}") '
328            f"VALUES (:key, {self.vector_value_sql.format(':vec')})"
329        )
330        params = [
331            {"key": row[self.key_column], "vec": row[_VECTOR_KEY]} for row in rows
332        ]
333
334        for chunk in chunkify(params, self._param_batch(2)):
335            await conn.execute(stmt, chunk)
336
337    async def reindex(self, batch_size: int = 1000) -> int:
338        """Rebuild the shadow indexes from the existing main-table rows.
339
340        Clears the ``vec0`` / FTS5 shadows and repopulates them page by page.
341        Use this once after pointing the storage at a host table that already
342        holds data (writes made *through* cbrkit keep the shadows in sync on
343        their own).  Until it returns, a search sees only the pages rebuilt so
344        far: an interrupted run is incomplete rather than inconsistent, so it
345        needs no atomic mode of its own, just a re-run.
346
347        Returns:
348            The number of rows indexed.
349        """
350        async with self._engine.begin() as conn:
351            await self._ensure_schema(conn)
352            if self.has_dense:
353                await conn.execute(sa.text(f'DELETE FROM "{self.vec_table_name}"'))
354            if self.has_sparse:
355                await conn.execute(sa.text(f'DELETE FROM "{self.fts_table_name}"'))
356
357        return await self._rebuild_pages(self._populate_shadows, batch_size)
358
359    async def _populate_shadows(
360        self, conn: AsyncConnection, rows: list[dict[str, Any]]
361    ) -> None:
362        """Insert prepared rows directly into the FTS5 / ``vec0`` shadows.
363
364        Used by :meth:`reindex` to backfill from existing rows; the normal
365        write path keeps FTS in sync via triggers instead.
366        """
367        assert self.value_column is not None
368
369        if self.has_sparse:
370            await conn.execute(
371                sa.insert(self.fts_table),
372                [
373                    {
374                        self.key_column: row[self.key_column],
375                        self.value_column: row[self.value_column],
376                    }
377                    for row in rows
378                ],
379            )
380        if self.has_dense:
381            await self._insert_vectors(conn, rows)
382
383
384@dataclass(slots=True)
385class sqlite_vec[K: int | str, V = Mapping[str, Any]](sqlalchemy[K, V]):
386    """Sync facade over :class:`sqlite_vec_async`.
387
388    Adds the ``sqlite-vec``-specific configuration on top of the
389    :class:`cbrkit.indexable.sqlalchemy` sync facade and overrides
390    :meth:`_build_async` to construct a :class:`sqlite_vec_async`.
391    """
392
393    vector_column: str = "embedding"
394    vector_dim: int | None = None
395    vector_type: Literal["float32", "int8"] = "float32"
396    index_type: Literal["dense", "sparse", "hybrid"] = "dense"
397    metric_type: Literal["cosine", "l2", "l1"] = "cosine"
398    fts_tokenizer: str | None = None
399    conversion_func: BatchConversionFunc[str, NumpyArray] | None = None
400
401    def _build_async(self) -> sqlite_vec_async[K, V]:
402        return sqlite_vec_async[K, V](
403            engine=self._engine, **forward_fields(self, exclude={"url"})
404        )
405
406    @property
407    def async_storage(self) -> sqlite_vec_async[K, V]:
408        """The wrapped async storage (used by sync retriever facades)."""
409        return cast("sqlite_vec_async[K, V]", self._async)
410
411    def reindex(self, batch_size: int = 1000) -> int:
412        """Rebuild the shadow indexes from existing main-table rows."""
413        return run_coroutine(self.async_storage.reindex(batch_size))
414
415
416__all__ = [
417    "sqlite_vec",
418    "sqlite_vec_async",
419]