98 lines
3.4 KiB
JavaScript
98 lines
3.4 KiB
JavaScript
import { RetrieveCommand, BedrockAgentRuntimeClient, } from "@aws-sdk/client-bedrock-agent-runtime";
|
|
import { BaseRetriever } from "@langchain/core/retrievers";
|
|
/**
|
|
* @deprecated The AmazonKnowledgeBaseRetriever integration has been moved to the `@langchain/aws` package. Import from `@langchain/aws` instead.
|
|
*
|
|
* Class for interacting with Amazon Bedrock Knowledge Bases, a RAG workflow oriented service
|
|
* provided by AWS. Extends the BaseRetriever class.
|
|
* @example
|
|
* ```typescript
|
|
* const retriever = new AmazonKnowledgeBaseRetriever({
|
|
* topK: 10,
|
|
* knowledgeBaseId: "YOUR_KNOWLEDGE_BASE_ID",
|
|
* region: "us-east-2",
|
|
* clientOptions: {
|
|
* credentials: {
|
|
* accessKeyId: "YOUR_ACCESS_KEY_ID",
|
|
* secretAccessKey: "YOUR_SECRET_ACCESS_KEY",
|
|
* },
|
|
* },
|
|
* });
|
|
*
|
|
* const docs = await retriever.getRelevantDocuments("How are clouds formed?");
|
|
* ```
|
|
*/
|
|
export class AmazonKnowledgeBaseRetriever extends BaseRetriever {
|
|
static lc_name() {
|
|
return "AmazonKnowledgeBaseRetriever";
|
|
}
|
|
constructor({ knowledgeBaseId, topK = 10, clientOptions, region, }) {
|
|
super();
|
|
Object.defineProperty(this, "lc_namespace", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: ["langchain", "retrievers", "amazon_bedrock_knowledge_base"]
|
|
});
|
|
Object.defineProperty(this, "knowledgeBaseId", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "topK", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "bedrockAgentRuntimeClient", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
this.topK = topK;
|
|
this.bedrockAgentRuntimeClient = new BedrockAgentRuntimeClient({
|
|
region,
|
|
...clientOptions,
|
|
});
|
|
this.knowledgeBaseId = knowledgeBaseId;
|
|
}
|
|
/**
|
|
* Cleans the result text by replacing sequences of whitespace with a
|
|
* single space and removing ellipses.
|
|
* @param resText The result text to clean.
|
|
* @returns The cleaned result text.
|
|
*/
|
|
cleanResult(resText) {
|
|
const res = resText.replace(/\s+/g, " ").replace(/\.\.\./g, "");
|
|
return res;
|
|
}
|
|
async queryKnowledgeBase(query, topK) {
|
|
const retrieveCommand = new RetrieveCommand({
|
|
knowledgeBaseId: this.knowledgeBaseId,
|
|
retrievalQuery: {
|
|
text: query,
|
|
},
|
|
retrievalConfiguration: {
|
|
vectorSearchConfiguration: {
|
|
numberOfResults: topK,
|
|
},
|
|
},
|
|
});
|
|
const retrieveResponse = await this.bedrockAgentRuntimeClient.send(retrieveCommand);
|
|
return (retrieveResponse.retrievalResults?.map((result) => ({
|
|
pageContent: this.cleanResult(result.content?.text || ""),
|
|
metadata: {
|
|
source: result.location?.s3Location?.uri,
|
|
score: result.score,
|
|
...result.metadata,
|
|
},
|
|
})) ?? []);
|
|
}
|
|
async _getRelevantDocuments(query) {
|
|
const docs = await this.queryKnowledgeBase(query, this.topK);
|
|
return docs;
|
|
}
|
|
}
|