quak/src/library/mlsearch.ts

// The content-similarity search surface over the CLIP index (issue #50).
//
// This is `lib.mldata`. It answers three questions against the ML-data cache
// (#49) without touching the network:
//
//   - `forFile` returns the whole stored payload (face boxes, landmarks,
//     embedding) for a file, read from disk on demand — the only method here
//     that touches the disk, and the only one that is async.
//   - `similar` and `searchByEmbedding` rank fileIDs by cosine similarity over
//     the packed `Float32Array` index alone. That index (~50k×512) already
//     lives in RAM, so each query is a plain loop over it and nothing else.
//
// quak bundles no text encoder (owner-deferred), so `searchByEmbedding` takes
// the query vector the caller has produced elsewhere; `similar` uses the
// query file's own indexed embedding.

import type { MLData } from "../mldata-fetch.js";
import type { MLDataStore, MLIndex } from "./mldata.js";

// How many nearest files a query returns when the caller names no limit.
const DEFAULT_LIMIT = 20;

// One ranked result: a fileID and its cosine similarity to the query, in
// [-1, 1]. Callers wanting only the ids read `.fileID`.
export interface SimilarResult {
    fileID: number;
    score: number;
}

export interface MLDataAPI {
    // The whole stored ML payload for a file, or undefined when it is not
    // cached. Reads the payload from disk, so it is async.
    forFile(args: { fileID: number }): Promise<MLData | undefined>;
    // The files nearest the given file by cosine over their CLIP embeddings,
    // most similar first, excluding the file itself. Empty when the file has
    // no indexed embedding.
    similar(args: { fileID: number; limit?: number }): SimilarResult[];
    // The files nearest a caller-supplied query embedding by cosine, most
    // similar first. Empty when the query is the wrong length for the index,
    // has zero magnitude, or the index is empty.
    searchByEmbedding(args: {
        embedding: ArrayLike<number>;
        limit?: number;
    }): SimilarResult[];
}

// Rank the packed index by cosine similarity to `query`, most similar first,
// and return the top `limit`. `skip` (a query file's own id) is left out. Both
// each row's magnitude and the query's are computed here rather than cached:
// the index mutates as ML data is fetched, and one plain pass over ~50k×512
// floats is fast enough that a norm cache would only add a staleness bug. A
// zero-magnitude vector has no direction, so it is dropped rather than divided
// by zero.
const topByCosine = (
    index: MLIndex,
    query: ArrayLike<number>,
    limit: number,
    skip?: number,
): SimilarResult[] => {
    const { fileIDs, embeddingLength, embeddings } = index;
    if (embeddingLength === 0 || query.length !== embeddingLength) return [];

    // Every indexed read below is in range: the inner loops run to
    // `embeddingLength`, the query is exactly that long (checked above), and
    // the packed buffer holds `fileIDs.length * embeddingLength` floats.
    // `noUncheckedIndexedAccess` still widens each read to `number | undefined`,
    // so they are asserted non-null rather than paying a per-element guard in
    // this hot ~50k×512 loop.
    let queryNorm = 0;
    for (let k = 0; k < embeddingLength; k++) {
        const q = query[k]!;
        queryNorm += q * q;
    }
    queryNorm = Math.sqrt(queryNorm);
    if (queryNorm === 0) return [];

    const results: SimilarResult[] = [];
    for (let i = 0; i < fileIDs.length; i++) {
        const id = fileIDs[i]!;
        if (id === skip) continue;
        const base = i * embeddingLength;
        let dot = 0;
        let norm = 0;
        for (let k = 0; k < embeddingLength; k++) {
            const v = embeddings[base + k]!;
            dot += query[k]! * v;
            norm += v * v;
        }
        if (norm === 0) continue;
        results.push({
            fileID: id,
            score: dot / (queryNorm * Math.sqrt(norm)),
        });
    }

    // Descending score, ties broken by ascending fileID for a stable order.
    results.sort((a, b) => b.score - a.score || a.fileID - b.fileID);
    return results.slice(0, Math.max(0, Math.trunc(limit)));
};

// Build the search surface over a store the library supplies lazily (the store
// is absent when the client cannot fetch ML data). Reading it per call keeps
// the surface current as the index grows.
export const makeMLDataAPI = (
    store: () => MLDataStore | undefined,
): MLDataAPI => ({
    forFile: ({ fileID }): Promise<MLData | undefined> => {
        const s = store();
        return s ? s.readPayload(fileID) : Promise.resolve(undefined);
    },
    similar: ({ fileID, limit }): SimilarResult[] => {
        const s = store();
        if (!s) return [];
        const index = s.getIndex();
        const pos = index.fileIDs.indexOf(fileID);
        if (pos < 0) return [];
        const base = pos * index.embeddingLength;
        const query = index.embeddings.subarray(
            base,
            base + index.embeddingLength,
        );
        return topByCosine(index, query, limit ?? DEFAULT_LIMIT, fileID);
    },
    searchByEmbedding: ({ embedding, limit }): SimilarResult[] => {
        const s = store();
        if (!s) return [];
        return topByCosine(s.getIndex(), embedding, limit ?? DEFAULT_LIMIT);
    },
});