[ENH] AmazonBedrockEmbeddingFunction for JS

clients/js/package.json

@@ -77,11 +77,15 @@
     "isomorphic-fetch": "^3.0.0"
   },
   "peerDependencies": {
+    "@aws-sdk/client-bedrock-runtime": "^3.496.0",
     "@google/generative-ai": "^0.1.1",
     "cohere-ai": "^5.0.0 || ^6.0.0 || ^7.0.0",
     "openai": "^3.0.0 || ^4.0.0"
   },
   "peerDependenciesMeta": {
+    "@aws-sdk/client-bedrock-runtime": {
+      "optional": true
+    },
     "@google/generative-ai": {
       "optional": true
     },

clients/js/src/embeddings/AmazonBedrockEmbeddingFunction.ts

@@ -0,0 +1,75 @@
+import { IEmbeddingFunction } from "./IEmbeddingFunction";
+
+let amazonBedrockApi: any;
+
+export class AmazonBedrockEmbeddingFunction implements IEmbeddingFunction {
+  private model: string;
+  private configuration: any; // Assume this is BedrockRuntimeClientConfig
+  private amazonBedrockApi?: any;
+  private invokeCommand?: any;
+
+  constructor({ config, model }: { config: any; model?: string }) {
+    this.configuration = config;
+    if (!this.configuration.region) {
+      this.configuration.region = "us-east-1";
+    }
+    this.model = model || "amazon.titan-embed-text-v1";
+  }
+
+  private async loadClient() {
+    if (this.amazonBedrockApi) return;
+    try {
+      // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+      const { amazonBedrock } = await AmazonBedrockEmbeddingFunction.import();
+      amazonBedrockApi = new amazonBedrock.BedrockRuntimeClient(
+        this.configuration,
+      );
+      this.invokeCommand = amazonBedrock.InvokeModelCommand;
+    } catch (_a) {
+      // @ts-ignore
+      if (_a.code === "MODULE_NOT_FOUND") {
+        throw new Error(
+          "Please install the @aws-sdk/client-bedrock-runtime package to use the AmazonBedrockEmbeddingFunction, `npm install -S @aws-sdk/client-bedrock-runtime`",
+        );
+      }
+      throw _a; // Re-throw other errors
+    }
+    this.amazonBedrockApi = amazonBedrockApi;
+  }
+
+  public async generate(texts: string[]): Promise<number[][]> {
+    await this.loadClient();
+    const td = new TextDecoder();
+    const embeddings: number[][] = [];
+    return Promise.all(
+      texts.map(async (text) => {
+        const input = {
+          modelId: this.model,
+          contentType: "application/json",
+          accept: "application/json",
+          body: JSON.stringify({ inputText: text }),
+        };
+        const command = new this.invokeCommand(input);
+        const response = await this.amazonBedrockApi.send(command);
+        const parsedBody = JSON.parse(td.decode(response.body));
+        return parsedBody.embedding;
+      }),
+    );
+  }
+
+  /** @ignore */
+  static async import(): Promise<{
+    // @ts-ignore
+    amazonBedrock: typeof import("@aws-sdk/client-bedrock-runtime");
+  }> {
+    try {
+      // @ts-ignore
+      const amazonBedrock = await import("@aws-sdk/client-bedrock-runtime");
+      return { amazonBedrock };
+    } catch (_a) {
+      throw new Error(
+        "Please install the @aws-sdk/client-bedrock-runtime package to use the AmazonBedrockEmbeddingFunction, `npm install -S @aws-sdk/client-bedrock-runtime`",
+      );
+    }
+  }
+}

clients/js/src/index.ts

@@ -10,6 +10,7 @@ export { DefaultEmbeddingFunction } from "./embeddings/DefaultEmbeddingFunction"
 export { HuggingFaceEmbeddingServerFunction } from "./embeddings/HuggingFaceEmbeddingServerFunction";
 export { JinaEmbeddingFunction } from "./embeddings/JinaEmbeddingFunction";
 export { GoogleGenerativeAiEmbeddingFunction } from "./embeddings/GoogleGeminiEmbeddingFunction";
+export { AmazonBedrockEmbeddingFunction } from "./embeddings/AmazonBedrockEmbeddingFunction";
 export { OllamaEmbeddingFunction } from "./embeddings/OllamaEmbeddingFunction";
 
 export {

clients/js/test/embeddings/aws.bedrock.test.ts

@@ -0,0 +1,36 @@
+import { expect, test } from "@jest/globals";
+import chroma from "../initClient";
+import { DOCUMENTS, IDS } from "../data";
+import { IncludeEnum } from "../../src/types";
+import { AmazonBedrockEmbeddingFunction } from "../../src/embeddings/AmazonBedrockEmbeddingFunction";
+
+if (!process.env.AWS_ACCESS_KEY_ID || !process.env.AWS_SECRET_ACCESS_KEY) {
+  test.skip("it should add Amazon Bedrock embeddings", async () => {});
+} else {
+  test("it should add Amazon Bedrock embeddings", async () => {
+    await chroma.reset();
+    const embedder = new AmazonBedrockEmbeddingFunction({
+      config: {
+        credentials: {
+          accessKeyId: process.env.AWS_ACCESS_KEY_ID,
+          secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY,
+          sessionToken: process.env.AWS_SESSION_TOKEN,
+        },
+        region: "us-east-1",
+      },
+    });
+    const collection = await chroma.createCollection({
+      name: "test",
+      embeddingFunction: embedder,
+    });
+    const embeddings = await embedder.generate(DOCUMENTS);
+    await collection.add({ ids: IDS, embeddings: embeddings });
+    const count = await collection.count();
+    expect(count).toBe(3);
+    var res = await collection.get({
+      ids: IDS,
+      include: [IncludeEnum.Embeddings],
+    });
+    expect(res.embeddings).toEqual(embeddings); // reverse because of the order of the ids
+  });
+}