diff --git a/src/core/embed-stale.ts b/src/core/embed-stale.ts index 474a3ea03..4652c0adb 100644 --- a/src/core/embed-stale.ts +++ b/src/core/embed-stale.ts @@ -20,6 +20,7 @@ import type { BrainEngine } from './engine.ts'; import type { ChunkInput } from './types.ts'; import { embedBatchWithBackoff } from '../commands/embed.ts'; +import { resolveEmbeddingModelLabel } from './embedding.ts'; import { type DbPacer, createNoopPacer, observed } from './db-pacer.ts'; import { AbortError } from './abort-check.ts'; @@ -200,11 +201,17 @@ export async function embedStaleForSource( for (let j = 0; j < stale.length; j++) { staleIdxToEmbedding.set(stale[j].chunk_index, embeddings[j]); } + // #1717: label re-embedded chunks with the model that produced the + // vector; preserved chunks keep their existing model. Without this, + // upsertChunks falls back to DEFAULT_EMBEDDING_MODEL for every chunk + // (the same mislabel the embed.ts paths fixed). + const embedModelLabel = resolveEmbeddingModelLabel(); const merged: ChunkInput[] = existing.map((c) => ({ chunk_index: c.chunk_index, chunk_text: c.chunk_text, chunk_source: c.chunk_source, embedding: staleIdxToEmbedding.get(c.chunk_index) ?? undefined, + model: staleIdxToEmbedding.has(c.chunk_index) && embedModelLabel ? embedModelLabel : c.model, token_count: c.token_count || Math.ceil(c.chunk_text.length / 4), // Carry through per-chunk metadata. upsertChunks writes these as // EXCLUDED. (not COALESCE), so omitting them here resets image diff --git a/test/embed-stale.serial.test.ts b/test/embed-stale.serial.test.ts index 9c3a88221..efe22b007 100644 --- a/test/embed-stale.serial.test.ts +++ b/test/embed-stale.serial.test.ts @@ -15,6 +15,7 @@ import { describe, test, expect, beforeAll, afterAll, beforeEach } from 'bun:tes import { PGLiteEngine } from '../src/core/pglite-engine.ts'; import { resetPgliteState } from './helpers/reset-pglite.ts'; import { embedStaleForSource } from '../src/core/embed-stale.ts'; +import { configureGateway, resetGateway } from '../src/core/ai/gateway.ts'; import type { ChunkInput } from '../src/core/types.ts'; let engine: PGLiteEngine; @@ -276,4 +277,49 @@ describe('embedStaleForSource', () => { // The stale text row actually got its embedding. expect(txtRow.embedded_at).not.toBeNull(); }); + + // #1717: the backfill path must label re-embedded chunks with the model + // that produced the vector, and preserve the existing label on chunks it + // did not touch (before the fix, both were reset to the engine default). + test('labels re-embedded chunks with the gateway model, preserves untouched labels (#1717)', async () => { + configureGateway({ + embedding_model: 'openai:text-embedding-3-large', + env: { OPENAI_API_KEY: 'sk-test-embed-stale-1717' }, + }); + try { + await engine.putPage('notes/model-label', { + type: 'note', + title: 'model-label', + compiled_truth: '# model-label\n\nseeded', + }); + await engine.upsertChunks('notes/model-label', [ + { + chunk_index: 0, + chunk_text: 'already embedded elsewhere', + chunk_source: 'compiled_truth', + embedding: new Float32Array(1536).fill(0.01), + model: 'voyage:voyage-3', + token_count: 4, + }, + { + chunk_index: 1, + chunk_text: 'stale chunk needing embed', + chunk_source: 'compiled_truth', + token_count: 5, + embedding: undefined, // stale + }, + ]); + + const result = await embedStaleForSource(engine, 'default', { embedFn: fakeEmbedFn }); + expect(result.embedded).toBe(1); + + const after = await engine.getChunks('notes/model-label'); + const preserved = after.find((c) => c.chunk_index === 0)!; + const reembedded = after.find((c) => c.chunk_index === 1)!; + expect(reembedded.model).toBe('openai:text-embedding-3-large'); + expect(preserved.model).toBe('voyage:voyage-3'); + } finally { + resetGateway(); + } + }); });