Support external/remote inference endpoints for embeddings and reranking (OpenAI-compatible), with local fallback #2

New Issue

2026-08-22T20:48:55+01:00

bobby commented

2026-08-22 20:48:55 +01:00

Summary

Request: allow the engine to call an external inference server for embeddings and reranking over an OpenAI-compatible HTTP API, falling back to the current in-process models when no URL is configured.

Today both models are loaded in-process and are hard-bound to a local torch device. _resolve_device() in engine/kb/embeddings.py only ever returns cuda or cpu — there is no third option and nothing in the engine takes a host or URL. Confirmed by grep: no references to base_url, api_base, or any HTTP model client anywhere under engine/ (httpx is a dev/test dependency only).

This is an enhancement request — please label accordingly (the repo currently has no labels defined).

Motivation

Reranking landed in the engine but is currently unusable on an 8 GB consumer GPU, because the reranker must share VRAM with the embedding model in the same process.

Observed on a host with an 8 GB NVIDIA card (<gpu-host>), engine v3.3.0:

Reranker model	`KB_RERANK_CANDIDATES`	Result
`BAAI/bge-reranker-v2-m3`	40	OOM — `Tried to allocate 3.98 GiB`, 1.07 GiB free
`tomaarsen/Qwen3-Reranker-0.6B-seq-cls`	20	OOM — `Tried to allocate 818.00 MiB`, 681 MiB free
`tomaarsen/Qwen3-Reranker-0.6B-seq-cls`	8	OOM — `Tried to allocate 144.00 MiB`, 45 MiB free

With both models resident the GPU sat at 7127 MiB / 8188 MiB at idle. Failing to allocate 144 MiB at 8 candidates shows candidate count is not the lever — the resident model is consuming the card. KB_RERANK_ENABLED is currently back to false in that deployment so search keeps working.

Being able to point the reranker at a separate inference host would sidestep this entirely, and would let a machine with a different/larger GPU serve models without moving the whole engine (and its SQLite data) to that machine.

Proposed change

Two new optional settings in engine/kb/config.py:

KB_EMBEDDINGS_URL — OpenAI-compatible base URL; when set, embed_texts() POSTs to /v1/embeddings instead of using the in-process SentenceTransformer.
KB_RERANKER_URL — when set, rerank_scores() POSTs to /v1/rerank instead of using the in-process CrossEncoder.

Plus KB_INFERENCE_API_KEY (optional bearer token) for servers that require auth.

When unset, behaviour is unchanged — models load locally exactly as now. This matches the graceful-degradation pattern engine/kb/reranker.py already documents ("Reranking is strictly optional — search degrades gracefully to plain hybrid retrieval when the model is absent").

/v1/embeddings and /v1/rerank are the de-facto standard: both llama.cpp's llama-server (with --pooling rank for rerank) and LocalAI expose them, as do several hosted providers. Implementing the interface rather than a specific vendor keeps this small.

Suggested scope split

Reranking is the low-risk half and could ship alone. It is stateless — the model can be swapped freely with no migration.

Embeddings are not symmetric. An existing corpus (~123k chunks in one deployment) is embedded with bge-base-en-v1.5 at 768 dims. Changing the embedding model requires a full reindex, and a model with a different dimension additionally requires a sqlite-vec schema change. Worth handling as a separate follow-up, ideally with a dimension check at startup that refuses to start rather than silently corrupting search.

Related bugs found while testing

These are worth fixing regardless of whether remote inference is adopted, since they are what makes local reranking fragile:

No batch_size or max_length in rerank_scores() (engine/kb/reranker.py). All candidates are passed to CrossEncoder.predict() in a single call, so peak VRAM scales with candidate count x longest chunk. bge-reranker-v2-m3 has an 8192-token limit, so a long chunk pads the whole batch — this is the direct cause of the 3.98 GiB allocation above.
Model loads in fp32. CrossEncoder(model_name, device=resolved_device) with no dtype. Passing model_kwargs={"torch_dtype": "float16"} would roughly halve resident VRAM and may be sufficient on its own for 8 GB cards.
The default reranker is heavy, and the proposal's sizing does not apply to it. docs/kb-enhancements-proposal.htm specifies qwen3-reranker-0.6b (~640MB GGUF), chosen explicitly as "small enough to sit alongside bge on the 4070 permanently". But CrossEncoder cannot load GGUF, and the default became BAAI/bge-reranker-v2-m3 — a 568M-param XLM-RoBERTa-large. The loadable HF safetensors port of qwen3-0.6b is 2.38 GB fp32, ~4x the budgeted size. Either the default should be something that fits (e.g. BAAI/bge-reranker-base, 278M, 512-token cap) or the docs should be corrected.
Qwen/Qwen3-Reranker-0.6B must not be used directly. It is Qwen3ForCausalLM — no sequence-classification head. Loading it via CrossEncoder would not error: transformers randomly initialises the score head, giving loaded: true in /status and silently meaningless rankings. A startup guard that checks config.architectures ends in ForSequenceClassification — or at minimum a documented warning — would prevent a silent quality regression. (tomaarsen/Qwen3-Reranker-0.6B-seq-cls is a correct port and loads cleanly with no random-init warning.)

Environment

Engine v3.3.0, client v3.3.0
Reranker host: 8 GB NVIDIA consumer GPU, nvidia image variant
Candidate external host: 16 GB AMD GPU, ROCm 7.2.3, currently running other models

## Summary Request: allow the engine to call an **external inference server** for embeddings and reranking over an OpenAI-compatible HTTP API, falling back to the current in-process models when no URL is configured. Today both models are loaded in-process and are hard-bound to a local torch device. `_resolve_device()` in `engine/kb/embeddings.py` only ever returns `cuda` or `cpu` — there is no third option and nothing in the engine takes a host or URL. Confirmed by grep: no references to `base_url`, `api_base`, or any HTTP model client anywhere under `engine/` (`httpx` is a dev/test dependency only). This is an **enhancement request** — please label accordingly (the repo currently has no labels defined). ## Motivation Reranking landed in the engine but is currently unusable on an 8 GB consumer GPU, because the reranker must share VRAM with the embedding model in the same process. Observed on a host with an 8 GB NVIDIA card (`<gpu-host>`), engine v3.3.0: | Reranker model | `KB_RERANK_CANDIDATES` | Result | |---|---|---| | `BAAI/bge-reranker-v2-m3` | 40 | OOM — `Tried to allocate 3.98 GiB`, 1.07 GiB free | | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls` | 20 | OOM — `Tried to allocate 818.00 MiB`, 681 MiB free | | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls` | 8 | OOM — `Tried to allocate 144.00 MiB`, **45 MiB free** | With both models resident the GPU sat at **7127 MiB / 8188 MiB at idle**. Failing to allocate 144 MiB at 8 candidates shows candidate count is not the lever — the resident model is consuming the card. `KB_RERANK_ENABLED` is currently back to `false` in that deployment so search keeps working. Being able to point the reranker at a separate inference host would sidestep this entirely, and would let a machine with a different/larger GPU serve models without moving the whole engine (and its SQLite data) to that machine. ## Proposed change Two new optional settings in `engine/kb/config.py`: - `KB_EMBEDDINGS_URL` — OpenAI-compatible base URL; when set, `embed_texts()` POSTs to `/v1/embeddings` instead of using the in-process `SentenceTransformer`. - `KB_RERANKER_URL` — when set, `rerank_scores()` POSTs to `/v1/rerank` instead of using the in-process `CrossEncoder`. Plus `KB_INFERENCE_API_KEY` (optional bearer token) for servers that require auth. When unset, behaviour is unchanged — models load locally exactly as now. This matches the graceful-degradation pattern `engine/kb/reranker.py` already documents ("Reranking is strictly optional — search degrades gracefully to plain hybrid retrieval when the model is absent"). `/v1/embeddings` and `/v1/rerank` are the de-facto standard: both llama.cpp's `llama-server` (with `--pooling rank` for rerank) and LocalAI expose them, as do several hosted providers. Implementing the interface rather than a specific vendor keeps this small. ### Suggested scope split **Reranking is the low-risk half and could ship alone.** It is stateless — the model can be swapped freely with no migration. **Embeddings are not symmetric.** An existing corpus (~123k chunks in one deployment) is embedded with `bge-base-en-v1.5` at 768 dims. Changing the embedding model requires a full `reindex`, and a model with a different dimension additionally requires a `sqlite-vec` schema change. Worth handling as a separate follow-up, ideally with a dimension check at startup that refuses to start rather than silently corrupting search. ## Related bugs found while testing These are worth fixing regardless of whether remote inference is adopted, since they are what makes local reranking fragile: 1. **No `batch_size` or `max_length` in `rerank_scores()`** (`engine/kb/reranker.py`). All candidates are passed to `CrossEncoder.predict()` in a single call, so peak VRAM scales with candidate count x longest chunk. `bge-reranker-v2-m3` has an 8192-token limit, so a long chunk pads the whole batch — this is the direct cause of the 3.98 GiB allocation above. 2. **Model loads in fp32.** `CrossEncoder(model_name, device=resolved_device)` with no dtype. Passing `model_kwargs={"torch_dtype": "float16"}` would roughly halve resident VRAM and may be sufficient on its own for 8 GB cards. 3. **The default reranker is heavy, and the proposal's sizing does not apply to it.** `docs/kb-enhancements-proposal.htm` specifies `qwen3-reranker-0.6b (~640MB GGUF)`, chosen explicitly as "small enough to sit alongside bge on the 4070 permanently". But `CrossEncoder` cannot load GGUF, and the default became `BAAI/bge-reranker-v2-m3` — a 568M-param XLM-RoBERTa-large. The loadable HF safetensors port of qwen3-0.6b is 2.38 GB fp32, ~4x the budgeted size. Either the default should be something that fits (e.g. `BAAI/bge-reranker-base`, 278M, 512-token cap) or the docs should be corrected. 4. **`Qwen/Qwen3-Reranker-0.6B` must not be used directly.** It is `Qwen3ForCausalLM` — no sequence-classification head. Loading it via `CrossEncoder` would **not** error: transformers randomly initialises the `score` head, giving `loaded: true` in `/status` and silently meaningless rankings. A startup guard that checks `config.architectures` ends in `ForSequenceClassification` — or at minimum a documented warning — would prevent a silent quality regression. (`tomaarsen/Qwen3-Reranker-0.6B-seq-cls` is a correct port and loads cleanly with no random-init warning.) ## Environment - Engine v3.3.0, client v3.3.0 - Reranker host: 8 GB NVIDIA consumer GPU, `nvidia` image variant - Candidate external host: 16 GB AMD GPU, ROCm 7.2.3, currently running other models

steve added the enhancement label 2026-08-22 21:07:55 +01:00

Sign in to join this conversation.

1 Participants

Notifications

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: steve/kb#2