Update for AWS Bedrock knowledge bases to support filters and overrid…

…eSearchType, Update the KB to support other locations as sources (langchain-ai#6189) Co-authored-by: Brace Sproul <[email protected]>
tylerdavis · Jul 24, 2024 · 63305a0 · 63305a0
1 parent 29c08a3
commit 63305a0
Show file tree

Hide file tree

Showing 4 changed files with 931 additions and 16 deletions.
diff --git a/libs/langchain-aws/package.json b/libs/langchain-aws/package.json
@@ -39,15 +39,15 @@
   "author": "LangChain",
   "license": "MIT",
   "dependencies": {
-    "@aws-sdk/client-bedrock-agent-runtime": "^3.583.0",
+    "@aws-sdk/client-bedrock-agent-runtime": "^3.616.0",
     "@aws-sdk/client-bedrock-runtime": "^3.602.0",
     "@aws-sdk/client-kendra": "^3.352.0",
     "@aws-sdk/credential-provider-node": "^3.600.0",
     "@langchain/core": ">=0.2.16 <0.3.0",
     "zod-to-json-schema": "^3.22.5"
   },
   "devDependencies": {
-    "@aws-sdk/types": "^3.598.0",
+    "@aws-sdk/types": "^3.609.0",
     "@jest/globals": "^29.5.0",
     "@langchain/scripts": "~0.0.14",
     "@langchain/standard-tests": "0.0.0",
@@ -97,4 +97,4 @@
     "index.d.ts",
     "index.d.cts"
   ]
-}
+}
diff --git a/libs/langchain-aws/src/retrievers/bedrock.ts b/libs/langchain-aws/src/retrievers/bedrock.ts
@@ -2,6 +2,8 @@ import {
   RetrieveCommand,
   BedrockAgentRuntimeClient,
   type BedrockAgentRuntimeClientConfig,
+  type SearchType,
+  type RetrievalFilter,
 } from "@aws-sdk/client-bedrock-agent-runtime";
 
 import { BaseRetriever } from "@langchain/core/retrievers";
@@ -16,6 +18,8 @@ export interface AmazonKnowledgeBaseRetrieverArgs {
   topK: number;
   region: string;
   clientOptions?: BedrockAgentRuntimeClientConfig;
+  filter?: RetrievalFilter;
+  overrideSearchType?: SearchType;
 }
 
 /**
@@ -51,15 +55,23 @@ export class AmazonKnowledgeBaseRetriever extends BaseRetriever {
 
   bedrockAgentRuntimeClient: BedrockAgentRuntimeClient;
 
+  filter?: RetrievalFilter;
+
+  overrideSearchType?: SearchType;
+
   constructor({
     knowledgeBaseId,
     topK = 10,
     clientOptions,
     region,
+    filter,
+    overrideSearchType,
   }: AmazonKnowledgeBaseRetrieverArgs) {
     super();
 
     this.topK = topK;
+    this.filter = filter;
+    this.overrideSearchType = overrideSearchType;
     this.bedrockAgentRuntimeClient = new BedrockAgentRuntimeClient({
       region,
       ...clientOptions,
@@ -78,7 +90,12 @@ export class AmazonKnowledgeBaseRetriever extends BaseRetriever {
     return res;
   }
 
-  async queryKnowledgeBase(query: string, topK: number) {
+  async queryKnowledgeBase(
+    query: string,
+    topK: number,
+    filter?: RetrievalFilter,
+    overrideSearchType?: SearchType
+  ) {
     const retrieveCommand = new RetrieveCommand({
       knowledgeBaseId: this.knowledgeBaseId,
       retrievalQuery: {
@@ -87,6 +104,8 @@ export class AmazonKnowledgeBaseRetriever extends BaseRetriever {
       retrievalConfiguration: {
         vectorSearchConfiguration: {
           numberOfResults: topK,
+          overrideSearchType,
+          filter,
         },
       },
     });
@@ -96,19 +115,48 @@ export class AmazonKnowledgeBaseRetriever extends BaseRetriever {
     );
 
     return (
-      retrieveResponse.retrievalResults?.map((result) => ({
-        pageContent: this.cleanResult(result.content?.text || ""),
-        metadata: {
-          source: result.location?.s3Location?.uri,
-          score: result.score,
-          ...result.metadata,
-        },
-      })) ?? ([] as Array<Document>)
+      retrieveResponse.retrievalResults?.map((result) => {
+        let source;
+        switch (result.location?.type) {
+          case "CONFLUENCE":
+            source = result.location?.confluenceLocation?.url;
+            break;
+          case "S3":
+            source = result.location?.s3Location?.uri;
+            break;
+          case "SALESFORCE":
+            source = result.location?.salesforceLocation?.url;
+            break;
+          case "SHAREPOINT":
+            source = result.location?.sharePointLocation?.url;
+            break;
+          case "WEB":
+            source = result.location?.webLocation?.url;
+            break;
+          default:
+            source = result.location?.s3Location?.uri;
+            break;
+        }
+
+        return {
+          pageContent: this.cleanResult(result.content?.text || ""),
+          metadata: {
+            source,
+            score: result.score,
+            ...result.metadata,
+          },
+        };
+      }) ?? ([] as Array<Document>)
     );
   }
 
   async _getRelevantDocuments(query: string): Promise<Document[]> {
-    const docs = await this.queryKnowledgeBase(query, this.topK);
+    const docs = await this.queryKnowledgeBase(
+      query,
+      this.topK,
+      this.filter,
+      this.overrideSearchType
+    );
     return docs;
   }
 }
diff --git a/libs/langchain-aws/src/retrievers/tests/bedrock.int.test.ts b/libs/langchain-aws/src/retrievers/tests/bedrock.int.test.ts
@@ -16,6 +16,8 @@ test("AmazonKnowledgeBaseRetriever", async () => {
     topK: 10,
     knowledgeBaseId: process.env.AMAZON_KNOWLEDGE_BASE_ID || "",
     region: process.env.BEDROCK_AWS_REGION,
+    overrideSearchType: "HYBRID",
+    filter: undefined,
     clientOptions: {
       credentials: {
         accessKeyId: process.env.BEDROCK_AWS_ACCESS_KEY_ID,