IndexedDbVectorStorage.ts

/**
 * @license
 * Copyright 2025 Steven Roussey <sroussey@gmail.com>
 * SPDX-License-Identifier: Apache-2.0
 */

import type {
  AnyVectorStorage,
  AutoGeneratedKeys,
  ClientProvidedKeysOption,
  InsertEntity,
  IVectorStorage,
  VectorSearchOptions,
} from "@workglow/storage";
import {
  assertVectorShape,
  getMetadataProperty,
  getVectorProperty,
  matchesFilter,
  safeEmit,
  validateVectorEntities,
} from "@workglow/storage";
import type { EventEmitter } from "@workglow/util";
import { createServiceToken } from "@workglow/util";
import type {
  DataPortSchemaObject,
  FromSchema,
  TypedArray,
  TypedArrayConstructor,
  TypedArraySchemaOptions,
} from "@workglow/util/schema";
import { cosineSimilarity } from "@workglow/util/schema";
import type { MigrationOptions } from "./IndexedDbTable";
import { IndexedDbTabularStorage } from "./IndexedDbTabularStorage";

export const IDB_VECTOR_REPOSITORY = createServiceToken<AnyVectorStorage>(
  "storage.vectorRepository.indexedDb"
);

/**
 * IndexedDB vector storage implementation.
 * Extends IndexedDbTabularStorage for storage.
 * Suitable for browser applications that need persistent vector storage.
 * No vector serialization needed since IndexedDB supports TypedArrays
 * natively via structured clone.
 *
 * @template Schema - The schema definition for the entity
 * @template PrimaryKeyNames - The primary key names
 * @template Metadata - The metadata type for the vector
 * @template VectorCtor - Constructor for stored vectors (default {@link typeof Float32Array})
 * @template Entity - The entity type
 */
export class IndexedDbVectorStorage<
  Schema extends DataPortSchemaObject,
  PrimaryKeyNames extends ReadonlyArray<keyof Schema["properties"]>,
  Metadata extends Record<string, unknown> = Record<string, unknown>,
  Entity = FromSchema<Schema, TypedArraySchemaOptions>,
  InsertType extends InsertEntity<Entity, AutoGeneratedKeys<Schema>> = InsertEntity<
    Entity,
    AutoGeneratedKeys<Schema>
  >,
>
  extends IndexedDbTabularStorage<Schema, PrimaryKeyNames, Entity>
  implements IVectorStorage<Metadata, Schema, Entity, PrimaryKeyNames>
{
  private vectorDimensions: number;
  private vectorPropertyName: keyof Entity;
  private metadataPropertyName: keyof Entity | undefined;

  constructor(
    table: string = "vectors",
    schema: Schema,
    primaryKeyNames: PrimaryKeyNames,
    indexes: readonly (keyof NoInfer<Entity> | readonly (keyof NoInfer<Entity>)[])[] = [],
    dimensions: number,
    _vectorCtor: TypedArrayConstructor = Float32Array,
    migrationOptions: MigrationOptions = {},
    clientProvidedKeys: ClientProvidedKeysOption = "if-missing"
  ) {
    super(table, schema, primaryKeyNames, indexes, migrationOptions, clientProvidedKeys);

    this.vectorDimensions = dimensions;

    const vectorProp = getVectorProperty(schema);
    if (!vectorProp) {
      throw new Error("Schema must have a property with type array and format TypedArray");
    }
    this.vectorPropertyName = vectorProp as keyof Entity;
    this.metadataPropertyName = getMetadataProperty(schema) as keyof Entity | undefined;
  }

  getVectorDimensions(): number {
    return this.vectorDimensions;
  }

  /**
   * Validate vector shape SYNCHRONOUSLY before any IDB transaction opens.
   * The throw propagates through the awaited Promise; no IDB request is ever
   * created, so a malformed put cannot leak a corrupt row into the object store.
   */
  override async put(record: InsertType): Promise<Entity> {
    assertVectorShape(
      (record as Record<string, unknown>)[this.vectorPropertyName as string],
      this.vectorDimensions,
      "write"
    );
    return super.put(record);
  }

  /**
   * Validate every entry up front so a single malformed vector aborts the
   * whole batch before any row is written, then route through a single IDB
   * `readwrite` transaction for atomicity. The inherited tabular `putBulk`
   * opens one transaction per row via `Promise.all`, so a duplicate-key or
   * quota error on the second or later record can leave earlier rows
   * committed — breaking the all-or-nothing guarantee vector ingestion (RAG
   * chunk upserts) depends on. The single-transaction path defers `put` event
   * emission to `tx.oncomplete`, which means subscribers see a successful
   * batch as a burst rather than interleaved with the per-record IDB requests.
   */
  override async putBulk(records: InsertType[]): Promise<Entity[]> {
    validateVectorEntities(
      records as ReadonlyArray<Record<string, unknown>>,
      this.vectorPropertyName as string,
      this.vectorDimensions
    );
    return this.putBulkInTransaction(records);
  }

  async similaritySearch(
    query: TypedArray,
    options: VectorSearchOptions<Record<string, unknown>> = {}
  ) {
    assertVectorShape(query, this.vectorDimensions, "query");
    // Default to no floor: cosine similarity ranges over [-1, 1], so a default
    // of 0 would silently drop negatively-correlated hits and return fewer than
    // `topK`. Callers opt into a relevance floor explicitly via `scoreThreshold`.
    const { topK = 10, filter, scoreThreshold = -Infinity } = options;
    const results: Array<Entity & { score: number }> = [];

    const allEntities = (await this.getAll()) || [];

    for (const entity of allEntities) {
      // IndexedDB stores TypedArrays natively via structured clone
      const vector = entity[this.vectorPropertyName] as TypedArray;
      // A present-but-null metadata column value coalesces to `{}` so a filtered
      // search treats the row as a non-match rather than dereferencing null.
      const rawMetadata = this.metadataPropertyName ? entity[this.metadataPropertyName] : undefined;
      const metadata = (rawMetadata ?? {}) as Metadata;

      if (filter && !matchesFilter(metadata, filter)) {
        continue;
      }

      const score = cosineSimilarity(query, vector);

      if (score < scoreThreshold) {
        continue;
      }

      results.push({
        ...entity,
        score,
      } as Entity & { score: number });
    }

    results.sort((a, b) => b.score - a.score);
    const topResults = results.slice(0, topK);
    // The inherited `events` emitter is typed for the tabular event surface;
    // `similaritySearch` lives on the vector extension of that surface. The
    // emitter instance is the same object, so widen the view to a record that
    // carries the event so it can be emitted type-safely.
    type SimilaritySearchEvents = {
      similaritySearch: (query: TypedArray, results: (Entity & { score: number })[]) => void;
    };
    safeEmit(
      this.events as unknown as EventEmitter<SimilaritySearchEvents>,
      "similaritySearch",
      query,
      topResults
    );

    return topResults;
  }
}