akm-cli/dist/storage/repositories/embeddings-repository.js

akm (Agent Knowledge Manager) — a portable, local-first capability library for AI agents. Discover, load, share, and improve reusable skills, scripts, workflows, and knowledge across any shell-capable coding agent, including Claude Code, OpenCode, and Cur

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
/**
 * Convert a `number[]` embedding vector to the `Float32Array` byte
 * representation stored in the `body_embeddings.embedding` BLOB column.
 */
export function embeddingToBlob(vec) {
    const f32 = new Float32Array(vec);
    return new Uint8Array(f32.buffer);
}
/**
 * Convert the raw `Uint8Array` bytes from the `body_embeddings.embedding`
 * BLOB column back to a `number[]` embedding vector.
 */
export function blobToEmbedding(blob) {
    // SQLite BLOB columns are returned as Uint8Array; re-interpret as Float32.
    const f32 = new Float32Array(blob.buffer, blob.byteOffset, blob.byteLength / 4);
    return Array.from(f32);
}
/**
 * Bulk-fetch cached body embeddings for a set of content hashes.
 * Returns a Map keyed by `content_hash` (embedding decoded to `number[]`).
 * Empty input → empty map (no query issued).
 *
 * If the stored `model_id` does not match `expectedModelId` the entire table
 * is cleared (drop-all on model mismatch) and an empty map is returned so
 * callers re-embed everything on this run.
 */
export function getBodyEmbeddings(db, contentHashes, expectedModelId) {
    const out = new Map();
    if (contentHashes.length === 0)
        return out;
    // Model-id mismatch: vectors are in the wrong metric space — drop all rows.
    const firstRow = db.prepare("SELECT model_id FROM body_embeddings LIMIT 1").get();
    if (firstRow && firstRow.model_id !== expectedModelId) {
        db.exec("DELETE FROM body_embeddings");
        return out;
    }
    // SQLite has a ~999 param ceiling; chunk if needed.
    const CHUNK = 500;
    for (let i = 0; i < contentHashes.length; i += CHUNK) {
        const chunk = contentHashes.slice(i, i + CHUNK);
        const placeholders = chunk.map(() => "?").join(",");
        const rows = db
            .prepare(`SELECT content_hash, embedding FROM body_embeddings WHERE content_hash IN (${placeholders})`)
            .all(...chunk);
        for (const row of rows) {
            out.set(row.content_hash, blobToEmbedding(row.embedding));
        }
    }
    return out;
}
/**
 * Upsert body-embedding rows in a single transaction.
 * Each entry maps a `cacheHash` → `number[]` vector. `model_id` is stored
 * so a future model change can trigger a drop-all purge.
 */
export function upsertBodyEmbeddings(db, entries) {
    if (entries.length === 0)
        return;
    const now = Date.now();
    const stmt = db.prepare(`
    INSERT OR REPLACE INTO body_embeddings (content_hash, embedding, model_id, created_at)
    VALUES (?, ?, ?, ?)
  `);
    db.transaction(() => {
        for (const { contentHash, embedding, modelId } of entries) {
            stmt.run(contentHash, embeddingToBlob(embedding), modelId, now);
        }
    })();
}