enhance embedding model support and update database schema for multi-model compatibility

This commit is contained in:
ImBenji 2026-04-18 17:04:09 +01:00
parent 3b8955c80c
commit 88f465e71f
2 changed files with 48 additions and 15 deletions

View file

@ -112,13 +112,6 @@ db.exec(`
CREATE INDEX IF NOT EXISTS idx_articles_normalized_title ON articles(normalized_title); CREATE INDEX IF NOT EXISTS idx_articles_normalized_title ON articles(normalized_title);
`); `);
db.exec(`
CREATE VIRTUAL TABLE IF NOT EXISTS article_embeddings USING vec0(
article_id INTEGER PRIMARY KEY,
embedding FLOAT[1024]
);
`);
db.exec(` db.exec(`
CREATE TABLE IF NOT EXISTS article_embedding_store ( CREATE TABLE IF NOT EXISTS article_embedding_store (
article_id INTEGER NOT NULL, article_id INTEGER NOT NULL,
@ -137,6 +130,30 @@ db.exec(`
); );
`); `);
// vec0 table — fixed at 8192 dims to cover any model on openrouter, shorter embeddings get zero-padded
{
const existing = db.prepare(`SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'article_embeddings'`).get();
const currentDim = existing && existing.sql && existing.sql.match(/FLOAT\[(\d+)\]/);
if (!existing) {
db.exec(`
CREATE VIRTUAL TABLE article_embeddings USING vec0(
article_id INTEGER PRIMARY KEY,
embedding FLOAT[8192]
);
`);
} else if (!currentDim || parseInt(currentDim[1], 10) !== 8192) {
db.exec(`DROP TABLE article_embeddings`);
db.exec(`DELETE FROM article_embedding_meta`);
db.exec(`
CREATE VIRTUAL TABLE article_embeddings USING vec0(
article_id INTEGER PRIMARY KEY,
embedding FLOAT[8192]
);
`);
}
}
// migrate query_embeddings to include model in primary key // migrate query_embeddings to include model in primary key
{ {
const cols = db.prepare(`PRAGMA table_info(query_embeddings)`).all(); const cols = db.prepare(`PRAGMA table_info(query_embeddings)`).all();

View file

@ -124,7 +124,10 @@ function rebuildVec0IfModelChanged() {
const insertMeta = db.prepare(`INSERT INTO article_embedding_meta (article_id, model) VALUES (?, ?)`); const insertMeta = db.prepare(`INSERT INTO article_embedding_meta (article_id, model) VALUES (?, ?)`);
for (const row of rows) { for (const row of rows) {
insertVec.run(BigInt(row.article_id), row.embedding); const vals = new Float32Array(row.embedding.buffer, row.embedding.byteOffset, row.embedding.byteLength / 4);
const padded = new Float32Array(VEC0_DIM);
padded.set(vals);
insertVec.run(BigInt(row.article_id), Buffer.from(padded.buffer));
insertMeta.run(row.article_id, EMBEDDING_MODEL); insertMeta.run(row.article_id, EMBEDDING_MODEL);
} }
@ -153,10 +156,19 @@ function buildEmbeddingInput(article) {
return [title, description, content].join('\n\n'); return [title, description, content].join('\n\n');
} }
const VEC0_DIM = 8192;
function serializeEmbedding(values) { function serializeEmbedding(values) {
return Buffer.from(new Float32Array(values).buffer); return Buffer.from(new Float32Array(values).buffer);
} }
function padEmbeddingForVec0(values) {
if (values.length === VEC0_DIM) return serializeEmbedding(values);
const padded = new Float32Array(VEC0_DIM);
padded.set(values);
return Buffer.from(padded.buffer);
}
function normalizeQuery(input) { function normalizeQuery(input) {
return String(input || '') return String(input || '')
.trim() .trim()
@ -197,8 +209,8 @@ async function requestEmbedding(input) {
const payload = await response.json(); const payload = await response.json();
const embedding = payload && payload.data && payload.data[0] && payload.data[0].embedding; const embedding = payload && payload.data && payload.data[0] && payload.data[0].embedding;
if (!Array.isArray(embedding) || embedding.length !== 1024) { if (!Array.isArray(embedding) || embedding.length === 0) {
throw new Error(`unexpected embedding length: ${Array.isArray(embedding) ? embedding.length : 'missing'}`); throw new Error(`invalid embedding in response: ${Array.isArray(embedding) ? 'empty' : 'missing'}`);
} }
return embedding; return embedding;
@ -221,8 +233,9 @@ async function generateAndStoreEmbedding(id) {
// already in store — make sure vec0 is also up to date // already in store — make sure vec0 is also up to date
if (!selectEmbeddingBuffer.get(id)) { if (!selectEmbeddingBuffer.get(id)) {
const row = selectEmbeddingFromStore.get(id, EMBEDDING_MODEL); const row = selectEmbeddingFromStore.get(id, EMBEDDING_MODEL);
const vals = new Float32Array(row.embedding.buffer, row.embedding.byteOffset, row.embedding.byteLength / 4);
deleteEmbedding.run(BigInt(id)); deleteEmbedding.run(BigInt(id));
insertEmbedding.run(BigInt(id), row.embedding); insertEmbedding.run(BigInt(id), padEmbeddingForVec0(vals));
upsertEmbeddingMeta.run(id, EMBEDDING_MODEL); upsertEmbeddingMeta.run(id, EMBEDDING_MODEL);
} }
@ -251,7 +264,7 @@ async function generateAndStoreEmbedding(id) {
upsertEmbeddingStore.run(id, EMBEDDING_MODEL, buffer); upsertEmbeddingStore.run(id, EMBEDDING_MODEL, buffer);
deleteEmbedding.run(BigInt(id)); deleteEmbedding.run(BigInt(id));
insertEmbedding.run(BigInt(id), buffer); insertEmbedding.run(BigInt(id), padEmbeddingForVec0(embedding));
upsertEmbeddingMeta.run(id, EMBEDDING_MODEL); upsertEmbeddingMeta.run(id, EMBEDDING_MODEL);
return { stored: true, shouldPauseBatch: false }; return { stored: true, shouldPauseBatch: false };
@ -289,7 +302,9 @@ async function backfillMissingEmbeddings(limit = 100) {
function getEmbeddingBuffer(articleId) { function getEmbeddingBuffer(articleId) {
const row = selectEmbeddingFromStore.get(articleId, EMBEDDING_MODEL); const row = selectEmbeddingFromStore.get(articleId, EMBEDDING_MODEL);
return row ? row.embedding : null; if (!row) return null;
const vals = new Float32Array(row.embedding.buffer, row.embedding.byteOffset, row.embedding.byteLength / 4);
return padEmbeddingForVec0(vals);
} }
function findArticlesByEmbedding(embedding, limit) { function findArticlesByEmbedding(embedding, limit) {
@ -312,13 +327,14 @@ async function getOrCreateQueryEmbedding(query) {
const cached = selectQueryEmbedding.get(normalizedQuery, EMBEDDING_MODEL); const cached = selectQueryEmbedding.get(normalizedQuery, EMBEDDING_MODEL);
if (cached) { if (cached) {
return cached.embedding; const vals = new Float32Array(cached.embedding.buffer, cached.embedding.byteOffset, cached.embedding.byteLength / 4);
return padEmbeddingForVec0(vals);
} }
const embedding = await requestEmbedding(normalizedQuery); const embedding = await requestEmbedding(normalizedQuery);
const buffer = serializeEmbedding(embedding); const buffer = serializeEmbedding(embedding);
upsertQueryEmbedding.run(normalizedQuery, EMBEDDING_MODEL, buffer); upsertQueryEmbedding.run(normalizedQuery, EMBEDDING_MODEL, buffer);
return buffer; return padEmbeddingForVec0(embedding);
} }
function findSimilarArticles(articleId, limit) { function findSimilarArticles(articleId, limit) {