src/library/index.ts

@@ -48,6 +48,7 @@ import {
     type EnsureOptions,
     type EnsureResult,
 } from "./content.js";
+import { makeMLDataAPI, type MLDataAPI } from "./mlsearch.js";
 
 export {
     Album,
@@ -71,6 +72,7 @@ export {
     type EnsureResult,
     type EnsureEvent,
 } from "./content.js";
+export { type MLDataAPI, type SimilarResult } from "./mlsearch.js";
 import type { CollectionsPage, FilesPage } from "../client.js";
 import { MLDATA_BATCH_SIZE, type MLData } from "../mldata-fetch.js";
 import type { Collection, EnteFile } from "../model/types.js";
@@ -189,6 +191,10 @@ export class Library {
     // The thumbnail-prefetch surface (issue #46): drives the thumbnail pool
     // with priority, dedup, and abort.
     readonly thumbnails: ThumbnailsAPI;
+    // The content-similarity search surface over the CLIP index (issue #50).
+    // Present whether or not ML fetching is enabled; with no ML store it
+    // returns empty results.
+    readonly mldata: MLDataAPI;
 
     private readonly client: LibraryClient;
     private readonly store: MetadataStore;
@@ -200,7 +206,7 @@ export class Library {
     private readonly onProgress?: RefreshProgressCallback;
     private readonly pools: RequestPools;
     // The ML-data cache, present only when the client can fetch ML data.
-    private readonly mldata?: MLDataStore;
+    private readonly mlStore?: MLDataStore;
 
     private timer?: ReturnType<typeof setTimeout>;
     private refreshing = false;
@@ -243,7 +249,7 @@ export class Library {
         this.intervalMs = args.intervalMs;
         this.onProgress = args.onProgress;
         this.pools = args.pools;
-        this.mldata = args.mldata;
+        this.mlStore = args.mldata;
         this.cache = args.cache;
         this.lastRecords = this.deriveNow();
 
@@ -265,6 +271,8 @@ export class Library {
                 return this.cache.ensureThumbnails(opts);
             },
         };
+        // Reads the ML store live so results grow as ML data is fetched.
+        this.mldata = makeMLDataAPI(() => this.mlStore);
     }
 
     // Load the cache and start the refresh loop. With an empty cache the first
@@ -386,7 +394,7 @@ export class Library {
         for (const c of collections) {
             files += this.store.listFiles(c.id).length;
         }
-        const ml = this.mldata?.stats();
+        const ml = this.mlStore?.stats();
         const originals = this.cache?.originalsStatus();
         return {
             userID: this.store.userID,
@@ -583,7 +591,7 @@ export class Library {
     // advanced), through the metadata pool, and update the CLIP index. Guarded
     // so passes never overlap; a failure is reported, not thrown.
     private async runMLFetch(): Promise<void> {
-        const mldata = this.mldata;
+        const mldata = this.mlStore;
         // Bind so the call keeps the client as its receiver when invoked
         // through the pool below.
         const fetchMLData = this.client.fetchMLData?.bind(this.client);

src/library/mlsearch.ts

@@ -0,0 +1,129 @@
+// 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);
+    },
+});

test/library/mlsearch.test.ts

@@ -0,0 +1,108 @@
+/**
+ * Tests for the content-similarity search surface over the CLIP index
+ * (issue #50).
+ *
+ * The surface is `lib.mldata`: `forFile` reads the full stored payload from
+ * disk, while `similar` and `searchByEmbedding` rank fileIDs by cosine
+ * similarity over the in-RAM `Float32Array` index alone (no disk, no network).
+ * The fixture uses axis-aligned vectors so the correct cosine ranking is
+ * obvious by inspection; cosine ignores magnitude, so `[2, 0, 0]` ranks above
+ * `[0.8, 0.6, 0]` for a `[1, 0, 0]` query.
+ */
+
+import { describe, it, expect, beforeEach, afterEach } from "vitest";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+import { MLDataStore } from "../../src/library/mldata.js";
+import { makeMLDataAPI, type MLDataAPI } from "../../src/library/mlsearch.js";
+import type { MLData } from "../../src/mldata-fetch.js";
+
+// A payload shaped like Ente's: a CLIP embedding plus face data that only the
+// on-disk payload carries (never the RAM index).
+const payload = (embedding: number[]): MLData => ({
+    face: { faces: [{ faceID: "f", detection: { box: { x: 0.5 } } }] },
+    clip: { embedding },
+});
+
+// A small fixture index. Directions are chosen so every cosine ranking below
+// is unambiguous.
+const fixture = (): Map<number, MLData> =>
+    new Map([
+        [10, payload([1, 0, 0])],
+        [20, payload([0.8, 0.6, 0])],
+        [30, payload([0, 1, 0])],
+        [40, payload([-1, 0, 0])],
+        [50, payload([2, 0, 0])],
+    ]);
+
+describe("lib.mldata content-similarity search", () => {
+    let dir: string;
+    let store: MLDataStore;
+    let api: MLDataAPI;
+
+    beforeEach(async () => {
+        dir = mkdtempSync(join(tmpdir(), "quak-mlsearch-"));
+        store = await MLDataStore.open(dir);
+        const updation = new Map([...fixture().keys()].map((id) => [id, 1]));
+        await store.storeFetched(fixture(), updation);
+        api = makeMLDataAPI(() => store);
+    });
+
+    afterEach(() => {
+        rmSync(dir, { recursive: true, force: true });
+    });
+
+    it("forFile returns the whole stored payload, or undefined when uncached", async () => {
+        const full = await api.forFile({ fileID: 20 });
+        expect(full).toBeDefined();
+        // Face data lives only in the payload, never in the RAM index.
+        expect(full?.face).toBeDefined();
+        expect(full?.clip).toEqual({ embedding: [0.8, 0.6, 0] });
+        expect(await api.forFile({ fileID: 999 })).toBeUndefined();
+    });
+
+    it("similar ranks other files by cosine and excludes the query itself", () => {
+        // Query is file 10 = [1, 0, 0]. By cosine: 50 (1.0) > 20 (0.8) >
+        // 30 (0) > 40 (-1); 10 itself is left out.
+        const ranked = api.similar({ fileID: 10 });
+        expect(ranked.map((r) => r.fileID)).toEqual([50, 20, 30, 40]);
+        // Cosine ignores magnitude: [2,0,0] is a perfect match for [1,0,0].
+        expect(ranked[0]).toMatchObject({ fileID: 50 });
+        expect(ranked[0].score).toBeCloseTo(1, 5);
+    });
+
+    it("similar honours limit and returns [] for an unindexed file", () => {
+        expect(
+            api.similar({ fileID: 10, limit: 2 }).map((r) => r.fileID),
+        ).toEqual([50, 20]);
+        expect(api.similar({ fileID: 999 })).toEqual([]);
+    });
+
+    it("searchByEmbedding ranks the index by cosine to the query vector", () => {
+        // Query [0, 1, 0]: 30 (1.0) > 20 (0.6) > {10, 40, 50} all 0, broken by
+        // ascending fileID.
+        const ranked = api.searchByEmbedding({ embedding: [0, 1, 0] });
+        expect(ranked.map((r) => r.fileID)).toEqual([30, 20, 10, 40, 50]);
+        expect(ranked[0].score).toBeCloseTo(1, 5);
+
+        expect(
+            api
+                .searchByEmbedding({ embedding: [0, 1, 0], limit: 2 })
+                .map((r) => r.fileID),
+        ).toEqual([30, 20]);
+    });
+
+    it("searchByEmbedding returns [] for a wrong-length or zero query", () => {
+        expect(api.searchByEmbedding({ embedding: [1, 0] })).toEqual([]);
+        expect(api.searchByEmbedding({ embedding: [0, 0, 0] })).toEqual([]);
+    });
+
+    it("degrades to empty results when no ML store is present", async () => {
+        const none = makeMLDataAPI(() => undefined);
+        expect(await none.forFile({ fileID: 10 })).toBeUndefined();
+        expect(none.similar({ fileID: 10 })).toEqual([]);
+        expect(none.searchByEmbedding({ embedding: [1, 0, 0] })).toEqual([]);
+    });
+});