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);
`);
db.exec(`
CREATE VIRTUAL TABLE IF NOT EXISTS article_embeddings USING vec0(
article_id INTEGER PRIMARY KEY,
embedding FLOAT[1024]
);
`);
db.exec(`
CREATE TABLE IF NOT EXISTS article_embedding_store (
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
{
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 (?, ?)`);
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);
}
@ -153,10 +156,19 @@ function buildEmbeddingInput(article) {
return [title, description, content].join('\n\n');
}
const VEC0_DIM = 8192;
function serializeEmbedding(values) {
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) {
return String(input || '')
.trim()
@ -197,8 +209,8 @@ async function requestEmbedding(input) {
const payload = await response.json();
const embedding = payload && payload.data && payload.data[0] && payload.data[0].embedding;
if (!Array.isArray(embedding) || embedding.length !== 1024) {
throw new Error(`unexpected embedding length: ${Array.isArray(embedding) ? embedding.length : 'missing'}`);
if (!Array.isArray(embedding) || embedding.length === 0) {
throw new Error(`invalid embedding in response: ${Array.isArray(embedding) ? 'empty' : 'missing'}`);
}
return embedding;
@ -221,8 +233,9 @@ async function generateAndStoreEmbedding(id) {
// already in store — make sure vec0 is also up to date
if (!selectEmbeddingBuffer.get(id)) {
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));
insertEmbedding.run(BigInt(id), row.embedding);
insertEmbedding.run(BigInt(id), padEmbeddingForVec0(vals));
upsertEmbeddingMeta.run(id, EMBEDDING_MODEL);
}
@ -251,7 +264,7 @@ async function generateAndStoreEmbedding(id) {
upsertEmbeddingStore.run(id, EMBEDDING_MODEL, buffer);
deleteEmbedding.run(BigInt(id));
insertEmbedding.run(BigInt(id), buffer);
insertEmbedding.run(BigInt(id), padEmbeddingForVec0(embedding));
upsertEmbeddingMeta.run(id, EMBEDDING_MODEL);
return { stored: true, shouldPauseBatch: false };
@ -289,7 +302,9 @@ async function backfillMissingEmbeddings(limit = 100) {
function getEmbeddingBuffer(articleId) {
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) {
@ -312,13 +327,14 @@ async function getOrCreateQueryEmbedding(query) {
const cached = selectQueryEmbedding.get(normalizedQuery, EMBEDDING_MODEL);
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 buffer = serializeEmbedding(embedding);
upsertQueryEmbedding.run(normalizedQuery, EMBEDDING_MODEL, buffer);
return buffer;
return padEmbeddingForVec0(embedding);
}
function findSimilarArticles(articleId, limit) {