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 ``connect`` 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
 51import asyncio
 52from collections.abc import Mapping
 53from dataclasses import dataclass, field
 54from typing import Any, Literal, cast
 55
 56import numpy as np
 57import sqlite_vec as sqlite_vec_ext
 58import sqlalchemy as sa
 59from sqlalchemy import event
 60from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 61
 62from ..helpers import forward_fields, run_coroutine
 63from ..typing import BatchConversionFunc, NumpyArray
 64from ._common import SQLITE_VEC_METRICS, SQLITE_VEC_TYPES
 65from .sqlalchemy import build_indexable_table, sqlalchemy, sqlalchemy_async
 66
 67
 68def _attach_sqlite_vec_loader(engine: AsyncEngine) -> None:
 69    """Load ``sqlite-vec`` on every new connection of a SQLite async engine.
 70
 71    The real ``sqlite3`` connection lives in ``aiosqlite``'s worker thread,
 72    so the extension must be loaded *in that thread* — reached here through
 73    the async driver connection's coroutine API driven by the adapter's
 74    ``await_`` bridge.  Attaching is idempotent per engine.
 75    """
 76    sync_engine = engine.sync_engine
 77    if sync_engine.dialect.name != "sqlite":
 78        return
 79    if getattr(sync_engine, "_cbrkit_sqlite_vec_loaded", False):
 80        return
 81    setattr(sync_engine, "_cbrkit_sqlite_vec_loaded", True)  # noqa: B010
 82    ext_path = sqlite_vec_ext.loadable_path()
 83
 84    @event.listens_for(sync_engine, "connect")
 85    def _load(dbapi_conn: Any, _: Any) -> None:
 86        driver = dbapi_conn.driver_connection
 87        dbapi_conn.await_(driver.enable_load_extension(True))
 88        dbapi_conn.await_(driver.load_extension(ext_path))
 89        dbapi_conn.await_(driver.enable_load_extension(False))
 90
 91
 92@dataclass(slots=True)
 93class sqlite_vec_async[K: int | str, V = Mapping[str, Any]](sqlalchemy_async[K, V]):
 94    """Async SQLite/``sqlite-vec`` tabular storage.
 95
 96    Extends :class:`sqlalchemy_async` with a ``vec0`` shadow table for dense
 97    KNN and an FTS5 shadow table for sparse search, both derived from the
 98    inherited :paramref:`value_column` (the embeddable text column).
 99
100    Args:
101        vector_column: Name of the vector column inside the ``vec0`` table.
102        vector_dim: Embedding dimension.  Required for *index_type* ∈
103            {``"dense"``, ``"hybrid"``} (``vec0`` declares the column as
104            ``<type>[dim]``).
105        vector_type: ``vec0`` element type — ``"float32"`` (exact) or
106            ``"int8"`` (quantized, ~4x smaller; embeddings are quantized by
107            ``sqlite-vec`` assuming unit-normalized vectors).
108        index_type: ``"dense"`` (vector KNN only), ``"sparse"`` (FTS5 only),
109            or ``"hybrid"`` (both).
110        metric_type: Distance metric for the ``vec0`` table (and mirrored at
111            search time by the retriever wrapper).
112        fts_tokenizer: Optional FTS5 ``tokenize=`` directive (e.g.
113            ``"porter unicode61"``); ``None`` uses the FTS5 default.
114        conversion_func: Embedding function.  Required for ``"dense"`` /
115            ``"hybrid"`` index types.
116    """
117
118    vector_column: str = "embedding"
119    vector_dim: int | None = None
120    vector_type: Literal["float32", "int8"] = "float32"
121    index_type: Literal["dense", "sparse", "hybrid"] = "dense"
122    metric_type: Literal["cosine", "l2", "l1"] = "cosine"
123    fts_tokenizer: str | None = None
124    conversion_func: BatchConversionFunc[str, NumpyArray] | None = None
125    _shadows_ready: bool = field(init=False, default=False, repr=False)
126
127    @property
128    def has_dense(self) -> bool:
129        """Whether this storage maintains a dense ``vec0`` index."""
130        return self.index_type in ("dense", "hybrid")
131
132    @property
133    def has_sparse(self) -> bool:
134        """Whether this storage maintains a sparse FTS5 index."""
135        return self.index_type in ("sparse", "hybrid")
136
137    @property
138    def vec_table_name(self) -> str:
139        """Name of the ``vec0`` shadow table (``<table>_vec``)."""
140        return f"{self.sa_table.name}_vec"
141
142    @property
143    def fts_table_name(self) -> str:
144        """Name of the FTS5 shadow table (``<table>_fts``)."""
145        return f"{self.sa_table.name}_fts"
146
147    @property
148    def vector_value_sql(self) -> str:
149        """SQL template wrapping a bound float32 BLOB for the element type."""
150        return SQLITE_VEC_TYPES[self.vector_type].value_template
151
152    @property
153    def fts_table(self) -> sa.Table:
154        """A lightweight :class:`sa.Table` over the FTS5 shadow for queries.
155
156        Lives on its own :class:`sa.MetaData` so it never participates in
157        the main table's DDL; the retriever joins it back to the main table
158        on the key column.
159        """
160        assert self.value_column is not None
161        return build_indexable_table(
162            self.fts_table_name,
163            metadata=sa.MetaData(),
164            key_column=self.key_column,
165            key_type=self.key_type,
166            columns={self.value_column: sa.Text()},
167        )
168
169    def _validate_init(self) -> None:
170        super(sqlite_vec_async, self)._validate_init()
171        if self.value_column is None:
172            raise ValueError("value_column is required for sqlite_vec")
173        if self.table is not None and self.value_column not in self.table.columns:
174            raise ValueError(
175                f"value_column={self.value_column!r} must be a column of the "
176                "model / table"
177            )
178        if self.has_dense and self.conversion_func is None:
179            raise ValueError(
180                f"conversion_func is required for index_type={self.index_type!r}"
181            )
182        if self.has_dense and self.vector_dim is None:
183            raise ValueError(
184                f"vector_dim is required for index_type={self.index_type!r} "
185                "(the vec0 shadow declares the column as <type>[dim])"
186            )
187        if self.vector_dim is not None and self.vector_dim <= 0:
188            raise ValueError(
189                f"vector_dim must be a positive int (got {self.vector_dim!r})"
190            )
191
192    def _init_engine(self) -> None:
193        super(sqlite_vec_async, self)._init_engine()
194        _attach_sqlite_vec_loader(self._engine)
195
196    async def _ensure_schema(self, conn: AsyncConnection) -> None:
197        # The base creates the main data table only when manage_schema=True;
198        # the shadow indexes are cbrkit-owned and created unconditionally.
199        await super(sqlite_vec_async, self)._ensure_schema(conn)
200        await self._ensure_shadows(conn)
201
202    async def _ensure_shadows(self, conn: AsyncConnection) -> None:
203        if self._shadows_ready:
204            return
205        if self.has_dense or self.has_sparse:
206            await conn.run_sync(self._create_shadows)
207        self._shadows_ready = True
208
209    def _create_shadows(self, sync_conn: sa.Connection) -> None:
210        assert self.value_column is not None
211        key, val, main = self.key_column, self.value_column, self._table.name
212
213        if self.has_dense:
214            assert self.vector_dim is not None
215            vec = self.vec_table_name
216            pk_type = "integer" if self.key_type == "int" else "text"
217            col_type = SQLITE_VEC_TYPES[self.vector_type].column_type
218            metric = SQLITE_VEC_METRICS[self.metric_type].distance_metric
219            # vec0's own DDL parser does not accept quoted column identifiers.
220            sync_conn.execute(
221                sa.text(
222                    f'CREATE VIRTUAL TABLE IF NOT EXISTS "{vec}" USING vec0('
223                    f"{key} {pk_type} primary key, "
224                    f"{self.vector_column} {col_type}[{self.vector_dim}] "
225                    f"distance_metric={metric})"
226                )
227            )
228            # The embedding is computed in Python, so only deletes are
229            # trigger-maintained; inserts happen in _do_upsert.
230            sync_conn.execute(
231                sa.text(
232                    f'CREATE TRIGGER IF NOT EXISTS "{vec}_ad" AFTER DELETE ON "{main}" '
233                    f'BEGIN DELETE FROM "{vec}" WHERE "{key}" = old."{key}"; END'
234                )
235            )
236
237        if self.has_sparse:
238            fts = self.fts_table_name
239            tok = f", tokenize='{self.fts_tokenizer}'" if self.fts_tokenizer else ""
240            sync_conn.execute(
241                sa.text(
242                    f'CREATE VIRTUAL TABLE IF NOT EXISTS "{fts}" '
243                    f'USING fts5("{key}" UNINDEXED, "{val}"{tok})'
244                )
245            )
246            # Triggers keep the FTS shadow in sync with every write path
247            # (the base upserts via delete-then-insert; UPDATE is covered too).
248            sync_conn.execute(
249                sa.text(
250                    f'CREATE TRIGGER IF NOT EXISTS "{fts}_ai" AFTER INSERT ON "{main}" '
251                    f'BEGIN INSERT INTO "{fts}"("{key}", "{val}") '
252                    f'VALUES (new."{key}", new."{val}"); END'
253                )
254            )
255            sync_conn.execute(
256                sa.text(
257                    f'CREATE TRIGGER IF NOT EXISTS "{fts}_ad" AFTER DELETE ON "{main}" '
258                    f'BEGIN DELETE FROM "{fts}" WHERE "{key}" = old."{key}"; END'
259                )
260            )
261            sync_conn.execute(
262                sa.text(
263                    f'CREATE TRIGGER IF NOT EXISTS "{fts}_au" AFTER UPDATE ON "{main}" '
264                    f'BEGIN DELETE FROM "{fts}" WHERE "{key}" = old."{key}"; '
265                    f'INSERT INTO "{fts}"("{key}", "{val}") '
266                    f'VALUES (new."{key}", new."{val}"); END'
267                )
268            )
269
270    async def _do_upsert(
271        self, conn: AsyncConnection, rows: list[dict[str, Any]]
272    ) -> None:
273        # Write the main table (triggers drop stale vec0/fts rows and refresh
274        # fts), then repopulate the vec0 shadow with freshly computed vectors.
275        await super(sqlite_vec_async, self)._do_upsert(conn, rows)
276        if self.has_dense and rows:
277            assert self.value_column is not None
278            await self._insert_vectors(
279                conn,
280                [row[self.key_column] for row in rows],
281                [row[self.value_column] for row in rows],
282            )
283
284    async def _insert_vectors(
285        self, conn: AsyncConnection, keys: list[Any], texts: list[Any]
286    ) -> None:
287        """Embed *texts* and insert the vectors into the ``vec0`` shadow."""
288        assert self.conversion_func is not None
289        # Off the event loop: an embedding batch would otherwise stall the
290        # host application's loop for its full duration.
291        vectors = await asyncio.to_thread(self.conversion_func, texts)
292        stmt = sa.text(
293            f'INSERT INTO "{self.vec_table_name}"'
294            f'("{self.key_column}", "{self.vector_column}") '
295            f"VALUES (:key, {self.vector_value_sql.format(':vec')})"
296        )
297        params = [
298            {
299                "key": key,
300                "vec": sqlite_vec_ext.serialize_float32(
301                    np.asarray(vec, dtype=np.float32).tolist()
302                ),
303            }
304            for key, vec in zip(keys, vectors, strict=True)
305        ]
306        batch_size = max(1, self._PARAM_LIMIT // 2)
307        for start in range(0, len(params), batch_size):
308            await conn.execute(stmt, params[start : start + batch_size])
309
310    async def reindex(self, batch_size: int = 1000) -> int:
311        """Rebuild the shadow indexes from the existing main-table rows.
312
313        Clears the ``vec0`` / FTS5 shadows and repopulates them by streaming
314        the main table.  Use this once after pointing the storage at a host
315        table that already holds data (writes made *through* cbrkit keep the
316        shadows in sync on their own).
317
318        Returns:
319            The number of rows indexed.
320        """
321        assert self.value_column is not None
322        total = 0
323
324        async with self._engine.begin() as conn:
325            await self._ensure_schema(conn)
326            kc = self._table.c[self.key_column]
327            vc = self._table.c[self.value_column]
328            if self.has_dense:
329                await conn.execute(sa.text(f'DELETE FROM "{self.vec_table_name}"'))
330            if self.has_sparse:
331                await conn.execute(sa.text(f'DELETE FROM "{self.fts_table_name}"'))
332
333            offset = 0
334            while True:
335                rows = (
336                    await conn.execute(
337                        sa.select(kc, vc).order_by(kc).limit(batch_size).offset(offset)
338                    )
339                ).all()
340                if not rows:
341                    break
342                await self._populate_shadows(
343                    conn, [r[0] for r in rows], [r[1] for r in rows]
344                )
345                total += len(rows)
346                offset += batch_size
347
348        return total
349
350    async def _populate_shadows(
351        self, conn: AsyncConnection, keys: list[Any], texts: list[Any]
352    ) -> None:
353        """Insert ``(key, text)`` pairs directly into the FTS5 / ``vec0`` shadows.
354
355        Used by :meth:`reindex` to backfill from existing rows; the normal
356        write path keeps FTS in sync via triggers instead.
357        """
358        if self.has_sparse:
359            await conn.execute(
360                sa.text(
361                    f'INSERT INTO "{self.fts_table_name}"'
362                    f'("{self.key_column}", "{self.value_column}") '
363                    f"VALUES (:key, :val)"
364                ),
365                [{"key": k, "val": t} for k, t in zip(keys, texts, strict=True)],
366            )
367        if self.has_dense:
368            await self._insert_vectors(conn, keys, texts)
369
370
371@dataclass(slots=True)
372class sqlite_vec[K: int | str, V = Mapping[str, Any]](sqlalchemy[K, V]):
373    """Sync facade over :class:`sqlite_vec_async`.
374
375    Adds the ``sqlite-vec``-specific configuration on top of the
376    :class:`cbrkit.indexable.sqlalchemy` sync facade and overrides
377    :meth:`_build_async` to construct a :class:`sqlite_vec_async`.
378    """
379
380    vector_column: str = "embedding"
381    vector_dim: int | None = None
382    vector_type: Literal["float32", "int8"] = "float32"
383    index_type: Literal["dense", "sparse", "hybrid"] = "dense"
384    metric_type: Literal["cosine", "l2", "l1"] = "cosine"
385    fts_tokenizer: str | None = None
386    conversion_func: BatchConversionFunc[str, NumpyArray] | None = None
387
388    def _build_async(self) -> sqlite_vec_async[K, V]:
389        return sqlite_vec_async[K, V](
390            engine=self._engine, **forward_fields(self, exclude={"url"})
391        )
392
393    @property
394    def async_storage(self) -> sqlite_vec_async[K, V]:
395        """The wrapped async storage (used by sync retriever facades)."""
396        return cast("sqlite_vec_async[K, V]", self._async)
397
398    def reindex(self, batch_size: int = 1000) -> int:
399        """Rebuild the shadow indexes from existing main-table rows."""
400        return run_coroutine(self.async_storage.reindex(batch_size))
401
402
403__all__ = [
404    "sqlite_vec",
405    "sqlite_vec_async",
406]