Skip to content

fix(embedding): enable RAG with the built-in local model and batch local embeddings - #134

Merged
mrsibe merged 2 commits into
mainfrom
fix/embedding-local-pipeline
Sep 25, 2026
Merged

mrsibe merged 2 commits into
mainfrom
fix/embedding-local-pipeline

Conversation

@mrsibe

@mrsibe mrsibe commented Sep 25, 2026

Copy link
Copy Markdown
Owner

What does this PR do?

Two fixes on the embedding path, both found from the same log line:

  1. RAG now runs with the built-in local model. chatHandlers decided whether to retrieve at all from ConnectionManager.getEmbeddingClient(), which only returns a client for a remote embedding connection. Local is the default, so RAG was skipped on every message.
  2. Local embeddings are batched. The remote path split work into batches; the local path sent the whole array through one ONNX forward pass.

It also stops logging the local fallback as a warning, and makes the import and test-loader changes needed to test any of this.

Why?

The RAG gate (the "chat does not retrieve" blocker)

const embeddingClient = await connectionManager.getEmbeddingClient()
if (embeddingClient) { /* retrieve + inject context */ }
else { Logger.debug('ChatHandlers', 'RAG disabled: No embedding model configured') }

getEmbeddingClient() is remote-only. With the built-in local model (the default) it is always null, so retrieval stayed none, no system context was injected, and no sources/citations were recorded — while the notebook was indexed correctly. EmbeddingService.resolveBackend() already falls back to the local backend; the search was never the problem, the gate was.

The same call produced [ConnectionManager] No embedding model configured on every index and query for local users. That is the normal fallback path, not an error.

Local batching (the "import hangs" blocker)

EmbeddingService.embedBatch() handled remote in batches of 20 but handed the local path the entire array in one call, which LocalEmbeddingBackend tokenizes and runs as a single forward pass. For the 694-chunk book in #133 that is one ~694 x 512 input: unbounded memory peak, a long synchronous tokenization, and onProgress firing once at the very end, so the UI progress sat at its start value.

Fixes #132. Fixes #133. Related to #125, #82.

What changed?

  • chatHandlers.ts: gate on knowledgeService.isEmbeddingAvailable() (remote configured or local model installed) instead of the remote-only client.
  • KnowledgeService.ts: add isEmbeddingAvailable(), delegating to EmbeddingService.isAvailable().
  • ConnectionManager.ts: the missing-remote-connection case is now a debug line that names the local fallback.
  • EmbeddingService.ts: new localBatchSize (default 16), embedLocalInBatches() — per-batch onProgress and an event-loop yield between batches. The ConnectionManager import becomes type-only (it was only used as a type and dragged electron/config into node --test).
  • test/ts-resolve.mjs: also resolve directory imports to index.ts, the other shape the sources use (../../shared/types).

How was this tested?

  • npm test — 281 pass, 0 fail (5 new).
    • test/embeddingBatching.test.ts: a fake pipeline records batch sizes — [16, 16, 8] for 40 texts with localBatchSize: 16, [16, 16, 1] for the default, per-batch progress [16,40] -> [32,40] -> [40,40], and the single-embed path.
    • test/chatRagGate.test.ts: fails if the gate goes back to connectionManager.getEmbeddingClient(). Verified by temporarily restoring that call.
  • npm run typecheck, npm run build — pass.
  • npm run build:unpack and smoke:packaged — 21/21 checks pass (run headless with an isolated --user-data-dir; CI uses xvfb-run).

Not verified: a real 694-chunk document end to end against the actual local model. The new tests pin the batching contract with a fake pipeline; the book-sized run is the acceptance criterion in #133, and the retrieval: 'used' path for a local-only profile is the acceptance criterion in #132.

Screenshots / recordings

Not applicable.

Checklist

  • I have reviewed my own changes.
  • npm run typecheck passes.
  • npm run build passes.
  • I have tested the affected user workflow.
  • I have not included unrelated changes.
  • I have updated documentation when necessary.

Desktop / build changes

  • npm run build:unpack passes.
  • npm run smoke:packaged passes.

`chatHandlers` decided whether to retrieve at all from
`ConnectionManager.getEmbeddingClient()`, which only returns a client for a
**remote** embedding connection. The built-in local model is the default
configuration, so this returned null on every message and the whole RAG block
was skipped: `retrieval` stayed 'none', no system context was injected, and no
sources or citations were recorded. Notebook chat answered from the model's own
knowledge while the notebook was indexed correctly.

Gate on `knowledgeService.isEmbeddingAvailable()` (remote configured, or the
local model installed) instead. `EmbeddingService.resolveBackend()` already
falls back to the local backend, so the search itself needs no change.

The same call produced the misleading warning
`[ConnectionManager] No embedding model configured` on every index and query for
local users. It is the normal fallback path, not an error, so it is now a debug
line that names the fallback.

Verification: `npm test` (the new test fails if the gate goes back to
`getEmbeddingClient()`), `npm run typecheck`, `npm run build`.
The remote path split work into batches; the local path handed the entire array
to the ONNX pipeline in one call. For a 694-chunk book that is one tokenization
and one forward pass over ~694 x 512 tokens: long stalls, an unbounded memory
peak, and `onProgress` firing once at the very end, so indexing progress sat at
its start value the whole time. Small fixtures never exercised it.

Split local work into batches of 16 (`localBatchSize`), report progress after
each batch, and yield to the event loop between batches so the synchronous
tokenization cannot starve IPC.

Two supporting changes so the behaviour is testable:
- `EmbeddingService`'s `ConnectionManager` import is type-only now. It was only
  used as a type, and the value import dragged the whole electron/config chain
  into `node --test`, which is why no test could import `EmbeddingService`.
- the test loader resolves directory imports to `index.ts`, the other shape the
  sources use (`../../shared/types`), alongside `./logger` -> `logger.ts`.

Verification: `npm test` covers batch boundaries ([16,16,8] for 40 texts), the
default of 16, per-batch progress, and the single-embed path. `npm run
typecheck`, `npm run build`.
@github-actions github-actions Bot added the bug Something isn't working label Sep 25, 2026
@mrsibe
mrsibe merged commit 925e27e into main Sep 25, 2026
4 checks passed
@mrsibe
mrsibe deleted the fix/embedding-local-pipeline branch September 25, 2026 17:59
mrsibe added a commit that referenced this pull request Sep 27, 2026
… either way (#147)

Fixes #146. A failed import used to produce nothing at all — no error, no row, and
a panel that still showed its empty state, which reads as "the upload worked and
the app is not showing it". #37 was closed with exactly that report and nothing in
the code could have told that user what went wrong.

Two halves, both needed:

**The failure was discarded before it could be shown.** `SourcePanel` calls the
store and ignores what it returns. The store reports failure by *returning*
`{ success: false, error }` — the IPC layer already catches and describes it — so
an unchecked return value makes a failure indistinguishable from a success.
`handleFileUpload`, `handleUrlImport`, `handleTextPaste` and `handleNoteImport`
all did this.

**The list was only re-read when the import succeeded.** `KnowledgeService` inserts
the source row and *then* indexes it, so a failure at the indexing step leaves a row
with `status: 'failed'` — and `DocumentList` already has the feedback for it, an
alert saying the document failed to embed. It never fired, because the list was
never re-read. A failure *before* the row exists (an unsupported type, a corrupt or
scanned file that extracts nothing) has no row and so needs the message to say so.

Now one rule covers every entry point: an import attempt always ends by re-reading
the library and recording the reason, and the reason reaches the reader as a toast.
The rejection path re-reads too, deliberately: an IPC call can fail after the main
process already wrote the source row, so the failure path cannot assume that
nothing changed.

Not in this PR: recording an attempt that failed before its source row existed
(there is still no row for an unparseable file — that is #95's ingestion record),
and any retry affordance (#95). The `hasEmbeddingModel` gate that disables `+` is
untouched; after #134 the built-in local backend makes it effectively always open.

One behaviour note: the empty-note case used to `alert()` from a branch that could
not be reached — the store returns `{ success: false }` rather than throwing, so the
`catch` never ran. It now shows the same message as a toast, like every other
import failure, and names the note instead of printing `Note note_123 not found`.

Verified: `npm test` 303 pass (6 new in `test/knowledgeStore.test.ts`, which drive all
four entry points through the store with a stubbed `window.api`), `npm run typecheck`,
`npm run check:design`, `npx eslint` on the changed files (0 problems), `npm run build`.

`tsconfig.test.json` now also includes `src/preload/index.d.ts`: importing renderer
code means typechecking it, and that file is where `window.api` is declared — without
it a test cannot import a store that legitimately uses the preload global.

Not run: `build:unpack` + `smoke:packaged`. Neither main nor preload changed; the
packaged checks are in CI for this PR and were run in #145 for the last main-process
change. Not hand-verified: the toast itself, which needs the GUI.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Local embedding processes every chunk in one batch [Bug] Notebook chat never retrieves when embeddings use the built-in local model

1 participant