mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-09-13 18:32:12 +03:00
fix(memory): measure the embedding width instead of waiting for a probe (#12180)
* fix(memory): measure the embedding width instead of waiting for a probe
resolveEmbeddingSource() reports dimensions: null for any source the
hard-coded registry does not describe, and a self-hosted endpoint is by
definition absent from it. Both write paths then deadlocked on that null:
- scheduleVectorUpsert called ensureReady() with the null resolution, which
declines to create vec_memories, and then ignored the {ready:false} answer
and upserted anyway -- straight into the catch, so every memory was stored,
marked needs_reindex, and never vectorized;
- reindexPending refused to embed until the width was known, and the width
could only ever come from an embedding.
Nothing surfaced it: POST /api/memory returned 200 and the health check
stayed green while rowCount stayed at 0.
The comment on EmbeddingResolution.dimensions already calls this a lazy
probe; nobody performed the probe. The upsert path holds a finished vector
when it calls ensureReady, so measure it there, and let reindex spend one
embedding up front to measure -- reusing that vector rather than paying for
it twice. withMeasuredDimensions rebuilds the signature the same way the
resolution did, identity first, so two endpoints serving the same model id
still reindex independently.
scheduleVectorUpsert now also honours a {ready:false} answer instead of
upserting into a table that is not there.
Fixes #12154
* chore(changelog): point the fragment at the real PR number
* fix(memory): extract reindex helpers so the complexity ratchet stays green
runReindexBatch grew past max-lines-per-function and cognitive-complexity
when the lazy-probe path landed. Split measure/ready/item helpers without
changing the #12154 behavior.
Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>
---------
Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>
Co-authored-by: Diego Rodrigues de Sa e Souza <diegosouza.pw@gmail.com>
This commit is contained in:
5
changelog.d/fixes/12180-embedding-lazy-probe.md
Normal file
5
changelog.d/fixes/12180-embedding-lazy-probe.md
Normal file
@@ -0,0 +1,5 @@
|
||||
- **fix(memory):** self-hosted embedding endpoints now vectorize — the vector width is
|
||||
measured from the first embedding that comes back instead of being read from a registry
|
||||
that cannot describe them, so `vec_memories` is created and memories stop piling up
|
||||
unvectorized behind a green health check
|
||||
([#12180](https://github.com/diegosouzapw/OmniRoute/pull/12180)) — thanks @kanade-hoshino
|
||||
@@ -52,6 +52,39 @@ function resolveRemoteDimensions(model: string): number | null {
|
||||
return typeof dim === "number" ? dim : null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Fill in the vector width the lazy probe was waiting for.
|
||||
*
|
||||
* `dimensions` is null for every source the hard-coded registry does not
|
||||
* describe — a self-hosted endpoint by definition — and the only thing that can
|
||||
* answer it is an embedding that has actually come back. Callers that hold one
|
||||
* pass its length here; the signature is rebuilt the same way the resolution
|
||||
* built it, so reindex detection still sees a model change as a change. (#12154)
|
||||
*/
|
||||
export function withMeasuredDimensions(
|
||||
resolution: EmbeddingResolution,
|
||||
dimensions: number
|
||||
): EmbeddingResolution {
|
||||
if (
|
||||
resolution.dimensions !== null ||
|
||||
!resolution.source ||
|
||||
!Number.isInteger(dimensions) ||
|
||||
dimensions <= 0
|
||||
) {
|
||||
return resolution;
|
||||
}
|
||||
return {
|
||||
...resolution,
|
||||
dimensions,
|
||||
signature: makeSignature(
|
||||
resolution.source,
|
||||
resolution.identity ?? resolution.model,
|
||||
dimensions
|
||||
),
|
||||
reason: `${resolution.reason} [dim=${dimensions} measured]`,
|
||||
};
|
||||
}
|
||||
|
||||
/** Build the remote EmbeddingResolution used by both explicit + auto paths. */
|
||||
function remoteResolution(model: string, reasonPrefix: string): EmbeddingResolution {
|
||||
const dimensions = resolveRemoteDimensions(model);
|
||||
|
||||
@@ -8,7 +8,8 @@ import {
|
||||
countMemoryReindexPending,
|
||||
markMemoryNeedsReindex,
|
||||
} from "@/lib/db/memoryVec";
|
||||
import { resolveEmbeddingSource, embed } from "./embedding";
|
||||
import { resolveEmbeddingSource, embed, withMeasuredDimensions } from "./embedding";
|
||||
import type { EmbeddingResolution } from "./embedding/types";
|
||||
import { getVectorStore } from "./vectorStore";
|
||||
import { getMemorySettings } from "./settings";
|
||||
import { logger } from "../../../open-sse/utils/logger.ts";
|
||||
@@ -16,6 +17,97 @@ import { sanitizeErrorMessage } from "../../../open-sse/utils/error.ts";
|
||||
|
||||
const log = logger("MEMORY_REINDEX");
|
||||
|
||||
type ReindexItem = { id: string; content: string; key: string };
|
||||
type MemorySettings = Awaited<ReturnType<typeof getMemorySettings>>;
|
||||
type VectorStore = NonNullable<ReturnType<typeof getVectorStore>>;
|
||||
|
||||
function errMsg(err: unknown): string {
|
||||
return sanitizeErrorMessage(err instanceof Error ? err.message : String(err));
|
||||
}
|
||||
|
||||
/**
|
||||
* Nothing but a returned embedding can supply the vector width for a source the
|
||||
* registry does not describe. Spend one embed to measure it and reuse that
|
||||
* vector rather than paying for it twice (#12154).
|
||||
*/
|
||||
async function measureUnknownWidth(
|
||||
resolution: EmbeddingResolution,
|
||||
probeItem: ReindexItem | undefined,
|
||||
settings: MemorySettings
|
||||
): Promise<{ effective: EmbeddingResolution; probed: Map<string, Float32Array> }> {
|
||||
const probed = new Map<string, Float32Array>();
|
||||
if (!probeItem) {
|
||||
return { effective: resolution, probed };
|
||||
}
|
||||
const probe = await embed(probeItem.content, settings);
|
||||
if (!("vector" in probe)) {
|
||||
return { effective: resolution, probed };
|
||||
}
|
||||
probed.set(probeItem.id, probe.vector);
|
||||
return {
|
||||
effective: withMeasuredDimensions(resolution, probe.vector.length),
|
||||
probed,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* ensureReady() returns `{ ready: false }` (without throwing) when dimensions
|
||||
* are still unknown — abort so we don't burn embed credits upserting into a
|
||||
* missing `vec_memories` table (#8074).
|
||||
*/
|
||||
async function ensureReindexStoreReady(
|
||||
vec: VectorStore,
|
||||
effective: EmbeddingResolution,
|
||||
resolution: EmbeddingResolution,
|
||||
pending: number
|
||||
): Promise<boolean> {
|
||||
try {
|
||||
const ready = await vec.ensureReady(effective);
|
||||
if (ready.ready) return true;
|
||||
log.warn("memory.reindex.ensure_ready.skipped", {
|
||||
reason: ready.reason,
|
||||
pending,
|
||||
model: resolution.model,
|
||||
dimensions: resolution.dimensions,
|
||||
});
|
||||
return false;
|
||||
} catch (err: unknown) {
|
||||
log.warn("memory.reindex.ensure_ready.fail", { error: errMsg(err) });
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
async function reindexOneItem(
|
||||
item: ReindexItem,
|
||||
settings: MemorySettings,
|
||||
vec: VectorStore,
|
||||
probed: Map<string, Float32Array>
|
||||
): Promise<"processed" | "error"> {
|
||||
try {
|
||||
const reusable = probed.get(item.id);
|
||||
const embeddingResult = reusable ? { vector: reusable } : await embed(item.content, settings);
|
||||
|
||||
if (!("vector" in embeddingResult)) {
|
||||
log.warn("memory.reindex.embed.fail", {
|
||||
id: item.id,
|
||||
reason: embeddingResult.reason,
|
||||
message: sanitizeErrorMessage(embeddingResult.message),
|
||||
});
|
||||
return "error";
|
||||
}
|
||||
|
||||
await vec.upsertVector(item.id, embeddingResult.vector);
|
||||
markMemoryNeedsReindex(item.id, false);
|
||||
return "processed";
|
||||
} catch (err: unknown) {
|
||||
log.warn("memory.reindex.item.fail", {
|
||||
id: item.id,
|
||||
error: errMsg(err),
|
||||
});
|
||||
return "error";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Process up to `limit` memories that are marked needs_reindex=1.
|
||||
* Generates embedding + upserts into sqlite-vec for each.
|
||||
@@ -30,7 +122,6 @@ export async function runReindexBatch(limit = 100): Promise<{ processed: number;
|
||||
return { processed: 0, errors: 0 };
|
||||
}
|
||||
|
||||
// Resolve embedding source and vector store once for the whole batch
|
||||
const settings = await getMemorySettings();
|
||||
const resolution = resolveEmbeddingSource(settings);
|
||||
|
||||
@@ -48,25 +139,11 @@ export async function runReindexBatch(limit = 100): Promise<{ processed: number;
|
||||
return { processed: 0, errors: 0 };
|
||||
}
|
||||
|
||||
// Ensure the vector table is ready before processing. ensureReady() returns
|
||||
// `{ ready: false }` (without throwing) when dimensions are still unknown —
|
||||
// abort the batch in that case so we don't burn embed credits upserting into
|
||||
// a missing `vec_memories` table (#8074).
|
||||
try {
|
||||
const ready = await vec.ensureReady(resolution);
|
||||
if (!ready.ready) {
|
||||
log.warn("memory.reindex.ensure_ready.skipped", {
|
||||
reason: ready.reason,
|
||||
pending: queue.length,
|
||||
model: resolution.model,
|
||||
dimensions: resolution.dimensions,
|
||||
});
|
||||
return { processed: 0, errors: 0 };
|
||||
}
|
||||
} catch (err: unknown) {
|
||||
log.warn("memory.reindex.ensure_ready.fail", {
|
||||
error: sanitizeErrorMessage(err instanceof Error ? err.message : String(err)),
|
||||
});
|
||||
const probeItem = resolution.dimensions === null ? queue[0] : undefined;
|
||||
const { effective, probed } = await measureUnknownWidth(resolution, probeItem, settings);
|
||||
|
||||
const ready = await ensureReindexStoreReady(vec, effective, resolution, queue.length);
|
||||
if (!ready) {
|
||||
return { processed: 0, errors: 0 };
|
||||
}
|
||||
|
||||
@@ -74,29 +151,9 @@ export async function runReindexBatch(limit = 100): Promise<{ processed: number;
|
||||
let errors = 0;
|
||||
|
||||
for (const item of queue) {
|
||||
try {
|
||||
const embeddingResult = await embed(item.content, settings);
|
||||
|
||||
if (!("vector" in embeddingResult)) {
|
||||
log.warn("memory.reindex.embed.fail", {
|
||||
id: item.id,
|
||||
reason: embeddingResult.reason,
|
||||
message: sanitizeErrorMessage(embeddingResult.message),
|
||||
});
|
||||
errors++;
|
||||
continue;
|
||||
}
|
||||
|
||||
await vec.upsertVector(item.id, embeddingResult.vector);
|
||||
markMemoryNeedsReindex(item.id, false);
|
||||
processed++;
|
||||
} catch (err: unknown) {
|
||||
log.warn("memory.reindex.item.fail", {
|
||||
id: item.id,
|
||||
error: sanitizeErrorMessage(err instanceof Error ? err.message : String(err)),
|
||||
});
|
||||
errors++;
|
||||
}
|
||||
const outcome = await reindexOneItem(item, settings, vec, probed);
|
||||
if (outcome === "processed") processed++;
|
||||
else errors++;
|
||||
}
|
||||
|
||||
log.info("memory.reindex.batch.complete", { processed, errors, batchSize: queue.length });
|
||||
|
||||
@@ -7,7 +7,7 @@ import { upsertSemanticMemoryPoint, deleteSemanticMemoryPoint } from "./qdrant";
|
||||
import { Memory, MemoryType } from "./types";
|
||||
import { logger } from "../../../open-sse/utils/logger.ts";
|
||||
import { sanitizeErrorMessage } from "../../../open-sse/utils/error.ts";
|
||||
import { resolveEmbeddingSource, embed } from "./embedding";
|
||||
import { resolveEmbeddingSource, embed, withMeasuredDimensions } from "./embedding";
|
||||
import { getVectorStore } from "./vectorStore";
|
||||
import { getMemorySettings } from "./settings";
|
||||
import { markMemoryNeedsReindex } from "@/lib/db/memoryVec";
|
||||
@@ -154,7 +154,18 @@ function scheduleVectorUpsert(id: string, content: string): void {
|
||||
return;
|
||||
}
|
||||
|
||||
await vec.ensureReady(resolution);
|
||||
// The vector in hand is the lazy probe the resolution is waiting for: the
|
||||
// registry has no width for a self-hosted endpoint, so without this
|
||||
// ensureReady() never creates vec_memories and every upsert below fails
|
||||
// into the catch, leaving the memory stored but never vectorized (#12154).
|
||||
const ready = await vec.ensureReady(
|
||||
withMeasuredDimensions(resolution, embeddingResult.vector.length)
|
||||
);
|
||||
if (!ready.ready) {
|
||||
log.warn("memory.vec.ensure_ready.skipped", { id, reason: ready.reason });
|
||||
safeMarkNeedsReindex(id, true);
|
||||
return;
|
||||
}
|
||||
await vec.upsertVector(id, embeddingResult.vector);
|
||||
safeMarkNeedsReindex(id, false);
|
||||
} catch (err: unknown) {
|
||||
|
||||
77
tests/unit/memory-vec-lazy-probe-12154.test.ts
Normal file
77
tests/unit/memory-vec-lazy-probe-12154.test.ts
Normal file
@@ -0,0 +1,77 @@
|
||||
/**
|
||||
* #12154 — a self-hosted embedding endpoint never got `vec_memories` created,
|
||||
* so memories were stored but never vectorized while health stayed green.
|
||||
*
|
||||
* `resolveEmbeddingSource` returns `dimensions: null` for any source the
|
||||
* hard-coded registry does not describe, and a self-hosted endpoint is by
|
||||
* definition absent from it. Both write paths then deadlocked: the vector store
|
||||
* refuses to create the table without a width, and the width can only come from
|
||||
* an embedding that has actually come back.
|
||||
*/
|
||||
import test from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
|
||||
const { withMeasuredDimensions } = await import("../../src/lib/memory/embedding/index.ts");
|
||||
|
||||
function resolution(overrides = {}) {
|
||||
return {
|
||||
source: "remote",
|
||||
model: "memory-custom/Qwen3-Embedding-0.6B",
|
||||
dimensions: null,
|
||||
identity: "http://tei.internal:8080|Qwen3-Embedding-0.6B",
|
||||
signature: "remote:http://tei.internal:8080|Qwen3-Embedding-0.6B:null",
|
||||
reason: "custom remote provider configured (dim=unknown, will probe at embed time)",
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
test("a measured width fills in the pending lazy probe", () => {
|
||||
const filled = withMeasuredDimensions(resolution(), 1024);
|
||||
assert.equal(filled.dimensions, 1024);
|
||||
assert.match(filled.reason, /dim=1024 measured/);
|
||||
});
|
||||
|
||||
test("the signature keeps its identity and gains the width", () => {
|
||||
const filled = withMeasuredDimensions(resolution(), 1024);
|
||||
// Identity, not model: the same model id can exist at several custom endpoints,
|
||||
// and the resolution built its signature that way too.
|
||||
assert.equal(filled.signature, "remote:http://tei.internal:8080|Qwen3-Embedding-0.6B:1024");
|
||||
});
|
||||
|
||||
test("a resolution with no identity signs by model", () => {
|
||||
const filled = withMeasuredDimensions(
|
||||
resolution({
|
||||
identity: undefined,
|
||||
model: "openai/text-embedding-3-small",
|
||||
signature: "remote:openai/text-embedding-3-small:null",
|
||||
}),
|
||||
1536
|
||||
);
|
||||
assert.equal(filled.signature, "remote:openai/text-embedding-3-small:1536");
|
||||
});
|
||||
|
||||
test("a width the registry already knows is never overwritten", () => {
|
||||
const known = resolution({ dimensions: 1536, signature: "remote:openai/x:1536" });
|
||||
assert.equal(withMeasuredDimensions(known, 1024), known);
|
||||
});
|
||||
|
||||
test("a nonsense measurement is ignored rather than written into the signature", () => {
|
||||
const pending = resolution();
|
||||
for (const bad of [0, -1, 1.5, Number.NaN]) {
|
||||
assert.equal(withMeasuredDimensions(pending, bad), pending, `width ${bad}`);
|
||||
}
|
||||
});
|
||||
|
||||
test("a resolution with no source stays unusable", () => {
|
||||
const none = resolution({ source: null, model: null, signature: "null:null:null" });
|
||||
assert.equal(withMeasuredDimensions(none, 1024), none);
|
||||
});
|
||||
|
||||
test("two endpoints serving the same model id do not share a signature", () => {
|
||||
const a = withMeasuredDimensions(resolution(), 1024);
|
||||
const b = withMeasuredDimensions(
|
||||
resolution({ identity: "http://other.internal:8080|Qwen3-Embedding-0.6B" }),
|
||||
1024
|
||||
);
|
||||
assert.notEqual(a.signature, b.signature);
|
||||
});
|
||||
Reference in New Issue
Block a user