test(database): fix returning ids

src/database.test.ts

@@ -1,9 +1,10 @@
-/* eslint-disable @typescript-eslint/no-non-null-assertion */
 import { afterAll, describe, expect, test } from "bun:test";
 import { Database } from "./database";
 import { Index } from "@upstash/vector";
 import { awaitUntilIndexed } from "./test-utils";
 
+const nanoidLength = 21;
+
 describe("Database", () => {
   const vector = new Index();
 
@@ -34,4 +35,48 @@ describe("Database", () => {
     });
     expect(result.map(({ data }) => data).join(" ")).toContain("330");
   });
+
+  test("should save and retrieve info using id field", async () => {
+    const database = new Database(vector);
+
+    await database.save({
+      type: "text",
+      id: "eiffel-tower",
+      data: "Paris, the capital of France, is renowned for its iconic landmark, the Eiffel Tower, which was completed in 1889 and stands at 330 meters tall.",
+    });
+
+    await awaitUntilIndexed(vector);
+
+    const result = await database.retrieve({
+      question:
+        "What year was the construction of the Eiffel Tower completed, and what is its height?",
+    });
+    // sadly we can't test the id here
+    // because the questions are not unique
+    // every time we run the test we will add the same question
+    // and the id will be different because the database find the similarity
+    // we could create a unique namespace for each test
+    expect(result[0].data).toContain("1889");
+    expect(result[0].data).toContain("330 meters");
+  });
+
+  test("should return generated id when saving without providing an id", async () => {
+    const database = new Database(vector);
+
+    const saveResult = await database.save({
+      type: "text",
+      data: "Paris, the capital of France, is renowned for its iconic landmark, the Eiffel Tower.",
+    });
+
+    expect(saveResult.success).toBe(true);
+    if (!saveResult.success) {
+      throw new Error(`Failed to save entry: ${saveResult.error}`);
+    }
+
+    expect(saveResult.success).toBe(true);
+    expect(saveResult.ids).toHaveLength(1);
+    expect(typeof saveResult.ids[0]).toBe("string");
+    expect(saveResult.ids[0]).toBeTruthy();
+    expect(saveResult.ids[0]).toHaveLength(nanoidLength);
+  });
 });

src/database.ts

@@ -153,32 +153,34 @@ export class Database {
     const { namespace } = input.options ?? {};
     if (input.type === "text") {
       try {
-        const vectorId = await this.index.upsert(
+        const returnId = input.id ?? nanoid();
+        await this.index.upsert(
           {
             data: input.data,
-            id: input.id ?? nanoid(),
+            id: returnId,
             metadata: input.options?.metadata,
           },
           { namespace }
         );
 
-        return { success: true, ids: [vectorId.toString()] };
+        return { success: true, ids: [returnId.toString()] };
       } catch (error) {
         return { success: false, error: JSON.stringify(error, Object.getOwnPropertyNames(error)) };
       }
     } else if (input.type === "embedding") {
       try {
-        const vectorId = await this.index.upsert(
+        const returnId = input.id ?? nanoid();
+        await this.index.upsert(
           {
             vector: input.data,
             data: input.text,
-            id: input.id ?? nanoid(),
+            id: returnId,
             metadata: input.options?.metadata,
           },
           { namespace }
         );
 
-        return { success: true, ids: [vectorId.toString()] };
+        return { success: true, ids: [returnId.toString()] };
       } catch (error) {
         return { success: false, error: JSON.stringify(error, Object.getOwnPropertyNames(error)) };
       }