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; } }