@galdor/memory-s3vectors 0.3.0

package/src/index.ts

@@ -0,0 +1,470 @@
+/**
+ * `@galdor/memory-s3vectors` — an Amazon S3 Vectors long-term memory store.
+ *
+ * Implements the core {@link Store} interface from `@galdor/core/memory` over the
+ * Amazon S3 Vectors API, so it drops in behind a {@link Retriever} exactly like
+ * the bundled `InMemoryStore` — but persists vectors in an S3 Vectors index.
+ *
+ * Credentials are resolved by the AWS SDK's default provider chain (environment
+ * variables → shared `~/.aws` config → ECS container credentials → EC2 IMDS /
+ * task role); no static keys are accepted here. The index is created on first
+ * use if it does not already exist.
+ *
+ * Construct one with {@link openS3Vectors}, which probes (and auto-creates) the
+ * index before returning a ready store.
+ *
+ * @example
+ * ```ts
+ * const store = await openS3Vectors({ bucket: "my-vectors", index: "galdor-chunks", region: "us-east-1", dim: 1024 });
+ * await store.add([{ id: "c1", documentId: "d1", index: 0, text: "…", embedding: vec, metadata: { lang: "es" } }]);
+ * const hits = await store.retrieve({ embedding: queryVec, k: 5, filter: { lang: "es" } });
+ * await store.close();
+ * ```
+ */
+
+import {
+  CreateIndexCommand,
+  DataType,
+  DeleteVectorsCommand,
+  DistanceMetric,
+  GetIndexCommand,
+  ListVectorsCommand,
+  PutVectorsCommand,
+  type PutVectorsCommandInput,
+  QueryVectorsCommand,
+  type QueryVectorsCommandInput,
+  S3VectorsClient,
+} from "@aws-sdk/client-s3vectors";
+import type { Chunk, Query, Result, Store } from "@galdor/core/memory";
+
+// ── Reserved metadata keys (stable names for cross-store compatibility) ──
+const META_DOCUMENT_ID = "__document_id";
+const META_INDEX = "__index";
+const META_TEXT = "__text";
+const RESERVED_PREFIX = "__";
+
+/** S3 Vectors batch/page limits. */
+const MAX_PUT_BATCH = 500;
+const MAX_DELETE_BATCH = 500;
+const LIST_PAGE_SIZE = 1000;
+
+const DEFAULT_INDEX = "galdor-chunks";
+const DEFAULT_K = 5;
+
+/** A JSON metadata document, as the S3 Vectors API models it. */
+type Metadata = Record<string, unknown>;
+
+/**
+ * The subset of the S3 Vectors API this store uses. Extracted as an interface so
+ * the store can be unit-tested against a fake, without a live AWS account.
+ */
+export interface S3VectorsAPI {
+  getIndex(input: { vectorBucketName: string; indexName: string }): Promise<{ index?: { dimension?: number } }>;
+  createIndex(input: {
+    vectorBucketName: string;
+    indexName: string;
+    dataType: DataType;
+    dimension: number;
+    distanceMetric: DistanceMetric;
+    metadataConfiguration?: { nonFilterableMetadataKeys: string[] };
+  }): Promise<unknown>;
+  putVectors(
+    input: {
+      vectorBucketName: string;
+      indexName: string;
+      vectors: { key: string; data: { float32: number[] }; metadata?: Metadata }[];
+    },
+    opts?: SendOptions,
+  ): Promise<unknown>;
+  queryVectors(
+    input: {
+      vectorBucketName: string;
+      indexName: string;
+      topK: number;
+      queryVector: { float32: number[] };
+      returnMetadata: boolean;
+      returnDistance: boolean;
+      filter?: Metadata;
+      nextToken?: string;
+    },
+    opts?: SendOptions,
+  ): Promise<{
+    vectors?: { key?: string; distance?: number; metadata?: Metadata }[];
+    distanceMetric?: DistanceMetric;
+    nextToken?: string;
+  }>;
+  listVectors(
+    input: {
+      vectorBucketName: string;
+      indexName: string;
+      returnMetadata: boolean;
+      maxResults: number;
+      nextToken?: string;
+    },
+    opts?: SendOptions,
+  ): Promise<{ vectors?: { key?: string; metadata?: Metadata }[]; nextToken?: string }>;
+  deleteVectors(input: { vectorBucketName: string; indexName: string; keys: string[] }, opts?: SendOptions): Promise<unknown>;
+}
+
+/** Per-call options forwarded to the AWS SDK; carries the caller's abort signal. */
+export interface SendOptions {
+  abortSignal?: AbortSignal;
+}
+
+/** Construction options for {@link openS3Vectors}. */
+export interface Config {
+  /** S3 Vectors bucket name. The bucket must already exist. Required. */
+  bucket: string;
+  /** Index name within the bucket. Defaults to `"galdor-chunks"`. Created if missing. */
+  index?: string;
+  /** AWS region. When empty, resolved from the default AWS chain. */
+  region?: string;
+  /** Embedding dimensionality, e.g. `1024`. Required and must be > 0. */
+  dim: number;
+  /** Distance metric used when creating the index. Defaults to `"cosine"`. Ignored if the index exists. */
+  distance?: DistanceMetric;
+  /** Inject a custom API implementation (for tests). Defaults to a real {@link S3VectorsClient}. */
+  api?: S3VectorsAPI;
+}
+
+/** Wrap a real {@link S3VectorsClient} as the {@link S3VectorsAPI} the store uses. */
+function clientAPI(client: S3VectorsClient): S3VectorsAPI {
+  // Our metadata/filter are plain JSON objects; the SDK models them as the
+  // recursive `DocumentType`. Cast at this single boundary rather than leak the
+  // SDK's document type into the store's own contract.
+  return {
+    getIndex: (i) => client.send(new GetIndexCommand(i)),
+    createIndex: (i) => client.send(new CreateIndexCommand(i)),
+    putVectors: (i, o) => client.send(new PutVectorsCommand(i as unknown as PutVectorsCommandInput), o),
+    queryVectors: (i, o) => client.send(new QueryVectorsCommand(i as unknown as QueryVectorsCommandInput), o),
+    listVectors: (i, o) => client.send(new ListVectorsCommand(i), o),
+    deleteVectors: (i, o) => client.send(new DeleteVectorsCommand(i), o),
+  } as S3VectorsAPI;
+}
+
+/**
+ * Open an S3 Vectors-backed {@link Store}, probing and auto-creating the index.
+ *
+ * @param cfg - The {@link Config}.
+ * @returns A ready store.
+ * @throws {Error} When `bucket` is empty, `dim` is not positive, or the index
+ *   exists with a different dimension.
+ */
+export async function openS3Vectors(cfg: Config): Promise<S3VectorsStore> {
+  if (!cfg.bucket || cfg.bucket.trim() === "") throw new Error("memory/s3vectors: bucket is required");
+  if (!Number.isInteger(cfg.dim) || cfg.dim <= 0) throw new Error("memory/s3vectors: dim must be a positive integer");
+  const index = cfg.index && cfg.index.trim() !== "" ? cfg.index : DEFAULT_INDEX;
+  validateIndexName(index);
+  const distance = cfg.distance ?? DistanceMetric.COSINE;
+
+  const ownsClient = cfg.api === undefined;
+  const client = ownsClient ? new S3VectorsClient(cfg.region ? { region: cfg.region } : {}) : undefined;
+  const api = cfg.api ?? clientAPI(client!);
+
+  const store = new S3VectorsStore(api, cfg.bucket, index, cfg.dim, distance, client);
+  await store.ensureIndex();
+  return store;
+}
+
+/**
+ * A {@link Store} backed by an Amazon S3 Vectors index.
+ *
+ * Prefer {@link openS3Vectors} to construct one; the constructor is exported for
+ * advanced use and testing.
+ */
+export class S3VectorsStore implements Store {
+  readonly #api: S3VectorsAPI;
+  readonly #bucket: string;
+  readonly #index: string;
+  readonly #dim: number;
+  readonly #distance: DistanceMetric;
+  readonly #client?: S3VectorsClient;
+
+  constructor(
+    api: S3VectorsAPI,
+    bucket: string,
+    index: string,
+    dim: number,
+    distance: DistanceMetric,
+    client?: S3VectorsClient,
+  ) {
+    this.#api = api;
+    this.#bucket = bucket;
+    this.#index = index;
+    this.#dim = dim;
+    this.#distance = distance;
+    if (client) this.#client = client;
+  }
+
+  /** Probe the index; create it (with this store's dim + metric) if it is missing. */
+  async ensureIndex(): Promise<void> {
+    try {
+      const out = await this.#api.getIndex({ vectorBucketName: this.#bucket, indexName: this.#index });
+      const got = out.index?.dimension;
+      if (typeof got === "number" && got !== this.#dim) {
+        throw new Error(`memory/s3vectors: index "${this.#index}" is ${got}-dim; Config.dim is ${this.#dim}`);
+      }
+      return;
+    } catch (err) {
+      if (!isNotFound(err)) throw err;
+    }
+    // Index missing: create it. `__text` is large and not useful as a filter, so
+    // it is declared non-filterable.
+    await this.#api.createIndex({
+      vectorBucketName: this.#bucket,
+      indexName: this.#index,
+      dataType: DataType.FLOAT32,
+      dimension: this.#dim,
+      distanceMetric: this.#distance,
+      metadataConfiguration: { nonFilterableMetadataKeys: [META_TEXT] },
+    });
+  }
+
+  /**
+   * Upsert chunks, keyed by their `id` (idempotent). Each chunk's embedding is
+   * stored as the vector; `documentId`, `index`, `text` and any chunk metadata
+   * are stored as queryable/returnable metadata.
+   *
+   * @throws {Error} When a chunk has an empty `id`, a missing/mismatched-dimension
+   *   embedding, or metadata using the reserved `__` prefix.
+   */
+  async add(chunks: Chunk[], signal?: AbortSignal): Promise<void> {
+    if (chunks.length === 0) return;
+    const vectors = chunks.map((c) => {
+      if (c.id === "") throw new Error("memory/s3vectors: Chunk.id is empty (caller must assign ids)");
+      if (c.embedding === undefined || c.embedding.length !== this.#dim) {
+        throw new Error(
+          `memory/s3vectors: chunk "${c.id}" has ${c.embedding?.length ?? 0}-dim embedding; index is ${this.#dim}-dim`,
+        );
+      }
+      return { key: c.id, data: { float32: [...c.embedding] }, metadata: buildMetadata(c) };
+    });
+    for (let start = 0; start < vectors.length; start += MAX_PUT_BATCH) {
+      await this.#api.putVectors(
+        {
+          vectorBucketName: this.#bucket,
+          indexName: this.#index,
+          vectors: vectors.slice(start, start + MAX_PUT_BATCH),
+        },
+        sendOpts(signal),
+      );
+    }
+  }
+
+  /**
+   * Return the top-K nearest chunks by the query embedding, in descending
+   * relevance. `query.filter` is applied as an exact-match AND over metadata.
+   *
+   * @throws {Error} When `query.embedding` is absent or its dimension mismatches.
+   */
+  async retrieve(query: Query, signal?: AbortSignal): Promise<Result[]> {
+    if (query.embedding === undefined || query.embedding.length === 0) {
+      throw new Error("memory/s3vectors: Query.embedding is required (this backend is vector-only)");
+    }
+    if (query.embedding.length !== this.#dim) {
+      throw new Error(`memory/s3vectors: query has ${query.embedding.length}-dim embedding; index is ${this.#dim}-dim`);
+    }
+    let k = query.k ?? 0;
+    if (k <= 0) k = DEFAULT_K;
+
+    const filter = buildFilter(query.filter);
+    const results: Result[] = [];
+    let nextToken: string | undefined;
+    for (;;) {
+      const out = await this.#api.queryVectors(
+        {
+          vectorBucketName: this.#bucket,
+          indexName: this.#index,
+          topK: k,
+          queryVector: { float32: [...query.embedding] },
+          returnMetadata: true,
+          returnDistance: true,
+          ...(filter ? { filter } : {}),
+          ...(nextToken ? { nextToken } : {}),
+        },
+        sendOpts(signal),
+      );
+      const metric = out.distanceMetric || this.#distance;
+      for (const v of out.vectors ?? []) {
+        const chunk = chunkFromVector(v.key, v.metadata);
+        const dist = v.distance ?? 0;
+        if (metric === DistanceMetric.EUCLIDEAN) {
+          // Euclidean (squared L2) ∈ [0, ∞): map to (0, 1], monotone-decreasing.
+          results.push({ chunk, score: 1 / (1 + dist) });
+          continue;
+        }
+        // Cosine distance ∈ [0, 2] → similarity = 1 - distance ∈ [-1, 1].
+        // Drop anti-correlated hits for parity with the other backends.
+        const score = 1 - dist;
+        if (score < 0) continue;
+        results.push({ chunk, score });
+      }
+      if (!out.nextToken || results.length >= k) break;
+      nextToken = out.nextToken;
+    }
+    if (results.length > k) results.length = k;
+    return results;
+  }
+
+  /**
+   * Remove every chunk whose `documentId` metadata matches. Scans the index for
+   * matching keys, then batch-deletes them. A no-op when nothing matches.
+   *
+   * @throws {Error} When `documentId` is empty.
+   */
+  async delete(documentId: string, signal?: AbortSignal): Promise<void> {
+    if (documentId === "") throw new Error("memory/s3vectors: delete called with empty documentId");
+    const keys: string[] = [];
+    let nextToken: string | undefined;
+    for (;;) {
+      const out = await this.#api.listVectors(
+        {
+          vectorBucketName: this.#bucket,
+          indexName: this.#index,
+          returnMetadata: true,
+          maxResults: LIST_PAGE_SIZE,
+          ...(nextToken ? { nextToken } : {}),
+        },
+        sendOpts(signal),
+      );
+      for (const v of out.vectors ?? []) {
+        if (v.key && metadataString(v.metadata, META_DOCUMENT_ID) === documentId) keys.push(v.key);
+      }
+      if (!out.nextToken) break;
+      nextToken = out.nextToken;
+    }
+    for (let start = 0; start < keys.length; start += MAX_DELETE_BATCH) {
+      await this.#api.deleteVectors(
+        {
+          vectorBucketName: this.#bucket,
+          indexName: this.#index,
+          keys: keys.slice(start, start + MAX_DELETE_BATCH),
+        },
+        sendOpts(signal),
+      );
+    }
+  }
+
+  /**
+   * Number of vectors in the index. Scans the index (paginated `ListVectors`), so
+   * it is O(index size). Useful for tests; not part of {@link Store}.
+   */
+  async len(): Promise<number> {
+    let count = 0;
+    let nextToken: string | undefined;
+    for (;;) {
+      const out = await this.#api.listVectors({
+        vectorBucketName: this.#bucket,
+        indexName: this.#index,
+        returnMetadata: false,
+        maxResults: LIST_PAGE_SIZE,
+        ...(nextToken ? { nextToken } : {}),
+      });
+      count += out.vectors?.length ?? 0;
+      if (!out.nextToken) break;
+      nextToken = out.nextToken;
+    }
+    return count;
+  }
+
+  /** Release the underlying client. The AWS client pools its own transport; safe to call repeatedly. */
+  async close(): Promise<void> {
+    this.#client?.destroy();
+  }
+}
+
+// ── helpers ──────────────────────────────────────────────────────────────────
+
+/** Build the metadata document for a chunk; rejects caller keys using the reserved prefix. */
+function buildMetadata(c: Chunk): Metadata {
+  const m: Metadata = { [META_DOCUMENT_ID]: c.documentId, [META_INDEX]: c.index, [META_TEXT]: c.text };
+  if (c.metadata) {
+    for (const [k, v] of Object.entries(c.metadata)) {
+      if (k.startsWith(RESERVED_PREFIX)) {
+        throw new Error(`memory/s3vectors: chunk "${c.id}" metadata key "${k}" uses the reserved "${RESERVED_PREFIX}" prefix`);
+      }
+      m[k] = v;
+    }
+  }
+  return m;
+}
+
+/**
+ * Build the S3 Vectors filter document from a metadata filter:
+ *  - 0 keys → no filter
+ *  - 1 key  → bare `{ k: v }` (implicit equals)
+ *  - >1 key → explicit `{ "$and": [{ k1: v1 }, …] }` (top-level implicit AND is rejected by the API)
+ */
+function buildFilter(filter: Record<string, string> | undefined): Metadata | undefined {
+  if (!filter) return undefined;
+  const keys = Object.keys(filter);
+  if (keys.length === 0) return undefined;
+  if (keys.length === 1) {
+    const k = keys[0]!;
+    return { [k]: filter[k]! };
+  }
+  return { $and: keys.map((k) => ({ [k]: filter[k]! })) };
+}
+
+/** Reconstruct a {@link Chunk} from a vector key + returned metadata. */
+function chunkFromVector(key: string | undefined, meta: Metadata | undefined): Chunk {
+  const chunk: Chunk = { id: key ?? "", documentId: "", index: 0, text: "" };
+  if (!meta) return chunk;
+  if (typeof meta[META_DOCUMENT_ID] === "string") chunk.documentId = meta[META_DOCUMENT_ID] as string;
+  if (typeof meta[META_TEXT] === "string") chunk.text = meta[META_TEXT] as string;
+  chunk.index = toInt(meta[META_INDEX]);
+  let userMeta: Record<string, string> | undefined;
+  for (const [k, v] of Object.entries(meta)) {
+    if (k.startsWith(RESERVED_PREFIX)) continue;
+    if (typeof v !== "string") continue;
+    (userMeta ??= {})[k] = v;
+  }
+  if (userMeta) chunk.metadata = userMeta;
+  return chunk;
+}
+
+/** Read a string-valued metadata field, or `""` when absent or non-string. */
+function metadataString(meta: Metadata | undefined, key: string): string {
+  const v = meta?.[key];
+  return typeof v === "string" ? v : "";
+}
+
+/** Coerce a metadata value (number, or numeric string) to an integer; 0 otherwise. */
+function toInt(v: unknown): number {
+  if (typeof v === "number") return Math.trunc(v);
+  if (typeof v === "string" && v !== "" && !Number.isNaN(Number(v))) return Math.trunc(Number(v));
+  return 0;
+}
+
+/** Whether an error is an S3 Vectors `NotFoundException` (by class name, robust across SDK copies). */
+function isNotFound(err: unknown): boolean {
+  return !!err && typeof err === "object" && (err as { name?: string }).name === "NotFoundException";
+}
+
+/** Options passed to the AWS SDK for a single call; `undefined` when there is no signal. */
+function sendOpts(signal: AbortSignal | undefined): SendOptions | undefined {
+  return signal ? { abortSignal: signal } : undefined;
+}
+
+/**
+ * Validate an index name against the S3 Vectors DNS-style rules: 3–63 chars;
+ * lowercase letters, digits, '-' and '.'; must begin and end with a letter or
+ * digit. No underscores or uppercase — a common slip when carrying over the
+ * pgvector/qdrant table/collection naming convention.
+ */
+function validateIndexName(name: string): void {
+  if (name.length < 3 || name.length > 63) {
+    throw new Error(`memory/s3vectors: index name "${name}" must be 3–63 characters`);
+  }
+  const isAlnum = (ch: string): boolean => /^[a-z0-9]$/.test(ch);
+  if (!isAlnum(name[0]!) || !isAlnum(name[name.length - 1]!)) {
+    throw new Error(`memory/s3vectors: index name "${name}" must begin and end with a lowercase letter or digit`);
+  }
+  for (const ch of name) {
+    if (!isAlnum(ch) && ch !== "-" && ch !== ".") {
+      throw new Error(
+        `memory/s3vectors: index name "${name}" has invalid character "${ch}" (allowed: a-z, 0-9, '-', '.'; no underscores or uppercase)`,
+      );
+    }
+  }
+}