Metadata-Version: 2.5
Name: mistralai-search-toolkit-plugins-postgres
Version: 0.0.13
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,
> navigation, and partial updates are implemented. Keyword-only retrieval
> (`KeywordStoreIndex`) is not — no backend in the toolkit implements it. 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`),
and `PatchableIndex` (partial updates). It does not implement the keyword-only protocol
(`KeywordStoreIndex`), which no backend does; a keyword query is expressed as the `query`
field of a `VectorSearchQuery`. Search honors `exclude_ids` 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`.

## 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`**, 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.

Only `content` is indexed for text. 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.

## 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, dim=1024)])
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
`exclude_ids`: 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.

## 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 extension 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 | `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 | 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`.

**`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.
