vrr/SurfSense
2
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-10-19 07:34:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

feat(local): choose the embedding model in onboarding

Browse source 
The embedding model is now chosen in the last onboarding step and fixed
for the whole library. Skipping keeps bge-small, which ships with the
app and can't be deleted. Changing the model later is not built.

- An embedding_indexes row records the model's spec and its vector
  table, and each document records which index embedded it. Migration
  0021 stamps existing libraries with bge-small. Routes that queue
  embedding work return 409 embedding_not_chosen until a model is
  chosen.
- Ingest, Studio and search embed through the active index; the old
  worker embedding module is removed. Workers wait for the schema, and
  documents interrupted when the app closed are marked failed.
- Curated embedders (bge-small, granite-97m, granite-311m) live in the
  local manifest, run by a new onnxruntime engine. EMBEDDING is a
  catalog type, never a selection slot.
- Hugging Face search finds ONNX embedders. Each download is checked: a
  probe sets its width, and a ten-question search check must pass.
  Small tokenizer.json files that Hugging Face lists without a checksum
  are hashed at the pinned commit. Results use the GGUF search's shapes.
- Onboarding renders the embedding step with the same component and
  shared list as the other steps; only its choice hook differs. Use
  marks the choice and Finish sends it. Bundled models have no Delete.
- Settings has a read-only Embedding section.
- ADRs 0036 and 0037, a new embedding architecture doc, and updates to
  the docs it touches. New strings are translated into every language.
This commit is contained in:
Anish Sarkar 2026-10-02 01:45:10 +05:30
parent 67e07d9949
commit 84c7ffad18
161 changed files with 5287 additions and 451 deletions
Download patch file Download diff file Expand all files Collapse all files

1
.agents/skills/translate/glossary.md
View file

@ -50,6 +50,7 @@ SurfSense, Studio, llama.cpp, GGUF, Hugging Face, API, MCP, and every model and
| report an issue | 問題を報告 | 문제 신고 | 报告问题 | समस्या रिपोर्ट करें | Problem melden | informar de un problema | signaler un problème | relatar um problema | сообщить о проблеме |
| session log | セッションログ | 세션 로그 | 会话日志 | सेशन लॉग | Sitzungsprotokoll | registro de la sesión | journal de session | log da sessão | журнал сеанса |
| clipboard | クリップボード | 클립보드 | 剪贴板 | क्लिपबोर्ड | Zwischenablage | portapapeles | presse-papiers | área de transferência | буфер обмена |
| embedding (the model that turns the library into vectors) | 埋め込み | 임베딩 | 嵌入 | एम्बेडिंग | Embedding | embedding | embedding | embedding | эмбеддинг |
## What the English means

4
docs/adr/0007-bundled-embeddings.md
View file

@ -1,6 +1,6 @@
# ADR 0007: Embeddings come from a bundled bge-small model run in process on the CPU
- **Status:** Accepted
- **Status:** Accepted; amended by [ADR 0036](0036-the-index-records-its-embedder.md): the width comes from the index's own record, and bge-small is the default rather than the only embedder
- **Date:** 2026-09-04
- **Source:** [Umbrella plan L106](https://github.com/MODSetter/SurfSense/blob/431914fae066e0c42a38b1fdbbab64e8f92d3d00/plans/community-local/00-umbrella-plan.md#L106), [Data model L220–224](https://github.com/MODSetter/SurfSense/blob/431914fae066e0c42a38b1fdbbab64e8f92d3d00/plans/community-local/00c-data-model.md#L220-L224)
@ -10,7 +10,7 @@ Ingest and search both need embeddings, and the app has to produce them with no
## Decision
- The embedder is bge-small-en-v1.5, int8 ONNX, 384 dimensions, about 66 MB. It is bundled with the app and runs in process on onnxruntime's CPU provider ([`worker/ingestion/embedding.py`](../../surfsense_local/backend/worker/ingestion/embedding.py)). Ingest and the query side use the same model.
- The embedder is bge-small-en-v1.5, int8 ONNX, 384 dimensions, about 66 MB. It is bundled with the app and runs in process on onnxruntime's CPU provider (now [`modules/embedding/encoder.py`](../../surfsense_local/backend/modules/embedding/encoder.py)). Ingest and the query side use the same model.
- `SURFSENSE_LOCAL_EMBEDDING_DIMENSION` declares the width, 384 by default. Vectors from another model are not the wrong shape but unrelated numbers, so startup compares the width `chunk_vectors` was created with against the setting and refuses to open a database that no longer matches (`_check_embedding_width()` in [`shared/migrations.py`](../../surfsense_local/backend/shared/migrations.py)).
- Remote embedding is a later opt-in, not a launch dependency.

30
docs/adr/0036-the-index-records-its-embedder.md Normal file
View file

@ -0,0 +1,30 @@
# ADR 0036: The index records which embedder built it, and the embedder is fixed once per install
- **Status:** Accepted
- **Date:** 2026-10-01
- **Amends:** [ADR 0007](0007-bundled-embeddings.md): the width comes from the index's own record rather than a setting, and bge-small is the default rather than the only embedder
- **Source:** [Choosing the embedding model once](../proposals/embedding-model-choice.md)
## Context
[ADR 0007](0007-bundled-embeddings.md) made one model, bge-small-en-v1.5, the embedder for every library, and declared its width in `SURFSENSE_LOCAL_EMBEDDING_DIMENSION`. Startup compared that setting with the width `chunk_vectors` was created at. Nothing recorded which model wrote the vectors, so a different model of the same width would have searched an index built by another one without anyone noticing, and the files were fetched from the repo's `main` branch, so a change upstream would have shipped a different vector space unannounced.
The [proposal](../proposals/embedding-model-choice.md) lets a new install choose its embedder during onboarding. Changing it afterwards means re-embedding the whole library, which is [deferred](../proposals/embedding-model-change.md). Whatever is chosen, the index has to say what it was built with.
## Decision
- **An index is a row.** `embedding_indexes` holds a snapshot of the embedder's spec (repo, pinned revision and file hashes, width, pooling, normalisation, prefixes, `max_tokens`, `semantic_weight`) and names the vec0 table holding its vectors. Search embeds the question with the active row's model and reads its table; ingest and Studio write to it; each document records the index it was embedded into.
- **The choice is fixed when onboarding finishes.** Until then a fresh install has no row, and the routes that queue embedding work answer `409 embedding_not_chosen`. Finishing onboarding with nothing chosen locks bge-small. Locking creates the vector table at the chosen width, and refuses once any chunk exists.
- **An existing library is recorded, not rebuilt.** Migration `0021` writes a bge-small row pointing at the `chunk_vectors` it already has, and stamps its documents. Nothing is re-embedded.
- **bge-small is pinned.** Its revision and the hashes of its weights and tokenizer are part of its spec, and the fetch script refuses a file that does not match.
- **Chunk boundaries do not follow the embedder.** The chunker measures with bge-small's tokenizer, pinned to the chunker, so a future change of embedder would not re-cut documents and move every chunk id that citations point at.
- **Rejected: the embedder as a `selected_models` slot.** Those rows are preferences, changed freely and cascaded away with their connection. The embedder is a property of the index.
- **Rejected: a singleton row.** A second row, `building`, is how changing the model would start; a `CHECK (id = 1)` would make that a schema change.
## Consequences
- `SURFSENSE_LOCAL_EMBEDDING_DIMENSION` is gone. `0001` creates the first table at 384, bge-small's width, and locking rebuilds it at the chosen one.
- A same-width replacement is now detected, because the row names the model, closing the gap ADR 0007 recorded.
- Locking is the one place outside migrations that emits DDL, safe only because it runs on an empty table.
- Workers wait until the API has migrated the database before taking a job. Electron starts them together, and a job queued before an update would otherwise run against the old schema.
- The bundled bge-small file can never change without a re-embed. Its FP16 weights stay, int8 or not ([search](../architecture/search.md), Known gaps).

22
docs/adr/0037-embedding-is-a-type-not-a-slot.md Normal file
View file

@ -0,0 +1,22 @@
# ADR 0037: Curated embedders are a fourth engine in the local catalog, and embedding is a type, never a slot
- **Status:** Accepted
- **Date:** 2026-10-01
- **Source:** [Choosing the embedding model once](../proposals/embedding-model-choice.md)
## Context
Onboarding offers a short list of embedders to download ([ADR 0036](0036-the-index-records-its-embedder.md)). Chat, image and audio models are already listed, pinned, licence-checked, downloaded and deleted by the local catalog. A second catalog for embedders would repeat all of it. But every model type the catalog knows is also a selection slot ([selection](../architecture/local-models/selection.md)), changed freely, and the embedder must not be: every vector in the library depends on it.
## Decision
- **One manifest.** Curated embedders live in the local manifest, run by a fourth engine, onnxruntime, whose `embedding` block holds what the index's spec needs and no file states, including the ranking weight the retrieval eval measured.
- **`EMBEDDING` is a catalog type and never a slot.** The classifier's embedder group answers it, so the manifest and the engine can describe one. `SLOTS` leaves it out: `selectable_for` never offers it, not even to a model nothing recognises, and the selection routes refuse it. The existing `selected_models` check lists the five slot types, so the database refuses it too.
- **A converted build must match the original.** A build pinned from someone else's conversion records its cosine to the original model's vectors, and the schema refuses one under 0.99 mean or 0.98 minimum. granite-97m's int8 build measured 0.961 and is not listed; its full-precision build is.
- **Rejected: a separate manifest for embedders.** It would keep `ModelType` purely a slot, at the price of a second refresh script, pinning and licence pipeline, and install and delete path.
## Consequences
- An embedder is downloaded under its catalog id, not its weights' file name: every ONNX repo calls its weights `model.onnx`.
- Downloads land in `embeddings/` under the data directory, not the models pack, which a packaged app ships read-only.
- The delete route refuses the embedder the active index names, as it refuses a model the app ships.

4
docs/adr/README.md
View file

@ -41,7 +41,7 @@ The Date is the day the decision was first written down. Source links are permal
| 0004 | [The desktop app is its own tree on FastAPI, SQLite and Huey, with no accounts](0004-desktop-app-is-its-own-tree.md) | Accepted, in part superseded by 0035 |
| 0005 | [Schema changes are hand-written Alembic revisions that the API applies at startup](0005-hand-written-migrations.md) | Accepted |
| 0006 | [Retrieval widens recall with FTS5 and sqlite-vec, then orders the union by cosine similarity](0006-hybrid-retrieval.md) | Accepted, the ordering decision superseded by 0031 |
| 0007 | [Embeddings come from a bundled bge-small model run in process on the CPU](0007-bundled-embeddings.md) | Accepted |
| 0007 | [Embeddings come from a bundled bge-small model run in process on the CPU](0007-bundled-embeddings.md) | Accepted, amended by 0036 |
| 0008 | [Ingest and Studio jobs run on separate Huey queues, each drained by its own worker](0008-two-job-queues.md) | Accepted |
| 0009 | [The UI stays fresh by invalidating queries on server-sent events, with no sync engine](0009-freshness-by-invalidation.md) | Accepted |
| 0010 | [Studio models emit structured content and trusted builders render it, so no model-written code runs](0010-studio-builders-not-sandboxes.md) | Accepted, in part superseded by 0028 |
@ -70,3 +70,5 @@ The Date is the day the decision was first written down. Source links are permal
| 0033 | [Every candidate is scored on its own cosine, because a chunk the vector leg did not reach is unmeasured rather than unrelated](0033-every-candidate-is-scored-on-its-own-cosine.md) | Accepted |
| 0034 | [Whether a chat model reads images is llama.cpp's own answer for a local model and the catalog's for a remote one, and it is stored nowhere](0034-vision-is-the-runtimes-answer-stored-nowhere.md) | Accepted |
| 0035 | [The desktop app's stack also ships as one Docker container, with Caddy as its only listener](0035-docker-compose-runs-the-desktop-stack.md) | Accepted |
| 0036 | [The index records which embedder built it, and the embedder is fixed once per install](0036-the-index-records-its-embedder.md) | Accepted |
| 0037 | [Curated embedders are a fourth engine in the local catalog, and embedding is a type, never a slot](0037-embedding-is-a-type-not-a-slot.md) | Accepted |

31
docs/architecture/data-model.md
View file

@ -3,7 +3,7 @@
Everything the desktop app knows lives in one SQLite file, `surfsense.db`, and in the files under `data/`; the job queues live apart in `huey.db`. The schema is a subset of the hosted SurfSense domain model, keeping the words users know (workspace, document, chunk, chat thread, artifact) and fixing the names that had drifted. The models are the source of truth, Alembic is the only thing that creates schema, and every revision is written by hand, because the database is one user's only copy.
**Code:** [`shared/db.py`](../../surfsense_local/backend/shared/db.py), [`shared/migrations.py`](../../surfsense_local/backend/shared/migrations.py), [`alembic/versions/`](../../surfsense_local/backend/alembic/versions/), the `models.py` in each folder of [`modules/`](../../surfsense_local/backend/modules/)
**Decisions:** [ADR 0005](../adr/0005-hand-written-migrations.md), [ADR 0003](../adr/0003-artifacts-as-documents.md), [ADR 0006](../adr/0006-hybrid-retrieval.md), [ADR 0007](../adr/0007-bundled-embeddings.md), [ADR 0018](../adr/0018-keychain-envelope-encryption.md)
**Decisions:** [ADR 0005](../adr/0005-hand-written-migrations.md), [ADR 0003](../adr/0003-artifacts-as-documents.md), [ADR 0006](../adr/0006-hybrid-retrieval.md), [ADR 0007](../adr/0007-bundled-embeddings.md), [ADR 0036](../adr/0036-the-index-records-its-embedder.md), [ADR 0018](../adr/0018-keychain-envelope-encryption.md)
## Principles
@ -62,6 +62,7 @@ The API seeds one workspace, "My Workspace", at startup when none exists. Deleti
| `content_hash` | declared but never written |
| `dedup_key` | SHA-256 of an uploaded or imported file's bytes; null for notes and artifacts |
| `document_metadata` | JSON: `mime_type`, `size_bytes`, `suffix` for files; import adds `folder_path`, `source` and the hosted ids |
| `embedding_index_id` | foreign key to `embedding_indexes`: the index its chunks were embedded into; null until its first ingest |
| `created_at`, `updated_at` | |
`documents_workspace` indexes `workspace_id`. `documents_workspace_dedup_key` is unique on `(workspace_id, dedup_key)` where `dedup_key IS NOT NULL`, so dedup is per workspace and rows without a key stay out of it. Routes and behaviour are in [`documents.md`](documents.md).
@ -81,11 +82,23 @@ The API seeds one workspace, "My Workspace", at startup when none exists. Deleti
| Table | Kind | Purpose |
|---|---|---|
| `chunks_fts` | FTS5, external content over `chunks` (`content_rowid='id'`) | the BM25 keyword leg; stores no text of its own |
| `chunk_vectors` | sqlite-vec `vec0(embedding float[D])`, rowid = `chunks.id` | the nearest-neighbour leg |
| `chunk_vectors` | sqlite-vec `vec0(embedding float[D])`, rowid = `chunks.id` | the nearest-neighbour leg; named by the active row of `embedding_indexes` |
Three triggers on `chunks` keep the keyword index in step. `chunks_after_insert` adds the text to `chunks_fts`; `chunks_after_update` deletes the old text and adds the new; `chunks_after_delete` deletes the text from `chunks_fts` and the row from `chunk_vectors`. An external-content table keeps no copy, so a delete must hand it the old text or it goes on matching. The triggers fire on cascade too, which is how every real delete arrives, since the user removes a document or a workspace, never a chunk. Nothing else could reach these tables: a virtual table takes no foreign key.
Ingest writes the `chunk_vectors` row itself, because only ingest holds the vector. `D` is `SURFSENSE_LOCAL_EMBEDDING_DIMENSION`, 384 for the bundled bge-small-en-v1.5, fixed when `0001` creates the table. `upgrade_to_head()` reads the declared width back from `sqlite_master` and refuses to start when it differs from the setting: vectors from another model are unrelated numbers, not merely the wrong shape, so the database has to be reindexed. How the two legs are queried is in [`search.md`](search.md).
Ingest writes the `chunk_vectors` row itself, because only ingest holds the vector. `0001` creates the table at 384, bge-small-en-v1.5's width; on a fresh install, finishing onboarding drops the still empty table and creates it again at the chosen model's width ([`lock.py`](../../surfsense_local/backend/modules/embedding/lock.py)), the one place outside migrations that emits DDL. `upgrade_to_head()` reads the declared width back from `sqlite_master` and refuses to start when it differs from the active index's spec: vectors from another model are unrelated numbers, not merely the wrong shape. How the two legs are queried is in [`search.md`](search.md).
### `embedding_indexes`
| Column | Notes |
|---|---|
| `id` | |
| `spec` | JSON snapshot of the embedder's spec: repo, pinned revision and file hashes, width, pooling, normalisation, prefixes, `max_tokens`, `semantic_weight`. A later edit to the app's copy of a spec cannot change what an existing index means |
| `vector_table` | the vec0 table holding this index's vectors; checked against a fixed name pattern before it reaches SQL |
| `state` | `active`; `building` and `retired` are reserved for changing the model, which is not built |
| `created_at` | |
One `active` row, enforced in code rather than by a singleton constraint, because a second row is how changing the model would start ([proposal](../proposals/embedding-model-change.md)). A fresh install has no row until onboarding finishes; `0021` writes a bge-small row for a library that already existed, pointing at its `chunk_vectors`, and stamps every document with chunks. Nothing is re-embedded.
### `chat_threads` and `chat_messages`
@ -180,6 +193,7 @@ erDiagram
documents ||--o| artifacts : "body of"
chunks ||--|| chunks_fts : "rowid, by trigger"
chunks ||--|| chunk_vectors : "rowid, by ingest"
embedding_indexes |o--o{ documents : "embedded"
chat_threads ||--o{ chat_messages : contains
chat_threads |o--o{ artifacts : "produced, set null"
artifacts ||--o{ artifact_files : stores
@ -203,9 +217,17 @@ erDiagram
text content_hash
text dedup_key
json document_metadata
int embedding_index_id FK
datetime created_at
datetime updated_at
}
embedding_indexes {
int id PK
json spec
text vector_table
text state "active, building, retired"
datetime created_at
}
chunks {
int id PK
int document_id FK
@ -219,7 +241,7 @@ erDiagram
text content "FTS5, external content"
}
chunk_vectors {
blob embedding "vec0, 384 floats by default"
blob embedding "vec0, the active spec's width"
}
chat_threads {
int id PK
@ -334,6 +356,7 @@ erDiagram
| `0016` | `0016_local_audio_provider.py` | `selected_models` rebuilt so `audiocpp` may hold `audio_gen`, and only that, without a connection; downgrading drops an `audiocpp` selection |
| `0019` | `0019_selection_settings.py` | `selected_models.settings`, a nullable JSON column added in place; its first entry is a server audio model's `voices` |
| `0020` | `0020_plugin_runs.py` | `plugin_runs` |
| `0021` | `0021_embedding_indexes.py` | `embedding_indexes` and `documents.embedding_index_id`; an existing library gets a bge-small row and its documents are stamped. The column is added by a plain `ALTER TABLE`, since a batch rebuild of `documents` would cascade to every chunk |
- Migrations run on every API start and are idempotent. Autogenerate is off: it renders a rename as a drop plus an add, which deletes a column's data silently, and `env.py` carries no `target_metadata`, so it cannot be used by accident.
- SQLite cannot alter a CHECK constraint in place, so `0004`, `0009`, `0012` and `0013` copy `selected_models` into a new table.

19
docs/architecture/documents.md
View file

@ -3,7 +3,7 @@
A workspace holds a library of sources: uploaded files, notes written in the app, and the bodies of Studio artifacts. Every source ends up the same way, as a markdown body on its `documents` row cut into passages that are embedded and indexed, so chat and search treat them alike. The API only stores bytes and enqueues a job; parsing, chunking and embedding happen in the ingest worker, one document at a time, and a document is `ready` only once it is searchable. Studio indexes an artifact's body itself, in its own worker, with the same chunking and embedding code.
**Code:** [`modules/workspaces/`](../../surfsense_local/backend/modules/workspaces/), [`modules/documents/`](../../surfsense_local/backend/modules/documents/), [`worker/ingestion/`](../../surfsense_local/backend/worker/ingestion/), [`worker/jobs.py`](../../surfsense_local/backend/worker/jobs.py), [`shared/queue.py`](../../surfsense_local/backend/shared/queue.py), [`frontend/src/features/sources/`](../../surfsense_local/frontend/src/features/sources/)
**Decisions:** [ADR 0003](../adr/0003-artifacts-as-documents.md), [ADR 0007](../adr/0007-bundled-embeddings.md), [ADR 0008](../adr/0008-two-job-queues.md), [ADR 0009](../adr/0009-freshness-by-invalidation.md)
**Decisions:** [ADR 0003](../adr/0003-artifacts-as-documents.md), [ADR 0007](../adr/0007-bundled-embeddings.md), [ADR 0036](../adr/0036-the-index-records-its-embedder.md), [ADR 0008](../adr/0008-two-job-queues.md), [ADR 0009](../adr/0009-freshness-by-invalidation.md)
## Workspaces
@ -84,25 +84,28 @@ The desktop app previews an original in its left rail when its MIME type has a s
`ingest_document(document_id)` is declared in [`modules/documents/tasks.py`](../../surfsense_local/backend/modules/documents/tasks.py), because Huey binds a task to its queue when it is decorated and the API is what enqueues it. The task body imports the pipeline lazily and Docling is imported only when the converter is first built, so the API never loads Docling or torch. Each job opens its own database engine.
1. **Start.** `begin_job` marks the row `processing` unless it was cancelled, and the worker notifies the API.
1. **Start.** `begin_job` marks the row `processing` unless it was cancelled, and the worker notifies the API. The job then reads the active embedding index; with none, because onboarding has not chosen an embedder, it fails with `no embedding model has been chosen yet`.
2. **Parse** ([`parsing.py`](../../surfsense_local/backend/worker/ingestion/parsing.py)). A `NOTE` or `ARTIFACT` already has its markdown in `content`. A `FILE` with a text extension (`.md`, `.markdown`, `.txt`, `.text`) is read as UTF-8; anything else goes to Docling. OCR and table structure are switched on, so a scanned PDF is readable and a table survives as a table. RapidOCR's per-line orientation classifier is off: on noisy scans it turned upright lines 180 degrees and they read back as garbage, and dropping it cut the character error rate on generated scans about eightfold. An image is first turned upright by its EXIF orientation, which Docling ignores, and given a DPI that reads it as one A4 page, never finer than the 216 dpi OCR renders at nor coarser than its own DPI ([`image_page.py`](../../surfsense_local/backend/worker/ingestion/image_page.py)); with none, Docling takes it as 72 dpi and OCR upsampled a photo threefold. Images and PDFs share one set of options, so they share one pipeline and its models. Inference gets one thread per physical core, at most 8, as audio.cpp does, instead of Docling's 4: on an 8-core CPU that parsed generated scans, photos and PDFs 1.4 times faster. The converter is built once per process and its models load on the first conversion, so the first document a worker parses pays for the import and the load. `HF_HOME` defaults to the models directory, because Docling would otherwise write weights into `site-packages`, which is read-only in a frozen bundle. When the bundled parser pack is present Docling uses it with RapidOCR, and `HF_HUB_OFFLINE=1` is set before Docling is imported, because huggingface_hub reads it only then, so a PDF or an image converts with networking off. The markdown is kept only in `content`; nothing is written beside the original.
3. **Chunk** ([`chunking.py`](../../surfsense_local/backend/worker/ingestion/chunking.py)). Chonkie's `RecursiveChunker` splits at the coarsest boundary that fits: a heading, a paragraph, a line, a sentence, a word, then a bare split so an over-long line still ends. Passages are at most 480 tokens by bge-small's own tokenizer, under its 512-token limit with room for the two tokens it adds. Pieces under 24 characters are merged into a neighbour while splitting, so only a document shorter than that yields a shorter passage. The tokenizer is passed as an object, because one named by a string would reach the network even offline. Each passage records the lines its first and last characters fall on.
4. **Embed** ([`embedding.py`](../../surfsense_local/backend/worker/ingestion/embedding.py)). bge-small-en-v1.5 as an int8 ONNX model on onnxruntime's CPU provider: tokenize (truncated at 512), CLS-pool, L2-normalise, in batches of 32. No network and no model server. A vector whose width is not the configured dimension is refused rather than stored.
5. **Index** ([`indexing.py`](../../surfsense_local/backend/worker/ingestion/indexing.py)). The document's old chunks are deleted, which clears both index tables through the delete trigger, and the new chunks are inserted with their float32 embeddings, followed by one `chunk_vectors` row per chunk. The keyword index follows by trigger; the vector cannot, because only ingest holds it. Editing a note therefore stops its old text being findable.
3. **Chunk** ([`chunking.py`](../../surfsense_local/backend/worker/ingestion/chunking.py)). Chonkie's `RecursiveChunker` splits at the coarsest boundary that fits: a heading, a paragraph, a line, a sentence, a word, then a bare split so an over-long line still ends. Passages are at most 480 tokens by bge-small's tokenizer, under its 512-token limit with room for the two tokens it adds. That tokenizer is pinned to the chunker rather than taken from the active embedder, so changing the embedder would not re-cut documents and move every chunk id. Pieces under 24 characters are merged into a neighbour while splitting, so only a document shorter than that yields a shorter passage. The tokenizer is passed as an object, because one named by a string would reach the network even offline. Each passage records the lines its first and last characters fall on.
4. **Embed** ([`encoder.py`](../../surfsense_local/backend/modules/embedding/encoder.py)). The active index's model, as its spec says, on onnxruntime's CPU provider: the document prefix, tokenize (truncated at the spec's `max_tokens`), pool, normalise, in batches of 32. For bge-small that is no prefix, 512 tokens, CLS pooling and L2 normalisation. No network and no model server. A vector whose width is not the spec's is refused rather than stored.
5. **Index** ([`indexing.py`](../../surfsense_local/backend/worker/ingestion/indexing.py)). The document's old chunks are deleted, which clears both index tables through the delete trigger, and the new chunks are inserted with their float32 embeddings, followed by one row per chunk in the active index's vector table, and the document records that index in `embedding_index_id`. The keyword index follows by trigger; the vector cannot, because only ingest holds it. Editing a note therefore stops its old text being findable.
6. **Finish.** The markdown is stored on `documents.content`, and `finish_job` writes `ready` unless a cancel won the race.
Cancellation is checked after parsing and after embedding. On any other failure the worker rolls back and re-reads the row. If the document was deleted or cancelled meanwhile it stops; otherwise it writes `failed` with `error_message` set to the exception type and message, cut to 500 characters, notifies, and re-raises so Huey runs the job again, up to two more times. A later success clears the message. The sources panel shows it and offers Retry.
**Embedding dimension.** `SURFSENSE_LOCAL_EMBEDDING_DIMENSION` (default 384) is schema, not preference: `chunk_vectors` is created at that width, and `upgrade_to_head()` refuses to start against a database built at another one ([`data-model.md`](data-model.md)). The per-batch check in `embed()` is the second line of defence.
**Embedding index.** Which model embeds is a property of the library, not a preference: a row in `embedding_indexes` holds the model's spec and names its vector table, which is built at the spec's width when onboarding locks the choice ([`data-model.md`](data-model.md)). `upgrade_to_head()` refuses to start when the table's width differs from the spec's. The per-batch check in `embed()` is the second line of defence. Studio indexes an artifact's body through the same index.
**Models on disk.** The bge-small files (`model_optimized.onnx`, `tokenizer.json` and `config.json`, from `Qdrant/bge-small-en-v1.5-onnx-Q`) live in `bge-small-en-v1.5/` under the models directory. Packaged builds ship them; a development checkout fetches them with `scripts/fetch_embedding_model.py`. Chat answers `503` while they are missing.
**Before the choice.** Onboarding gates the UI but not the API, so the routes that queue work which ends up embedded (writing, uploading, editing or retrying a document, a cloud import, and creating or regenerating a Studio artifact) answer `409` with the code `embedding_not_chosen` until an index exists.
**Models on disk.** The bge-small files (`model_optimized.onnx` and `tokenizer.json`, from `Qdrant/bge-small-en-v1.5-onnx-Q`) live in `bge-small-en-v1.5/` under the models directory. They are pinned to one revision and checked against their hashes ([`bundled.py`](../../surfsense_local/backend/modules/embedding/bundled.py)): every existing library was built with that exact file, so a different build would be a different vector space. Packaged builds ship them; a development checkout fetches them with `scripts/fetch_embedding_model.py`, which refuses a file that does not match. Any other curated embedder downloads into `embeddings/<id>/` under the data directory, since the models pack is read-only in a packaged app. Chat answers `503` while the active model's files are missing.
## The queue
- `ingest` is a `SqliteHuey` queue in `huey.db`, a file of its own so the consumer's constant polling never contends for the write lock on `surfsense.db`.
- `worker-ingest` drains it with one thread. Ingest saturates a CPU and writes to the database the API is serving from, so a second thread would spend its time behind the first one's lock. Studio has its own queue so an import never sits in front of a summary ([ADR 0008](../adr/0008-two-job-queues.md)).
- The queue survives restarts. A job carries only its task's name, so each consumer imports every task before it starts (`import_tasks()`), and [`tests/integration/test_registration.py`](../../surfsense_local/backend/tests/integration/test_registration.py) fails when that list falls behind the task modules on disk.
- The queue survives restarts. Electron starts the workers beside the API and only the API migrates, so each consumer waits until the database is at the latest revision before taking a job ([`wait_for_schema.py`](../../surfsense_local/backend/worker/wait_for_schema.py)); a job queued before an update would otherwise run against the old schema and fail. A job carries only its task's name, so each consumer imports every task before it starts (`import_tasks()`), and [`tests/integration/test_registration.py`](../../surfsense_local/backend/tests/integration/test_registration.py) fails when that list falls behind the task modules on disk.
- Jobs are enqueued by upload, note creation, a note content edit and retry, and by import ([`import.md`](import.md)).
- A source the app quit in the middle of ingesting is failed with `interrupted when the app closed` when the ingest worker next starts, before it takes a job ([`interrupted_documents.py`](../../surfsense_local/backend/worker/interrupted_documents.py)). Huey drops a job as it starts it, so nothing else would end it, and Retry accepts only a failed or cancelled document. A source still `pending` is still queued and is left alone.
- Each transition the ingest worker makes, to `processing`, `ready` or `failed`, sends a `documents` event keyed by document id, and so does each change the API makes itself: a note written, an upload, a rename or edit, a retry, a cancel, and a delete, whose status is `deleted`. The sources panel reloads its list on each event and whenever a dropped stream is back, so a note a plugin writes shows as it is written; it still refetches every 1.5 seconds while any row is `pending` or `processing` ([`overview.md`](overview.md#freshness)).
## Known gaps

41
docs/architecture/embedding.md Normal file
View file

@ -0,0 +1,41 @@
# Embedding model
Which model turns passages and questions into vectors. It is chosen once, in onboarding, and fixed for the whole library: every vector in it was made by that model, and changing it means embedding everything again, which is not built ([proposal](../proposals/embedding-model-change.md)).
**Code:** [`modules/embedding/`](../../surfsense_local/backend/modules/embedding/), [`engines/onnxruntime/`](../../surfsense_local/backend/modules/llm/catalog/local/engines/onnxruntime/), [`frontend/src/features/onboarding/model-step/use-embedding-step.ts`](../../surfsense_local/frontend/src/features/onboarding/model-step/use-embedding-step.ts), [`frontend/src/features/embedding/`](../../surfsense_local/frontend/src/features/embedding/)
**Decisions:** [ADR 0007](../adr/0007-bundled-embeddings.md), [ADR 0036](../adr/0036-the-index-records-its-embedder.md), [ADR 0037](../adr/0037-embedding-is-a-type-not-a-slot.md)
## The index
A row in `embedding_indexes` holds a snapshot of the model's spec and names its vector table ([data model](data-model.md)). Ingest and Studio write through it, search reads through it, and each document records which index it was embedded into ([documents](documents.md), [search](search.md)). Until onboarding finishes there is no row, and the routes that queue embedding work answer `409 embedding_not_chosen`.
## Choosing
Onboarding's last step lists the choices; finishing onboarding locks the one chosen ([selection](local-models/selection.md#onboarding)), and skipping means bge-small. Settings › Embedding shows it, its vector size, and whether SurfSense tested it, and offers no way to change it.
| Source | Where it comes from | Label |
|---|---|---|
| bge-small | bundled in the read-only models pack, pinned by hash | measured |
| curated | the local manifest, run by onnxruntime ([catalog](local-models/catalog.md)); ranking weight measured by the retrieval eval | measured |
| Hugging Face | any ONNX embedder, found by search | declared, or inferred |
## Hugging Face
`GET /embedding/huggingface/search?q=` lists repos tagged `sentence-similarity` or `feature-extraction`, most downloaded first. `GET /embedding/huggingface/repo/{repo}` opens one, without downloading weights, as a catalog row with one build, or none and why. Both answer in the GGUF search's shapes (`SearchRead`, `RepoRead`), so onboarding renders them through the same search as the chat step:
- **Checksums:** a small file Hugging Face keeps out of LFS, such as most `tokenizer.json` files, lists no sha256, so it is hashed from its bytes at the pinned commit.
- **Refused** when gated, when Hugging Face's own security scan flags a file it would take or anything as `unsafe`, when it has no ONNX build or no `tokenizer.json`, or when a weights file lists no checksum.
- **Which file:** a generic int8 build (`model_int8.onnx`, `model_quantized.onnx`), else full precision; never one tuned for a single CPU, an `O1`–`O4` variant, or an fp16, q4 or bnb4 build. External data downloads with it under its own name ([`pick.py`](../../surfsense_local/backend/modules/embedding/huggingface/pick.py)).
- **The spec** comes from the repo's sentence-transformers files: pooling, prompts for questions and passages, maximum length, capped at 2,048 tokens. Vectors are always normalised, since search compares by cosine. A repo with those files is `declared`; one with only the tag is `inferred` and takes bge's defaults. The ranking weight is 0.65, bge's, since nobody measured it ([`spec_from_repo.py`](../../surfsense_local/backend/modules/embedding/huggingface/spec_from_repo.py)).
An open repo is installed through the same install jobs as every local model, under the name `hf--<owner>--<name>`. After the download, two checks run before it can be chosen ([`verify.py`](../../surfsense_local/backend/modules/embedding/verify.py)):
1. **The probe** embeds one passage; its width is the one kept, over whatever the config said.
2. **The search check** asks ten questions over twenty passages, each answer beside a decoy on the same topic, and every answer must rank first ([`search_check.py`](../../surfsense_local/backend/modules/embedding/search_check.py)). bge-small and both curated granite models found 10 of 10; chat models served as embedders found 5 to 8, except one at 10. It shows a model can find answers, not that it was built to embed.
A model that fails is deleted and the install says why. One that passes keeps its settled spec beside its files and is listed as a downloaded row, labelled not tested by SurfSense.
## Known gaps
- A Hugging Face repo deleted or made private after it was chosen cannot be downloaded again; nothing can repair the library until changing the model is built.
- A remote embedder is not offered yet ([proposal](../proposals/embedding-model-choice.md)).

44
docs/architecture/local-models/catalog.md
View file

@ -36,12 +36,17 @@ sd.cpp nor audio.cpp has a fit estimate, so the row states the download size and
nothing about this machine, and nothing blocks its install. The screen renders these fields and computes none
of them.
`GET /llm/catalog/local` returns llama.cpp's rows, then sd.cpp's and audio.cpp's.
`GET /llm/catalog/local` returns llama.cpp's rows, then sd.cpp's, audio.cpp's and onnxruntime's.
sd.cpp's are the six curated image models and two video models, only when Electron handed the API an
images folder, which it does only when it staged sd-server, in dev or packaged
([`index.ts`](../../../surfsense_local/electron/src/main/index.ts); [packaging](../packaging.md)).
audio.cpp's are the three curated audio models, only when Electron handed the API
an audio folder, which it does only when it staged audio.cpp's server.
onnxruntime's are the curated embedders, then any Hugging Face embedder
downloaded and checked ([embedding](../embedding.md)), always: they run in the API and the
workers, so no runtime has to be staged. An embedder row adds `description`, what
it is for, which onboarding's search step shows; it has no fit, and can be
selected for nothing ([selection](selection.md)).
## One slice per engine
@ -63,6 +68,7 @@ catalog/local/
builds/, rows/, images_folder/
audiocpp/ engine.py, manifest_fields.py, evidence.py,
builds/, rows/, audio_folder/
onnxruntime/ engine.py, manifest_fields.py, rows.py, spec.py
service.py router.py schemas.py dependencies.py
```
@ -80,14 +86,14 @@ that the preset does not name, it rewrites the preset once; later reads with no
new model leave it alone, so Electron does not restart the router repeatedly.
The selection an install fills is the engine's `model_type` and `provider`.
| | [`engines/llamacpp/`](../../../surfsense_local/backend/modules/llm/catalog/local/engines/llamacpp/) | [`engines/sdcpp/`](../../../surfsense_local/backend/modules/llm/catalog/local/engines/sdcpp/) | [`engines/audiocpp/`](../../../surfsense_local/backend/modules/llm/catalog/local/engines/audiocpp/) |
|---|---|---|---|
| Evidence | GGUF metadata keys | tensor names, as sd.cpp dispatches on them (`evidence.py`) | the family key, read from the header's front (`evidence.py`) |
| Manifest fields | `context` (required), `shape`, `template`, `sampling` | `image` defaults, or `video` ones for a video model: one of them (required) | `audio`: voices, languages, sample rate, measured memory (required) |
| Which files make a build | every build a repo offers, each with the repo's projector (`builds/in_repo.py`) | a GGUF at the repo's root, with the VAE and text encoder the entry names from their own repos (`builds/in_repo.py`) | every GGUF in the model's own folder of a shared repo, alone (`builds/in_repo.py`) |
| Default build | the first in a preference order led by `UD-Q4_K_XL` (`builds/choice/`) | the first the entry pins, in its reviewed order (`builds/choice.py`) | the first the entry pins, in its reviewed order (`builds/choice.py`) |
| Its folder | `models_folder/`: the scan, the preset, readiness | `images_folder/`: its files, where each lands, legacy downloads, the installed image; `launch.py`: its flags | `audio_folder/`: its files, the installed models, `server.json` |
| Also | support, pricing, the recommended build, the lead build, the star, search | | |
| | [`engines/llamacpp/`](../../../surfsense_local/backend/modules/llm/catalog/local/engines/llamacpp/) | [`engines/sdcpp/`](../../../surfsense_local/backend/modules/llm/catalog/local/engines/sdcpp/) | [`engines/audiocpp/`](../../../surfsense_local/backend/modules/llm/catalog/local/engines/audiocpp/) | [`engines/onnxruntime/`](../../../surfsense_local/backend/modules/llm/catalog/local/engines/onnxruntime/) |
|---|---|---|---|---|
| Evidence | GGUF metadata keys | tensor names, as sd.cpp dispatches on them (`evidence.py`) | the family key, read from the header's front (`evidence.py`) | the repo config's `model_type`, and its pipeline tag |
| Manifest fields | `context` (required), `shape`, `template`, `sampling` | `image` defaults, or `video` ones for a video model: one of them (required) | `audio`: voices, languages, sample rate, measured memory (required) | `embedding`: width, pooling, normalisation, prefixes, `max_tokens`, `semantic_weight`, `batch`, and `parity` for a converted build (required) |
| Which files make a build | every build a repo offers, each with the repo's projector (`builds/in_repo.py`) | a GGUF at the repo's root, with the VAE and text encoder the entry names from their own repos (`builds/in_repo.py`) | every GGUF in the model's own folder of a shared repo, alone (`builds/in_repo.py`) | the ONNX weights the entry names, with the repo's `tokenizer.json`, which every build must carry |
| Default build | the first in a preference order led by `UD-Q4_K_XL` (`builds/choice/`) | the first the entry pins, in its reviewed order (`builds/choice.py`) | the first the entry pins, in its reviewed order (`builds/choice.py`) | its only build: another would be another vector space |
| Its folder | `models_folder/`: the scan, the preset, readiness | `images_folder/`: its files, where each lands, legacy downloads, the installed image; `launch.py`: its flags | `audio_folder/`: its files, the installed models, `server.json` | `embeddings/<id>/` under the data directory, not the models pack, which a packaged app ships read-only; bge-small is read in place from the pack |
| Also | support, pricing, the recommended build, the lead build, the star, search | | | `spec_for`: an entry as the spec an index stores |
A slice holds what the catalog knows about a runtime. Running it stays outside
`catalog/`: the clients in `providers/`, and `fit/`, which prices llama.cpp.
@ -203,9 +209,9 @@ engine's `refresh.py`, is the other. For each repo the script reads the
commit sha, the file listing at that commit (sizes and `lfs.oid` hashes) and the
repo's `params` file, then the default build's header and its projector's header
over HTTP range requests. No weights are downloaded. Each engine reads its own
entries: `local_manifest/llamacpp/`, `local_manifest/sdcpp/` and
`local_manifest/audiocpp/` each hold a `refresh.py` (the network) and an
`assemble.py` (pure, the written entry).
entries: `local_manifest/llamacpp/`, `local_manifest/sdcpp/`,
`local_manifest/audiocpp/` and `local_manifest/onnxruntime/` each hold a
`refresh.py` (the network) and an `assemble.py` (pure, the written entry).
- **An image entry's evidence comes from its tensors**, the only thing an sd.cpp
file reliably states, and what a person reviews comes from the entry: the
@ -222,6 +228,12 @@ entries: `local_manifest/llamacpp/`, `local_manifest/sdcpp/` and
64 KiB: audio.cpp writes it before the files it embeds, which make Kokoro's
header 38.6 MB and Supertonic's 57 MB. The entry names the model's folder in
the shared repo, its builds most preferred first, and the reviewed `audio` block.
- **An embedder entry pins the weights it names with the repo's tokenizer**, and
the reviewed `embedding` block: the ranking weight the retrieval eval measured
for that model, and for a build converted by someone else, how close its
vectors are to the original's. The schema refuses a build under 0.99 mean or
0.98 minimum cosine. A file kept in git rather than LFS lists no sha256, so the
refresh downloads and hashes it when it is small; a tokenizer is.
- **It pins every chat build in the quantization preference order**
([`builds/choice/preference.py`](../../../surfsense_local/backend/modules/llm/catalog/local/engines/llamacpp/builds/choice/preference.py)),
and nothing outside it: no imatrix files, drafters, big endian builds, or
@ -244,7 +256,8 @@ entries: `local_manifest/llamacpp/`, `local_manifest/sdcpp/` and
The script's assembly is tested over recorded input, with no network
([`test_local_manifest.py`](../../../surfsense_local/backend/tests/unit/scripts/test_local_manifest.py),
[`test_local_manifest_sdcpp.py`](../../../surfsense_local/backend/tests/unit/scripts/test_local_manifest_sdcpp.py),
[`test_local_manifest_audiocpp.py`](../../../surfsense_local/backend/tests/unit/scripts/test_local_manifest_audiocpp.py)).
[`test_local_manifest_audiocpp.py`](../../../surfsense_local/backend/tests/unit/scripts/test_local_manifest_audiocpp.py),
[`test_local_manifest_onnxruntime.py`](../../../surfsense_local/backend/tests/unit/scripts/test_local_manifest_onnxruntime.py)).
## Which files make a build
@ -318,7 +331,8 @@ recommended one. A light spill is described quietly in the reason line instead.
turns an architecture and a pipeline tag into a type. It names what is not a chat
model and calls everything else `TEXT_GEN`, because llama.cpp gains
architectures faster than any list is updated: embedders, rerankers, speech
recognisers, labellers, OCR, drafters and projectors are known and typeless;
recognisers, labellers, OCR, drafters and projectors are known and typeless,
except embedders, which are `EMBEDDING`, a type onnxruntime runs and no slot holds;
diffusion architectures are `IMAGE_GEN`, video `VIDEO_GEN`, text to speech
`AUDIO_GEN`, audio.cpp's voice families among them (`kokoro_tts`, `supertonic`,
`kitten_tts`). Every audio.cpp file declares `general.architecture = audiocpp`,
@ -559,7 +573,7 @@ delete is in `modules/llm/router.py` beside the selection routes.
| `GET /llm/installs/{id}` | one job | none |
| `GET /llm/installs/events` | NDJSON: every job, now and on each change | none |
| `DELETE /llm/installs/{id}` | `204`; `404` when it is unknown or already over | none |
| `DELETE /llm/models/{model_name:path}` | `ModelDeleteRead`, with `selection_cleared` | none |
| `DELETE /llm/models/{model_name:path}` | `ModelDeleteRead`, with `selection_cleared`; `409` for a model the app ships, or the embedder that made the library's vectors | none |
Search, repo reads and downloads share one consent, `host:huggingface.co`
([`../egress.md`](../egress.md)). A destination that is off is a `403` with

37
docs/architecture/local-models/selection.md
View file

@ -28,6 +28,14 @@ choosing or clearing a model never touches.
A type no feature reads can still be chosen; the feature that first reads one
brings the client that calls it.
`embedding` is a `ModelType` too, and the one exception: a catalog type, so the
manifest and an engine can describe an embedder, never a slot. The embedder
belongs to the library's index, fixed when onboarding finishes
([search](../search.md), [ADR 0037](../../adr/0037-embedding-is-a-type-not-a-slot.md)). `SLOTS` in
[`selectable.py`](../../../surfsense_local/backend/modules/llm/selectable.py) is every type but it:
`selectable_for` never offers `embedding`, not even to a model nothing
recognises, and both `/llm/selection/{model_type}` routes answer `422` for it.
A row stores the provider, the connection when remote, the exact model id, and
three fingerprint facts. A check constraint requires a `connection_id` exactly
when the provider is `openai_compatible`, a second (`local_runtime_type`) lets
@ -163,7 +171,11 @@ three, and `worker.spec` takes those plus every `*.md` under `worker.studio`.
`GET /llm/onboarding` returns `{"completed": bool}`, true once the singleton
`onboarding_completion` row exists. `POST /llm/onboarding` writes that row and
requires a persisted `text_gen` selection, answering `422 chat model required`
otherwise; image, image editing, video and audio models are optional. The marker means the user finished
otherwise; image, image editing, video and audio models are optional. Its body
may name `embedding_model`, a curated embedder already downloaded; the route
locks it as the library's embedder before writing the marker, refusing `409` one
not yet downloaded and leaving onboarding unfinished. No name locks the bundled
bge-small ([`choose.py`](../../../surfsense_local/backend/modules/embedding/choose.py)). The marker means the user finished
choosing, and it is the one thing that must not become true early.
Two invariants, both easy to break from the frontend: selecting or clearing a
@ -172,12 +184,29 @@ route. Only the onboarding page's last step does, once a chat model is
persisted. Once the marker exists the app never shows onboarding again, and a
missing selection is fixed from Settings' Chat section.
The onboarding page opens on a welcome screen, then five steps: chat, image, image editing, audio and video model. The welcome is not counted as a step, but it is part of onboarding and gated by the same marker, so it is never shown again once onboarding is done. The five
model steps are one component for any slot
The onboarding page opens on a welcome screen, then six steps: chat, image, image editing, audio, video and search model. The welcome is not counted as a step, but it is part of onboarding and gated by the same marker, so it is never shown again once onboarding is done.
The embedding step comes last and is not a slot, but it is the same component
as the model steps below, with one more entry in their tables
([`model-step/`](../../../surfsense_local/frontend/src/features/onboarding/model-step/)).
What a step does on Use is its own hook: a slot's saves the selection, and the
embedding step's ([`use-embedding-step.ts`](../../../surfsense_local/frontend/src/features/onboarding/model-step/use-embedding-step.ts))
only marks a choice, In use until another is used, bge-small by default and again
if the chosen one is deleted. Its downloads install with `select: false`; a
Hugging Face pick is labelled not tested by SurfSense; the bundled bge-small has
no Delete. Its search is `ModelSearch` given the embedding endpoints, which answer
in the GGUF search's shapes, with a note that larger models are slower and the
choice is permanent, and a notice below the list says the choice is kept for the
whole library. It offers no server until remote embedders exist. Being last, its
Finish sends the choice with the call that ends onboarding; Skip and finish, or
Finish with the choice untouched, sends none, which means bge-small
([embedding](../embedding.md)).
The model steps are one component for any slot
([`frontend/src/features/onboarding/model-step/`](../../../surfsense_local/frontend/src/features/onboarding/model-step/)),
built on the same hooks as Settings but with its own screens. Each lists every
model this computer can run at once, the catalog's starred row first, with
Download, Use and Delete as in Settings; a download's progress shows under its
Download, Use and Delete as in Settings, and no Delete on a model the app ships; a download's progress shows under its
row and never moves the page. The chat step also offers Settings' Hugging Face
search, closed until asked for; the image, image editing, audio and video steps have none, since
sd.cpp's and audio.cpp's models are the few the catalog ships. The image editing step lists first the model chosen for images earlier when it edits too, so FLUX.2 klein is one Use away, and its downloads fill `image_edit`. A server sits one line below the list and names

1
docs/architecture/overview.md
View file

@ -110,6 +110,7 @@ The rules that keep the processes out of each other's way:
- [Data model](data-model.md): tables, the search index, revisions, files on disk.
- [Documents](documents.md): workspaces, upload, notes, the ingest pipeline.
- [Search](search.md): `retrieve()`.
- [Embedding model](embedding.md): the model that embeds the library, chosen once at onboarding.
- [Chat](chat.md): grounding, the stream, citations.
- [Studio](studio.md): artifact formats, jobs, viewers.
- [Connections](connections.md): OpenAI-compatible endpoints and where keys live.

2
docs/architecture/packaging.md
View file

@ -52,7 +52,7 @@ Three packs are staged into `backend/models` before packaging and ship as `resou
| Pack | Staged by | Holds |
|---|---|---|
| Embedding | `build:model`, `scripts/fetch_embedding_model.py` | `bge-small-en-v1.5`: the ONNX model, tokenizer and config |
| Embedding | `build:model`, `scripts/fetch_embedding_model.py` | `bge-small-en-v1.5`: the ONNX model and its tokenizer, from a pinned revision, sha256-checked |
| Voice | `build:voice`, `scripts/fetch_bundled_voice.py` | `audio/`: the manifest's first audio model in its default build, Kokoro 82M `Q8_0` (190 MB), with its install record, fetched as a catalog install is: from its pinned commit, sha256-checked |
| Parser | `build:parser`, `scripts/fetch_docling_models.py` | Docling's layout, table and RapidOCR weights, pruned of the variants ingest never loads |

13
docs/architecture/search.md
View file

@ -2,8 +2,9 @@
`retrieve()` finds the passages in a workspace that best answer a query. Two legs widen recall, a keyword match and a nearest-neighbour search over embeddings, and a weighted blend of the two decides the order, meaning counting for rather more than words. It runs in the calling process against the same SQLite file as everything else. Chat calls it, and so does Studio's grounding ([`studio.md`](studio.md), Grounding).
**Code:** [`shared/search.py`](../../surfsense_local/backend/shared/search.py), [`shared/tokenizer.py`](../../surfsense_local/backend/shared/tokenizer.py), [`worker/ingestion/embedding.py`](../../surfsense_local/backend/worker/ingestion/embedding.py)
**Decisions:** [ADR 0006](../adr/0006-hybrid-retrieval.md), [ADR 0007](../adr/0007-bundled-embeddings.md), [ADR 0031](../adr/0031-ranking-blends-absolute-leg-scores.md), [ADR 0032](../adr/0032-one-tokenizer-for-index-and-question.md), [ADR 0033](../adr/0033-every-candidate-is-scored-on-its-own-cosine.md)
**Code:** [`shared/search.py`](../../surfsense_local/backend/shared/search.py), [`shared/tokenizer.py`](../../surfsense_local/backend/shared/tokenizer.py), [`modules/embedding/`](../../surfsense_local/backend/modules/embedding/)
**Embedder:** [embedding](embedding.md)
**Decisions:** [ADR 0006](../adr/0006-hybrid-retrieval.md), [ADR 0007](../adr/0007-bundled-embeddings.md), [ADR 0036](../adr/0036-the-index-records-its-embedder.md), [ADR 0031](../adr/0031-ranking-blends-absolute-leg-scores.md), [ADR 0032](../adr/0032-one-tokenizer-for-index-and-question.md), [ADR 0033](../adr/0033-every-candidate-is-scored-on-its-own-cosine.md)
## Interface
@ -18,10 +19,10 @@ retrieve(session, workspace_id, query, top_k=5, document_ids=None) -> list[Hit]
## How it ranks
1. **Embed the query** with the bundled bge-small model at the same width as ingest. The import is lazy, so onnxruntime and the model load on the first query rather than when the API starts.
1. **Embed the query** with the active embedding index's model, the one ingest wrote with, adding the spec's query prefix if it has one (bge-small has none). The import is lazy, so onnxruntime and the model load on the first query rather than when the API starts. With no index chosen yet, `retrieve()` raises the same `embedding_not_chosen` error the routes answer `409` with.
2. **Keyword leg.** The query is split into terms by the index's own tokenizer, each quoted against FTS5's query grammar and joined with `OR` to keep recall wide. `chunks_fts` is matched, joined to `chunks` and `documents` for the workspace and document filters, and the best 20 by BM25 are kept. Each candidate then scores the fraction of the query's distinct terms it matched, one further lookup per term.
3. **Vector leg.** A sqlite-vec nearest-neighbour search over `chunk_vectors` takes the 20 closest chunks. It sits in its own CTE, so only `MATCH` and `k` constrain it, as vec0 requires; the workspace and document filters then apply to that set. It proposes those chunks and does not score them.
4. **Blend.** `0.65 × semantic + 0.35 × keyword` over the union, cut to `top_k` with cosine breaking a tie. Semantic strength is each candidate's own cosine similarity, floored at zero, measured here for the whole union rather than taken from the vector leg ([ADR 0033](../adr/0033-every-candidate-is-scored-on-its-own-cosine.md)). A candidate the keyword leg is missing contributes nothing from it, that absence being a measurement.
3. **Vector leg.** A sqlite-vec nearest-neighbour search over the active index's vector table, `chunk_vectors`, takes the 20 closest chunks. It sits in its own CTE, so only `MATCH` and `k` constrain it, as vec0 requires; the workspace and document filters then apply to that set. It proposes those chunks and does not score them.
4. **Blend.** `w × semantic + (1 − w) × keyword` over the union, where `w` is the active spec's `semantic_weight`, 0.65 for bge-small, cut to `top_k` with cosine breaking a tie. Semantic strength is each candidate's own cosine similarity, floored at zero, measured here for the whole union rather than taken from the vector leg ([ADR 0033](../adr/0033-every-candidate-is-scored-on-its-own-cosine.md)). A candidate the keyword leg is missing contributes nothing from it, that absence being a measurement.
Both legs decide the order, and both score on a scale that means the same thing from one query to the next, so a leg with nothing to say adds nothing. That is what keeps a keyword match honest: a chunk matching one term of five is weak whatever BM25 says about it, so no stopword list is needed. It holds because a term is a word: [`shared/tokenizer.py`](../../surfsense_local/backend/shared/tokenizer.py) states the one rule that splits both the index and the question, and keeps a combining mark inside its word so Devanagari is not cut into letters ([ADR 0032](../adr/0032-one-tokenizer-for-index-and-question.md)). A paraphrase that shares no words with its passage still arrives through the vector leg. There is no reranker.
@ -48,5 +49,5 @@ Two things it cannot currently show. LIMIT-small is 200 of its 266 queries, so t
## Known gaps
- A question in one language does not find its answer in another. Cosine alone puts the answering passage in the top 5 for 1 of 8 such queries, and the blend for 2 of 8: bge-small is English-only, so no ranking recovers it. Same-language retrieval is unaffected, at 100% in all four non-English slices. A multilingual embedder is the fix, designed in the [retrieval proposal](../proposals/retrieval.md).
- The bundled model is FP16, not the int8 [ADR 0007](../adr/0007-bundled-embeddings.md) specifies: `Qdrant/bge-small-en-v1.5-onnx-Q` is named for quantization but its only ONNX file stores FLOAT16 weights. An int8 export would roughly halve the 63 MB and speed embedding, and needs checking for whether its vectors are close enough to skip a re-embed.
- The bundled model is FP16, not the int8 [ADR 0007](../adr/0007-bundled-embeddings.md) specifies: `Qdrant/bge-small-en-v1.5-onnx-Q` is named for quantization but its only ONNX file stores FLOAT16 weights. An int8 export would roughly halve the 63 MB and speed embedding, but every library was built with this exact file, which is now pinned by hash, so a different build is a different vector space: it can replace this one only through a re-embed ([proposal](../proposals/embedding-model-change.md)).
- FTS5's tokenizer keeps a Japanese or Chinese clause as one token, so the keyword leg finds nothing in those scripts and the blend runs on meaning alone there. `trigram` would segment them at the cost of what BM25 means everywhere else ([ADR 0032](../adr/0032-one-tokenizer-for-index-and-question.md)).

1
docs/architecture/studio.md
View file

@ -124,6 +124,7 @@ Every pipeline returns a `Built`: a `title`, the `markdown` that is always the i
- A failure rolls back and writes `failed` with a reason cut to 500 characters: the error's first line, or for an HTTP error its whole message after "The model could not be reached: ". The job is then re-raised for Huey's one retry, except an image error, since the endpoint may already have generated, and billed, an image, and a podcast refused for memory, since a retry would draft the episode again and refuse again.
- **Cancel** marks the document `cancelled` and revokes a queued copy. A running job stops at its next check. During a model call it checks every second, while tokens stream and while the model is still reading the prompt, and hangs up; closing the request stops llama-server within 1.5 seconds, measured through the router ([`generate.py`](../../surfsense_local/backend/worker/studio/shared/generate.py)). A podcast checks before voicing each turn, so the rest of the episode is never sent, and an office job checks before each attempt, so a cancelled job asks the model for no retry. The check is a context variable ([`shared/cancellation.py`](../../surfsense_local/backend/shared/cancellation.py)), in `shared/` because the audio.cpp provider reads it too.
- Every Studio model call turns thinking off. Measured on Qwen3 1.7B: with it on, a mindmap over a 12,000-token prompt thought past 15,000 tokens without answering, holding the runtime's only slot so chat queued behind it.
- **A job the app quit in the middle of** is failed with `interrupted when the app closed` when the Studio worker next starts, before it takes a job ([`interrupted_documents.py`](../../surfsense_local/backend/worker/interrupted_documents.py)): Huey dropped the job as it started it, so nothing else would end it, and Regenerate would refuse it forever.
- **Regenerate** refuses a job still `pending` or `processing`, rechecks availability, resets the document to `pending`, clears the error, increments `generation` and re-enqueues with the same sources, prompt and options. The new run replaces the files and the indexed body; the artifact and its document keep their ids.
- Each transition the Studio worker makes sends an `artifacts` event keyed by artifact id; the API's own changes, to `pending` and `cancelled`, send none. The frontend does not listen yet; the artifact list refetches every 1.5 seconds while one is running ([`overview.md`](overview.md#freshness)).

4
docs/proposals/embedding-model-choice.md
View file

@ -1,5 +1,5 @@
---
status: proposed
status: in-progress
code:
- surfsense_local/backend/modules/embedding/
- surfsense_local/backend/modules/llm/catalog/local/
@ -21,7 +21,7 @@ code:
> The user picks the embedding model during onboarding: bge-small or another curated model, an ONNX embedder found on Hugging Face, or a remote model through a connection. Every pick SurfSense did not measure is tested before it is accepted, and the choice is fixed for the install. Changing it later means re-embedding the library, which is [future work](embedding-model-change.md). This proposal builds the data shape that work needs, so it adds code without changing a schema.
Today one model, bge-small-en-v1.5, is hard-coded in [`embedding.py`](../../surfsense_local/backend/worker/ingestion/embedding.py) ([ADR 0007](../adr/0007-bundled-embeddings.md)). It is English-only, so a question in one language does not find its answer in another ([search](../architecture/search.md), Known gaps). A multilingual or hosted model fixes that for the users who need it, and costs everyone else disk, memory, speed or privacy for nothing. So it is a choice, not a swap.
Today one model, bge-small-en-v1.5, is the only embedder ([`bundled.py`](../../surfsense_local/backend/modules/embedding/bundled.py), [ADR 0007](../adr/0007-bundled-embeddings.md)). It is English-only, so a question in one language does not find its answer in another ([search](../architecture/search.md), Known gaps). A multilingual or hosted model fixes that for the users who need it, and costs everyone else disk, memory, speed or privacy for nothing. So it is a choice, not a swap.
## Decisions

4
docs/proposals/retrieval.md
View file

@ -1,7 +1,7 @@
---
status: proposed
code:
- surfsense_local/backend/worker/ingestion/embedding.py
- surfsense_local/backend/modules/embedding/
- surfsense_local/backend/shared/search.py
- surfsense_local/backend/alembic/versions/
---
@ -32,7 +32,7 @@ bge scores Hindi **backwards**: an unrelated Hindi sentence sits closer to the E
## The candidate
`ibm-granite/granite-embedding-97m-multilingual-r2`. Apache-2.0, ungated, 52 languages, int8 ONNX. It is a drop-in: `hidden_size` 384 so `vec0` keeps its width and [ADR 0007](../adr/0007-bundled-embeddings.md)'s width check still passes, CLS pooling like bge, and no query prefix — [`embedding.py`](../../surfsense_local/backend/worker/ingestion/embedding.py) needs no change beyond the repo constants.
`ibm-granite/granite-embedding-97m-multilingual-r2`. Apache-2.0, ungated, 52 languages, int8 ONNX. It is a drop-in: `hidden_size` 384 so `vec0` keeps its width and [ADR 0007](../adr/0007-bundled-embeddings.md)'s width check still passes, CLS pooling like bge, and no query prefix — it needs a spec of its own and no change to [`encoder.py`](../../surfsense_local/backend/modules/embedding/encoder.py).
Swept on `dev`, because `SEMANTIC_WEIGHT` is defined as a plateau measured per embedder and judging a candidate at the incumbent's tuned weight is not a fair test:

18
plugins/core/cli/tests/conftest.py
View file

@ -19,6 +19,22 @@ import pytest
BACKEND_DIR = Path(__file__).resolve().parents[4] / "surfsense_local" / "backend"
# The app as a user has it once onboarding is done: plugins run only after it,
# so the embedder is fixed (bge-small, the default) before the API starts.
SERVE_ONBOARDED = (
"from modules.embedding.lock import lock_default_if_unchosen; "
"from shared.config import get_storage_settings; "
"from shared.db import create_db_engine, create_session_factory, import_models; "
"from shared.migrations import upgrade_to_head; "
"import_models(); "
"engine = create_db_engine(get_storage_settings().database_path); "
"upgrade_to_head(engine); "
"session = create_session_factory(engine)(); "
"lock_default_if_unchosen(session); session.commit(); session.close(); "
"engine.dispose(); "
"from api.server import serve; serve()"
)
# A first `uv run` may still be installing the backend's own packages.
STARTUP_SECONDS = 180
@ -164,7 +180,7 @@ def _started_app(folder: Path) -> Iterator[RealApp]:
str(BACKEND_DIR),
"python",
"-c",
"from api.server import serve; serve()",
SERVE_ONBOARDED,
],
env=environment,
stdout=output,

18
plugins/core/sdk/tests/contract/conftest.py
View file

@ -14,6 +14,22 @@ import pytest
BACKEND_DIR = Path(__file__).resolve().parents[5] / "surfsense_local" / "backend"
# The app as a user has it once onboarding is done: plugins run only after it,
# so the embedder is fixed (bge-small, the default) before the API starts.
SERVE_ONBOARDED = (
"from modules.embedding.lock import lock_default_if_unchosen; "
"from shared.config import get_storage_settings; "
"from shared.db import create_db_engine, create_session_factory, import_models; "
"from shared.migrations import upgrade_to_head; "
"import_models(); "
"engine = create_db_engine(get_storage_settings().database_path); "
"upgrade_to_head(engine); "
"session = create_session_factory(engine)(); "
"lock_default_if_unchosen(session); session.commit(); session.close(); "
"engine.dispose(); "
"from api.server import serve; serve()"
)
# A first `uv run` may still be installing the backend's own packages.
STARTUP_SECONDS = 180
@ -94,7 +110,7 @@ def real_app(tmp_path_factory: pytest.TempPathFactory) -> Iterator[RealApp]:
str(BACKEND_DIR),
"python",
"-c",
"from api.server import serve; serve()",
SERVE_ONBOARDED,
],
env=environment,
stdout=output,

10
surfsense_local/backend/alembic/versions/0001_initial_schema.py
View file

@ -11,7 +11,6 @@ from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from shared.config import get_search_settings
revision: str = "0001"
down_revision: str | Sequence[str] | None = None
@ -43,12 +42,9 @@ def _create_search_index() -> None:
"CREATE VIRTUAL TABLE chunks_fts USING fts5("
"content, content='chunks', content_rowid='id')"
)
# Fixed at this width forever; upgrade_to_head refuses a database that
# no longer matches the setting.
op.execute(
"CREATE VIRTUAL TABLE chunk_vectors USING vec0("
f"embedding float[{get_search_settings().embedding_dimension}])"
)
# bge-small's width. A fresh install's table is rebuilt at the chosen
# model's width when onboarding locks it (modules/embedding/lock.py).
op.execute("CREATE VIRTUAL TABLE chunk_vectors USING vec0(embedding float[384])")
op.execute("""
CREATE TRIGGER chunks_after_insert AFTER INSERT ON chunks BEGIN

116
surfsense_local/backend/alembic/versions/0021_embedding_indexes.py Normal file
View file

@ -0,0 +1,116 @@
"""Record which embedding model built the index.
An existing library was built by the bundled bge-small, so it gets that row,
pointing at the vectors it already has; nothing is re-embedded. A fresh install
gets no row: onboarding offers the choice and writes it.
The spec is written out here rather than imported, so this revision means the
same thing whatever the app's own copy of bge's spec later becomes.
`documents` gains its column by a plain ALTER TABLE, never batch mode: a batch
rebuild drops the table, and with foreign keys on that cascades to every chunk.
Revision ID: 0021
Revises: 0020
"""
import json
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
revision: str = "0021"
down_revision: str | Sequence[str] | None = "0020"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_BGE = {
"id": "bge-small-en-v1.5",
"source": "curated",
"identified": "measured",
"repo": "Qdrant/bge-small-en-v1.5-onnx-Q",
"revision": "aa8f8b060edb00e03bfdd08813a2949946c8ba55",
"weights": {
"path": "model_optimized.onnx",
"sha256": "51f1bd0addd6e859e42c2c8021a5e5461385bb676a649f4b269aa445449f2431",
},
"tokenizer": {
"path": "tokenizer.json",
"sha256": "d241a60d5e8f04cc1b2b3e9ef7a4921b27bf526d9f6050ab90f9267a1f9e5c66",
},
"dimension": 384,
"pooling": "cls",
"normalize": True,
"query_prefix": "",
"document_prefix": "",
"max_tokens": 512,
"semantic_weight": 0.65,
}
def upgrade() -> None:
op.create_table(
"embedding_indexes",
sa.Column("id", sa.Integer(), nullable=False),
sa.Column("spec", sa.JSON(), nullable=False),
sa.Column("vector_table", sa.String(), nullable=False),
sa.Column(
"state",
sa.Enum(
"active",
"building",
"retired",
name="indexstate",
native_enum=False,
create_constraint=True,
),
nullable=False,
),
sa.Column(
"created_at",
sa.DateTime(),
server_default=sa.text("(CURRENT_TIMESTAMP)"),
nullable=False,
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_embedding_indexes")),
)
# Inline, because Alembic will not add a constraint to a SQLite table and
# SQLite accepts a reference on an added column only in the column itself.
op.execute(
"ALTER TABLE documents ADD COLUMN embedding_index_id INTEGER "
"CONSTRAINT fk_documents_embedding_index_id_embedding_indexes "
"REFERENCES embedding_indexes (id)"
)
bind = op.get_bind()
built = bind.execute(
sa.text(
"SELECT EXISTS (SELECT 1 FROM onboarding_completion) "
"OR EXISTS (SELECT 1 FROM chunks)"
)
).scalar_one()
if not built:
return
index_id = bind.execute(
sa.text(
"INSERT INTO embedding_indexes(spec, vector_table, state) "
"VALUES (:spec, 'chunk_vectors', 'active') RETURNING id"
),
{"spec": json.dumps(_BGE)},
).scalar_one()
# Only what is indexed; a document still waiting is stamped by its ingest.
bind.execute(
sa.text(
"UPDATE documents SET embedding_index_id = :id "
"WHERE id IN (SELECT DISTINCT document_id FROM chunks)"
),
{"id": index_id},
)
def downgrade() -> None:
op.drop_column("documents", "embedding_index_id")
op.drop_table("embedding_indexes")

18
surfsense_local/backend/api/main.py
View file

@ -17,6 +17,11 @@ from modules.chat.router import router as chat_router
from modules.documents.router import router as documents_router
from modules.egress.router import router as egress_router
from modules.egress.service import EgressDeniedError
from modules.embedding.active import EmbeddingNotChosenError
from modules.embedding.huggingface.router import (
router as embedding_huggingface_router,
)
from modules.embedding.router import router as embedding_router
from modules.events.broker import EventBroker
from modules.events.router import router as events_router
from modules.health.router import router as health_router
@ -157,8 +162,11 @@ def create_app() -> FastAPI:
app.include_router(migration_router)
app.include_router(license_router)
app.include_router(egress_router)
app.include_router(embedding_router)
app.include_router(embedding_huggingface_router)
app.include_router(resource_usage_router)
app.add_exception_handler(EgressDeniedError, egress_denied)
app.add_exception_handler(EmbeddingNotChosenError, embedding_not_chosen)
app.add_exception_handler(UnreadableSecretError, unreadable_secret)
return app
@ -176,6 +184,16 @@ def unreadable_secret(_request: Request, error: UnreadableSecretError) -> JSONRe
)
def embedding_not_chosen(
_request: Request, error: EmbeddingNotChosenError
) -> JSONResponse:
"""Before onboarding fixes the embedder, nothing that would be embedded is taken."""
return JSONResponse(
{"detail": {"code": "embedding_not_chosen", "message": str(error)}},
status.HTTP_409_CONFLICT,
)
def egress_denied(_request: Request, error: EgressDeniedError) -> JSONResponse:
return JSONResponse(
{

6
surfsense_local/backend/main.py
View file

@ -12,10 +12,8 @@ def check_retrieval_runtime() -> None:
import onnxruntime
from tokenizers import Tokenizer
# The first-query entry point, imported the way retrieval reaches it: this
# walks worker.ingestion's package init through chonkie, whose optional
# transformers and pandas imports are what api.spec excludes.
from worker.ingestion.embedding import embed
# The first-query entry point, imported the way retrieval reaches it.
from modules.embedding.encoder import embed
sys.stdout.write(
f"retrieval imports OK (onnxruntime {onnxruntime.__version__}, "

3
surfsense_local/backend/modules/artifacts/router.py
View file

@ -40,6 +40,7 @@ from modules.artifacts.service import (
regenerate_artifact,
)
from modules.documents.models import Document, DocumentType
from modules.embedding.dependencies import EMBEDDER_CHOSEN
from modules.workspaces.dependencies import WorkspaceDep
from shared.config import get_storage_settings
@ -61,6 +62,7 @@ def studio_formats(workspace: WorkspaceDep, session: SessionDep) -> list[FormatR
@router.post(
"/workspaces/{workspace_id}/studio/jobs",
dependencies=[EMBEDDER_CHOSEN],
response_model=ArtifactRead,
status_code=status.HTTP_201_CREATED,
summary="Generate an artifact from documents",
@ -103,6 +105,7 @@ def read_artifact(artifact: ArtifactDep) -> ArtifactDetail:
@router.post(
"/artifacts/{artifact_id}/regenerate",
dependencies=[EMBEDDER_CHOSEN],
response_model=ArtifactRead,
status_code=status.HTTP_202_ACCEPTED,
summary="Generate a finished or failed artifact again",

5
surfsense_local/backend/modules/chat/router.py
View file

@ -38,6 +38,7 @@ from modules.chat.schemas import (
)
from modules.chat.title import generate_title
from modules.documents.sources import load_selected_sources
from modules.embedding.active import require_active_index
from modules.llm.activity import ModelBusyError, model_activity, model_key
from modules.llm.providers.protocols import Generator
from modules.llm.resolution import (
@ -407,9 +408,9 @@ def _ground(
if payload.document_ids is not None:
load_selected_sources(session, thread.workspace_id, payload.document_ids)
# Keep numpy/onnxruntime lazy: only chat and ingestion need this module.
from worker.ingestion.embedding import missing_embedding_files
from modules.embedding.encoder import missing_files
if missing_embedding_files():
if missing_files(require_active_index(session).spec):
raise HTTPException(
status.HTTP_503_SERVICE_UNAVAILABLE,
"local embedding model is not installed; "

4
surfsense_local/backend/modules/documents/models.py
View file

@ -55,6 +55,10 @@ class Document(Base):
content_hash: Mapped[str | None]
dedup_key: Mapped[str | None]
document_metadata: Mapped[dict[str, Any] | None] = mapped_column(JSON)
# The index its chunks were embedded into; what a re-embed has not reached.
embedding_index_id: Mapped[int | None] = mapped_column(
ForeignKey("embedding_indexes.id")
)
created_at: Mapped[datetime] = mapped_column(server_default=func.now())
updated_at: Mapped[datetime] = mapped_column(
server_default=func.now(), onupdate=func.now()

5
surfsense_local/backend/modules/documents/router.py
View file

@ -31,6 +31,7 @@ from modules.documents.storage import (
validate_upload,
)
from modules.documents.tasks import ingest_document
from modules.embedding.dependencies import EMBEDDER_CHOSEN
from modules.events.dependencies import EventBrokerDep
from modules.workspaces.dependencies import WorkspaceDep
from shared.config import get_storage_settings
@ -149,6 +150,7 @@ def read_document(document: DocumentDep) -> Document:
@router.post(
"",
response_model=DocumentDetail,
dependencies=[EMBEDDER_CHOSEN],
status_code=status.HTTP_201_CREATED,
summary="Write a note",
)
@ -177,6 +179,7 @@ def create_note(
@router.post(
"/upload",
dependencies=[EMBEDDER_CHOSEN],
response_model=UploadOutcome,
status_code=status.HTTP_201_CREATED,
summary="Upload files",
@ -274,6 +277,7 @@ def upload_documents(
@router.patch(
"/{document_id}",
dependencies=[EMBEDDER_CHOSEN],
response_model=DocumentRead,
summary="Edit a document",
)
@ -311,6 +315,7 @@ def update_document(
@router.post(
"/{document_id}/retry",
dependencies=[EMBEDDER_CHOSEN],
response_model=DocumentRead,
summary="Requeue a failed or cancelled document",
)

0
surfsense_local/backend/modules/embedding/__init__.py Normal file
View file

38
surfsense_local/backend/modules/embedding/active.py Normal file
View file

@ -0,0 +1,38 @@
from dataclasses import dataclass
from sqlalchemy import select
from sqlalchemy.orm import Session
from modules.embedding.models import EmbeddingIndex, IndexState
from modules.embedding.spec import EmbedderSpec
class EmbeddingNotChosenError(Exception):
"""Nothing may be embedded before onboarding fixes the model."""
def __init__(self) -> None:
super().__init__("no embedding model has been chosen yet")
@dataclass(frozen=True)
class ActiveIndex:
id: int
spec: EmbedderSpec
vector_table: str
def active_index(session: Session) -> ActiveIndex | None:
"""The index search reads and ingest writes. None before onboarding chooses."""
row = session.scalars(
select(EmbeddingIndex).where(EmbeddingIndex.state == IndexState.ACTIVE)
).one_or_none()
if row is None:
return None
return ActiveIndex(row.id, EmbedderSpec.model_validate(row.spec), row.vector_table)
def require_active_index(session: Session) -> ActiveIndex:
index = active_index(session)
if index is None:
raise EmbeddingNotChosenError
return index

49
surfsense_local/backend/modules/embedding/bundled.py Normal file
View file

@ -0,0 +1,49 @@
"""bge-small, the embedder the app ships: the default, and never deletable.
Pinned by revision and hash because every existing library was built with this
exact file. A different build of the same model is a different vector space.
"""
from pathlib import Path
from modules.embedding.spec import (
EmbedderSpec,
Identified,
PinnedFile,
Pooling,
Source,
)
from shared.config import get_storage_settings
BGE = EmbedderSpec(
id="bge-small-en-v1.5",
source=Source.CURATED,
identified=Identified.MEASURED,
repo="Qdrant/bge-small-en-v1.5-onnx-Q",
# A commit and checksums, public by nature: not secrets.
# pragma: allowlist nextline secret
revision="aa8f8b060edb00e03bfdd08813a2949946c8ba55",
weights=PinnedFile(
path="model_optimized.onnx",
# pragma: allowlist nextline secret
sha256="51f1bd0addd6e859e42c2c8021a5e5461385bb676a649f4b269aa445449f2431",
),
tokenizer=PinnedFile(
path="tokenizer.json",
# pragma: allowlist nextline secret
sha256="d241a60d5e8f04cc1b2b3e9ef7a4921b27bf526d9f6050ab90f9267a1f9e5c66",
),
dimension=384,
pooling=Pooling.CLS,
normalize=True,
max_tokens=512,
# How much of the order meaning owns. The eval corpus is flat across 0.6 to
# 0.65 for bge and falls away either side: the middle of a measured plateau,
# not a default. Reciprocal rank fusion was tried first and rejected at 88%
# of answers in the top 5, against 98% here.
semantic_weight=0.65,
)
def bundled_dir() -> Path:
return get_storage_settings().models_dir / BGE.id

53
surfsense_local/backend/modules/embedding/choose.py Normal file
View file

@ -0,0 +1,53 @@
"""What finishing onboarding fixes as the library's embedder."""
from fastapi import HTTPException, status
from sqlalchemy.orm import Session
from modules.embedding.active import active_index
from modules.embedding.bundled import BGE
from modules.embedding.lock import lock_index
from modules.llm.catalog.local.dependencies import get_local_catalog
from modules.llm.catalog.local.engines.onnxruntime.spec import spec_for
from modules.llm.catalog.local.manifest import load_local_manifest
def lock_chosen(session: Session, model_id: str | None) -> None:
"""The curated embedder onboarding chose, or bge when it chose none.
Once locked, the same choice again is a no-op and any other is refused:
changing the embedder is a re-embed, which is not built.
"""
chosen = model_id or BGE.id
active = active_index(session)
if active is not None:
if model_id is not None and active.spec.id != chosen:
raise HTTPException(
status.HTTP_409_CONFLICT,
f"{active.spec.id} is already this library's embedding model",
)
return
if chosen == BGE.id:
lock_index(session, BGE)
return
picked = get_local_catalog().onnxruntime.installed_spec(chosen)
if picked is not None:
lock_index(session, picked)
return
model = next(
(
m
for m in load_local_manifest().models
if m.id == chosen and m.embedding is not None
),
None,
)
if model is None:
raise HTTPException(
status.HTTP_422_UNPROCESSABLE_CONTENT,
f"no curated embedding model named {chosen}",
)
if not get_local_catalog().onnxruntime.holds(chosen):
raise HTTPException(
status.HTTP_409_CONFLICT, f"{chosen} has not finished downloading"
)
lock_index(session, spec_for(model))

12
surfsense_local/backend/modules/embedding/dependencies.py Normal file
View file

@ -0,0 +1,12 @@
from fastapi import Depends
from api.dependencies import SessionDep
from modules.embedding.active import require_active_index
def _require_embedder_chosen(session: SessionDep) -> None:
"""Refuse before queueing work that would be embedded with no model chosen."""
require_active_index(session)
EMBEDDER_CHOSEN = Depends(_require_embedder_chosen)

104
surfsense_local/backend/modules/embedding/encoder.py Normal file
View file

@ -0,0 +1,104 @@
"""Text to vectors, run as the spec says, in process on the CPU."""
from enum import StrEnum
from functools import lru_cache
from pathlib import Path
from typing import Any
import numpy as np
from modules.embedding.bundled import BGE, bundled_dir
from modules.embedding.pooling import pool
from modules.embedding.spec import EmbedderSpec
from shared.config import get_storage_settings
class Purpose(StrEnum):
"""Asymmetric models embed a question and a passage differently."""
QUERY = "query"
DOCUMENT = "document"
def model_dir(spec: EmbedderSpec) -> Path:
"""bge in the read-only pack the app ships; anything else where it downloaded."""
if (spec.repo, spec.revision, spec.weights) == (
BGE.repo,
BGE.revision,
BGE.weights,
):
return bundled_dir()
return get_storage_settings().embedding_models_dir / spec.id
def missing_files(spec: EmbedderSpec) -> list[str]:
"""Files the model still needs before anything can be embedded."""
directory = model_dir(spec)
return [
pinned.path
for pinned in (spec.weights, spec.tokenizer)
if not (directory / pinned.path).is_file()
]
def embed(spec: EmbedderSpec, texts: list[str], purpose: Purpose) -> list[list[float]]:
"""Embed every text with the spec's model, in batches."""
prefix = spec.query_prefix if purpose is Purpose.QUERY else spec.document_prefix
prefixed = [prefix + text for text in texts]
vectors: list[list[float]] = []
for start in range(0, len(prefixed), spec.batch):
vectors.extend(_embed_batch(spec, prefixed[start : start + spec.batch]))
wrong = next((len(v) for v in vectors if len(v) != spec.dimension), None)
if wrong is not None:
raise ValueError(
f"{spec.id} returned {wrong} dimensions where its spec says "
f"{spec.dimension}"
)
return vectors
def width(spec: EmbedderSpec, text: str) -> int:
"""How wide the model's vectors really are, whatever its spec says."""
return len(_embed_batch(spec, [text])[0])
def _embed_batch(spec: EmbedderSpec, texts: list[str]) -> list[list[float]]:
session, encoder = _loaded(spec, model_dir(spec))
encoded = encoder.encode_batch(texts)
feed = {
"input_ids": np.array([e.ids for e in encoded], dtype=np.int64),
"attention_mask": np.array([e.attention_mask for e in encoded], dtype=np.int64),
}
# Some exports drop token_type_ids; feed it only if the graph asks.
inputs = {i.name for i in session.get_inputs()}
if "token_type_ids" in inputs:
feed["token_type_ids"] = np.array([e.type_ids for e in encoded], dtype=np.int64)
# Decoder-style exports, Qwen3-Embedding's among them, take positions too.
if "position_ids" in inputs:
ids = feed["input_ids"]
feed["position_ids"] = np.broadcast_to(
np.arange(ids.shape[1]), ids.shape
).copy()
output = session.run(None, feed)[0]
pooled = pool(spec.pooling, output, feed["attention_mask"])
if spec.normalize:
pooled = pooled / np.linalg.norm(pooled, axis=1, keepdims=True)
return pooled.astype(np.float32).tolist()
@lru_cache(maxsize=2)
def _loaded(spec: EmbedderSpec, directory: Path) -> tuple[Any, Any]:
"""One session and tokenizer per model, so a process follows the active index
rather than whatever it loaded first."""
import onnxruntime as ort
from tokenizers import Tokenizer
encoder = Tokenizer.from_file(str(directory / spec.tokenizer.path))
encoder.enable_truncation(max_length=spec.max_tokens)
encoder.enable_padding()
session = ort.InferenceSession(
str(directory / spec.weights.path), providers=["CPUExecutionProvider"]
)
return session, encoder

0
surfsense_local/backend/modules/embedding/huggingface/__init__.py Normal file
View file

17
surfsense_local/backend/modules/embedding/huggingface/hub.py Normal file
View file

@ -0,0 +1,17 @@
"""The one way this slice reaches Hugging Face."""
import httpx
BASE = "https://huggingface.co"
TIMEOUT = httpx.Timeout(15.0, connect=10.0)
def transport() -> httpx.AsyncBaseTransport | None:
"""The network itself; tests answer in its place."""
return None
def client() -> httpx.AsyncClient:
return httpx.AsyncClient(
base_url=BASE, transport=transport(), timeout=TIMEOUT, follow_redirects=True
)

73
surfsense_local/backend/modules/embedding/huggingface/pick.py Normal file
View file

@ -0,0 +1,73 @@
"""Which files of a Hugging Face repo an embedder is downloaded as.
A generic int8 build first, which measured smaller and faster on a CPU than full
precision, then full precision; nothing tuned for one CPU, optimised for a GPU,
or in fp16, which runs slowly on a CPU. A bad quantisation is left to the probe
and the search check after the download.
"""
import re
from collections.abc import Collection, Iterable, Mapping
from dataclasses import dataclass
from typing import Any
from modules.llm.catalog.local.listed_file import ListedFile
_PREFERENCE = (
"onnx/model_int8.onnx",
"onnx/model_quantized.onnx",
"model_int8.onnx",
"model_quantized.onnx",
"onnx/model.onnx",
"model.onnx",
)
TOKENIZER = "tokenizer.json"
class NotRunnableError(Exception):
"""Why this repo cannot be an embedder here, in words the screen shows."""
@dataclass(frozen=True)
class PickedFiles:
weights: ListedFile
# The graph's external weights, which must land beside it under their names.
data: tuple[ListedFile, ...]
tokenizer: ListedFile
@property
def all(self) -> tuple[ListedFile, ...]:
return (self.weights, *self.data, self.tokenizer)
def pick_files(listing: Iterable[ListedFile]) -> PickedFiles:
files = {f.path: f for f in listing}
weights = next((files[p] for p in _PREFERENCE if p in files), None)
if weights is None:
raise NotRunnableError("This repo has no ONNX build SurfSense can run.")
tokenizer = files.get(TOKENIZER)
if tokenizer is None:
raise NotRunnableError("This repo has no tokenizer.json.")
data = tuple(
f
for path, f in sorted(files.items())
if re.fullmatch(re.escape(weights.path) + r"(_data|\.data)(_\d+)?", path)
)
picked = PickedFiles(weights, data, tokenizer)
if unhashed := [f.path for f in picked.all if not f.sha256]:
raise NotRunnableError(f"Hugging Face lists no checksum for {unhashed[0]}.")
return picked
def scan_verdict(status: Mapping[str, Any] | None, ours: Collection[str]) -> str | None:
"""Why Hugging Face's own scan refuses the repo, or None.
Unsafe anywhere refuses it, and so does any flag on a file we would take.
"""
for issue in (status or {}).get("filesWithIssues") or ():
level = issue.get("level")
if level == "safe":
continue
if level == "unsafe" or issue.get("path") in ours:
return "Hugging Face's security scan flags this repo."
return None

140
surfsense_local/backend/modules/embedding/huggingface/resolve.py Normal file
View file

@ -0,0 +1,140 @@
"""One Hugging Face repo, opened: whether it can run here, and as what.
Reads its summary and security scan, its file list at that commit, and the few
small files that say how it embeds. No weights are downloaded.
"""
import asyncio
import hashlib
import json
from dataclasses import dataclass
from typing import Any
import httpx
from modules.embedding.huggingface import hub
from modules.embedding.huggingface.pick import (
TOKENIZER,
NotRunnableError,
PickedFiles,
pick_files,
scan_verdict,
)
from modules.embedding.huggingface.spec_from_repo import RepoConfig, spec_from_repo
from modules.embedding.spec import EmbedderSpec
from modules.llm.catalog.local.listed_file import ListedFile
_CONFIG_FILES = (
"config.json",
"1_Pooling/config.json",
"modules.json",
"config_sentence_transformers.json",
"sentence_bert_config.json",
)
# Hugging Face keeps files under 10 MB out of LFS, so lists no sha256 for them.
_SMALL_FILE_BYTES = 10 * 1024 * 1024
@dataclass(frozen=True)
class Resolved:
repo: str
revision: str | None
picked: PickedFiles | None
spec: EmbedderSpec | None
# Why it cannot be an embedder here; None when it can.
reason: str | None
gated: bool = False
async def resolve(repo: str) -> Resolved:
async with hub.client() as client:
info = await _json(
client,
f"/api/models/{repo}",
{"securityStatus": "true", "expand[]": ["sha", "pipeline_tag", "gated"]},
)
revision = info["sha"]
if info.get("gated"):
return Resolved(
repo,
revision,
None,
None,
"This repo needs an account to download.",
gated=True,
)
tree = await _json(
client, f"/api/models/{repo}/tree/{revision}", {"recursive": "true"}
)
listing = [
ListedFile(
row["path"], int(row.get("size", 0)), (row.get("lfs") or {}).get("oid")
)
for row in tree
if row.get("type") == "file"
]
listing = [
await _hashed(client, repo, revision, f) if f.path == TOKENIZER else f
for f in listing
]
try:
picked = pick_files(listing)
except NotRunnableError as error:
return Resolved(repo, revision, None, None, str(error))
ours = {f.path for f in picked.all}
if refused := scan_verdict(info.get("securityRepoStatus"), ours):
return Resolved(repo, revision, None, None, refused)
listed = {f.path for f in listing}
configs = await asyncio.gather(
*(
_optional(client, f"/{repo}/resolve/{revision}/{path}")
if path in listed
else _none()
for path in _CONFIG_FILES
)
)
config, pooling, modules, sentence, bert = configs
found = RepoConfig(
config=config,
pooling=pooling,
modules=modules,
prompts=(sentence or {}).get("prompts"),
max_seq_length=(bert or {}).get("max_seq_length"),
)
spec = spec_from_repo(repo, revision, picked, found, info.get("pipeline_tag"))
return Resolved(repo, revision, picked, spec, None)
async def _json(client: httpx.AsyncClient, path: str, params: dict[str, Any]) -> Any:
reply = await client.get(path, params=params)
reply.raise_for_status()
return reply.json()
async def _hashed(
client: httpx.AsyncClient, repo: str, revision: str, file: ListedFile
) -> ListedFile:
"""A small file without an LFS hash, hashed from its bytes at the pinned
commit, so its download is checked like any other."""
if file.sha256 or file.size_bytes > _SMALL_FILE_BYTES:
return file
reply = await client.get(f"/{repo}/resolve/{revision}/{file.path}")
if reply.status_code != 200:
return file
return ListedFile(
file.path, file.size_bytes, hashlib.sha256(reply.content).hexdigest()
)
async def _optional(client: httpx.AsyncClient, path: str) -> Any:
reply = await client.get(path)
if reply.status_code != 200:
return None
try:
return json.loads(reply.content)
except ValueError:
return None
async def _none() -> None:
return None

72
surfsense_local/backend/modules/embedding/huggingface/router.py Normal file
View file

@ -0,0 +1,72 @@
"""Embedders on Hugging Face: search, and one repo opened as an install.
Answered in the GGUF search's shapes, so onboarding renders both through one
list. Egress gated like that search: what the user types goes to a third party.
Downloading goes through the same install jobs as every local model.
"""
from typing import Annotated
from fastapi import APIRouter, Query
from api.dependencies import SessionDep, transact
from modules.egress import service as egress
from modules.embedding.huggingface.resolve import resolve
from modules.embedding.huggingface.search import search
from modules.llm.catalog.local.build import Build, BuildFile, FileRole
from modules.llm.catalog.local.dependencies import LocalCatalogDep
from modules.llm.catalog.local.engines.onnxruntime.rows import searched_row
from modules.llm.catalog.local.router import row_read
from modules.llm.catalog.local.rows import BuildRow
from modules.llm.catalog.local.schemas import RepoRead, SearchRead
router = APIRouter(prefix="/embedding/huggingface", tags=["embedding"])
@router.get("/search", response_model=SearchRead, summary="Search for embedders")
async def search_embedders(
q: Annotated[str, Query(min_length=1, max_length=200)], session: SessionDep
) -> dict:
await transact(session, egress.require, egress.HUGGINGFACE)
return {"results": [hit.__dict__ for hit in await search(q)]}
@router.get(
"/repo/{repo:path}", response_model=RepoRead, summary="Open one embedder repo"
)
async def open_repo(repo: str, service: LocalCatalogDep, session: SessionDep) -> dict:
await transact(session, egress.require, egress.HUGGINGFACE)
resolved = await resolve(repo)
if resolved.picked is None or resolved.spec is None:
row = searched_row(repo, None, resolved.reason)
return {"repo": repo, "gated": resolved.gated, "row": row_read(row, {})}
build = Build(
"ONNX",
tuple(
BuildFile(
FileRole.TOKENIZER
if f is resolved.picked.tokenizer
else FileRole.WEIGHTS,
f.path,
f.size_bytes,
f.sha256,
repo,
resolved.revision or "main",
)
for f in resolved.picked.all
),
)
# Installed means checked: files without a settled spec are not a pick.
held = service.onnxruntime.installed_spec(resolved.spec.id) is not None
offered = BuildRow(
catalog_id=service.offer_embedder(build, resolved.spec),
build=build,
fit=None,
badge=None,
installed_as=resolved.spec.id if held else None,
recommended=False,
reads_images=False,
projector_checked=False,
)
row = searched_row(repo, offered, None)
return {"repo": repo, "gated": False, "row": row_read(row, {})}

58
surfsense_local/backend/modules/embedding/huggingface/search.py Normal file
View file

@ -0,0 +1,58 @@
"""Embedders on Hugging Face by name, described and not yet judged.
Whether one can run here needs its file list, read when it is opened. Hits take
the GGUF search's shape, so both searches render through one list.
"""
import asyncio
from modules.embedding.huggingface import hub
from modules.llm.catalog.local.engines.llamacpp.search.hits import SearchHit
# The two tags an embedder carries; the search API takes one at a time.
TAGS = ("sentence-similarity", "feature-extraction")
LIMIT = 20
async def search(query: str) -> list[SearchHit]:
async with hub.client() as client:
replies = await asyncio.gather(
*(
client.get(
"/api/models",
params={
"search": query,
"pipeline_tag": tag,
"limit": LIMIT,
"expand[]": [
"downloads",
"likes",
"lastModified",
"gated",
"tags",
],
},
)
for tag in TAGS
)
)
found: dict[str, SearchHit] = {}
for reply in replies:
reply.raise_for_status()
for row in reply.json():
found.setdefault(row["id"], _hit(row))
return sorted(found.values(), key=lambda hit: -hit.downloads)[:LIMIT]
def _hit(row: dict) -> SearchHit:
tags = tuple(row.get("tags") or ())
return SearchHit(
repo=row["id"],
downloads=int(row.get("downloads") or 0),
likes=int(row.get("likes") or 0),
last_modified=row.get("lastModified"),
license=next(
(t.removeprefix("license:") for t in tags if t.startswith("license:")), None
),
gated=bool(row.get("gated")),
)

82
surfsense_local/backend/modules/embedding/huggingface/spec_from_repo.py Normal file
View file

@ -0,0 +1,82 @@
"""A Hugging Face embedder's spec from the files its repo ships.
sentence-transformers states pooling, normalisation, prompts and length; a repo
without those files gets bge's defaults and is labelled inferred. The width here
is the config's; the probe after download has the last word.
"""
from dataclasses import dataclass
from typing import Any
from modules.embedding.huggingface.pick import PickedFiles
from modules.embedding.spec import EmbedderSpec, Identified, PinnedFile, Pooling, Source
# The ranking weight is measured per model; one not measured takes bge's.
UNMEASURED_WEIGHT = 0.65
# Passages are cut at 480 of bge's tokens; another tokenizer may count more,
# never four times as many, and a longer window only costs memory.
MAX_TOKENS_CAP = 2048
_POOLING = {
"pooling_mode_cls_token": Pooling.CLS,
"pooling_mode_mean_tokens": Pooling.MEAN,
"pooling_mode_lasttoken": Pooling.LAST,
}
@dataclass(frozen=True)
class RepoConfig:
"""The repo's own files, each None where it ships none."""
config: dict[str, Any] | None
pooling: dict[str, Any] | None
modules: list[dict[str, Any]] | None
prompts: dict[str, str] | None
max_seq_length: int | None
def spec_from_repo(
repo: str,
revision: str,
picked: PickedFiles,
found: RepoConfig,
pipeline_tag: str | None,
) -> EmbedderSpec:
del pipeline_tag # the search admits only embedding tags; recorded for reading
declared = found.pooling is not None or found.modules is not None
prompts = found.prompts or {}
config = found.config or {}
length = found.max_seq_length or config.get("max_position_embeddings") or 512
return EmbedderSpec(
id=install_id(repo),
source=Source.HUGGINGFACE,
identified=Identified.DECLARED if declared else Identified.INFERRED,
repo=repo,
revision=revision,
weights=PinnedFile(
path=_name(picked.weights.path), sha256=picked.weights.sha256
),
tokenizer=PinnedFile(
path=_name(picked.tokenizer.path), sha256=picked.tokenizer.sha256
),
dimension=int(config.get("hidden_size") or 384),
pooling=next(
(mode for key, mode in _POOLING.items() if (found.pooling or {}).get(key)),
Pooling.CLS,
),
# Search compares by cosine, so a unit vector loses nothing and keeps the
# nearest-neighbour leg, which measures distance, in agreement with it.
normalize=True,
query_prefix=prompts.get("query", ""),
document_prefix=prompts.get("document") or prompts.get("passage") or "",
max_tokens=min(int(length), MAX_TOKENS_CAP),
semantic_weight=UNMEASURED_WEIGHT,
)
def install_id(repo: str) -> str:
"""What a Hugging Face pick is installed as: never a curated id."""
return "hf--" + repo.replace("/", "--").lower()
def _name(path: str) -> str:
return path.rsplit("/", 1)[-1]

53
surfsense_local/backend/modules/embedding/lock.py Normal file
View file

@ -0,0 +1,53 @@
"""Fixing the library's embedder, once, when onboarding finishes.
The one place outside migrations that emits DDL: the vector table's width is the
chosen model's, which no migration can know. Safe only while the table is empty,
which is what the refusals guarantee.
"""
from sqlalchemy import exists, select, text
from sqlalchemy.orm import Session
from modules.chunks.models import Chunk
from modules.embedding.active import active_index
from modules.embedding.bundled import BGE
from modules.embedding.models import EmbeddingIndex, IndexState
from modules.embedding.spec import EmbedderSpec
# The table migration 0001 created, which the chunk delete trigger names.
FIRST_TABLE = "chunk_vectors"
class IndexAlreadyBuiltError(Exception):
"""The library already has vectors from some model, or a locked choice."""
def lock_index(session: Session, spec: EmbedderSpec) -> EmbeddingIndex:
if session.scalar(
select(exists().where(EmbeddingIndex.state == IndexState.ACTIVE))
):
raise IndexAlreadyBuiltError("an embedding model is already chosen")
if session.scalar(select(exists().select_from(Chunk))):
raise IndexAlreadyBuiltError("the library already holds embedded passages")
session.execute(text(f"DROP TABLE IF EXISTS {FIRST_TABLE}"))
session.execute(
text(
f"CREATE VIRTUAL TABLE {FIRST_TABLE} "
f"USING vec0(embedding float[{int(spec.dimension)}])"
)
)
index = EmbeddingIndex(
spec=spec.model_dump(mode="json"),
vector_table=FIRST_TABLE,
state=IndexState.ACTIVE,
)
session.add(index)
session.flush()
return index
def lock_default_if_unchosen(session: Session) -> None:
"""Finishing onboarding with no embedder chosen means the bundled one."""
if active_index(session) is None:
lock_index(session, BGE)

31
surfsense_local/backend/modules/embedding/models.py Normal file
View file

@ -0,0 +1,31 @@
import enum
from datetime import datetime
from typing import Any
from sqlalchemy import JSON, func
from sqlalchemy.orm import Mapped, mapped_column
from shared.db import Base, text_enum
class IndexState(enum.StrEnum):
ACTIVE = "active"
# Written by changing the model, which a later release adds.
BUILDING = "building"
RETIRED = "retired"
class EmbeddingIndex(Base):
"""Which model built a set of vectors, and where they are.
One active row, enforced in code rather than by a singleton constraint: a
second row is how changing the model starts.
"""
__tablename__ = "embedding_indexes"
id: Mapped[int] = mapped_column(primary_key=True)
spec: Mapped[dict[str, Any]] = mapped_column(JSON)
vector_table: Mapped[str]
state: Mapped[IndexState] = mapped_column(text_enum(IndexState))
created_at: Mapped[datetime] = mapped_column(server_default=func.now())

19
surfsense_local/backend/modules/embedding/pooling.py Normal file
View file

@ -0,0 +1,19 @@
"""One vector per text from a model's per-token output, as its spec says."""
import numpy as np
from modules.embedding.spec import Pooling
def pool(pooling: Pooling, output: np.ndarray, mask: np.ndarray) -> np.ndarray:
"""`output` is (texts, tokens, width), or (texts, width) for a build that
pools inside its own graph. Padding is on the right, as the tokenizer pads."""
if pooling is Pooling.IN_MODEL:
return output
if pooling is Pooling.CLS:
return output[:, 0]
if pooling is Pooling.MEAN:
weights = mask[:, :, None].astype(output.dtype)
return (output * weights).sum(axis=1) / np.maximum(weights.sum(axis=1), 1)
last = mask.sum(axis=1) - 1
return output[np.arange(output.shape[0]), last]

31
surfsense_local/backend/modules/embedding/router.py Normal file
View file

@ -0,0 +1,31 @@
from fastapi import APIRouter
from sqlalchemy.orm import Session
from api.dependencies import SessionDep, transact
from modules.embedding.active import active_index
from modules.embedding.schemas import EmbeddingIndexRead, IndexRead
from modules.llm.catalog.local.manifest import load_local_manifest
router = APIRouter(prefix="/embedding", tags=["embedding"])
@router.get(
"/index",
response_model=EmbeddingIndexRead,
summary="Read which embedding model the library is built with",
)
async def read_index(session: SessionDep) -> EmbeddingIndexRead:
return await transact(session, _read)
def _read(session: Session) -> EmbeddingIndexRead:
index = active_index(session)
if index is None:
return EmbeddingIndexRead(active=None)
names = {m.id: m.name for m in load_local_manifest().models}
return EmbeddingIndexRead(
# A Hugging Face pick is named by its repo; its install name is ours.
active=IndexRead(
name=names.get(index.spec.id, index.spec.repo), spec=index.spec
)
)

16
surfsense_local/backend/modules/embedding/schemas.py Normal file
View file

@ -0,0 +1,16 @@
from pydantic import BaseModel
from modules.embedding.spec import EmbedderSpec
class IndexRead(BaseModel):
"""One index: the spec its vectors were made with, and what to call it."""
name: str
spec: EmbedderSpec
class EmbeddingIndexRead(BaseModel):
active: IndexRead | None
# Always null until changing the model is built; carries its progress then.
building: IndexRead | None = None

104
surfsense_local/backend/modules/embedding/search_check.py Normal file
View file

@ -0,0 +1,104 @@
"""The bar a model SurfSense did not measure must clear before it is locked.
Ten questions over twenty passages, each answer beside a decoy on the same
topic, and every answer must rank first. Measured on 1 Oct 2026: bge-small and
both curated granite models found 10 of 10. Chat models served as embedders
found 5 to 8, except one at 10, so this proves a model can find answers, which
is what search needs, not that it was built to embed.
A probe alone was not enough: easy paraphrase pairs separated a chat model as
well as bge did.
"""
from collections.abc import Callable
from dataclasses import dataclass
import numpy as np
# (question, the passage that answers it, a same-topic passage that does not)
QUESTIONS = (
(
"How long does the scanner battery last?",
"The scanner battery lasts about ten hours on a full charge.",
"The scanner battery charges in the four-slot cradle.",
),
(
"When is the invoice due?",
"Invoices are due thirty days after the date of issue.",
"Invoices are sent by email as PDF attachments.",
),
(
"Who approves travel expenses?",
"Travel expenses are approved by your line manager.",
"Travel expenses include flights, hotels and meals.",
),
(
"What temperature should the server room be kept at?",
"Keep the server room between 18 and 24 degrees Celsius.",
"The server room is on the second floor next to the lift.",
),
(
"How many vacation days do new employees get?",
"New employees receive twenty days of paid vacation per year.",
"Vacation requests are submitted in the HR portal.",
),
(
"Where is the spare key for the office?",
"The spare office key is kept in the reception safe.",
"The office is locked automatically at 8 pm.",
),
(
"What does the red warning light mean?",
"A red warning light means the filter must be replaced.",
"The warning light is next to the power switch.",
),
(
"Which port does the database listen on?",
"The database listens on port 5432 by default.",
"The database is backed up every night at midnight.",
),
(
"How do I reset my password?",
"To reset your password, click Forgot password on the sign-in page.",
"Passwords must be at least twelve characters long.",
),
(
"When does the cafeteria close?",
"The cafeteria closes at three in the afternoon.",
"The cafeteria serves vegetarian options every day.",
),
)
@dataclass(frozen=True)
class SearchCheck:
first: int
asked: int
@property
def passed(self) -> bool:
return self.first == self.asked
def search_check(
embed_queries: Callable[[list[str]], list[list[float]]],
embed_passages: Callable[[list[str]], list[list[float]]] | None = None,
) -> SearchCheck:
"""How many answers rank first. Questions and passages embed apart, so an
asymmetric model is checked with its own prefixes."""
passages = [text for _, answer, decoy in QUESTIONS for text in (answer, decoy)]
queries = _unit(embed_queries([q for q, _, _ in QUESTIONS]))
stored = _unit((embed_passages or embed_queries)(passages))
similarity = queries @ stored.T
# First means nothing scores above the answer; a tie does not demote it.
first = sum(
int(not (similarity[i] > similarity[i, 2 * i]).any())
for i in range(len(QUESTIONS))
)
return SearchCheck(first, len(QUESTIONS))
def _unit(vectors: list[list[float]]) -> np.ndarray:
array = np.asarray(vectors, dtype=np.float64)
norms = np.linalg.norm(array, axis=1, keepdims=True)
return array / np.where(norms == 0, 1, norms)

60
surfsense_local/backend/modules/embedding/spec.py Normal file
View file

@ -0,0 +1,60 @@
"""Everything that decides what a vector means.
An index stores a snapshot of this, so a later edit to a manifest entry cannot
change what an existing index's vectors are.
"""
from enum import StrEnum
from pydantic import BaseModel, ConfigDict, Field
class Source(StrEnum):
CURATED = "curated"
HUGGINGFACE = "huggingface"
REMOTE = "remote"
class Identified(StrEnum):
"""How SurfSense knows the model is built to embed."""
MEASURED = "measured"
DECLARED = "declared"
INFERRED = "inferred"
UNVERIFIED = "unverified"
class Pooling(StrEnum):
CLS = "cls"
MEAN = "mean"
LAST = "last"
# The build pools inside its own graph and returns one vector per text.
IN_MODEL = "in_model"
class PinnedFile(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
path: str = Field(min_length=1)
sha256: str = Field(pattern=r"^[0-9a-f]{64}$")
class EmbedderSpec(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
id: str = Field(min_length=1)
source: Source
identified: Identified
repo: str = Field(min_length=1)
revision: str = Field(pattern=r"^[0-9a-f]{40}$")
weights: PinnedFile
tokenizer: PinnedFile
dimension: int = Field(gt=0)
pooling: Pooling
normalize: bool
query_prefix: str = ""
document_prefix: str = ""
max_tokens: int = Field(gt=0)
semantic_weight: float = Field(gt=0, lt=1)
# Passages embedded at once; the batch, not the weights, sets peak memory.
batch: int = Field(default=32, gt=0)

27
surfsense_local/backend/modules/embedding/vector_table.py Normal file
View file

@ -0,0 +1,27 @@
"""The vec0 table an index's vectors live in.
Its name comes from a database row and is spliced into SQL, so it is checked
against a fixed shape first.
"""
import re
from sqlalchemy import Connection, text
_NAME = re.compile(r"^chunk_vectors(_[0-9]+)?$")
_DECLARED_WIDTH = re.compile(r"float\[(\d+)\]")
def checked(name: str) -> str:
if not _NAME.fullmatch(name):
raise ValueError(f"not a vector table name: {name!r}")
return name
def declared_width(connection: Connection, name: str) -> int | None:
"""The width the table was created with, or None if it does not exist."""
schema = connection.execute(
text("SELECT sql FROM sqlite_master WHERE name = :name"),
{"name": checked(name)},
).scalar_one_or_none()
return None if schema is None else int(_DECLARED_WIDTH.search(schema).group(1))

27
surfsense_local/backend/modules/embedding/verify.py Normal file
View file

@ -0,0 +1,27 @@
"""Whether a model SurfSense did not measure works, after its download."""
from modules.embedding.encoder import Purpose, embed, width
from modules.embedding.search_check import search_check
from modules.embedding.spec import EmbedderSpec
PROBE = "The scanner battery lasts about ten hours on a full charge."
def verify(spec: EmbedderSpec) -> tuple[int, str | None]:
"""Its real width, and why it is refused, or None. The probe's width wins
over the config's, which a model with a projection layer gets wrong."""
probed = spec.model_copy(update={"dimension": width(spec, PROBE)})
result = search_check(
lambda texts: _embed(probed, texts, Purpose.QUERY),
lambda texts: _embed(probed, texts, Purpose.DOCUMENT),
)
if not result.passed:
return probed.dimension, (
f"It found {result.first} of {result.asked} answers first; a search "
f"model has to find all {result.asked}."
)
return probed.dimension, None
def _embed(spec: EmbedderSpec, texts: list[str], purpose: Purpose) -> list[list[float]]:
return embed(spec, texts, purpose)

2
surfsense_local/backend/modules/llm/catalog/local/build.py
View file

@ -13,6 +13,8 @@ class FileRole(StrEnum):
# diffusion model, and several models use the same one.
VAE = "vae"
TEXT_ENCODER = "text_encoder"
# An ONNX embedder reads text through its own tokenizer, which no header carries.
TOKENIZER = "tokenizer"
@dataclass(frozen=True)

2
surfsense_local/backend/modules/llm/catalog/local/classifier.py
View file

@ -145,7 +145,7 @@ GROUPS: dict[str, Group] = {
**dict.fromkeys(
_EMBEDDING,
Group(
_NONE,
(ModelType.EMBEDDING,),
(
"This model turns text into numbers for search. It cannot answer "
"questions, and SurfSense already has its own."

3
surfsense_local/backend/modules/llm/catalog/local/engines/onnxruntime/__init__.py Normal file
View file

@ -0,0 +1,3 @@
"""onnxruntime: what the local catalog knows about embedding models."""
ENGINE = "onnxruntime"

129
surfsense_local/backend/modules/llm/catalog/local/engines/onnxruntime/engine.py Normal file
View file

@ -0,0 +1,129 @@
"""onnxruntime behind the engine seam: embedders, one folder each, run in process.
There is no server to configure: the encoder loads whichever model the active
index names. bge is shipped in the read-only models pack and never downloaded.
"""
import asyncio
import shutil
from collections.abc import AsyncIterator, Callable, Sequence
from pathlib import Path
from modules.embedding.bundled import BGE, bundled_dir
from modules.embedding.encoder import missing_files
from modules.embedding.spec import EmbedderSpec
from modules.embedding.verify import verify
from modules.llm.catalog.local.build import Build, BuildFile
from modules.llm.catalog.local.engines.engine import InstallStep
from modules.llm.catalog.local.engines.onnxruntime import ENGINE
from modules.llm.catalog.local.engines.onnxruntime.rows import (
downloaded_rows,
embedding_catalog,
)
from modules.llm.catalog.local.install.plan import InstallPlan
from modules.llm.catalog.local.installs import forget_install, read_installs
from modules.llm.catalog.local.manifest import CuratedModel
from modules.llm.catalog.local.rows import LocalRow
from modules.llm.model_type import ModelType
# Beside a Hugging Face pick's files: the spec its checks settled.
SPEC_FILE = "spec.json"
class OnnxRuntimeEngine:
name = ENGINE
model_types = (ModelType.EMBEDDING,)
# Never a selection: the index names the embedder, not `selected_models`.
provider = ENGINE
def __init__(self, folder: Path) -> None:
self._folder = folder
# Hugging Face picks opened this session, by install name: their checks
# run after the download, before the spec is kept beside the files.
self._offered: dict[str, EmbedderSpec] = {}
def offer(self, spec: EmbedderSpec) -> None:
self._offered[spec.id] = spec
def installed_spec(self, model_id: str) -> EmbedderSpec | None:
"""A downloaded Hugging Face pick's spec, as its checks left it."""
path = self._folder / model_id / SPEC_FILE
if not self.holds(model_id) or not path.is_file():
return None
return EmbedderSpec.model_validate_json(path.read_text())
@property
def folder(self) -> Path:
return self._folder
def rows(
self,
models: Sequence[CuratedModel],
mint: Callable[[Build], str],
*,
selected: str | None,
) -> tuple[LocalRow, ...]:
installs = read_installs(self._folder) if self._folder.is_dir() else {}
picked = [
spec
for model_id in installs
if (spec := self.installed_spec(model_id)) is not None
]
return tuple(
embedding_catalog(
models,
installs,
mint,
bundled={BGE.id} if self._ships_bge() else set(),
)
) + tuple(downloaded_rows(picked))
def landing(self, file: BuildFile, model_id: str) -> str:
# Every ONNX repo calls its weights model.onnx: the folder tells them apart.
return f"{model_id}/{file.name}"
def holds(self, model_id: str) -> bool:
if model_id == BGE.id and self._ships_bge():
return True
installs = read_installs(self._folder) if self._folder.is_dir() else {}
return model_id in installs and not installs[model_id].pending
def bundled(self, model_id: str) -> bool:
return model_id == BGE.id and self._ships_bge()
def _ships_bge(self) -> bool:
return bundled_dir().is_dir() and not missing_files(BGE)
async def check(self, plan: InstallPlan) -> InstallPlan:
return plan # curated only: checked when the manifest was refreshed
async def after_install(self, model_id: str) -> AsyncIterator[InstallStep]:
# Nothing to start: the encoder loads it when the index names it. A
# Hugging Face pick is checked first, since nobody measured it.
offered = self._offered.pop(model_id, None)
if offered is not None:
yield InstallStep("verifying", "Checking that it finds answers")
width, refusal = await asyncio.to_thread(verify, offered)
if refusal is not None:
self._discard(model_id)
yield InstallStep("error", refusal)
return
checked = offered.model_copy(update={"dimension": width})
(self._folder / model_id / SPEC_FILE).write_text(checked.model_dump_json())
yield InstallStep("complete", "Model is ready")
def _discard(self, model_id: str) -> None:
for name in forget_install(self._folder, model_id):
(self._folder / name).unlink(missing_ok=True)
shutil.rmtree(self._folder / model_id, ignore_errors=True)
def after_remove(self) -> None:
"""A removed model's folder still holds its spec: drop the folders no
install names any more."""
installs = read_installs(self._folder) if self._folder.is_dir() else {}
for folder in self._folder.iterdir() if self._folder.is_dir() else ():
if folder.is_dir() and folder.name not in installs:
shutil.rmtree(folder, ignore_errors=True)
def on_startup(self) -> None:
pass

54
surfsense_local/backend/modules/llm/catalog/local/engines/onnxruntime/manifest_fields.py Normal file
View file

@ -0,0 +1,54 @@
"""What a curated embedding model commits beyond the shared entry: what its
index's spec needs and no file states, the weight measured by the retrieval eval
among it."""
from pydantic import BaseModel, Field, model_validator
from modules.embedding.spec import Pooling
from modules.llm.catalog.local.manifest.strict import STRICT
# Cosine against the original model, per passage. granite-97m's int8 build
# measured 0.961 mean and failed; granite-311m's measured 0.990 and passed.
PARITY_MEAN = 0.99
PARITY_MIN = 0.98
class Parity(BaseModel):
"""How close a converted build's vectors are to the original model's.
Absent for a build that is the original, which matches it by definition.
"""
model_config = STRICT
origin: str = Field(min_length=1)
mean: float = Field(gt=0, le=1)
min: float = Field(gt=0, le=1)
@model_validator(mode="after")
def _close_enough(self) -> "Parity":
if self.mean < PARITY_MEAN or self.min < PARITY_MIN:
raise ValueError(
f"parity {self.mean:.3f} mean, {self.min:.3f} min is under "
f"{PARITY_MEAN} and {PARITY_MIN}: another vector space"
)
return self
class EmbeddingDefaults(BaseModel):
model_config = STRICT
dimension: int = Field(gt=0)
pooling: Pooling
normalize: bool
query_prefix: str = ""
document_prefix: str = ""
max_tokens: int = Field(gt=0)
semantic_weight: float = Field(gt=0, lt=1)
# Peak memory is mostly the batch, not the weights, so a larger model keeps it smaller.
batch: int = Field(gt=0)
parity: Parity | None = None
ENTRY_OWNS = frozenset({"embedding"})
ENTRY_REQUIRES = frozenset({"embedding"})

151
surfsense_local/backend/modules/llm/catalog/local/engines/onnxruntime/rows.py Normal file
View file

@ -0,0 +1,151 @@
"""The manifest's embedders as rows, unpriced: an embedder is small beside a chat
model, and its size is what the screen shows."""
from collections.abc import Callable, Collection, Mapping, Sequence
from dataclasses import replace
from modules.embedding.spec import EmbedderSpec
from modules.llm.catalog.local.build import Build, BuildFile, FileRole
from modules.llm.catalog.local.classifier import classify
from modules.llm.catalog.local.engines.onnxruntime import ENGINE
from modules.llm.catalog.local.engines.registry import engine_for
from modules.llm.catalog.local.installs import InstalledBuild
from modules.llm.catalog.local.manifest import CuratedModel
from modules.llm.catalog.local.rows import (
BuildRow,
Lead,
LeadReason,
LocalRow,
LocalSupport,
Origin,
)
def embedding_catalog(
models: Sequence[CuratedModel],
installs: Mapping[str, InstalledBuild],
catalog_id: Callable[[Build], str],
*,
bundled: Collection[str] = (),
) -> list[LocalRow]:
"""`installs` is keyed by catalog id, which is what a download is named;
`bundled` names the ones the app ships, installed from the first start."""
rows = []
for model in models:
classification = classify(
model.evidence.architecture, model.evidence.pipeline_tag
)
engine = engine_for(classification.types)
if engine is None or engine.name != ENGINE:
continue
installed = model.id in bundled or (
model.id in installs and not installs[model.id].pending
)
build_rows = tuple(
BuildRow(
catalog_id=catalog_id(build),
build=build,
fit=None,
badge=None,
installed_as=model.id if installed else None,
recommended=False,
reads_images=False,
projector_checked=False,
bundled=model.id in bundled,
)
for build in model.as_builds()
)
lead = build_rows[0].build.quantization
rows.append(
LocalRow(
id=model.id,
origin=Origin.CURATED,
name=model.name,
family=model.family,
classification=classification,
support=LocalSupport(
context=None, reads_images=False, tools=None, reasoning=None
),
builds=build_rows,
default_quantization=lead,
recommended=False,
engine=ENGINE,
lead=Lead(
lead, LeadReason.INSTALLED if installed else LeadReason.DEFAULT
),
description=model.description,
)
)
return rows
def downloaded_rows(specs: Sequence[EmbedderSpec]) -> list[LocalRow]:
"""Hugging Face picks on disk: installed, with nothing left to download."""
classification = classify("", "feature-extraction")
rows = []
for spec in specs:
build = Build(
"ONNX",
(
BuildFile(FileRole.WEIGHTS, spec.weights.path, 0, spec.weights.sha256),
BuildFile(
FileRole.TOKENIZER, spec.tokenizer.path, 0, spec.tokenizer.sha256
),
),
)
rows.append(
LocalRow(
id=spec.id,
origin=Origin.DOWNLOADED,
name=spec.repo,
family="Hugging Face",
classification=classification,
support=LocalSupport(
context=None, reads_images=False, tools=None, reasoning=None
),
builds=(
BuildRow(
catalog_id="",
build=build,
fit=None,
badge=None,
installed_as=spec.id,
recommended=False,
reads_images=False,
projector_checked=False,
),
),
default_quantization="ONNX",
recommended=False,
engine=ENGINE,
lead=Lead("ONNX", LeadReason.INSTALLED),
)
)
return rows
def searched_row(repo: str, build: BuildRow | None, reason: str | None) -> LocalRow:
"""One opened Hugging Face repo, in the shape the GGUF search answers: its
one build, or none and why."""
classification = classify("", "feature-extraction")
if build is None:
# No build to offer: classified as nothing, so the row says why.
classification = replace(
classification,
types=(),
reason=reason or "SurfSense cannot run this model.",
)
return LocalRow(
id=repo,
origin=Origin.SEARCH,
name=repo,
family="Hugging Face",
classification=classification,
support=LocalSupport(
context=None, reads_images=False, tools=None, reasoning=None
),
builds=(build,) if build else (),
default_quantization=build.build.quantization if build else None,
recommended=False,
engine=ENGINE,
)

33
surfsense_local/backend/modules/llm/catalog/local/engines/onnxruntime/spec.py Normal file
View file

@ -0,0 +1,33 @@
"""A curated embedder as the spec an index stores a snapshot of."""
from modules.embedding.spec import EmbedderSpec, Identified, PinnedFile, Source
from modules.llm.catalog.local.build import FileRole
from modules.llm.catalog.local.manifest import CuratedModel
def spec_for(model: CuratedModel) -> EmbedderSpec:
"""Its default build: the first, which is the only one a curated embedder
lists, since another build would be another vector space."""
if model.embedding is None:
raise ValueError(f"{model.id} is not an embedding model")
build = model.as_builds()[0]
tokenizer = next(f for f in build.files if f.role is FileRole.TOKENIZER)
embedding = model.embedding
return EmbedderSpec(
id=model.id,
source=Source.CURATED,
identified=Identified.MEASURED,
repo=build.weights.repo,
revision=build.weights.revision,
# Named as they land: one folder per model, the upstream path dropped.
weights=PinnedFile(path=build.weights.name, sha256=build.weights.sha256),
tokenizer=PinnedFile(path=tokenizer.name, sha256=tokenizer.sha256),
dimension=embedding.dimension,
pooling=embedding.pooling,
normalize=embedding.normalize,
query_prefix=embedding.query_prefix,
document_prefix=embedding.document_prefix,
max_tokens=embedding.max_tokens,
semantic_weight=embedding.semantic_weight,
batch=embedding.batch,
)

15
surfsense_local/backend/modules/llm/catalog/local/engines/registry.py
View file

@ -4,13 +4,17 @@ manifest entry."""
from collections.abc import Iterable
from dataclasses import dataclass
from modules.llm.catalog.local.engines import audiocpp, sdcpp
from modules.llm.catalog.local.build import FileRole
from modules.llm.catalog.local.engines import audiocpp, onnxruntime, sdcpp
from modules.llm.catalog.local.engines.audiocpp import (
manifest_fields as audiocpp_fields,
)
from modules.llm.catalog.local.engines.llamacpp import (
manifest_fields as llamacpp_fields,
)
from modules.llm.catalog.local.engines.onnxruntime import (
manifest_fields as onnxruntime_fields,
)
from modules.llm.catalog.local.engines.sdcpp import manifest_fields as sdcpp_fields
from modules.llm.model_type import ModelType
@ -24,6 +28,8 @@ class Engine:
entry_requires: frozenset[str]
# Exactly one of these must be present, where the engine has such a choice.
entry_requires_one_of: frozenset[str] = frozenset()
# File roles every build must carry beside its weights.
build_requires: frozenset[FileRole] = frozenset()
ENGINES = (
@ -46,6 +52,13 @@ ENGINES = (
audiocpp_fields.ENTRY_OWNS,
audiocpp_fields.ENTRY_REQUIRES,
),
Engine(
onnxruntime.ENGINE,
(ModelType.EMBEDDING,),
onnxruntime_fields.ENTRY_OWNS,
onnxruntime_fields.ENTRY_REQUIRES,
build_requires=frozenset({FileRole.TOKENIZER}),
),
)
ENGINE_ENTRY_FIELDS = frozenset().union(*(e.entry_owns for e in ENGINES))

18
surfsense_local/backend/modules/llm/catalog/local/install/tickets.py
View file

@ -21,6 +21,10 @@ class Ticket:
# The repo's own tag, which the exact check reads beside the header.
pipeline_tag: str | None
issued_at: float
# Search reaches llama.cpp's catalog, and onnxruntime's for an embedder,
# which installs under a name of its own rather than its weights' file.
engine: str = "llamacpp"
model_id: str | None = None
class TicketStore:
@ -29,12 +33,22 @@ class TicketStore:
self._tickets: dict[str, Ticket] = {}
def mint(
self, build: Build, *, pipeline_tag: str | None = None, now: float | None = None
self,
build: Build,
*,
pipeline_tag: str | None = None,
now: float | None = None,
engine: str = "llamacpp",
model_id: str | None = None,
) -> str:
self._expire(now)
token = secrets.token_urlsafe(18)
self._tickets[token] = Ticket(
build, pipeline_tag, now if now is not None else time.monotonic()
build,
pipeline_tag,
now if now is not None else time.monotonic(),
engine,
model_id,
)
return token

4
surfsense_local/backend/modules/llm/catalog/local/install_jobs/steps.py
View file

@ -40,6 +40,10 @@ async def install_steps(
if step.kind == "complete":
ready = step.message
continue
# The model is on disk and the engine refused it: say why, and stop.
if step.kind == "error":
yield {"type": "error", "message": step.message}
return
yield {"type": step.kind, "message": step.message, "progress": step.progress}
selection = None
if select:

163
surfsense_local/backend/modules/llm/catalog/local/manifest/models.json
View file

@ -1,6 +1,6 @@
{
"schema_version": 1,
"refreshed_at": "2026-09-25",
"refreshed_at": "2026-10-01",
"models": [
{
"id": "qwen3-32b",
@ -4549,6 +4549,167 @@
}
}
]
},
{
"id": "bge-small-en-v1.5",
"name": "BGE Small (English)",
"family": "BGE",
"publisher": "BAAI",
"description": "Fast search in English. Comes with SurfSense.",
"license": "mit",
"source_repo": "BAAI/bge-small-en-v1.5",
"aliases": [
"BAAI/bge-small-en-v1.5"
],
"evidence": {
"architecture": "bert",
"pipeline_tag": "sentence-similarity"
},
"embedding": {
"dimension": 384,
"pooling": "cls",
"normalize": true,
"max_tokens": 512,
"semantic_weight": 0.65,
"batch": 32
},
"builds": [
{
"quantization": "F16",
"files": [
{
"role": "weights",
"repo": "Qdrant/bge-small-en-v1.5-onnx-Q",
"upstream_repo": null,
"revision": "aa8f8b060edb00e03bfdd08813a2949946c8ba55",
"path": "model_optimized.onnx",
"size_bytes": 66465124,
"sha256": "51f1bd0addd6e859e42c2c8021a5e5461385bb676a649f4b269aa445449f2431"
},
{
"role": "tokenizer",
"repo": "Qdrant/bge-small-en-v1.5-onnx-Q",
"upstream_repo": null,
"revision": "aa8f8b060edb00e03bfdd08813a2949946c8ba55",
"path": "tokenizer.json",
"size_bytes": 711396,
"sha256": "d241a60d5e8f04cc1b2b3e9ef7a4921b27bf526d9f6050ab90f9267a1f9e5c66"
}
],
"validated": {
"onnxruntime": "1.29.0"
}
}
]
},
{
"id": "granite-embedding-97m-multilingual-r2",
"name": "Granite Embedding 97M (Multilingual)",
"family": "Granite Embedding",
"publisher": "IBM",
"description": "Search across more than 50 languages.",
"license": "apache-2.0",
"source_repo": "ibm-granite/granite-embedding-97m-multilingual-r2",
"aliases": [
"ibm-granite/granite-embedding-97m-multilingual-r2"
],
"evidence": {
"architecture": "modernbert",
"pipeline_tag": "feature-extraction"
},
"embedding": {
"dimension": 384,
"pooling": "cls",
"normalize": true,
"max_tokens": 512,
"semantic_weight": 0.85,
"batch": 32
},
"builds": [
{
"quantization": "F32",
"files": [
{
"role": "weights",
"repo": "onnx-community/granite-embedding-97m-multilingual-r2-ONNX",
"upstream_repo": "ibm-granite/granite-embedding-97m-multilingual-r2",
"revision": "536a9f241cb3f02a9c5995a1e708c784bd274859",
"path": "onnx/model.onnx",
"size_bytes": 390004608,
"sha256": "68e592b160673d30250824c1116bc6ab33f70efb22b97c9e1d7ce1e69c1c9d70"
},
{
"role": "tokenizer",
"repo": "onnx-community/granite-embedding-97m-multilingual-r2-ONNX",
"upstream_repo": "ibm-granite/granite-embedding-97m-multilingual-r2",
"revision": "536a9f241cb3f02a9c5995a1e708c784bd274859",
"path": "tokenizer.json",
"size_bytes": 25301671,
"sha256": "51947676cae1f991fa51c6b9a24e14ee5460e5f0b9f692f13bb3159829d1592a"
}
],
"validated": {
"onnxruntime": "1.29.0"
}
}
]
},
{
"id": "granite-embedding-311m-multilingual-r2",
"name": "Granite Embedding 311M (Multilingual)",
"family": "Granite Embedding",
"publisher": "IBM",
"description": "The larger multilingual model, for machines with more memory.",
"license": "apache-2.0",
"source_repo": "ibm-granite/granite-embedding-311m-multilingual-r2",
"aliases": [
"ibm-granite/granite-embedding-311m-multilingual-r2"
],
"evidence": {
"architecture": "modernbert",
"pipeline_tag": "feature-extraction"
},
"embedding": {
"dimension": 768,
"pooling": "cls",
"normalize": true,
"max_tokens": 512,
"semantic_weight": 0.85,
"batch": 16,
"parity": {
"origin": "retrieval eval corpus, 23 passages, against onnx/model.onnx",
"mean": 0.992,
"min": 0.987
}
},
"builds": [
{
"quantization": "INT8",
"files": [
{
"role": "weights",
"repo": "onnx-community/granite-embedding-311m-multilingual-r2-ONNX",
"upstream_repo": "ibm-granite/granite-embedding-311m-multilingual-r2",
"revision": "8f039f21d4181327268271bea4b11ddcc7eef88d",
"path": "onnx/model_int8.onnx",
"size_bytes": 312556945,
"sha256": "54d33d10f865516eda6770d7b0bbf5eaece861058907da0ad095c62e52958c3d"
},
{
"role": "tokenizer",
"repo": "onnx-community/granite-embedding-311m-multilingual-r2-ONNX",
"upstream_repo": "ibm-granite/granite-embedding-311m-multilingual-r2",
"revision": "8f039f21d4181327268271bea4b11ddcc7eef88d",
"path": "tokenizer.json",
"size_bytes": 33384821,
"sha256": "0087c868b33bad550a78a08d19798cfd7f713cde4f020803b8f51f405503e15f"
}
],
"validated": {
"onnxruntime": "1.29.0"
}
}
]
}
]
}

12
surfsense_local/backend/modules/llm/catalog/local/manifest/schema.py
View file

@ -18,6 +18,9 @@ from modules.llm.catalog.local.engines.llamacpp.manifest_fields import (
ShapeSpec,
Template,
)
from modules.llm.catalog.local.engines.onnxruntime.manifest_fields import (
EmbeddingDefaults,
)
from modules.llm.catalog.local.engines.registry import ENGINE_ENTRY_FIELDS, engine_for
from modules.llm.catalog.local.engines.sdcpp.manifest_fields import (
ImageDefaults,
@ -70,6 +73,7 @@ class Validated(BaseModel):
llama_cpp: str | None = None
sd_cpp: str | None = None
audio_cpp: str | None = None
onnxruntime: str | None = None
class ManifestBuild(BaseModel):
@ -122,6 +126,7 @@ class CuratedModel(BaseModel):
image: ImageDefaults | None = None
video: VideoDefaults | None = None
audio: AudioDefaults | None = None
embedding: EmbeddingDefaults | None = None
# Text models only: an image model is not priced by the llama.cpp estimator.
shape: ShapeSpec | None = None
builds: list[ManifestBuild] = Field(min_length=1)
@ -146,6 +151,13 @@ class CuratedModel(BaseModel):
choice = engine.entry_requires_one_of
if choice and len(present & choice) != 1:
raise ValueError(f"{self.id}: {engine.name} needs one of {sorted(choice)}")
for build in self.builds:
roles = {f.role for f in build.files}
if missing_roles := engine.build_requires - roles:
raise ValueError(
f"{self.id} {build.quantization}: {engine.name} needs "
f"{sorted(missing_roles)}"
)
return self
@model_validator(mode="after")

8
surfsense_local/backend/modules/llm/catalog/local/router.py
View file

@ -68,7 +68,7 @@ def read_local_catalog(service: LocalCatalogDep, session: SessionDep) -> dict:
return {
"budget": _budget(catalog.budget),
"gpu_status": catalog.gpu_status.value,
"rows": [_row(row, in_use) for row in catalog.rows],
"rows": [row_read(row, in_use) for row in catalog.rows],
"recommended_id": catalog.recommended_id,
"projector_notices": [
{
@ -123,7 +123,7 @@ async def read_repo(repo: str, service: LocalCatalogDep, session: SessionDep) ->
row, gated = await service.repo(repo)
except httpx.HTTPError as error:
raise HTTPException(status.HTTP_503_SERVICE_UNAVAILABLE, UNREACHABLE) from error
return {"repo": repo, "gated": gated, "row": _row(row, in_use={})}
return {"repo": repo, "gated": gated, "row": row_read(row, in_use={})}
def _selected_local(
@ -147,7 +147,8 @@ def _budget(budget) -> dict:
}
def _row(row: LocalRow, in_use: Mapping[str, list[ModelType]]) -> dict:
def row_read(row: LocalRow, in_use: Mapping[str, list[ModelType]]) -> dict:
"""A row as every catalog route answers it, searched repos included."""
classification = row.classification
return {
"id": row.id,
@ -170,6 +171,7 @@ def _row(row: LocalRow, in_use: Mapping[str, list[ModelType]]) -> dict:
"default_quantization": row.default_quantization,
"recommended": row.recommended,
"engine": row.engine,
"description": row.description,
"lead": (
{"quantization": row.lead.quantization, "why": row.lead.why.value}
if row.lead

3
surfsense_local/backend/modules/llm/catalog/local/rows.py
View file

@ -105,6 +105,9 @@ class LocalRow:
lead: Lead | None = None
# An audio model's only.
voicing: Voicing | None = None
# What the model is for, as its entry says; an embedder's only, which
# onboarding lists without a chat model's fit and badges to tell them apart.
description: str | None = None
@property
def runnable(self) -> bool:

5
surfsense_local/backend/modules/llm/catalog/local/schemas.py
View file

@ -125,8 +125,11 @@ class LocalRowRead(BaseModel):
builds: list[BuildRead]
default_quantization: str | None
recommended: bool
# The engine that offered the row and would run it: llamacpp, sdcpp or audiocpp.
# The engine that offered the row and would run it: llamacpp, sdcpp,
# audiocpp or onnxruntime.
engine: str
# What the model is for; an embedder's only.
description: str | None = None
# Absent for a searched repo, which lists every build and leads with none.
lead: LeadRead | None
# An audio model's memory while voicing, voices and languages; null otherwise.

35
surfsense_local/backend/modules/llm/catalog/local/service.py
View file

@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from modules.embedding.spec import EmbedderSpec
from modules.llm.catalog.local.build import Build
from modules.llm.catalog.local.engines.audiocpp.audio_folder.espeak import Espeak
from modules.llm.catalog.local.engines.audiocpp.engine import AudioCppEngine
@ -17,6 +18,7 @@ from modules.llm.catalog.local.engines.llamacpp.models_folder.scan import (
ProjectorNotice,
)
from modules.llm.catalog.local.engines.llamacpp.sampling import publisher_temperature
from modules.llm.catalog.local.engines.onnxruntime.engine import OnnxRuntimeEngine
from modules.llm.catalog.local.engines.sdcpp.engine import SdCppEngine
from modules.llm.catalog.local.install import download
from modules.llm.catalog.local.install.disk_room import refuse_without_room
@ -41,6 +43,7 @@ from modules.llm.hardware import (
from modules.llm.hardware.inventory import OsGpu, Probe
from modules.llm.model_type import ModelType
from modules.llm.providers.types import DownloadProgress
from shared.config import get_storage_settings
@dataclass(frozen=True)
@ -67,6 +70,7 @@ class LocalCatalogService:
audio_dir: Path | None = None,
audio_bundled_dir: Path | None = None,
audio_espeak: Espeak | None = None,
embeddings_dir: Path | None = None,
probe: Probe = probe_devices,
os_gpu: OsGpu = os_reports_gpu,
) -> None:
@ -87,6 +91,14 @@ class LocalCatalogService:
for model in manifest.models
for build in model.as_builds()
}
# An embedder installs under its catalog id: every ONNX repo names its
# weights model.onnx, so the file cannot name the install as a GGUF does.
self._curated_entries: dict[tuple[str, str], str] = {
(build.weights.repo, build.weights.path): model.id
for model in manifest.models
if model.embedding is not None
for build in model.as_builds()
}
# One pull at a time: two downloads compete for one disk and one bar.
self._install_lock = asyncio.Lock()
self._install_jobs = InstallJobs(self._install_lock)
@ -98,10 +110,14 @@ class LocalCatalogService:
bundled_dir=audio_bundled_dir,
espeak=audio_espeak,
)
self.onnxruntime = OnnxRuntimeEngine(
embeddings_dir or get_storage_settings().embedding_models_dir
)
self._engines: tuple[LocalEngine, ...] = (
self.llamacpp,
self.sdcpp,
self.audiocpp,
self.onnxruntime,
)
def install_lock(self) -> asyncio.Lock:
@ -203,20 +219,27 @@ class LocalCatalogService:
curated = self._curated_ids.get(catalog_id)
if curated is not None:
build, engine = curated
shape = self._curated_shapes.get((build.weights.repo, build.weights.path))
return InstallPlan(build.runtime_name, build, engine, shape=shape)
key = (build.weights.repo, build.weights.path)
shape = self._curated_shapes.get(key)
name = self._curated_entries.get(key, build.runtime_name)
return InstallPlan(name, build, engine, shape=shape)
ticket = self._tickets.resolve(catalog_id)
if ticket is None:
return None
# Search reaches llama.cpp's catalog only.
return InstallPlan(
ticket.build.runtime_name,
ticket.model_id or ticket.build.runtime_name,
ticket.build,
self.llamacpp.name,
needs_check=True,
ticket.engine,
needs_check=ticket.engine == self.llamacpp.name,
pipeline_tag=ticket.pipeline_tag,
)
def offer_embedder(self, build: Build, spec: EmbedderSpec) -> str:
"""An install id for a Hugging Face embedder, its spec held for the
checks after its download."""
self.onnxruntime.offer(spec)
return self._tickets.mint(build, engine=self.onnxruntime.name, model_id=spec.id)
async def check(self, plan: InstallPlan) -> InstallPlan:
"""The engine's refusals first, then the disk's."""
engine = self.engine(plan.engine)

5
surfsense_local/backend/modules/llm/model_type.py
View file

@ -7,6 +7,10 @@ class ModelType(StrEnum):
The user picks one model per type, so a type is also a slot. A type no
feature reads yet can still be chosen; the feature that first reads it
brings the client that calls it.
EMBEDDING is the exception: a catalog type, so the manifest and an engine can
describe it, never a slot. The embedder belongs to the index, which onboarding
fixes once; a selection would offer the swap that breaks search.
"""
TEXT_GEN = "text_gen"
@ -14,3 +18,4 @@ class ModelType(StrEnum):
IMAGE_EDIT = "image_edit"
VIDEO_GEN = "video_gen"
AUDIO_GEN = "audio_gen"
EMBEDDING = "embedding"

40
surfsense_local/backend/modules/llm/router.py
View file

@ -1,12 +1,14 @@
from datetime import UTC, datetime
from typing import Annotated
from fastapi import APIRouter, BackgroundTasks, HTTPException, status
from fastapi import APIRouter, BackgroundTasks, Body, Depends, HTTPException, status
from sqlalchemy import select
from sqlalchemy.orm import Session
from api.dependencies import SessionDep, transact
from modules.artifacts.local_image_demand import local_image_demand
from modules.documents.models import Document, DocumentStatus, DocumentType
from modules.embedding.active import active_index
from modules.llm.activity import ModelBusyError, model_activity, model_key
from modules.llm.catalog.local.dependencies import LocalCatalogDep
from modules.llm.catalog.local.install_jobs.router import router as install_jobs_router
@ -27,13 +29,14 @@ from modules.llm.schemas import (
LocalImageRuntimeRead,
ModelDeleteRead,
ModelRead,
OnboardingComplete,
OnboardingStatusRead,
ProviderRead,
RuntimeFileRead,
SelectionRead,
SelectionWrite,
)
from modules.llm.selectable import selectable_for
from modules.llm.selectable import SLOTS, selectable_for
from modules.llm.selection import choose_model, complete_onboarding
from modules.llm.voices.router import router as voices_router
from shared.config import get_llm_settings
@ -62,8 +65,10 @@ def read_onboarding_status(session: SessionDep) -> OnboardingStatusRead:
response_model=OnboardingStatusRead,
summary="Complete model onboarding",
)
def mark_onboarding_complete(session: SessionDep) -> OnboardingStatusRead:
complete_onboarding(session)
def mark_onboarding_complete(
session: SessionDep, payload: Annotated[OnboardingComplete | None, Body()] = None
) -> OnboardingStatusRead:
complete_onboarding(session, payload.embedding_model if payload else None)
return OnboardingStatusRead(completed=True)
@ -129,6 +134,11 @@ async def delete_model(
status.HTTP_409_CONFLICT,
f"{model_name} comes with SurfSense and cannot be deleted",
)
if await transact(session, _embeds_the_library, model_name):
raise HTTPException(
status.HTTP_409_CONFLICT,
f"{model_name} made every vector in the library and cannot be deleted",
)
if await transact(session, _studio_running):
raise HTTPException(
status.HTTP_409_CONFLICT,
@ -166,6 +176,11 @@ async def delete_model(
return ModelDeleteRead(name=model_name, selection_cleared=selection_cleared)
def _embeds_the_library(session: Session, model_name: str) -> bool:
active = active_index(session)
return active is not None and active.spec.id == model_name
def _studio_running(session: Session) -> bool:
# ponytail: Studio does not persist the model used by each job, so block all
# local deletes while one runs. Record provider/model per job to narrow this.
@ -223,12 +238,25 @@ def read_local_image_runtime(
)
def _slot(model_type: ModelType) -> ModelType:
"""The embedder is fixed with the index at onboarding, never selected."""
if model_type not in SLOTS:
raise HTTPException(
status.HTTP_422_UNPROCESSABLE_CONTENT,
"the embedding model is chosen during onboarding, not selected",
)
return model_type
SlotDep = Annotated[ModelType, Depends(_slot)]
@router.get(
"/selection/{model_type}",
response_model=SelectionRead,
summary="Read the model chosen for a model type",
)
async def read_selection(model_type: ModelType, session: SessionDep) -> SelectionRead:
async def read_selection(model_type: SlotDep, session: SessionDep) -> SelectionRead:
return await _selection_read(
session, await transact(session, _chosen_or_404, model_type)
)
@ -249,7 +277,7 @@ def _chosen_or_404(session: Session, model_type: ModelType) -> SelectedModel:
summary="Choose the model for a model type",
)
async def set_selection(
model_type: ModelType,
model_type: SlotDep,
payload: SelectionWrite,
session: SessionDep,
background: BackgroundTasks,

5
surfsense_local/backend/modules/llm/schemas.py
View file

@ -127,6 +127,11 @@ class OnboardingStatusRead(BaseModel):
completed: bool
class OnboardingComplete(BaseModel):
# A curated embedder already downloaded; none means the bundled one.
embedding_model: str | None = Field(default=None, min_length=1, max_length=200)
class InstallRequest(BaseModel):
catalog_id: str = Field(min_length=1, max_length=128)
select: bool = True

7
surfsense_local/backend/modules/llm/selectable.py
View file

@ -2,6 +2,9 @@ from collections.abc import Iterable
from modules.llm.model_type import ModelType
# Every type but the embedder, which belongs to the index rather than a selection.
SLOTS = tuple(t for t in ModelType if t is not ModelType.EMBEDDING)
def selectable_for(types: Iterable[ModelType], known: bool) -> list[ModelType]:
"""The slots a model can fill, and the one rule selection and every picker share.
@ -11,6 +14,6 @@ def selectable_for(types: Iterable[ModelType], known: bool) -> list[ModelType]:
trusting it with a slot.
"""
if not known:
return list(ModelType)
return list(SLOTS)
held = set(types)
return [model_type for model_type in ModelType if model_type in held]
return [slot for slot in SLOTS if slot in held]

4
surfsense_local/backend/modules/llm/selection.py
View file

@ -5,6 +5,7 @@ from fastapi import HTTPException, status
from sqlalchemy.orm import Session
from api.dependencies import transact
from modules.embedding.choose import lock_chosen
from modules.llm.catalog.local.dependencies import get_local_catalog
from modules.llm.connections import discover_models
from modules.llm.connections.router import allowed_connection
@ -129,12 +130,13 @@ def _store(
return selected
def complete_onboarding(session: Session) -> bool:
def complete_onboarding(session: Session, embedding_model: str | None = None) -> bool:
if session.get(SelectedModel, ModelType.TEXT_GEN) is None:
raise HTTPException(
status.HTTP_422_UNPROCESSABLE_CONTENT,
"chat model required",
)
lock_chosen(session, embedding_model)
if session.get(OnboardingCompletion, 1) is None:
session.add(OnboardingCompletion())
session.flush()

2
surfsense_local/backend/modules/migration/router.py
View file

@ -12,6 +12,7 @@ from pydantic import ValidationError
from api.dependencies import SessionDep
from modules.documents.storage import stream_upload
from modules.embedding.dependencies import EMBEDDER_CHOSEN
from modules.migration.schemas import ImportAccepted, ImportedWorkspace, Manifest
from modules.migration.service import find_or_create_workspaces, import_bundle
from shared.config import get_storage_settings
@ -27,6 +28,7 @@ MAX_BUNDLE_UNPACKED_BYTES = 5 * 1024 * 1024 * 1024
@router.post(
"/import",
dependencies=[EMBEDDER_CHOSEN],
response_model=ImportAccepted,
status_code=status.HTTP_202_ACCEPTED,
summary="Import a SurfSense cloud export bundle",

49
surfsense_local/backend/scripts/fetch_embedding_model.py
View file

@ -3,43 +3,52 @@
`uv run scripts/fetch_embedding_model.py` places it where the app reads it in
development; pass a models root (`... models`) to stage it for an installer,
which electron-builder then copies into resources/models.
Fetched at bge's pinned revision and checked against its hashes: every existing
library was built with that exact file, so a different one is refused, not used.
"""
import hashlib
import sys
from pathlib import Path
import httpx
from worker.ingestion.embedding import (
EMBEDDING_FILES,
EMBEDDING_REPO,
MODEL_DIR_NAME,
embedding_dir,
)
from modules.embedding.bundled import BGE, bundled_dir
BASE = "https://huggingface.co"
def fetch(repo: str, files: list[str], into: Path) -> None:
def fetch(into: Path) -> None:
into.mkdir(parents=True, exist_ok=True)
with httpx.Client(follow_redirects=True, timeout=120.0) as client:
for name in files:
target = into / name
if target.exists():
print(f"have {target}")
continue
print(f"get {name}")
with client.stream("GET", f"{BASE}/{repo}/resolve/main/{name}") as reply:
reply.raise_for_status()
with target.open("wb") as sink:
for block in reply.iter_bytes():
sink.write(block)
for pinned in (BGE.weights, BGE.tokenizer):
target = into / pinned.path
if not target.exists():
print(f"get {pinned.path}")
url = f"{BASE}/{BGE.repo}/resolve/{BGE.revision}/{pinned.path}"
with client.stream("GET", url) as reply:
reply.raise_for_status()
with target.open("wb") as sink:
for block in reply.iter_bytes():
sink.write(block)
if _sha256(target) != pinned.sha256:
target.unlink()
raise SystemExit(f"{pinned.path} does not match bge's pinned hash")
print(f"have {target}")
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for block in iter(lambda: source.read(1 << 20), b""):
digest.update(block)
return digest.hexdigest()
def main() -> int:
# Optional models root for staging an installer; defaults to the dev location.
into = Path(sys.argv[1]) / MODEL_DIR_NAME if len(sys.argv) > 1 else embedding_dir()
fetch(EMBEDDING_REPO, EMBEDDING_FILES, into)
fetch(Path(sys.argv[1]) / BGE.id if len(sys.argv) > 1 else bundled_dir())
return 0

74
surfsense_local/backend/scripts/local_manifest/entries.py
View file

@ -9,6 +9,7 @@ in the written manifest is read from the files.
from local_manifest.audiocpp.entry import AudioEntry
from local_manifest.entry import Entry
from local_manifest.onnxruntime.entry import EmbeddingEntry
from local_manifest.sdcpp.entry import Companion, ImageEntry, VideoEntry
@ -446,4 +447,77 @@ ENTRIES: tuple[Entry, ...] = (
],
},
),
# Embedders: bge ships with the app and is the default; the rest are chosen
# once, at onboarding. Each semantic_weight is the middle of the retrieval
# eval's plateau for that model (1 Oct 2026, 266 queries): both granites
# peak at 0.85, where cross-lingual is 8/8 and 7/8 and LIMIT stays 200/200;
# 0.90 starts costing LIMIT.
EmbeddingEntry(
id="bge-small-en-v1.5",
name="BGE Small (English)",
family="BGE",
publisher="BAAI",
description="Fast search in English. Comes with SurfSense.",
license="mit",
source_repo="BAAI/bge-small-en-v1.5",
repo="Qdrant/bge-small-en-v1.5-onnx-Q",
aliases=("BAAI/bge-small-en-v1.5",),
builds=(("F16", "model_optimized.onnx"),),
embedding={
"dimension": 384,
"pooling": "cls",
"normalize": True,
"max_tokens": 512,
"semantic_weight": 0.65,
"batch": 32,
},
),
EmbeddingEntry(
id="granite-embedding-97m-multilingual-r2",
name="Granite Embedding 97M (Multilingual)",
family="Granite Embedding",
publisher="IBM",
description="Search across more than 50 languages.",
license="apache-2.0",
source_repo="ibm-granite/granite-embedding-97m-multilingual-r2",
repo="onnx-community/granite-embedding-97m-multilingual-r2-ONNX",
upstream_repo="ibm-granite/granite-embedding-97m-multilingual-r2",
aliases=("ibm-granite/granite-embedding-97m-multilingual-r2",),
# Full precision: its int8 build measured 0.961 against it and failed parity.
builds=(("F32", "onnx/model.onnx"),),
embedding={
"dimension": 384,
"pooling": "cls",
"normalize": True,
"max_tokens": 512,
"semantic_weight": 0.85,
"batch": 32,
},
),
EmbeddingEntry(
id="granite-embedding-311m-multilingual-r2",
name="Granite Embedding 311M (Multilingual)",
family="Granite Embedding",
publisher="IBM",
description="The larger multilingual model, for machines with more memory.",
license="apache-2.0",
source_repo="ibm-granite/granite-embedding-311m-multilingual-r2",
repo="onnx-community/granite-embedding-311m-multilingual-r2-ONNX",
upstream_repo="ibm-granite/granite-embedding-311m-multilingual-r2",
aliases=("ibm-granite/granite-embedding-311m-multilingual-r2",),
builds=(("INT8", "onnx/model_int8.onnx"),),
embedding={
"dimension": 768,
"pooling": "cls",
"normalize": True,
"max_tokens": 512,
"semantic_weight": 0.85,
"batch": 16,
"parity": {
"origin": "retrieval eval corpus, 23 passages, against onnx/model.onnx",
"mean": 0.992,
"min": 0.987,
},
},
),
)

0
surfsense_local/backend/scripts/local_manifest/onnxruntime/__init__.py Normal file
View file

86
surfsense_local/backend/scripts/local_manifest/onnxruntime/assemble.py Normal file
View file

@ -0,0 +1,86 @@
"""One embedding model's manifest entry from its repo at a commit.
Pure, like the other engines' assemblies: every rule is a unit test over
recorded input.
"""
from typing import Any
from local_manifest.onnxruntime.entry import EmbeddingEntry
from local_manifest.recorded import RepoAtRevision
from local_manifest.unreadable import UnreadableBuildError
from modules.llm.catalog.local.build import Build, BuildFile, FileRole
def pinned_builds(entry: EmbeddingEntry, snapshot: RepoAtRevision) -> list[Build]:
"""Each build's weights with the repo's tokenizer, every file hashed."""
listed = {f.path: f for f in snapshot.listing}
builds = []
for label, weights in entry.builds:
files = []
for role, path in (
(FileRole.WEIGHTS, weights),
(FileRole.TOKENIZER, entry.tokenizer),
):
found = listed.get(path)
if found is None or not found.sha256 or found.size_bytes <= 0:
raise UnreadableBuildError(
f"{snapshot.repo}: {path} has no hash or size"
)
files.append(
BuildFile(
role,
path,
found.size_bytes,
found.sha256,
snapshot.repo,
snapshot.revision,
)
)
builds.append(Build(label, tuple(files)))
return builds
def entry_for(
entry: EmbeddingEntry,
snapshot: RepoAtRevision,
builds: list[Build],
architecture: str,
validated: dict[str, str] | None = None,
) -> dict[str, Any]:
"""The written entry. `architecture` is the repo config's `model_type`."""
validated = validated or {}
return {
"id": entry.id,
"name": entry.name,
"family": entry.family,
"publisher": entry.publisher,
"description": entry.description,
"license": entry.license,
"source_repo": entry.source_repo,
"aliases": list(entry.aliases),
"evidence": {
"architecture": architecture,
"pipeline_tag": snapshot.pipeline_tag,
},
"embedding": dict(entry.embedding),
"builds": [
{
"quantization": build.quantization,
"files": [
{
"role": file.role.value,
"repo": file.repo,
"upstream_repo": entry.upstream_repo,
"revision": file.revision,
"path": file.path,
"size_bytes": file.size_bytes,
"sha256": file.sha256,
}
for file in build.files
],
"validated": {"onnxruntime": validated.get(build.weights.path)},
}
for build in builds
],
}

20
surfsense_local/backend/scripts/local_manifest/onnxruntime/entry.py Normal file
View file

@ -0,0 +1,20 @@
"""The hand-authored part of an embedding model's entry.
Beyond what every entry names, a person reviews what no file states: the width,
pooling and prefixes the index's spec needs, the ranking weight the retrieval
eval measured, and how close a converted build's vectors are to the original.
"""
from dataclasses import dataclass, field
from typing import Any
from local_manifest.entry import Entry
@dataclass(frozen=True)
class EmbeddingEntry(Entry):
# (label, weights path) to pin, most preferred first: the first is the default.
builds: tuple[tuple[str, str], ...] = ()
tokenizer: str = "tokenizer.json"
# `EmbeddingDefaults` as the manifest writes it.
embedding: dict[str, Any] = field(default_factory=dict)

61
surfsense_local/backend/scripts/local_manifest/onnxruntime/refresh.py Normal file
View file

@ -0,0 +1,61 @@
"""Read one embedding model's repo at its commit and write its entry."""
import hashlib
from dataclasses import replace
from typing import Any
import httpx
from local_manifest.hub import RESOLVE, repo_at_revision
from local_manifest.onnxruntime.assemble import entry_for, pinned_builds
from local_manifest.onnxruntime.entry import EmbeddingEntry
from local_manifest.unreadable import UnreadableBuildError
# weights path -> the onnxruntime a person ran the retrieval eval with
VALIDATED: dict[str, str] = {
"model_optimized.onnx": "1.29.0",
"onnx/model.onnx": "1.29.0",
"onnx/model_int8.onnx": "1.29.0",
}
# A file kept in git rather than LFS lists no sha256, so it is hashed here: a
# tokenizer is a few MB, and a pinned file must be checkable on download.
_HASHED_HERE_UP_TO = 64 * 1024 * 1024
async def refresh(client: httpx.AsyncClient, entry: EmbeddingEntry) -> dict[str, Any]:
snapshot = await repo_at_revision(client, entry.repo)
snapshot = await _hash_small_files(client, snapshot, entry)
builds = pinned_builds(entry, snapshot)
reply = await client.get(
RESOLVE.format(repo=entry.repo, revision=snapshot.revision, path="config.json")
)
reply.raise_for_status()
architecture = reply.json().get("model_type")
if not architecture:
raise UnreadableBuildError(f"{entry.repo}: config.json names no model_type")
written = entry_for(entry, snapshot, builds, architecture, VALIDATED)
print(
f" {entry.name:34s} {snapshot.revision[:8]} {len(builds):2d} builds, "
f"default {builds[0].quantization}, {architecture}"
)
return written
async def _hash_small_files(client, snapshot, entry: EmbeddingEntry):
wanted = {path for _, path in entry.builds} | {entry.tokenizer}
listing = []
for listed in snapshot.listing:
if listed.path in wanted and not listed.sha256:
if listed.size_bytes > _HASHED_HERE_UP_TO:
raise UnreadableBuildError(f"{entry.repo}: {listed.path} is unhashed")
reply = await client.get(
RESOLVE.format(
repo=entry.repo, revision=snapshot.revision, path=listed.path
)
)
reply.raise_for_status()
listed = replace(listed, sha256=hashlib.sha256(reply.content).hexdigest())
listing.append(listed)
return replace(snapshot, listing=tuple(listing))

4
surfsense_local/backend/scripts/refresh_local_manifest.py
View file

@ -27,6 +27,8 @@ from local_manifest.entries import ENTRIES
from local_manifest.guard import losses
from local_manifest.licence import refused
from local_manifest.llamacpp import refresh as llamacpp
from local_manifest.onnxruntime import refresh as onnxruntime
from local_manifest.onnxruntime.entry import EmbeddingEntry
from local_manifest.sdcpp import refresh as sdcpp
from local_manifest.sdcpp.entry import ImageEntry
from local_manifest.unreadable import UnreadableBuildError
@ -41,6 +43,8 @@ from modules.llm.catalog.local.manifest import (
def _engine(entry):
if isinstance(entry, AudioEntry):
return audiocpp
if isinstance(entry, EmbeddingEntry):
return onnxruntime
return sdcpp if isinstance(entry, ImageEntry) else llamacpp

15
surfsense_local/backend/scripts/retrieval_eval/embedder.py Normal file
View file

@ -0,0 +1,15 @@
"""The embedder a run indexes with: bge, or a curated one by manifest id."""
from modules.embedding.bundled import BGE
from modules.embedding.spec import EmbedderSpec
from modules.llm.catalog.local.engines.onnxruntime.spec import spec_for
from modules.llm.catalog.local.manifest import load_local_manifest
def embedder_spec(model_id: str) -> EmbedderSpec:
if model_id == BGE.id:
return BGE
for model in load_local_manifest().models:
if model.id == model_id and model.embedding is not None:
return spec_for(model)
raise ValueError(f"no curated embedder named {model_id}")

14
surfsense_local/backend/scripts/retrieval_eval/fingerprint.py
View file

@ -6,6 +6,7 @@ in seconds and makes only an embedder or corpus change pay to index again.
"""
import hashlib
import json
from pathlib import Path
@ -38,3 +39,16 @@ def embedder_identity(models_dir: Path, model_dir_name: str) -> bytes:
for path in sorted(directory.glob("*")) if directory.is_dir() else []:
parts.append(f"{path.name}:{path.stat().st_size}")
return "|".join(parts).encode("utf-8")
def curated_identity(model_id: str) -> bytes:
"""A curated embedder by what pins it: its id and its weights' hash, read
from the manifest file without importing the app."""
manifest = Path(__file__).resolve().parents[2] / (
"modules/llm/catalog/local/manifest/models.json"
)
for model in json.loads(manifest.read_text())["models"]:
if model["id"] == model_id:
weights = model["builds"][0]["files"][0]
return f"{model_id}|{weights['sha256']}".encode()
raise ValueError(f"no curated embedder named {model_id}")

18
surfsense_local/backend/scripts/retrieval_eval/run.py
View file

@ -6,7 +6,12 @@ so chunking, embedding and both retrieval legs are the app's own.
from dataclasses import asdict, dataclass
from sqlalchemy import update
from modules.documents.models import Document, DocumentType
from modules.embedding.lock import lock_index
from modules.embedding.models import EmbeddingIndex
from modules.embedding.spec import EmbedderSpec
from modules.workspaces.models import Workspace
from retrieval_eval.cases import Corpus
from retrieval_eval.score import Ranking, score
@ -29,13 +34,19 @@ class Result:
titles: list[str]
def open_index():
"""The workspace in an index built by an earlier run, or None."""
def open_index(spec: EmbedderSpec):
"""The workspace in an index built by an earlier run, or None.
The stored spec is replaced by `spec`, which differs from it at most in its
ranking weight: the index is keyed by everything else.
"""
database = get_storage_settings().database_path
if not database.exists():
return None
engine = create_db_engine(database)
session = create_session_factory(engine)()
session.execute(update(EmbeddingIndex).values(spec=spec.model_dump(mode="json")))
session.commit()
workspace = session.query(Workspace).first()
if workspace is None:
session.close()
@ -44,7 +55,7 @@ def open_index():
return session, workspace.id, engine
def index(corpus: Corpus):
def index(corpus: Corpus, spec: EmbedderSpec):
"""A workspace holding the corpus, every document ingested to ready.
The database is the configured one, which the entry script points at a
@ -58,6 +69,7 @@ def index(corpus: Corpus):
engine = create_db_engine(get_storage_settings().database_path)
upgrade_to_head(engine)
session = create_session_factory(engine)()
lock_index(session, spec) # as finishing onboarding would
workspace = Workspace(name="Retrieval eval")
session.add(workspace)
session.flush()

49
surfsense_local/backend/scripts/run_retrieval_eval.py
View file

@ -17,17 +17,22 @@ import sys
from pathlib import Path
from retrieval_eval.cases import CORPUS_DIR, LOCAL_DIR
from retrieval_eval.fingerprint import embedder_identity, fingerprint
from retrieval_eval.fingerprint import curated_identity, embedder_identity, fingerprint
# Safe this high, unlike the embedder below: it reads no settings, only sqlite3.
from shared.tokenizer import TOKENIZER
# Deliberately not imported from worker.ingestion.embedding, even though that
# is where it is defined: importing anything under `worker` pulls in
# shared.queue, which reads the settings at import time and caches them. That
# would freeze the data dir below before it is set, and the run would read the
# user's own library instead. Measured twice.
# Deliberately not imported from modules.embedding.bundled, even though that
# is where it is defined: app imports stay below the line that sets the data
# dir, because anything under `worker` pulls in shared.queue, which reads the
# settings at import time and caches them. That would freeze the data dir
# before it is set, and the run would read the user's own library instead.
# Measured twice.
MODEL_DIR_NAME = "bge-small-en-v1.5"
# A curated embedder to index with instead of bge, by manifest id. From the
# environment rather than a flag, because the data dir is keyed by it and is
# set before argparse runs.
EMBEDDER = os.environ.get("SURFSENSE_EVAL_EMBEDDER", MODEL_DIR_NAME)
# Before the first import that reads settings, as tests/conftest.py does. The
# ingest job opens its own engine from the configured data dir, so a run that
@ -48,20 +53,31 @@ CACHE_ROOT = Path.home() / ".surfsense-retrieval-eval"
_KEY = fingerprint(
CORPUS_DIR,
LOCAL_DIR / "corpus",
embedder=embedder_identity(
Path(os.environ["SURFSENSE_LOCAL_MODELS_DIR"]), MODEL_DIR_NAME
embedder=(
embedder_identity(
Path(os.environ["SURFSENSE_LOCAL_MODELS_DIR"]), MODEL_DIR_NAME
)
if EMBEDDER == MODEL_DIR_NAME
else curated_identity(EMBEDDER)
),
tokenizer=TOKENIZER,
)
os.environ["SURFSENSE_LOCAL_DATA_DIR"] = str(CACHE_ROOT / _KEY)
# A downloaded embedder is read where the user's app downloaded it, as the
# models are: only the database and its documents belong to this run.
_RUN_EMBEDDINGS = CACHE_ROOT / _KEY / "embeddings"
if not _RUN_EMBEDDINGS.exists() and (_DATA / "embeddings").is_dir():
_RUN_EMBEDDINGS.parent.mkdir(parents=True, exist_ok=True)
_RUN_EMBEDDINGS.symlink_to(_DATA / "embeddings", target_is_directory=True)
from retrieval_eval.cases import load # noqa: E402
from retrieval_eval.embedder import embedder_spec # noqa: E402
from retrieval_eval.run import as_record, ask, index, open_index # noqa: E402
from retrieval_eval.summary import summarize # noqa: E402
from modules.embedding.encoder import missing_files # noqa: E402
from shared.config import get_storage_settings # noqa: E402
from shared.db import import_models # noqa: E402
from worker.ingestion.embedding import missing_embedding_files # noqa: E402
def guard_data_dir() -> None:
@ -83,18 +99,23 @@ def guard_data_dir() -> None:
def run(args: argparse.Namespace) -> None:
guard_data_dir()
missing = missing_embedding_files()
spec = embedder_spec(EMBEDDER)
missing = missing_files(spec)
if missing:
raise ValueError(
f"the embedding model is missing {missing}; "
"run scripts/fetch_embedding_model.py"
f"{spec.id} is missing {missing}: fetch bge with "
"scripts/fetch_embedding_model.py, or install another from the app"
)
weight = os.environ.get("SURFSENSE_EVAL_WEIGHT")
if weight is not None:
spec = spec.model_copy(update={"semantic_weight": float(weight)})
print(f"{spec.id} at semantic weight {spec.semantic_weight}")
corpus = load()
print(f"{len(corpus.documents)} documents, {len(corpus.queries)} queries")
opened = None if args.rebuild else open_index()
opened = None if args.rebuild else open_index(spec)
if opened is None:
print(f"indexing into {CACHE_ROOT / _KEY}")
opened = index(corpus)
opened = index(corpus, spec)
else:
print(f"reusing the index in {CACHE_ROOT / _KEY}")
session, workspace_id, engine = opened

21
surfsense_local/backend/shared/config.py
View file

@ -18,6 +18,12 @@ class StorageSettings(BaseSettings):
override = os.environ.get("SURFSENSE_LOCAL_MODELS_DIR")
return Path(override) if override else self.data_dir / "models"
@property
def embedding_models_dir(self) -> Path:
"""Downloaded embedders, one folder each. Not `models_dir`, which a
packaged app ships read-only."""
return self.data_dir / "embeddings"
@property
def database_path(self) -> Path:
return self.data_dir / "surfsense.db"
@ -57,16 +63,6 @@ class StorageSettings(BaseSettings):
return self.data_dir / "plugins" / plugin_id / "data"
class SearchSettings(BaseSettings):
"""The index's shape, which both ingest and search have to agree on."""
model_config = SettingsConfigDict(env_prefix="SURFSENSE_LOCAL_")
# Schema, not preference: a vec0 table declares its width at creation.
# 384 is bge-small-en-v1.5, the bundled default.
embedding_dimension: int = 384
class LLMSettings(BaseSettings):
"""The generation runtime. Electron starts llama-server and passes its address."""
@ -103,11 +99,6 @@ def get_storage_settings() -> StorageSettings:
return StorageSettings()
@lru_cache
def get_search_settings() -> SearchSettings:
return SearchSettings()
@lru_cache
def get_llm_settings() -> LLMSettings:
return LLMSettings()

1
surfsense_local/backend/shared/db.py
View file

@ -69,6 +69,7 @@ def import_models() -> None:
import modules.chunks.models
import modules.documents.models
import modules.egress.models
import modules.embedding.models
import modules.license.models
import modules.llm.models
import modules.plugins.models

60
surfsense_local/backend/shared/migrations.py
View file

@ -1,43 +1,55 @@
import re
from pathlib import Path
from alembic.config import Config
from sqlalchemy import Engine, text
from alembic.migration import MigrationContext
from alembic.script import ScriptDirectory
from sqlalchemy import Engine
from sqlalchemy.orm import Session
from alembic import command
from shared.config import get_search_settings
from modules.embedding.active import active_index
from modules.embedding.vector_table import declared_width
DECLARED_WIDTH = re.compile(r"float\[(\d+)\]")
_REVISIONS = Path(__file__).resolve().parent.parent / "alembic"
def upgrade_to_head(engine: Engine) -> None:
"""Apply pending migrations. The only thing in the app allowed to emit DDL."""
revisions = Path(__file__).resolve().parent.parent / "alembic"
config = Config()
config.set_main_option("script_location", str(revisions))
"""Apply pending migrations. The API runs this; besides locking the embedder,
nothing else in the app emits DDL."""
config = _config()
config.attributes["engine"] = engine
command.upgrade(config, "head")
_check_embedding_width(engine)
def _check_embedding_width(engine: Engine) -> None:
"""Refuse a database whose vectors were written by a different model.
Nothing here can be migrated: embeddings of another width are not merely the
wrong shape, they are unrelated numbers. Search would return whatever
survived, so the app stops and asks to be reindexed instead.
"""
expected = get_search_settings().embedding_dimension
def is_migrated(engine: Engine) -> bool:
"""Whether the database is at this code's latest revision."""
head = ScriptDirectory.from_config(_config()).get_current_head()
with engine.connect() as connection:
schema = connection.execute(
text("SELECT sql FROM sqlite_master WHERE name = 'chunk_vectors'")
).scalar_one()
return MigrationContext.configure(connection).get_current_revision() == head
found = int(DECLARED_WIDTH.search(schema).group(1))
if found != expected:
def _config() -> Config:
config = Config()
config.set_main_option("script_location", str(_REVISIONS))
return config
def _check_embedding_width(engine: Engine) -> None:
"""Refuse an index whose table cannot hold the vectors its model makes.
Embeddings of another width are not merely the wrong shape, they are
unrelated numbers, so the app stops rather than search what survived.
"""
with Session(engine) as session:
index = active_index(session)
if index is None:
return # Onboarding has not chosen yet.
found = declared_width(session.connection(), index.vector_table)
if found != index.spec.dimension:
raise RuntimeError(
f"chunk_vectors holds {found}-dimension vectors but the configured "
f"embedding model produces {expected}. Reindex before starting."
f"{index.vector_table} holds {found}-dimension vectors but "
f"{index.spec.id} makes {index.spec.dimension}."
)

41
surfsense_local/backend/shared/search.py
View file

@ -4,6 +4,8 @@ from dataclasses import dataclass
from sqlalchemy import bindparam, text
from sqlalchemy.orm import Session
from modules.embedding.active import ActiveIndex, require_active_index
from modules.embedding.vector_table import checked
from shared.tokenizer import terms
# Each leg proposes this many for recall; the blend below orders the union.
@ -50,27 +52,24 @@ def retrieve(
# Lazy: pulls onnxruntime and the model, which only chat and ingest need.
from sqlite_vec import serialize_float32
from worker.ingestion.embedding import embed
from modules.embedding import encoder
vector = serialize_float32(embed([query])[0])
# The question is embedded by the index it is searched against, never by a
# model configured on its own.
index = require_active_index(session)
vector = serialize_float32(
encoder.embed(index.spec, [query], encoder.Purpose.QUERY)[0]
)
keyword = dict(_keyword_leg(session, workspace_id, query, selected_document_ids))
candidates = keyword.keys() | set(
_vector_leg(session, workspace_id, vector, selected_document_ids)
_vector_leg(
session, index.vector_table, workspace_id, vector, selected_document_ids
)
)
if not candidates:
return []
return _best(session, keyword, candidates, vector, top_k)
# How much of the order meaning owns. Our corpus is flat across 0.6 to 0.65 and
# falls away either side, so this is the middle of a measured plateau rather
# than a default: 98% of answers reach the top 5, against 50% for meaning alone.
# Reciprocal rank fusion was tried first and rejected at 88%; throwing away each
# leg's strength let a passage sharing only "a" and "in" with the query outrank
# the one that meant the same thing, because appearing in both legs beats
# topping one.
SEMANTIC_WEIGHT = 0.65
return _best(session, index, keyword, candidates, vector, top_k)
def _keyword_leg(
@ -142,6 +141,7 @@ def _matching(
def _vector_leg(
session: Session,
vector_table: str,
workspace_id: int,
vector: bytes,
document_ids: Sequence[int] | None,
@ -156,7 +156,7 @@ def _vector_leg(
# a workspace's hits start falling outside the global top CANDIDATES.
sql = (
"WITH knn AS ("
" SELECT rowid, distance FROM chunk_vectors "
f" SELECT rowid, distance FROM {checked(vector_table)} "
" WHERE embedding MATCH :vector AND k = :k"
") "
"SELECT c.id, knn.distance FROM knn "
@ -180,6 +180,7 @@ def _vector_leg(
def _best(
session: Session,
index: ActiveIndex,
keyword: dict[int, float],
candidates: set[int],
vector: bytes,
@ -203,18 +204,20 @@ def _best(
stmt = text(
f"SELECT {_COLUMNS}, "
"vec_distance_cosine(v.embedding, :vector) AS distance "
"FROM chunks c JOIN chunk_vectors v ON v.rowid = c.id "
f"FROM chunks c JOIN {checked(index.vector_table)} v ON v.rowid = c.id "
"JOIN documents d ON d.id = c.document_id "
"WHERE c.id IN :ids"
).bindparams(bindparam("ids", expanding=True))
rows = session.execute(stmt, {"vector": vector, "ids": list(candidates)}).all()
weight = index.spec.semantic_weight
def blended(row) -> float:
# A passage pointing the other way is no evidence rather than evidence
# against, so cosine floors at nothing.
return SEMANTIC_WEIGHT * max(0.0, 1.0 - row.distance) + (
1.0 - SEMANTIC_WEIGHT
) * keyword.get(row.id, 0.0)
return weight * max(0.0, 1.0 - row.distance) + (1.0 - weight) * keyword.get(
row.id, 0.0
)
ordered = sorted(rows, key=lambda row: (-blended(row), row.distance))
return [

37
surfsense_local/backend/tests/conftest.py
View file

@ -15,6 +15,8 @@ from httpx import ASGITransport, AsyncClient
from sqlalchemy import Engine
from api.main import create_app
from modules.embedding.bundled import BGE
from modules.embedding.lock import lock_index
from shared.config import get_storage_settings
from shared.db import (
create_db_engine,
@ -46,22 +48,43 @@ REAL_MODELS = next(
@pytest.fixture
async def client(engine: Engine) -> AsyncGenerator[AsyncClient, None]:
"""Drive a fresh app in-process, so tests never bind a port."""
app = create_app()
app.state.session_factory = create_session_factory(engine)
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as async_client:
async with _client_over(engine) as async_client:
yield async_client
@pytest.fixture
def engine(tmp_path: Path) -> Engine:
"""A database built the way a user's is: by migrations, never create_all."""
async def unlocked_client(
unlocked_engine: Engine,
) -> AsyncGenerator[AsyncClient, None]:
"""The app before onboarding has chosen an embedder."""
async with _client_over(unlocked_engine) as async_client:
yield async_client
def _client_over(engine: Engine) -> AsyncClient:
app = create_app()
app.state.session_factory = create_session_factory(engine)
return AsyncClient(transport=ASGITransport(app=app), base_url="http://test")
@pytest.fixture
def unlocked_engine(tmp_path: Path) -> Engine:
"""A database built the way a user's is, by migrations, before onboarding
has chosen an embedder."""
engine = create_db_engine(tmp_path / "surfsense.db")
upgrade_to_head(engine)
return engine
@pytest.fixture
def engine(unlocked_engine: Engine) -> Engine:
"""A user's database once onboarding has fixed the bundled embedder."""
with create_session_factory(unlocked_engine)() as session:
lock_index(session, BGE)
session.commit()
return unlocked_engine
@pytest.fixture(autouse=True)
def data_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[Path]:
"""Point uploads at this test's own tree, and empty the queue it enqueues to.

10
surfsense_local/backend/tests/integration/artifacts/test_flashcard_state_routes.py
View file

@ -1,5 +1,7 @@
"""The flashcard-state routes: a real deck artifact on disk, driven over HTTP."""
from typing import Any
import pytest
from httpx import AsyncClient
from sqlalchemy import Engine
@ -7,7 +9,6 @@ from sqlalchemy.orm import Session
from modules.artifacts.models import Artifact
from modules.documents.models import Document, DocumentStatus, DocumentType
from shared.config import get_search_settings
from shared.db import create_session_factory
from worker.studio.content.flashcards import pipeline as flashcards
from worker.studio.shared import persist
@ -20,12 +21,11 @@ def stub_model(monkeypatch: pytest.MonkeyPatch) -> None:
"""persist() indexes the deck's markdown; stand in for the embedder so
these tests don't need the bundled model on disk (see
tests/integration/worker/conftest.py's identical fixture)."""
width = get_search_settings().embedding_dimension
def embed(texts: list[str]) -> list[list[float]]:
return [[float(len(text) % 97)] * width for text in texts]
def embed(spec: Any, texts: list[str], _purpose: Any) -> list[list[float]]:
return [[float(len(text) % 97)] * spec.dimension for text in texts]
monkeypatch.setattr("worker.ingestion.embedding.embed", embed)
monkeypatch.setattr("modules.embedding.encoder.embed", embed)
monkeypatch.setattr(
"worker.ingestion.chunking._default_tokenizer", lambda: "character"
)

10
surfsense_local/backend/tests/integration/artifacts/test_quiz_state_routes.py
View file

@ -1,5 +1,7 @@
"""The quiz-state routes: a real quiz artifact on disk, driven over HTTP."""
from typing import Any
import pytest
from httpx import AsyncClient
from sqlalchemy import Engine
@ -7,7 +9,6 @@ from sqlalchemy.orm import Session
from modules.artifacts.models import Artifact
from modules.documents.models import Document, DocumentStatus, DocumentType
from shared.config import get_search_settings
from shared.db import create_session_factory
from worker.studio.content.quiz import pipeline as quiz
from worker.studio.shared import persist
@ -20,12 +21,11 @@ def stub_model(monkeypatch: pytest.MonkeyPatch) -> None:
"""persist() indexes the quiz's markdown; stand in for the embedder so
these tests don't need the bundled model on disk (see
tests/integration/worker/conftest.py's identical fixture)."""
width = get_search_settings().embedding_dimension
def embed(texts: list[str]) -> list[list[float]]:
return [[float(len(text) % 97)] * width for text in texts]
def embed(spec: Any, texts: list[str], _purpose: Any) -> list[list[float]]:
return [[float(len(text) % 97)] * spec.dimension for text in texts]
monkeypatch.setattr("worker.ingestion.embedding.embed", embed)
monkeypatch.setattr("modules.embedding.encoder.embed", embed)
monkeypatch.setattr(
"worker.ingestion.chunking._default_tokenizer", lambda: "character"
)

36
surfsense_local/backend/tests/integration/chat/test_chat.py
View file

@ -52,7 +52,9 @@ def _seed(
ids[title] = doc.id
session.add(
SelectedModel(
model_type=ModelType.TEXT_GEN, provider="llamacpp", name="Qwen3-1.7B-Q4_K_M"
model_type=ModelType.TEXT_GEN,
provider="llamacpp",
name="Qwen3-1.7B-Q4_K_M",
)
)
session.commit()
@ -67,7 +69,9 @@ def _choose_generation_model(engine: Engine) -> None:
with create_session_factory(engine)() as session:
session.add(
SelectedModel(
model_type=ModelType.TEXT_GEN, provider="llamacpp", name="Qwen3-1.7B-Q4_K_M"
model_type=ModelType.TEXT_GEN,
provider="llamacpp",
name="Qwen3-1.7B-Q4_K_M",
)
)
session.commit()
@ -521,7 +525,9 @@ async def test_missing_embedding_assets_are_an_actionable_503(
with create_session_factory(engine)() as session:
session.add(
SelectedModel(
model_type=ModelType.TEXT_GEN, provider="llamacpp", name="Qwen3-1.7B-Q4_K_M"
model_type=ModelType.TEXT_GEN,
provider="llamacpp",
name="Qwen3-1.7B-Q4_K_M",
)
)
session.commit()
@ -539,6 +545,30 @@ async def test_missing_embedding_assets_are_an_actionable_503(
)
async def test_chat_before_an_embedder_is_chosen_is_a_409(
unlocked_client: AsyncClient, unlocked_engine: Engine
) -> None:
"""The API reached before onboarding finishes: say why, rather than guess a model."""
with create_session_factory(unlocked_engine)() as session:
session.add(
SelectedModel(
model_type=ModelType.TEXT_GEN,
provider="llamacpp",
name="Qwen3-1.7B-Q4_K_M",
)
)
session.commit()
workspace = (await unlocked_client.post("/workspaces", json={"name": "w"})).json()
thread_id = await _open_thread(unlocked_client, workspace["id"])
reply = await unlocked_client.post(
f"/chat/threads/{thread_id}/messages", json={"text": "hi"}
)
assert reply.status_code == 409
assert reply.json()["detail"]["code"] == "embedding_not_chosen"
async def test_the_answer_reserves_room_instead_of_taking_the_whole_window(
client: AsyncClient, engine: Engine, real_model: object, llamacpp_server: list[dict]
) -> None:

4
surfsense_local/backend/tests/integration/chunks/test_search_index.py
View file

@ -8,8 +8,8 @@ from sqlalchemy.exc import OperationalError
from modules.chunks.models import Chunk
from modules.documents.models import Document, DocumentType
from modules.embedding.bundled import BGE
from modules.workspaces.models import Workspace
from shared.config import get_search_settings
from shared.tokenizer import terms
pytestmark = pytest.mark.integration
@ -18,7 +18,7 @@ pytestmark = pytest.mark.integration
@pytest.fixture
def chunked(engine: Engine) -> Engine:
"""A workspace holding one document of one indexed chunk."""
width = get_search_settings().embedding_dimension
width = BGE.dimension
with engine.begin() as connection:
connection.execute(insert(Workspace).values(id=1, name="Research"))

6
surfsense_local/backend/tests/integration/conftest.py
View file

@ -5,6 +5,7 @@ from collections.abc import AsyncIterator
import pytest
import uvicorn
from sqlalchemy import Engine
from api.main import create_app
@ -16,8 +17,9 @@ def _free_port() -> int:
@pytest.fixture
async def base_url(data_dir: object) -> AsyncIterator[str]:
"""Serve a migrated app on a free port; data_dir points it at this test's DB."""
async def base_url(data_dir: object, engine: Engine) -> AsyncIterator[str]:
"""Serve a migrated app on a free port; data_dir points it at this test's DB,
which `engine` has already built with the embedder fixed, as onboarding would."""
port = _free_port()
server = uvicorn.Server(
uvicorn.Config(create_app(), host="127.0.0.1", port=port, log_level="warning")

40
surfsense_local/backend/tests/integration/embedding/test_encoder.py Normal file
View file

@ -0,0 +1,40 @@
"""The encoder, run on the real bundled model."""
from pathlib import Path
import pytest
from modules.embedding.bundled import BGE
from modules.embedding.encoder import Purpose, embed
from modules.embedding.verify import verify
pytestmark = pytest.mark.integration
ASYMMETRIC = BGE.model_copy(update={"query_prefix": "query: "})
def test_a_query_prefix_applies_to_questions_only(real_model: Path) -> None:
"""Asymmetric models embed a question and a passage differently."""
question = embed(ASYMMETRIC, ["when does the ship sail"], Purpose.QUERY)
by_hand = embed(BGE, ["query: when does the ship sail"], Purpose.DOCUMENT)
passage = embed(ASYMMETRIC, ["when does the ship sail"], Purpose.DOCUMENT)
assert question == by_hand
assert passage != question
def test_vectors_have_the_specs_width(real_model: Path) -> None:
"""The width the index's table is built for."""
[vector] = embed(BGE, ["the ship sails at dawn"], Purpose.DOCUMENT)
assert len(vector) == BGE.dimension
def test_a_real_search_model_passes_its_checks(real_model: Path) -> None:
"""What every Hugging Face pick runs after download, on a model that works.
A config can misstate the width, so the probe's is the one kept."""
misstated = BGE.model_copy(update={"dimension": 768})
width, refusal = verify(misstated)
assert (width, refusal) == (384, None)

91
surfsense_local/backend/tests/integration/embedding/test_lock.py Normal file
View file

@ -0,0 +1,91 @@
"""Locking: the moment the library's embedder is chosen for good."""
import struct
from pathlib import Path
import pytest
from sqlalchemy import Engine, text
from sqlalchemy.exc import OperationalError
from modules.embedding.active import active_index
from modules.embedding.bundled import BGE
from modules.embedding.lock import IndexAlreadyBuiltError, lock_index
from shared.db import create_db_engine, create_session_factory
from shared.migrations import upgrade_to_head
pytestmark = pytest.mark.integration
NARROW = BGE.model_copy(update={"id": "narrow", "dimension": 8})
@pytest.fixture
def fresh(tmp_path: Path) -> Engine:
"""An install whose onboarding has not chosen an embedder."""
engine = create_db_engine(tmp_path / "fresh.db")
upgrade_to_head(engine)
return engine
def _store(engine: Engine, table: str, width: int) -> None:
with engine.begin() as connection:
connection.execute(
text(f"INSERT INTO {table}(rowid, embedding) VALUES (1, :vector)"),
{"vector": struct.pack(f"{width}f", *([0.1] * width))},
)
def test_the_chosen_model_becomes_the_active_index(fresh: Engine) -> None:
"""Search reads it and ingest writes it from here on."""
with create_session_factory(fresh)() as session:
lock_index(session, NARROW)
session.commit()
index = active_index(session)
assert index is not None
assert index.spec == NARROW
def test_its_table_holds_vectors_of_the_models_width(fresh: Engine) -> None:
"""Built at the chosen width, not the width the first migration guessed."""
with create_session_factory(fresh)() as session:
lock_index(session, NARROW)
session.commit()
table = active_index(session).vector_table
_store(fresh, table, 8)
with pytest.raises(OperationalError):
_store(fresh, table, 384)
def test_a_library_with_chunks_cannot_be_locked_again(fresh: Engine) -> None:
"""Its vectors were made by some model already; a new one would not match."""
with fresh.begin() as connection:
connection.execute(text("INSERT INTO workspaces(id, name) VALUES (1, 'w')"))
connection.execute(
text(
"INSERT INTO documents(id, workspace_id, title, document_type, "
"status) VALUES (1, 1, 'n', 'NOTE', 'ready')"
)
)
connection.execute(
text(
"INSERT INTO chunks(document_id, position, content) "
"VALUES (1, 0, 'text')"
)
)
with (
create_session_factory(fresh)() as session,
pytest.raises(IndexAlreadyBuiltError),
):
lock_index(session, NARROW)
def test_a_locked_library_cannot_be_locked_again(fresh: Engine) -> None:
"""The choice is made once; changing it is a re-embed, not a second lock."""
with create_session_factory(fresh)() as session:
lock_index(session, BGE)
session.commit()
with pytest.raises(IndexAlreadyBuiltError):
lock_index(session, NARROW)

54
surfsense_local/backend/tests/integration/embedding/test_not_chosen.py Normal file
View file

@ -0,0 +1,54 @@
"""Before onboarding fixes the embedder, nothing that would be embedded is taken.
Onboarding gates the UI only; the API is reached directly too, by scripts and
plugins, and a document accepted now would have no model to be embedded with.
"""
import pytest
from httpx import AsyncClient
from shared.queue import ingest_queue
pytestmark = pytest.mark.integration
NOT_CHOSEN = {
"code": "embedding_not_chosen",
"message": "no embedding model has been chosen yet",
}
@pytest.fixture
async def workspace_id(unlocked_client: AsyncClient) -> int:
"""Workspaces need no embedder, so one exists before onboarding finishes."""
created = await unlocked_client.post("/workspaces", json={"name": "Research"})
return int(created.json()["id"])
async def test_a_note_is_refused(
unlocked_client: AsyncClient, workspace_id: int
) -> None:
"""Refused before a row is written or a job is queued."""
reply = await unlocked_client.post(
f"/workspaces/{workspace_id}/documents",
json={"title": "Kickoff", "content": "agreed to ship"},
)
assert reply.status_code == 409
assert reply.json()["detail"] == NOT_CHOSEN
listed = await unlocked_client.get(f"/workspaces/{workspace_id}/documents")
assert listed.json() == []
assert ingest_queue.pending_count() == 0
async def test_an_upload_is_refused(
unlocked_client: AsyncClient, workspace_id: int
) -> None:
"""The bytes are not kept either: nothing would ever embed them."""
reply = await unlocked_client.post(
f"/workspaces/{workspace_id}/documents/upload",
files={"files": ("report.pdf", b"%PDF-1.7 fake", "application/pdf")},
)
assert reply.status_code == 409
assert reply.json()["detail"] == NOT_CHOSEN
assert ingest_queue.pending_count() == 0

11
surfsense_local/backend/tests/integration/llm/test_catalog_routes.py
View file

@ -14,7 +14,7 @@ async def test_the_catalog_renders_with_no_network_and_no_scan(
assert reply.status_code == 200
body = reply.json()
curated = [row for row in body["rows"] if row["origin"] == "curated"]
curated = [row for row in body["rows"] if row["origin"] == "curated" and row["engine"] == "llamacpp"]
assert curated
# A badge names a verdict exactly when it warns.
for row in curated:
@ -59,7 +59,7 @@ async def test_each_curated_model_names_its_default_and_at_most_one_recommended_
body = (await client.get("/llm/catalog/local")).json()
for row in body["rows"]:
if row["origin"] == "curated":
if row["origin"] == "curated" and row["engine"] == "llamacpp":
assert row["default_quantization"] == "UD-Q4_K_XL"
assert sum(b["recommended"] for b in row["builds"]) <= 1
@ -71,7 +71,10 @@ async def test_the_offload_fraction_survives_to_the_renderer(
body = (await client.get("/llm/catalog/local")).json()
assert all(
"offload_fraction" in b["fit"] for row in body["rows"] for b in row["builds"]
"offload_fraction" in b["fit"]
for row in body["rows"]
if row["engine"] == "llamacpp"
for b in row["builds"]
)
@ -128,7 +131,7 @@ async def test_every_curated_row_says_which_build_it_leads_with_and_why(
body = (await client.get("/llm/catalog/local")).json()
for row in body["rows"]:
if row["origin"] == "curated":
if row["origin"] == "curated" and row["engine"] == "llamacpp":
lead = row["lead"]
assert lead["why"] in {
"in_use",

69
surfsense_local/backend/tests/integration/llm/test_embedding_catalog_routes.py Normal file
View file

@ -0,0 +1,69 @@
"""Embedding models in the one local catalog, offered by onnxruntime: listed and
downloaded like any model, and never a slot."""
import pytest
from httpx import AsyncClient
from modules.embedding.bundled import BGE
from modules.embedding.encoder import missing_files
from modules.llm.catalog.local.engines.onnxruntime.spec import spec_for
from modules.llm.catalog.local.manifest import load_local_manifest
from tests.integration.llm.installs import install_to_end
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
GRANITE = "granite-embedding-97m-multilingual-r2"
def embedding_rows(body: dict) -> dict[str, dict]:
"""The rows onnxruntime offered, by id."""
return {r["id"]: r for r in body["rows"] if r["engine"] == "onnxruntime"}
@pytest.fixture
def bundled_bge(data_dir):
"""The models pack as the build script leaves it: bge's files, read-only."""
pack = data_dir / "models" / BGE.id
pack.mkdir(parents=True)
for pinned in (BGE.weights, BGE.tokenizer):
(pack / pinned.path).write_bytes(b"onnx")
return pack
async def test_embedders_are_rows_that_fill_no_slot(client: AsyncClient) -> None:
"""Downloadable from the catalog, chosen only at onboarding."""
rows = embedding_rows((await client.get("/llm/catalog/local")).json())
assert {BGE.id, GRANITE} <= set(rows)
granite = rows[GRANITE]
assert granite["types"] == ["embedding"]
assert granite["selectable_for"] == []
assert granite["runnable"]
# Onboarding says what each is for; the manifest is where a person wrote it.
assert granite["description"] == "Search across more than 50 languages."
async def test_the_embedder_the_app_ships_reads_installed_and_bundled(
client: AsyncClient, bundled_bge
) -> None:
"""In the installer, so on this computer from the first start."""
rows = embedding_rows((await client.get("/llm/catalog/local")).json())
(build,) = rows[BGE.id]["builds"]
assert (build["installed_as"], build["bundled"]) == (BGE.id, True)
async def test_an_embedder_installs_where_the_encoder_reads_it(
client: AsyncClient, fake_hub
) -> None:
"""Named by its catalog id: every ONNX repo calls its weights model.onnx."""
rows = embedding_rows((await client.get("/llm/catalog/local")).json())
(build,) = rows[GRANITE]["builds"]
job = await install_to_end(client, catalog_id=build["catalog_id"], select=False)
assert job["event"]["type"] == "complete", job
(model,) = [m for m in load_local_manifest().models if m.id == GRANITE]
assert missing_files(spec_for(model)) == []
rows = embedding_rows((await client.get("/llm/catalog/local")).json())
assert rows[GRANITE]["builds"][0]["installed_as"] == GRANITE

118
surfsense_local/backend/tests/integration/llm/test_embedding_choice_routes.py Normal file
View file

@ -0,0 +1,118 @@
"""Choosing the embedder at onboarding, reading it back, and keeping it.
The choice is made once: finishing onboarding locks it, the index reports it,
and nothing it depends on can be deleted while it is active.
"""
import pytest
from httpx import AsyncClient
from modules.embedding.bundled import BGE
from modules.llm.catalog.local.engines.onnxruntime.spec import spec_for
from modules.llm.catalog.local.manifest import load_local_manifest
from tests.integration.llm.installs import install_to_end
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
GRANITE = "granite-embedding-97m-multilingual-r2"
CHAT = {"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"}
async def _install(client: AsyncClient, model_id: str) -> None:
rows = (await client.get("/llm/catalog/local")).json()["rows"]
(row,) = [r for r in rows if r["id"] == model_id]
job = await install_to_end(
client, catalog_id=row["builds"][0]["catalog_id"], select=False
)
assert job["event"]["type"] == "complete", job
def _granite_spec() -> dict:
(model,) = [m for m in load_local_manifest().models if m.id == GRANITE]
return spec_for(model).model_dump(mode="json")
async def test_onboarding_locks_the_embedder_it_was_given(
unlocked_client: AsyncClient, llamacpp_server: str, fake_hub
) -> None:
"""Downloaded during the step, fixed when the user finishes."""
client = unlocked_client
await client.put("/llm/selection/text_gen", json=CHAT)
await _install(client, GRANITE)
finished = await client.post("/llm/onboarding", json={"embedding_model": GRANITE})
assert finished.status_code == 200, finished.text
index = (await client.get("/embedding/index")).json()
assert index["building"] is None
assert index["active"]["spec"] == _granite_spec()
assert index["active"]["name"] == "Granite Embedding 97M (Multilingual)"
async def test_an_embedder_not_yet_downloaded_cannot_be_locked(
unlocked_client: AsyncClient, llamacpp_server: str
) -> None:
"""Locked without its files, nothing could ever be embedded."""
client = unlocked_client
await client.put("/llm/selection/text_gen", json=CHAT)
refused = await client.post("/llm/onboarding", json={"embedding_model": GRANITE})
assert refused.status_code == 409
assert (await client.get("/embedding/index")).json()["active"] is None
assert (await client.get("/llm/onboarding")).json() == {"completed": False}
async def test_the_index_is_empty_until_onboarding_chooses(
unlocked_client: AsyncClient,
) -> None:
"""Settings reads this; before onboarding there is nothing to show."""
index = await unlocked_client.get("/embedding/index")
assert index.status_code == 200
assert index.json() == {"active": None, "building": None}
async def test_the_active_embedder_cannot_be_deleted(
unlocked_client: AsyncClient, llamacpp_server: str, fake_hub
) -> None:
"""Every vector in the library was made by it."""
client = unlocked_client
await client.put("/llm/selection/text_gen", json=CHAT)
await _install(client, GRANITE)
await client.post("/llm/onboarding", json={"embedding_model": GRANITE})
refused = await client.delete(f"/llm/models/{GRANITE}")
assert refused.status_code == 409
rows = (await client.get("/llm/catalog/local")).json()["rows"]
(row,) = [r for r in rows if r["id"] == GRANITE]
assert row["builds"][0]["installed_as"] == GRANITE
async def test_an_embedder_downloaded_but_not_chosen_can_be_deleted(
client: AsyncClient, fake_hub
) -> None:
"""bge was locked; the download is spare."""
await _install(client, GRANITE)
deleted = await client.delete(f"/llm/models/{GRANITE}")
assert deleted.status_code == 200, deleted.text
rows = (await client.get("/llm/catalog/local")).json()["rows"]
(row,) = [r for r in rows if r["id"] == GRANITE]
assert row["builds"][0]["installed_as"] is None
async def test_the_bundled_embedder_cannot_be_deleted(
client: AsyncClient, data_dir
) -> None:
"""It is the way back once changing the model exists."""
pack = data_dir / "models" / BGE.id
pack.mkdir(parents=True)
for pinned in (BGE.weights, BGE.tokenizer):
(pack / pinned.path).write_bytes(b"onnx")
refused = await client.delete(f"/llm/models/{BGE.id}")
assert refused.status_code == 409

272
surfsense_local/backend/tests/integration/llm/test_embedding_huggingface_routes.py Normal file
View file

@ -0,0 +1,272 @@
"""Finding an embedder on Hugging Face, downloading it, and locking it.
Hugging Face is faked at its HTTP boundary; the model's own checks after the
download are faked at `verify`, which the unit tests and the real-model tests
cover.
"""
import json
import httpx
import pytest
from httpx import AsyncClient
from modules.embedding.huggingface import hub
from modules.llm.catalog.local.dependencies import get_local_catalog
from tests.integration.llm.installs import wait_for_end
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
REPO = "intfloat/multilingual-e5-small"
REV = "c" * 40
INSTALLED_AS = "hf--intfloat--multilingual-e5-small"
TREE = [
{
"type": "file",
"path": "onnx/model.onnx",
"size": 470_000_000,
"lfs": {"oid": "1" * 64},
},
{
"type": "file",
"path": "onnx/model_int8.onnx",
"size": 118_000_000,
"lfs": {"oid": "2" * 64},
},
{
"type": "file",
"path": "onnx/model_O4.onnx",
"size": 235_000_000,
"lfs": {"oid": "3" * 64},
},
{
"type": "file",
"path": "tokenizer.json",
"size": 17_000_000,
"lfs": {"oid": "4" * 64},
},
{"type": "file", "path": "config.json", "size": 700},
{"type": "file", "path": "1_Pooling/config.json", "size": 200},
{"type": "file", "path": "modules.json", "size": 300},
{"type": "file", "path": "config_sentence_transformers.json", "size": 200},
{"type": "file", "path": "sentence_bert_config.json", "size": 60},
]
CONFIGS = {
"config.json": {"hidden_size": 384, "max_position_embeddings": 512},
"1_Pooling/config.json": {"pooling_mode_mean_tokens": True},
"modules.json": [{"type": "sentence_transformers.models.Normalize"}],
"config_sentence_transformers.json": {
"prompts": {"query": "query: ", "passage": "passage: "}
},
"sentence_bert_config.json": {"max_seq_length": 512},
}
def build_of(body: dict) -> dict:
"""The one build an opened repo offers."""
(build,) = body["row"]["builds"]
return build
def fake_hub(
*, scan: list | None = None, tree: list | None = None
) -> httpx.MockTransport:
"""Hugging Face's API and file server, answering for one repo."""
def handle(request: httpx.Request) -> httpx.Response:
path = request.url.path
if path == "/api/models" and request.url.params.get("pipeline_tag"):
tag = request.url.params["pipeline_tag"]
found = {
"sentence-similarity": [
{"id": REPO, "downloads": 900, "gated": False, "pipeline_tag": tag}
],
"feature-extraction": [
{
"id": "BAAI/bge-m3",
"downloads": 5000,
"gated": False,
"pipeline_tag": tag,
}
],
}
return httpx.Response(200, json=found.get(tag, []))
if path == f"/api/models/{REPO}":
return httpx.Response(
200,
json={
"sha": REV,
"pipeline_tag": "sentence-similarity",
"gated": False,
"securityRepoStatus": {
"scansDone": True,
"filesWithIssues": scan or [],
},
},
)
if path == f"/api/models/{REPO}/tree/{REV}":
return httpx.Response(200, json=tree if tree is not None else TREE)
prefix = f"/{REPO}/resolve/{REV}/"
if path == f"{prefix}tokenizer.json":
return httpx.Response(200, content=b'{"model": {}}')
if path.startswith(prefix) and path[len(prefix) :] in CONFIGS:
return httpx.Response(200, content=json.dumps(CONFIGS[path[len(prefix) :]]))
return httpx.Response(404)
return httpx.MockTransport(handle)
@pytest.fixture
def huggingface(monkeypatch: pytest.MonkeyPatch):
"""Hugging Face as the fake above, with consent already given."""
monkeypatch.setattr("modules.egress.service.require", lambda *a, **k: None)
def use(**kwargs):
monkeypatch.setattr(hub, "transport", lambda: fake_hub(**kwargs))
use()
return use
@pytest.fixture
def downloads(monkeypatch: pytest.MonkeyPatch):
"""Downloads that write a few bytes, and checks that answer as told."""
from modules.llm.catalog.local.install import download as download_module
from modules.llm.providers.types import DownloadProgress
async def download(url, destination, *, sha256=None, transport=None):
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(b"onnx")
yield DownloadProgress("complete", 4, 4)
monkeypatch.setattr(download_module, "download_gguf", download)
outcome = {"width": 384, "refusal": None}
monkeypatch.setattr(
"modules.llm.catalog.local.engines.onnxruntime.engine.verify",
lambda spec: (outcome["width"], outcome["refusal"]),
)
get_local_catalog.cache_clear()
return outcome
async def test_search_lists_embedders_most_downloaded_first(
client: AsyncClient, huggingface
) -> None:
"""Both tags an embedder carries, merged."""
reply = await client.get("/embedding/huggingface/search", params={"q": "e5"})
assert reply.status_code == 200
assert [hit["repo"] for hit in reply.json()["results"]] == ["BAAI/bge-m3", REPO]
async def test_a_runnable_repo_resolves_to_a_build_and_what_it_will_be(
client: AsyncClient, huggingface
) -> None:
"""A catalog row with one build: what Download fetches, and how big it is."""
reply = await client.get(f"/embedding/huggingface/repo/{REPO}")
body = reply.json()
assert reply.status_code == 200, body
assert body["row"]["runnable"] and body["row"]["not_runnable_reason"] is None
assert body["row"]["engine"] == "onnxruntime"
build = build_of(body)
assert build["footprint_bytes"] == 118_000_000 + 17_000_000
assert build["catalog_id"] and build["installed_as"] is None
async def test_a_repo_without_onnx_says_why_it_cannot_run(
client: AsyncClient, huggingface
) -> None:
"""Safetensors alone needs torch, which the API does not ship."""
huggingface(tree=[f for f in TREE if not f["path"].endswith(".onnx")])
body = (await client.get(f"/embedding/huggingface/repo/{REPO}")).json()
assert not body["row"]["runnable"]
assert "ONNX" in body["row"]["not_runnable_reason"]
assert body["row"]["builds"] == []
async def test_a_tokenizer_stored_without_lfs_is_hashed_from_its_bytes(
client: AsyncClient, huggingface
) -> None:
"""Small files skip LFS, so Hugging Face lists no sha256 for them; the
sentence-transformers repos store tokenizer.json that way."""
small = [
{"type": "file", "path": "tokenizer.json", "size": 700_000}
if f["path"] == "tokenizer.json"
else f
for f in TREE
]
huggingface(tree=small)
body = (await client.get(f"/embedding/huggingface/repo/{REPO}")).json()
assert body["row"]["runnable"], body
assert build_of(body)["footprint_bytes"] == 118_000_000 + 700_000
async def test_a_repo_hugging_face_flags_is_refused(
client: AsyncClient, huggingface
) -> None:
"""Hugging Face's own scan found malware: no download, and the screen says why."""
huggingface(scan=[{"path": "danger.pkl", "level": "unsafe"}])
body = (await client.get(f"/embedding/huggingface/repo/{REPO}")).json()
assert not body["row"]["runnable"]
assert "security scan" in body["row"]["not_runnable_reason"]
async def test_a_pick_that_passes_its_checks_installs_and_can_be_locked(
unlocked_client: AsyncClient, huggingface, downloads, llamacpp_server
) -> None:
"""Downloaded, checked, listed, then fixed as the library's embedder."""
client = unlocked_client
resolved = (await client.get(f"/embedding/huggingface/repo/{REPO}")).json()
started = await client.post(
"/llm/installs",
json={"catalog_id": build_of(resolved)["catalog_id"], "select": False},
)
job = await wait_for_end(client, started.json()["id"])
assert job["event"]["type"] == "complete", job
# Opened again, the search offers Use rather than a second download.
reopened = (await client.get(f"/embedding/huggingface/repo/{REPO}")).json()
assert build_of(reopened)["installed_as"] == INSTALLED_AS
rows = (await client.get("/llm/catalog/local")).json()["rows"]
(row,) = [r for r in rows if r["id"] == INSTALLED_AS]
assert row["builds"][0]["installed_as"] == INSTALLED_AS
await client.put(
"/llm/selection/text_gen",
json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"},
)
finished = await client.post(
"/llm/onboarding", json={"embedding_model": INSTALLED_AS}
)
assert finished.status_code == 200, finished.text
active = (await client.get("/embedding/index")).json()["active"]
spec = active["spec"]
# Named by its repo: the install name is ours, not something a person chose.
assert active["name"] == REPO
assert (spec["source"], spec["identified"]) == ("huggingface", "declared")
assert spec["query_prefix"] == "query: "
async def test_a_pick_that_fails_its_checks_is_removed(
client: AsyncClient, huggingface, downloads, data_dir
) -> None:
"""A chat model served as an embedder ranks decoys first."""
downloads["refusal"] = "It found 6 of 10 answers first."
resolved = (await client.get(f"/embedding/huggingface/repo/{REPO}")).json()
started = await client.post(
"/llm/installs",
json={"catalog_id": build_of(resolved)["catalog_id"], "select": False},
)
job = await wait_for_end(client, started.json()["id"])
assert job["event"]["type"] == "error"
assert "6 of 10" in json.dumps(job["event"])
assert not (data_dir / "embeddings" / INSTALLED_AS).exists()

54
surfsense_local/backend/tests/integration/llm/test_routes.py
View file

@ -3,6 +3,8 @@ from httpx import AsyncClient
from sqlalchemy import Engine
from modules.documents.models import Document, DocumentStatus, DocumentType
from modules.embedding.active import active_index
from modules.embedding.bundled import BGE
from modules.llm import providers as registry
from modules.llm.activity import model_activity, model_key
from modules.llm.catalog.local.dependencies import get_local_catalog
@ -88,6 +90,44 @@ async def test_selecting_a_chat_model_does_not_complete_onboarding(
assert completed.json() == {"completed": True}
async def test_the_embedder_cannot_be_selected_like_a_model(
client: AsyncClient,
) -> None:
"""It is fixed with the index at onboarding; a selection would swap it."""
reply = await client.put(
"/llm/selection/embedding",
json={
"provider": "onnxruntime",
"name": "granite-embedding-97m-multilingual-r2",
},
)
assert reply.status_code == 422
assert (await client.get("/llm/selection/embedding")).status_code == 422
async def test_finishing_onboarding_fixes_bge_as_the_embedder(
unlocked_client: AsyncClient, unlocked_engine: Engine, llamacpp_server: str
) -> None:
"""Nothing was chosen, so the bundled model is; finishing twice changes nothing."""
client = unlocked_client
await client.put(
"/llm/selection/text_gen",
json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"},
)
await client.post("/llm/onboarding")
with create_session_factory(unlocked_engine)() as session:
first = active_index(session)
assert (await client.post("/llm/onboarding")).status_code == 200
with create_session_factory(unlocked_engine)() as session:
again = active_index(session)
assert first is not None
assert first.spec == BGE
assert again == first
async def test_onboarding_cannot_complete_without_a_chat_model(
client: AsyncClient,
) -> None:
@ -120,7 +160,9 @@ async def test_deleting_the_selected_local_model_clears_only_the_selection(
# lists the model until its next restart, which the preset rewrite triggers:
# it scans its directory once at startup and has no way to be told otherwise.
rows = (await client.get("/llm/catalog/local")).json()["rows"]
installed = [b["installed_as"] for row in rows for b in row["builds"] if b["installed_as"]]
installed = [
b["installed_as"] for row in rows for b in row["builds"] if b["installed_as"]
]
assert installed == ["Qwen3-4B-Q4_K_M"]
@ -199,13 +241,17 @@ async def test_choosing_again_updates_in_place(
) -> None:
"""One row per role: the second choice replaces the first, not adds to it."""
await client.put(
"/llm/selection/text_gen", json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"}
"/llm/selection/text_gen",
json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"},
)
await client.put(
"/llm/selection/text_gen", json={"provider": "llamacpp", "name": "Qwen3-4B-Q4_K_M"}
"/llm/selection/text_gen",
json={"provider": "llamacpp", "name": "Qwen3-4B-Q4_K_M"},
)
assert (await client.get("/llm/selection/text_gen")).json()["name"] == "Qwen3-4B-Q4_K_M"
assert (await client.get("/llm/selection/text_gen")).json()[
"name"
] == "Qwen3-4B-Q4_K_M"
async def test_a_selection_names_a_known_provider(client: AsyncClient) -> None:

117
surfsense_local/backend/tests/integration/test_migration_0021.py Normal file
View file

@ -0,0 +1,117 @@
"""Revision 0021: the index records which model built it."""
import struct
from pathlib import Path
import pytest
from alembic.config import Config
from sqlalchemy import Engine, text
from alembic import command
from modules.embedding.active import active_index
from modules.embedding.bundled import BGE
from shared.db import create_db_engine, create_session_factory
from shared.migrations import upgrade_to_head
pytestmark = pytest.mark.integration
def _at_0020(tmp_path: Path) -> Engine:
engine = create_db_engine(tmp_path / "surfsense.db")
config = Config()
config.set_main_option(
"script_location", str(Path(__file__).parents[2] / "alembic")
)
config.attributes["engine"] = engine
command.upgrade(config, "0020")
return engine
def _library(engine: Engine) -> None:
"""One indexed document, and one still waiting for the worker."""
with engine.begin() as connection:
connection.execute(text("INSERT INTO workspaces(id, name) VALUES (1, 'Notes')"))
for document_id, status in ((1, "ready"), (2, "pending")):
connection.execute(
text(
"INSERT INTO documents(id, workspace_id, title, document_type, "
"status) VALUES (:id, 1, 'note', 'NOTE', :status)"
),
{"id": document_id, "status": status},
)
connection.execute(
text(
"INSERT INTO chunks(id, document_id, position, content) "
"VALUES (1, 1, 0, 'the ship sails at dawn')"
)
)
connection.execute(
text("INSERT INTO chunk_vectors(rowid, embedding) VALUES (1, :vector)"),
{"vector": struct.pack("384f", *([0.1] * 384))},
)
def _stamps(engine: Engine) -> dict[int, int | None]:
with engine.connect() as connection:
rows = connection.execute(text("SELECT id, embedding_index_id FROM documents"))
return {row.id: row.embedding_index_id for row in rows}
def test_an_existing_library_is_recorded_as_built_by_bge(tmp_path: Path) -> None:
"""Its vectors are bge's, so the row says so; nothing is re-embedded."""
engine = _at_0020(tmp_path)
_library(engine)
upgrade_to_head(engine)
with create_session_factory(engine)() as session:
index = active_index(session)
assert index is not None
assert index.spec == BGE
assert index.vector_table == "chunk_vectors"
assert _stamps(engine) == {1: index.id, 2: None}
def test_a_finished_onboarding_with_no_documents_is_still_bge(
tmp_path: Path,
) -> None:
"""The choice is offered only before onboarding finishes."""
engine = _at_0020(tmp_path)
with engine.begin() as connection:
connection.execute(text("INSERT INTO onboarding_completion(id) VALUES (1)"))
upgrade_to_head(engine)
with create_session_factory(engine)() as session:
index = active_index(session)
assert index is not None
assert index.spec == BGE
def test_startup_refuses_an_index_whose_table_is_another_width(
tmp_path: Path,
) -> None:
"""The row says what wrote the vectors; a table of another width cannot hold them."""
engine = create_db_engine(tmp_path / "surfsense.db")
upgrade_to_head(engine)
with engine.begin() as connection:
connection.execute(
text(
"INSERT INTO embedding_indexes(spec, vector_table, state) "
"VALUES (:spec, 'chunk_vectors', 'active')"
),
{"spec": BGE.model_copy(update={"dimension": 768}).model_dump_json()},
)
with pytest.raises(RuntimeError, match="768"):
upgrade_to_head(engine)
def test_a_fresh_install_has_not_chosen_yet(tmp_path: Path) -> None:
"""Onboarding offers the choice, so nothing is recorded for it."""
engine = create_db_engine(tmp_path / "surfsense.db")
upgrade_to_head(engine)
with create_session_factory(engine)() as session:
assert active_index(session) is None

Some files were not shown because too many files have changed in this diff Show more