enhance embedding model support and update database schema for multi-model compatibility
This commit is contained in:
parent
3b8955c80c
commit
88f465e71f
2 changed files with 48 additions and 15 deletions
31
src/db.js
31
src/db.js
|
|
@ -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();
|
||||||
|
|
|
||||||
|
|
@ -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) {
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue