@galdor/memory-s3vectors 0.3.0

package/src/s3vectors.test.ts

@@ -0,0 +1,241 @@
+/**
+ * Behavioral tests for the S3 Vectors store, driven by an injected fake API so
+ * they run without AWS. They mirror the reference adapter's coverage: index
+ * auto-create, round-trip with metadata, the filter shape, distance→score
+ * conversion, delete-by-documentId across pages, and idempotent add by id.
+ */
+import { describe, expect, test } from "bun:test";
+import type { Chunk } from "@galdor/core/memory";
+import { Retriever } from "@galdor/core/memory";
+import { openS3Vectors, type S3VectorsAPI } from "./index.ts";
+
+type Vec = { key: string; data: { float32: number[] }; metadata?: Record<string, unknown> };
+
+/** An in-memory fake of the S3 Vectors API, recording calls for assertions. */
+class FakeAPI implements S3VectorsAPI {
+  vectors = new Map<string, Vec>();
+  created: { dimension: number; distanceMetric: string; nonFilterable?: string[] } | undefined;
+  indexExists: boolean;
+  existingDim: number | undefined;
+  // Override query behavior per-test (default returns all stored, distance 0).
+  onQuery?: (input: { topK: number; filter?: Record<string, unknown> }) => {
+    vectors?: { key?: string; distance?: number; metadata?: Record<string, unknown> }[];
+    distanceMetric?: string;
+    nextToken?: string;
+  };
+
+  constructor(opts: { indexExists?: boolean; existingDim?: number } = {}) {
+    this.indexExists = opts.indexExists ?? false;
+    if (opts.existingDim !== undefined) this.existingDim = opts.existingDim;
+  }
+
+  async getIndex(_i: { vectorBucketName: string; indexName: string }) {
+    if (!this.indexExists) {
+      const e = new Error("not found");
+      e.name = "NotFoundException";
+      throw e;
+    }
+    return this.existingDim !== undefined ? { index: { dimension: this.existingDim } } : {};
+  }
+  async createIndex(i: {
+    dimension: number;
+    distanceMetric: string;
+    metadataConfiguration?: { nonFilterableMetadataKeys: string[] };
+  }) {
+    const c: { dimension: number; distanceMetric: string; nonFilterable?: string[] } = {
+      dimension: i.dimension,
+      distanceMetric: i.distanceMetric,
+    };
+    if (i.metadataConfiguration) c.nonFilterable = i.metadataConfiguration.nonFilterableMetadataKeys;
+    this.created = c;
+    this.indexExists = true;
+    return {};
+  }
+  async putVectors(i: { vectors: Vec[] }) {
+    for (const v of i.vectors) this.vectors.set(v.key, v); // upsert by key = idempotent
+    return {};
+  }
+  async queryVectors(i: { topK: number; filter?: Record<string, unknown> }) {
+    if (this.onQuery) return this.onQuery(i);
+    const vectors = [...this.vectors.values()].map((v) => ({ key: v.key, distance: 0, metadata: v.metadata }));
+    return { vectors, distanceMetric: "cosine" as const };
+  }
+  async listVectors(_i: { maxResults: number; nextToken?: string }) {
+    const vectors = [...this.vectors.values()].map((v) => ({ key: v.key, metadata: v.metadata }));
+    return { vectors };
+  }
+  async deleteVectors(i: { keys: string[] }) {
+    for (const k of i.keys) this.vectors.delete(k);
+    return {};
+  }
+}
+
+const vec = (...xs: number[]): number[] => xs;
+const chunk = (over: Partial<Chunk> & Pick<Chunk, "id" | "embedding">): Chunk => ({
+  documentId: "d1",
+  index: 0,
+  text: "t",
+  ...over,
+});
+
+describe("openS3Vectors / ensureIndex", () => {
+  test("creates the index when missing, with dim, cosine, and __text non-filterable", async () => {
+    const api = new FakeAPI({ indexExists: false });
+    await openS3Vectors({ bucket: "b", index: "galdor-chunks", dim: 4, api });
+    expect(api.created).toBeDefined();
+    expect(api.created!.dimension).toBe(4);
+    expect(api.created!.distanceMetric).toBe("cosine");
+    expect(api.created!.nonFilterable).toEqual(["__text"]);
+  });
+
+  test("does not create when the index already exists at the same dim", async () => {
+    const api = new FakeAPI({ indexExists: true, existingDim: 4 });
+    await openS3Vectors({ bucket: "b", dim: 4, api });
+    expect(api.created).toBeUndefined();
+  });
+
+  test("throws when the existing index has a different dimension", async () => {
+    const api = new FakeAPI({ indexExists: true, existingDim: 8 });
+    await expect(openS3Vectors({ bucket: "b", dim: 4, api })).rejects.toThrow(/8-dim/);
+  });
+
+  test("validates bucket, dim, and index name", async () => {
+    await expect(openS3Vectors({ bucket: "", dim: 4, api: new FakeAPI() })).rejects.toThrow(/bucket/);
+    await expect(openS3Vectors({ bucket: "b", dim: 0, api: new FakeAPI() })).rejects.toThrow(/dim/);
+    await expect(openS3Vectors({ bucket: "b", dim: 4, index: "AB", api: new FakeAPI() })).rejects.toThrow(/index name/);
+  });
+});
+
+describe("add / retrieve round-trip", () => {
+  test("stores embedding + metadata and reconstructs the chunk on retrieve", async () => {
+    const api = new FakeAPI();
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api });
+    await store.add([
+      chunk({ id: "c1", documentId: "doc1", index: 3, text: "hello", embedding: vec(1, 0, 0, 0), metadata: { lang: "es" } }),
+    ]);
+    const hits = await store.retrieve({ embedding: vec(1, 0, 0, 0), k: 5 });
+    expect(hits).toHaveLength(1);
+    const c = hits[0]!.chunk;
+    expect(c.id).toBe("c1");
+    expect(c.documentId).toBe("doc1");
+    expect(c.index).toBe(3);
+    expect(c.text).toBe("hello");
+    expect(c.metadata).toEqual({ lang: "es" });
+    expect(hits[0]!.score).toBeCloseTo(1, 6); // 1 - distance(0)
+  });
+
+  test("add is idempotent by chunk id (re-add overwrites, no duplicate)", async () => {
+    const api = new FakeAPI();
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api });
+    await store.add([chunk({ id: "c1", text: "v1", embedding: vec(1, 0, 0, 0) })]);
+    await store.add([chunk({ id: "c1", text: "v2", embedding: vec(1, 0, 0, 0) })]);
+    expect(api.vectors.size).toBe(1);
+    const hits = await store.retrieve({ embedding: vec(1, 0, 0, 0) });
+    expect(hits).toHaveLength(1);
+    expect(hits[0]!.chunk.text).toBe("v2");
+  });
+
+  test("rejects empty id, mismatched dimension, and reserved metadata keys", async () => {
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api: new FakeAPI() });
+    await expect(store.add([chunk({ id: "", embedding: vec(1, 0, 0, 0) })])).rejects.toThrow(/id is empty/);
+    await expect(store.add([chunk({ id: "c", embedding: vec(1, 0) })])).rejects.toThrow(/2-dim/);
+    await expect(
+      store.add([chunk({ id: "c", embedding: vec(1, 0, 0, 0), metadata: { __document_id: "x" } })]),
+    ).rejects.toThrow(/reserved/);
+  });
+});
+
+describe("retrieve scoring + filter", () => {
+  test("cosine: converts distance to similarity and drops anti-correlated hits", async () => {
+    const api = new FakeAPI();
+    api.onQuery = () => ({
+      distanceMetric: "cosine",
+      vectors: [
+        { key: "a", distance: 0.1, metadata: { __text: "near" } },
+        { key: "b", distance: 1.8, metadata: { __text: "anti" } }, // score 1-1.8 = -0.8 → dropped
+      ],
+    });
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api });
+    const hits = await store.retrieve({ embedding: vec(1, 0, 0, 0), k: 5 });
+    expect(hits).toHaveLength(1);
+    expect(hits[0]!.chunk.id).toBe("a");
+    expect(hits[0]!.score).toBeCloseTo(0.9, 6);
+  });
+
+  test("euclidean: keeps all, score is 1/(1+distance), monotone-decreasing", async () => {
+    const api = new FakeAPI();
+    api.onQuery = () => ({
+      distanceMetric: "euclidean",
+      vectors: [
+        { key: "near", distance: 0, metadata: { __text: "n" } },
+        { key: "far", distance: 9, metadata: { __text: "f" } },
+      ],
+    });
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api });
+    const hits = await store.retrieve({ embedding: vec(1, 0, 0, 0), k: 5 });
+    expect(hits).toHaveLength(2);
+    expect(hits[0]!.score).toBeCloseTo(1, 6);
+    expect(hits[1]!.score).toBeCloseTo(0.1, 6);
+    expect(hits[0]!.score).toBeGreaterThan(hits[1]!.score);
+  });
+
+  test("filter: single key is bare, multiple keys use $and", async () => {
+    const api = new FakeAPI();
+    let seen: Record<string, unknown> | undefined;
+    api.onQuery = (i) => {
+      seen = i.filter;
+      return { vectors: [], distanceMetric: "cosine" };
+    };
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api });
+
+    await store.retrieve({ embedding: vec(1, 0, 0, 0), filter: { lang: "es" } });
+    expect(seen).toEqual({ lang: "es" });
+
+    await store.retrieve({ embedding: vec(1, 0, 0, 0), filter: { type: "skill", category: "backend" } });
+    expect(seen).toHaveProperty("$and");
+    const conds = (seen as { $and: Record<string, string>[] }).$and;
+    expect(conds).toHaveLength(2);
+    expect(Object.assign({}, ...conds)).toEqual({ type: "skill", category: "backend" });
+  });
+
+  test("requires an embedding and a matching dimension", async () => {
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api: new FakeAPI() });
+    await expect(store.retrieve({ text: "x" })).rejects.toThrow(/embedding is required/);
+    await expect(store.retrieve({ embedding: vec(1, 0) })).rejects.toThrow(/2-dim/);
+  });
+});
+
+describe("delete", () => {
+  test("removes exactly the chunks of the given documentId", async () => {
+    const api = new FakeAPI();
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api });
+    await store.add([
+      chunk({ id: "d1#0", documentId: "d1", embedding: vec(1, 0, 0, 0) }),
+      chunk({ id: "d1#1", documentId: "d1", embedding: vec(0, 1, 0, 0) }),
+      chunk({ id: "d2#0", documentId: "d2", embedding: vec(0, 0, 1, 0) }),
+    ]);
+    await store.delete("d1");
+    expect([...api.vectors.keys()].sort()).toEqual(["d2#0"]);
+  });
+
+  test("delete with no matches is a no-op; empty id throws", async () => {
+    const api = new FakeAPI();
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api });
+    await store.add([chunk({ id: "d2#0", documentId: "d2", embedding: vec(1, 0, 0, 0) })]);
+    await store.delete("missing");
+    expect(api.vectors.size).toBe(1);
+    await expect(store.delete("")).rejects.toThrow(/empty documentId/);
+  });
+});
+
+describe("behind a Retriever (same path as InMemoryStore)", () => {
+  test("Retriever.retrieve delegates to the store and returns hits", async () => {
+    const api = new FakeAPI();
+    const store = await openS3Vectors({ bucket: "b", dim: 4, api });
+    await store.add([chunk({ id: "c1", text: "hola", embedding: vec(1, 0, 0, 0) })]);
+    const retriever = new Retriever({ store, defaultK: 5 });
+    const hits = await retriever.retrieve({ embedding: vec(1, 0, 0, 0) });
+    expect(hits).toHaveLength(1);
+    expect(hits[0]!.chunk.text).toBe("hola");
+  });
+});