Metadata-Version: 2.5
Name: mistralai-search-toolkit-plugins-postgres
Version: 0.2.0
Summary: PostgreSQL + pgvector backend for mistralai-search-toolkit
Author-email: Mistral AI <support@mistral.ai>
License: Apache-2.0
License-File: LICENSE
Keywords: ai,information-retrieval,mistral,pgvector,postgres,search
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: <3.15,>=3.12
Requires-Dist: alembic>=1.13
Requires-Dist: asyncpg>=0.30.0
Requires-Dist: mistralai-search-toolkit
Requires-Dist: pgvector>=0.3.6
Requires-Dist: pydantic<3,>=2.7
Requires-Dist: sqlalchemy[asyncio]<3,>=2.0
Requires-Dist: structlog<26,>=24
Description-Content-Type: text/markdown

# mistralai-search-toolkit-postgres

PostgreSQL (with the [pgvector](https://github.com/pgvector/pgvector) extension)
backend for `mistralai-search-toolkit`.

> **Status:** baseline. Dense (vector) search, hybrid (vector + full text) search,
> keyword-only collections (`KeywordStoreIndex`), navigation, partial updates, and the
> backend-agnostic Retriever DSL (`RetrievableIndex`) are implemented. See
> [Design Proposal: a Postgres backend for the Search Toolkit](https://app.notion.com/p/4536ba59a7fe827b98ec014fc2342700)
> (Mistral internal) for the interface design and open questions.

## Capabilities

`PostgresStoreIndex` implements the toolkit protocols `VectorStoreIndex` (dense and
hybrid search, index, delete), `NavigableIndex` (`navigate`/`read`/`grep`/`get_chunk`),
`PatchableIndex` (partial updates), `VersionedIndex`, and `RetrievableIndex` (the
backend-agnostic Retriever DSL via `retrieve`). On it, a keyword query is expressed as the
`query` field of a `VectorSearchQuery` or, through the Retriever DSL, as a
`KeywordRetriever`. A collection without embeddings is served by
`PostgresKeywordSearchIndex` instead — see [Keyword-only collections](#keyword-only-collections).
Search honors `exclude_ids`, `exclude_source_ids` (whole
documents, by the `source_id` column) and the `max_candidates`
dial (mapped to `hnsw.ef_search` per query, accepted range `1..1000`). The dial is
raised to `top_k` when it would otherwise sit below it: HNSW's candidate list bounds
how many rows can come back, so a smaller list silently returns fewer than `top_k`.

`retrieve(retriever)` answers the Retriever DSL — `KeywordRetriever` (BM25),
`NearestNeighbourRetriever` (HNSW over a precomputed `query_embedding`; the protocol
never embeds), and `RRFRetriever` (N-way reciprocal rank fusion of any number of leaf
arms). Filters are a backend-agnostic `Filter` tree (`Equal`/`In`/`Range`/`And`/`Or`/`Not`)
resolved against the table's columns. Keyword retrieval needs a usable BM25 index and
raises when it is missing or invalid, unlike `search`, which degrades a hybrid query to
dense.

## Hybrid search

A `VectorSearchQuery` that carries `query` as well as `embedding` is answered by both retrievers
at once. Nothing has to be declared for it:

```python
results = await store.search(VectorSearchQuery(query="quarterly revenue", embedding=vec, top_k=10))
```

Every collection's table carries a **BM25 index over `content`** (or over a column it declares, see
[Ranking another column](#ranking-another-column)), built by
[`pg_textsearch`](https://github.com/timescale/pg_textsearch) under `text_search_config`
(`"english"` by default; set it to the language of the corpus, or `"simple"` for a corpus that is
not natural language). There is no switch to turn it off — what is tunable is how much the lexical
half counts, at query time. Nothing is materialised beside the text: the posting lists live in the
index, so unlike a generated `tsvector` column there is no second copy of every document.

**This requires the extension.** `pg_textsearch` must be installed *and* listed in
`shared_preload_libraries` before the server starts — a `CREATE EXTENSION` alone is not enough.
See "Local Postgres for integration tests" below.

Both retrievers run over the same candidate depth (`max_candidates`, as for the dense path) and
are combined by **weighted reciprocal rank fusion**:

```
score = w_vector / (k + rank_vector) + w_text / (k + rank_text)
```

A chunk matching only the text, or only the vector, is still a candidate — the same union recall
Vespa gets from `userInput(@query) OR nearestNeighbor(embedding)`. The fusion is over ranks rather
than raw scores because BM25 is unbounded and runs to double digits where cosine similarity sits
around 0.6–0.9: weighting the raw values leaves the BM25 term deciding everything until each
deployment discovers its own scaling factor.

Two things change in the results. `score` is a fusion score, comparable only within one result
set — not `1 - distance` as on the dense path. And `distance` is `None` for a hit the vector half
never ranked.

One column is indexed for text: `content`, or the custom text column the collection declares
instead (see [Ranking another column](#ranking-another-column)). Vespa can weight bm25 per field;
matching that would need a per-column weight map and is not implemented.

### When the index is not usable

The declaration describes the table a migration should build, and nothing guarantees that
migration ran against the database the store is pointed at. So the store asks, once, on its first
hybrid query, and if it cannot rank lexically it answers with the vector half alone rather than
failing — a hybrid statement naming an index that cannot answer takes the dense half down with
it, and every search against that collection then returns nothing.

The question is whether the index is *usable*, not whether the name resolves: an index left
`INVALID` by a build that never finished exists, and Postgres refuses to use it. It counts as
absent here, and gets its own remedy — a rebuild, not a migration that would report itself
already applied.

Every downgraded query logs at **ERROR**, naming the collection, the index it looked for, whether
that index exists, whether `pg_textsearch` is installed, and which of the three remedies applies. Once per process would be tidier and
wrong: the symptom is a quality drop with no failure attached, so a line written at startup is
long gone by the time anyone looks. `PostgresStoreIndex.lexical_ranking_resolved` exposes the same
answer as a value (`True`, `False`, or `None` before the first hybrid query).

The result is cached, so a database that gains the index later is treated as though it had not
until the process restarts.

### Why BM25

A generated `tsvector` column ranked with `ts_rank` reaches the same accuracy. It is roughly an
order of magnitude slower. Measured through `PostgresStoreIndex.search` over a 78,341-chunk
corpus, 150 questions, both warm
([#48502](https://github.com/mistralai/dashboard/pull/48502)):

| | BM25 | `tsvector` + `ts_rank` |
|---|---|---|
| `page_recall@10` | 0.3333 | 0.3600 |
| `doc_recall@10` | 0.6533 | 0.6533 |
| latency, median | **14.6 ms** | 194 ms |
| latency, p90 | **19.6 ms** | 317 ms |

The two return identical page recall on 144 of the 150 questions and identical `doc_recall` to
four decimals, so accuracy is not what separates them. Latency is, and the gap is structural: a
GIN index stores nothing to rank by, so every row an ORed question matches has to be scored
before the top *k* can be taken. Block-Max WAND instead skips posting blocks that cannot reach
the running top-*k* threshold, and answers `ORDER BY … LIMIT k` from the index.

`tsvector` also holds the lexemes a second time — the column is `GENERATED ... STORED`, roughly
+17% on the table, recomputed on every write — where BM25 keeps its posting lists in the index.

So BM25, for two reasons, neither of them accuracy:

- it is an order of magnitude faster at the same recall;
- it is what deployments already running hybrid search on Postgres use.

`tsvector` is a poor compromise rather than a wrong answer, and there is one case it fits: a
managed Postgres that cannot have `pg_textsearch` — it has to be in `shared_preload_libraries`
before the server starts, so no migration can install it — and that still wants a lexical arm
rather than degrading to dense. That fallback is a possible follow-up, not part of this; see
[#48502](https://github.com/mistralai/dashboard/pull/48502) for what it costs.

### The fusion defaults

`rrf_k=4`, with the two weights left equal — one change from the 1.0/1.0/60 that reciprocal rank
fusion is usually quoted with, not two.

**`rrf_k` is the part that generalises.** Cormack et al.'s 60 assumes two retrievers of comparable
strength, and a BM25 arm and a vector arm rarely are. Two corpora that disagree about almost
everything else both wanted a sharp discount:

| corpus | shape | best lexical share | best `rrf_k` |
|---|---|---|---|
| BrowseComp-Plus | 830 long multi-constraint questions, 100k docs | 0.6–0.7 | 1–3 |
| FinanceBench | short keyword-shaped questions about figures | 0.0–0.2 | 5 |

On BrowseComp, holding the weights equal and moving `rrf_k` from 60 to 4 alone is worth
**+0.0118 nDCG@10, 95% CI [+0.0054, +0.0180]** — about half the total tuning gain, with no
weighting decision attached.

**The weighting does not generalise, and the default is a compromise.** Those two corpora want
opposite things. Equal weights are optimal for neither; they are the setting whose *worst case* is
smallest:

| lexical share | loss vs BrowseComp best | loss vs FinanceBench best | worst case |
|---|---|---|---|
| 0.3 | 31.8% | 3.2% | 31.8% |
| 0.4 | 22.8% | 9.0% | 22.8% |
| **0.5 (equal)** | **12.4%** | **10.6%** | **12.4%** |
| 0.6 | 0.6% | 26.8% | 26.8% |
| 0.7 | 1.5% | 35.8% | 35.8% |

So `text_weight` is a per-query field for a reason: a deployment that knows its corpus will beat
this default. Note the weights are a *ratio* — only their relative size affects the ranking, so
1.0/1.0 and 0.5/0.5 are the same setting.

`PostgresSearchQuery` carries the dials, in the way `VespaSearchQuery` carries that backend's:

```python
from mistralai.search.toolkit.plugins.postgres import PostgresSearchQuery

await store.search(
    PostgresSearchQuery(
        query="quarterly revenue",
        embedding=vec,
        vector_weight=1.0,
        text_weight=2.0,  # favour the lexical half
        rrf_k=60,  # lower sharpens the top of each list
    )
)
```

### No context-derived filtering

`IngestContext` and `RetrievalContext` are accepted on every entrypoint and carried for
propagation only. **This backend derives no scoping from them.** The base contexts hold
no fields — they exist to be subclassed — so there is nothing here to enforce, and no
equivalent of the Vespa backend's context protocols has been built yet.

That matters when switching backends. Vespa reads `group_id` off the context to scope a
streaming-mode query (and refuses the query without one), and binds `query_params` into a
schema's `yql_filter`. A deployment leaning on either for tenant isolation **does not get
it here** — every query sees the whole collection. Until this backend grows an equivalent,
isolate tenants with a collection (and therefore a table) per tenant, or filter above the
store.

A caller that writes its own SQL can also rank inside its own statement, with its own predicates,
joins included, in the same `WHERE` and so applied before the cut:

```python
from mistralai.search.toolkit.plugins.postgres import keyword_rank

table = store.table
rank = keyword_rank(table, query)
rows = await conn.execute(
    select(table.c.id, (-rank).label("score"))
    .join(allowed, allowed.c.source_id == table.c.source_id)
    .where(rank < 0)  # only rows carrying a query term
    .order_by(rank)  # ascending is best
    .limit(10)
)
```

`keyword_rank` ranks the column the collection's BM25 index covers and names the index the way
`to_bm25query` resolves it; the query is a bound parameter. Which plan answers the statement is the
planner's call: an ordered scan of the index yields only rows carrying a query term, but behind a
selective predicate Postgres may score the surviving rows on their own and sort them, and a row
carrying none of the terms then ranks 0 instead of dropping out — hence `rank < 0`.

The write side has the same seam. `PostgresCollectionWriter` replaces and deletes documents on a
connection the caller owns, inside its transaction, so the rows land or vanish with the caller's
own writes and its next read sees them:

```python
writer = PostgresCollectionWriter(COLLECTION)
async with engine.begin() as conn:
    await conn.execute(...)  # the caller's own writes
    await writer.replace(conn, documents)  # one DELETE, one INSERT, however many documents
    await writer.delete(conn, removed_ids)
```

It never begins, commits or rolls back a transaction. A document that cannot be written — one
without chunks, one with no embedded chunk in a collection that embeds, one passed twice — raises
`IndexingError` before any statement runs and leaves the transaction as it was. A statement that
fails raises its SQLAlchemy error unchanged, SQLSTATE included, since whether to retry is the
transaction owner's call, and leaves the transaction aborted for the caller to roll back. The
stores write through one, each call in a transaction of its own, and keep reporting failures as
`IndexingError`.

## Keyword-only collections

A collection declared with `PostgresKeywordCollectionSchema` ranks by BM25 alone:

```python
WORDS = PostgresKeywordCollectionSchema(collection_name="words", document_type=MyDoc)
app = PostgresApp([WORDS])
await app.create_schema(config, "words")
store = app.get_keyword_index(config, "words")  # -> PostgresKeywordSearchIndex
results = await store.search(KeywordSearchQuery(query="quarterly revenue", top_k=10))
```

Its table is the same as any other collection's minus the `embedding` column and its HNSW index,
so it needs `pg_textsearch` but not pgvector, and nothing has to embed a document to index it.
`PostgresKeywordSearchIndex` implements `KeywordStoreIndex` and the same navigation, patch,
versioning and Retriever DSL protocols as `PostgresStoreIndex`; a `NearestNeighbourRetriever` is
refused, an RRF of keyword arms is not. An embedding a chunk carries is not stored.

`get_keyword_index` serves a keyword-only collection and refuses one that embeds;
`get_search_index` does the reverse. Where `search` on
`PostgresStoreIndex` degrades a hybrid query to dense when the BM25 index is unusable, a
keyword-only search has nothing to degrade to and raises, as `retrieve` does.

## Ranking another column

`bm25_column` moves the BM25 index off `content` onto a custom text column of the document model:

```python
class RankedChunk(DocumentChunk):
    rank_terms: str = ""


RANKED = PostgresKeywordCollectionSchema(
    collection_name="ranked",
    document_type=RankedDocument,  # chunks: Sequence[RankedChunk]
    text_search_config="simple",
    bm25_column="rank_terms",
)
```

That column is then what `KeywordRetriever` ranks and the `and` operator filters, and a retriever's
`field` must name it or be left unset. It is for ranking text that is not the stored chunk text —
field-boosted or normalised text, or terms a caller analysed itself — while `content` stays what
`read`, `grep` and rerankers see. `"simple"` neither stems nor drops stop words, but it still
lowercases and breaks text up with Postgres's default parser, which splits at `_` and keeps `a-b`,
`1.5` or `a/b` whole besides: such terms survive as written only if they are lowercase letters and
digits separated by spaces, and a query has to be analysed the same way.

`bm25_column` names the column, not the field: a document field's column is `document_`-prefixed,
and `PostgresColumn(name=...)` renames either. Patching `content` changes no ranking then; patch the
declared column's field. The column has no length CHECK unless `PostgresColumn(max_length=...)`
gives it one, and the `and` operator, which lexes it per row, fails on text past Postgres's 1 MB
`tsvector` limit. Over `content` the index keeps its name; over another column the name carries
the column after `bm25`, so moving it is a drop plus a create in a migration, as the lexemes have
to be rebuilt.

## Construction

```python
from mistralai.search.toolkit.plugins.postgres import (
    PostgresCollectionSchema,
    PostgresApp,
    PostgresConnectionConfig,
)
from mistralai.search.toolkit.search import ApproximateQueryOptions, VectorSearchQuery

app = PostgresApp([PostgresCollectionSchema(collection_name="docs", document_type=MyDoc, embedding_model=MODEL)])
config = PostgresConnectionConfig(dsn="postgresql://postgres:postgres@localhost:5432/postgres")

await app.create_schema(config, "docs")  # verify extension + create table (idempotent)
store = app.get_search_index(config, "docs")  # -> PostgresStoreIndex
try:
    await store.index_document(doc)
    results = await store.search(
        VectorSearchQuery(
            embedding=vec,
            top_k=10,
            approximate_options=ApproximateQueryOptions(max_candidates=160),
        )
    )
finally:
    await store.aclose()  # disposes the app-created engine
```

The app object mirrors the Vespa backend's `get_search_index(config, collection)`
shape, so application code and the starter app swap backends by changing
construction and configuration only. A caller that manages its own engine can pass
engine directly instead of config as well (and the returned store will not dispose it),
or build the store directly with `PostgresStoreIndex(engine, schema)`.

`PostgresCollectionSchema` is the only thing the store is told about the collection: it builds
the table, declares the metric to rank with, and resolves the document model's custom
fields to their columns. Nothing is reflected from the database, and no column name is
inferred by convention — so a field mapped with `PostgresColumn(name=...)` works on the
ingest, search, and patch paths alike.

## Connecting over TLS

Managed Postgres hands out a DSN with `?sslmode=require` (or stricter). Pass it as-is:

```python
PostgresConnectionConfig(dsn="postgresql://user:pw@host:5432/db?sslmode=require")
```

`sslmode` is libpq's spelling; asyncpg takes the same values under `ssl` and rejects
`sslmode` as a connect argument, so the config lifts it out of the URL and hands it to
the driver correctly. Other query parameters are left on the URL untouched.

The parts form has no query string, so it takes the mode directly — and an explicit
`ssl=` also overrides one embedded in a DSN:

```python
PostgresConnectionConfig(host="host", database="db", user="user", ssl="verify-full")
```

## Requirements

pgvector **0.8 or later** is recommended. Everything works on older versions except the
exclusion filters: HNSW post-filters, so excluded rows are not replaced and a filtered
search can return fewer than `top_k`. 0.8 added `hnsw.iterative_scan`, which the store
enables for those queries. On an older server it logs a warning once and falls back to
the short-result behaviour rather than failing.

pg_textsearch **1.5 or later**. Earlier releases end an ordered BM25 scan after 100,000 index
entries, counting the ones a filter rejects and dead ones not yet vacuumed, so a keyword query over
a large collection returns fewer than `top_k` rows, or none, without an error
([pg_textsearch#516](https://github.com/timescale/pg_textsearch/issues/516)). The store logs a
warning on its first keyword or hybrid query against an older server.

## Creating the table

The table has to exist before the store can use it, and how it gets created is the
application's business. `PostgresApp.create_schema` is there for a deployment with no
migrations at all; anything else — Alembic, Flyway, Liquibase, a hand-written `.sql` file —
needs the statements themselves:

```console
$ mistral-postgres ddl myapp.search:DOCS_COLLECTION
-- docs_chunks: the table `mistralai-search-toolkit-plugins-postgres` expects.
-- Generated from the collection declaration; apply it with whatever migration tool you use.
-- Requires the pgvector and pg_textsearch extensions to be installed in the target database already.

CREATE TABLE docs_chunks (
        id TEXT NOT NULL,
        ...
);

CREATE INDEX docs_chunks_hnsw_m16_efc64_hv_cos ON docs_chunks USING hnsw (embedding halfvec_cosine_ops) WITH (m = 16, ef_construction = 64);
...
```

The argument is `module:attribute` naming a declared `PostgresCollectionSchema`; the module has
to be importable. `render_ddl(collection)` is the same thing as a function, for a project that
would rather generate the file from its own script.

The output is derived from the same `to_table()` the store binds to, so it cannot describe a
different table than the one the store will query. It does not include `CREATE EXTENSION
vector` — see below.

## Table shape

One denormalized row per chunk. The indexes follow the store's two disjoint access
patterns rather than covering columns one by one:

| Index | Serves |
| --- | --- |
| `(embedding)` HNSW, in a collection that embeds | `search` |
| `(source_id, chunk_type, start_offset)` | `navigate`/`read`/`grep` — the equality predicates lead, so the index also supplies the `start_offset` ordering |
| `(source_id, locator)` UNIQUE | the natural key: a chunk id is `uuid5(f"{source_id}:{locator}")` |
| `(document_id)` | the write side — reindex, `delete_document`, `patch_document` |
| `(metadata)` GIN | metadata containment |
| `(content)` BM25, or the declared `bm25_column` | the lexical half of `search`, via `pg_textsearch` |

There is deliberately no standalone index on `start_offset`, `end_offset` or
`chunk_type`. Those are only selective within a single source, and a lone `start_offset`
index is worse than none: the planner picks it for `navigate(PREVIOUS)` and scans it
backwards across every source in the table.

Text columns carry `CHECK (length(col) <= n)` ceilings — 255 for the ids and
`parent_ref`, 1024 for `source_id`, 512 for `locator`, 64 for `chunk_type`, and 100 000
for `content`. They are sized to catch a runaway rather than to constrain real data
(`content` is an embedded chunk, so the embedding model's input limit binds long before
this does). A CHECK rather than `varchar(n)` because widening one later is a catalogue
change plus an optional validation scan, where shrinking a `varchar` rewrites the table.

Custom columns are unbounded unless the model asks for a limit:

```python
section: Annotated[str | None, PostgresColumn(max_length=120)] = None
```

## Provisioning: two extensions must already be installed

The plugin never runs `CREATE EXTENSION`. On managed Postgres (RDS, Cloud SQL) that
statement needs privileges a least-privilege application role is not expected to hold, so
installing them belongs to whoever provisions the database:

```sql
CREATE EXTENSION vector;         -- once per database, as a privileged role
CREATE EXTENSION pg_textsearch;  -- likewise; see the caveat below
```

Both are preconditions of the DDL, because every collection's table declares an HNSW index
over the embedding *and* a BM25 index over `content` or its declared column — except a keyword-only one, which has no
embedding and needs `pg_textsearch` alone.

**`pg_textsearch` additionally needs the server.** It has to be listed in
`shared_preload_libraries` before startup, so unlike pgvector no privilege level can add it to a
running instance — `CREATE EXTENSION` fails until the server is restarted with it preloaded. If
that is not possible on your Postgres, the collection cannot be created today; a `tsvector`
fallback for exactly that case is a possible follow-up
([#48502](https://github.com/mistralai/dashboard/pull/48502)).

`PostgresApp.create_schema` verifies both before attempting any DDL, raising
`MissingVectorExtensionError` or `MissingBM25AccessMethodError`. Without those checks the same
misconfiguration surfaces as an opaque `type "halfvec" does not exist` or
`access method "bm25" does not exist`. The BM25 check looks for the access method rather than the
extension row, since that is what the DDL needs and it is registered per database.

## Local Postgres for integration tests

This plugin needs a Postgres carrying **both** `vector` and `pg_textsearch`. The shared instance
from `infra/local` has only the first, and `pg_textsearch` cannot be added to a running server:
it has to be listed in `shared_preload_libraries` before startup. So this directory ships a
compose file that builds the image and sets that flag:

```bash
docker compose -f dev/docker-compose.yaml up -d
POSTGRES_TEST_DSN=postgresql://pg:1234@localhost:5442/postgres uv run pytest tests/
```

It binds 5442, not 5432, so it runs *beside* the shared instance rather than in place of it. If
`pg_textsearch` reaches the shared image, delete `dev/` and point the tests back at 5432.

`POSTGRES_TEST_DSN` defaults to `postgresql://pg:1234@localhost:5432/postgres`, so it has to be
set for the compose server above. Either way it is an *admin* connection. The tests do not write
into that database: `tests/conftest.py` creates a
`test-search-toolkit-postgres-<pid>` database for the run, enables `vector` on it, and drops
it on teardown. That keeps runs isolated from each other and from your dev data.

The admin DSN is only used for that setup. Everything under test connects as a separate,
unprivileged `search_toolkit_test_app` role that is granted `CONNECT` plus `USAGE, CREATE`
on `public` — enough to create the collection's table, not enough to install an extension.
So the suite exercises the same privilege split as a real deployment rather than assuming
it, and `test_the_application_role_cannot_install_pgvector` asserts the role really is
denied. Without that, `MissingVectorExtensionError` would be guarding a situation the tests
never produce.

Unit tests run without a database; integration tests skip cleanly when Postgres is
unreachable. They do *not* skip when the server is reachable but the test database has no
pgvector: that is a setup error, and the failure says so.
