feat: add Bedrock provider integration with embedding functionality

packages/llamaindex/package.json

@@ -30,6 +30,7 @@
     "@google/generative-ai": "0.12.0",
     "@grpc/grpc-js": "^1.11.1",
     "@llamaindex/anthropic": "workspace:*",
+    "@llamaindex/bedrock": "workspace:*",
     "@llamaindex/clip": "workspace:*",
     "@llamaindex/cloud": "workspace:*",
     "@llamaindex/core": "workspace:*",

packages/llamaindex/src/embeddings/BedrockEmbedding.ts

@@ -0,0 +1 @@
+export * from "@llamaindex/bedrock";

packages/llamaindex/src/embeddings/index.ts

@@ -1,4 +1,5 @@
 export * from "@llamaindex/core/embeddings";
+export * from "./BedrockEmbedding.js";
 export { ClipEmbedding, ClipEmbeddingModelType } from "./ClipEmbedding.js";
 export { DeepInfraEmbedding } from "./DeepInfraEmbedding.js";
 export { FireworksEmbedding } from "./fireworks.js";

packages/providers/bedrock/CHANGELOG.md

@@ -0,0 +1,72 @@
+# @llamaindex/bedrock
+
+## 0.0.10
+
+### Patch Changes
+
+- Updated dependencies [9c73f0a]
+  - @llamaindex/[email protected]
+
+## 0.0.9
+
+### Patch Changes
+
+- Updated dependencies [359fd33]
+- Updated dependencies [efb7e1b]
+- Updated dependencies [98ba1e7]
+- Updated dependencies [620c63c]
+  - @llamaindex/[email protected]
+
+## 0.0.8
+
+### Patch Changes
+
+- Updated dependencies [60b185f]
+  - @llamaindex/[email protected]
+
+## 0.0.7
+
+### Patch Changes
+
+- Updated dependencies [691c5bc]
+  - @llamaindex/[email protected]
+
+## 0.0.6
+
+### Patch Changes
+
+- Updated dependencies [fa60fc6]
+  - @llamaindex/[email protected]
+  - @llamaindex/[email protected]
+
+## 0.0.5
+
+### Patch Changes
+
+- Updated dependencies [e2a0876]
+  - @llamaindex/[email protected]
+
+## 0.0.4
+
+### Patch Changes
+
+- Updated dependencies [0493f67]
+  - @llamaindex/[email protected]
+
+## 0.0.3
+
+### Patch Changes
+
+- Updated dependencies [4ba2cfe]
+  - @llamaindex/[email protected]
+  - @llamaindex/[email protected]
+
+## 0.0.2
+
+### Patch Changes
+
+- a75af83: refactor: move some llm and embedding to single package
+- Updated dependencies [ae49ff4]
+- Updated dependencies [a75af83]
+  - @llamaindex/[email protected]
+  - @llamaindex/[email protected]

packages/providers/bedrock/package.json

@@ -0,0 +1,43 @@
+{
+  "name": "@llamaindex/bedrock",
+  "description": "Bedrock Adapter for LlamaIndex",
+  "version": "0.0.1",
+  "type": "module",
+  "main": "./dist/index.cjs",
+  "module": "./dist/index.js",
+  "exports": {
+    ".": {
+      "require": {
+        "types": "./dist/index.d.cts",
+        "default": "./dist/index.cjs"
+      },
+      "import": {
+        "types": "./dist/index.d.ts",
+        "default": "./dist/index.js"
+      }
+    }
+  },
+  "files": [
+    "dist"
+  ],
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/run-llama/LlamaIndexTS.git",
+    "directory": "packages/providers/bedrock"
+  },
+  "scripts": {
+    "build": "bunchee",
+    "dev": "bunchee --watch"
+  },
+  "devDependencies": {
+    "bunchee": "5.5.1"
+  },
+  "dependencies": {
+    "@aws-sdk/client-bedrock": "^3.682.0",
+    "@aws-sdk/client-bedrock-runtime": "^3.642.0",
+    "@llamaindex/core": "workspace:*",
+    "@llamaindex/env": "workspace:*",
+    "aws-sdk": "^2.1691.0",
+    "remeda": "^2.12.0"
+  }
+}

packages/providers/bedrock/src/embedding.ts

@@ -0,0 +1,127 @@
+import {
+  BedrockRuntimeClient,
+  InvokeModelCommand,
+} from "@aws-sdk/client-bedrock-runtime";
+import { BaseEmbedding } from "@llamaindex/core/embeddings";
+import type { MessageContentDetail } from "@llamaindex/core/llms";
+import { extractSingleText } from "@llamaindex/core/utils";
+
+export enum PROVIDERS {
+  AMAZON = "amazon",
+}
+
+export enum ALL_BEDROCK_EMBEDDING_MODELS {
+  TITAN_EMBEDDING = "amazon.titan-embed-text-v1",
+  TITAN_EMBEDDING_V2_0 = "amazon.titan-embed-text-v2:0",
+  TITAN_EMBEDDING_G1_TEXT_02 = "amazon.titan-embed-g1-text-02",
+}
+
+const PROVIDER_SPECIFIC_IDENTIFIERS = {
+  [PROVIDERS.AMAZON]: {
+    getEmbeddingsFunc: (response: any) => response.embedding,
+  },
+};
+
+export class BedrockEmbedding extends BaseEmbedding {
+  private client: BedrockRuntimeClient;
+  private modelName: string;
+  private region: string;
+  public embedBatchSize: number = 10;
+
+  constructor({
+    modelName = ALL_BEDROCK_EMBEDDING_MODELS.TITAN_EMBEDDING,
+    awsAccessKeyId,
+    awsSecretAccessKey,
+    awsSessionToken,
+    region = "us-east-1",
+  }: {
+    modelName?: string;
+    awsAccessKeyId?: string;
+    awsSecretAccessKey?: string;
+    awsSessionToken?: string;
+    region?: string;
+  }) {
+    super();
+    this.modelName = modelName;
+    this.region = region;
+
+    this.client = new BedrockRuntimeClient({
+      region: this.region,
+      credentials: {
+        accessKeyId: awsAccessKeyId || "",
+        secretAccessKey: awsSecretAccessKey || "",
+        sessionToken: awsSessionToken || "",
+      },
+    });
+  }
+
+  public static listSupportedModels(): { [key: string]: string[] } {
+    const listModels: { [key: string]: string[] } = {};
+    for (const provider of Object.values(PROVIDERS)) {
+      listModels[provider] = Object.values(ALL_BEDROCK_EMBEDDING_MODELS).filter(
+        (model) => model.startsWith(provider),
+      );
+    }
+    return listModels;
+  }
+
+  private async getEmbedding(payload: string | string[]): Promise<any> {
+    if (!this.client) throw new Error("Client not set");
+
+    const provider = this.modelName.split(".")[0] as PROVIDERS;
+    const requestBody = this.getRequestBody(provider, payload);
+
+    const command = new InvokeModelCommand({
+      modelId: this.modelName,
+      body: requestBody,
+      accept: "application/json",
+      contentType: "application/json",
+    });
+
+    const response = await this.client.send(command);
+    const resp = JSON.parse(new TextDecoder().decode(response.body));
+    const identifiers = PROVIDER_SPECIFIC_IDENTIFIERS[provider];
+    if (!identifiers) throw new Error("Provider not supported");
+
+    return identifiers.getEmbeddingsFunc(resp);
+  }
+
+  private getRequestBody(provider: PROVIDERS, payload: string | string[]): any {
+    if (provider === PROVIDERS.AMAZON) {
+      if (Array.isArray(payload))
+        throw new Error("Amazon provider does not support list of texts");
+      return JSON.stringify({ inputText: payload });
+    } else {
+      throw new Error("Provider not supported");
+    }
+  }
+
+  public async getTextEmbedding(text: string): Promise<any> {
+    return this.getEmbedding(text);
+  }
+
+  getTextEmbeddings = async (texts: string[]): Promise<number[][]> => {
+    return this.getEmbeddingBatch(texts);
+  };
+
+  public async getQueryEmbedding(
+    query: MessageContentDetail,
+  ): Promise<number[] | null> {
+    const text = extractSingleText(query);
+    if (text) {
+      return this.getEmbedding(text);
+    } else {
+      return null;
+    }
+  }
+  public async getEmbeddingBatch(texts: string[]): Promise<number[][]> {
+    const embeddings: number[][] = [];
+
+    for (const text of texts) {
+      const embedding = await this.getEmbedding(text);
+      embeddings.push(embedding);
+    }
+
+    return embeddings;
+  }
+}

packages/providers/bedrock/src/index.ts

@@ -0,0 +1 @@
+export { ALL_BEDROCK_EMBEDDING_MODELS, BedrockEmbedding } from "./embedding";

packages/providers/bedrock/tsconfig.json

@@ -0,0 +1,20 @@
+{
+  "extends": "../../../tsconfig.json",
+  "compilerOptions": {
+    "target": "ESNext",
+    "module": "ESNext",
+    "moduleResolution": "bundler",
+    "outDir": "./lib",
+    "tsBuildInfoFile": "./lib/.tsbuildinfo",
+    "lib": ["DOM", "ES2022"]
+  },
+  "include": ["./src"],
+  "references": [
+    {
+      "path": "../../core/tsconfig.json"
+    },
+    {
+      "path": "../../env/tsconfig.json"
+    }
+  ]
+}