110 lines
4.2 KiB
JavaScript
110 lines
4.2 KiB
JavaScript
|
import { BedrockRuntimeClient, InvokeModelCommand, } from "@aws-sdk/client-bedrock-runtime";
|
||
|
|
import { Embeddings } from "@langchain/core/embeddings";
|
||
|
|
/**
|
||
|
|
* @deprecated The BedrockEmbeddings integration has been moved to the `@langchain/aws` package. Import from `@langchain/aws` instead.
|
||
|
|
*
|
||
|
|
* Class that extends the Embeddings class and provides methods for
|
||
|
|
* generating embeddings using the Bedrock API.
|
||
|
|
* @example
|
||
|
|
* ```typescript
|
||
|
|
* const embeddings = new BedrockEmbeddings({
|
||
|
|
* region: "your-aws-region",
|
||
|
|
* credentials: {
|
||
|
|
* accessKeyId: "your-access-key-id",
|
||
|
|
* secretAccessKey: "your-secret-access-key",
|
||
|
|
* },
|
||
|
|
* model: "amazon.titan-embed-text-v1",
|
||
|
|
* });
|
||
|
|
*
|
||
|
|
* // Embed a query and log the result
|
||
|
|
* const res = await embeddings.embedQuery(
|
||
|
|
* "What would be a good company name for a company that makes colorful socks?"
|
||
|
|
* );
|
||
|
|
* console.log({ res });
|
||
|
|
* ```
|
||
|
|
*/
|
||
|
|
export class BedrockEmbeddings extends Embeddings {
|
||
|
|
constructor(fields) {
|
||
|
|
super(fields ?? {});
|
||
|
|
Object.defineProperty(this, "model", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "client", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "batchSize", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: 512
|
||
|
|
});
|
||
|
|
this.model = fields?.model ?? "amazon.titan-embed-text-v1";
|
||
|
|
this.client =
|
||
|
|
fields?.client ??
|
||
|
|
new BedrockRuntimeClient({
|
||
|
|
region: fields?.region,
|
||
|
|
credentials: fields?.credentials,
|
||
|
|
});
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Protected method to make a request to the Bedrock API to generate
|
||
|
|
* embeddings. Handles the retry logic and returns the response from the
|
||
|
|
* API.
|
||
|
|
* @param request Request to send to the Bedrock API.
|
||
|
|
* @returns Promise that resolves to the response from the API.
|
||
|
|
*/
|
||
|
|
async _embedText(text) {
|
||
|
|
return this.caller.call(async () => {
|
||
|
|
try {
|
||
|
|
// replace newlines, which can negatively affect performance.
|
||
|
|
const cleanedText = text.replace(/\n/g, " ");
|
||
|
|
const res = await this.client.send(new InvokeModelCommand({
|
||
|
|
modelId: this.model,
|
||
|
|
body: JSON.stringify({
|
||
|
|
inputText: cleanedText,
|
||
|
|
}),
|
||
|
|
contentType: "application/json",
|
||
|
|
accept: "application/json",
|
||
|
|
}));
|
||
|
|
const body = new TextDecoder().decode(res.body);
|
||
|
|
return JSON.parse(body).embedding;
|
||
|
|
}
|
||
|
|
catch (e) {
|
||
|
|
console.error({
|
||
|
|
error: e,
|
||
|
|
});
|
||
|
|
// eslint-disable-next-line no-instanceof/no-instanceof
|
||
|
|
if (e instanceof Error) {
|
||
|
|
throw new Error(`An error occurred while embedding documents with Bedrock: ${e.message}`);
|
||
|
|
}
|
||
|
|
throw new Error("An error occurred while embedding documents with Bedrock");
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Method that takes a document as input and returns a promise that
|
||
|
|
* resolves to an embedding for the document. It calls the _embedText
|
||
|
|
* method with the document as the input.
|
||
|
|
* @param document Document for which to generate an embedding.
|
||
|
|
* @returns Promise that resolves to an embedding for the input document.
|
||
|
|
*/
|
||
|
|
embedQuery(document) {
|
||
|
|
return this.caller.callWithOptions({}, this._embedText.bind(this), document);
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Method to generate embeddings for an array of texts. Calls _embedText
|
||
|
|
* method which batches and handles retry logic when calling the AWS Bedrock API.
|
||
|
|
* @param documents Array of texts for which to generate embeddings.
|
||
|
|
* @returns Promise that resolves to a 2D array of embeddings for each input document.
|
||
|
|
*/
|
||
|
|
async embedDocuments(documents) {
|
||
|
|
return Promise.all(documents.map((document) => this._embedText(document)));
|
||
|
|
}
|
||
|
|
}
|