feat(client-bedrock-agent-runtime): This release adds support to override search strategy performed by the Retrieve and RetrieveAndGenerate APIs for Amazon Bedrock Agents

Author: awstools

clients/client-bedrock-agent-runtime/src/commands/RetrieveAndGenerateCommand.ts

@@ -54,6 +54,12 @@ export interface RetrieveAndGenerateCommandOutput extends RetrieveAndGenerateRes
  *     knowledgeBaseConfiguration: { // KnowledgeBaseRetrieveAndGenerateConfiguration
  *       knowledgeBaseId: "STRING_VALUE", // required
  *       modelArn: "STRING_VALUE", // required
+ *       retrievalConfiguration: { // KnowledgeBaseRetrievalConfiguration
+ *         vectorSearchConfiguration: { // KnowledgeBaseVectorSearchConfiguration
+ *           numberOfResults: Number("int"),
+ *           overrideSearchType: "HYBRID" || "SEMANTIC",
+ *         },
+ *       },
  *     },
  *   },
  *   sessionConfiguration: { // RetrieveAndGenerateSessionConfiguration

clients/client-bedrock-agent-runtime/src/commands/RetrieveCommand.ts

@@ -51,7 +51,8 @@ export interface RetrieveCommandOutput extends RetrieveResponse, __MetadataBeare
  *   },
  *   retrievalConfiguration: { // KnowledgeBaseRetrievalConfiguration
  *     vectorSearchConfiguration: { // KnowledgeBaseVectorSearchConfiguration
- *       numberOfResults: Number("int"), // required
+ *       numberOfResults: Number("int"),
+ *       overrideSearchType: "HYBRID" || "SEMANTIC",
  *     },
  *   },
  *   nextToken: "STRING_VALUE",

clients/client-bedrock-agent-runtime/src/models/models_0.ts

@@ -1554,6 +1554,50 @@ export interface RetrieveAndGenerateInput {
   text: string | undefined;
 }
 
+/**
+ * @public
+ * @enum
+ */
+export const SearchType = {
+  HYBRID: "HYBRID",
+  SEMANTIC: "SEMANTIC",
+} as const;
+
+/**
+ * @public
+ */
+export type SearchType = (typeof SearchType)[keyof typeof SearchType];
+
+/**
+ * @public
+ * Knowledge base vector search configuration
+ */
+export interface KnowledgeBaseVectorSearchConfiguration {
+  /**
+   * @public
+   * Top-K results to retrieve from knowledge base.
+   */
+  numberOfResults?: number;
+
+  /**
+   * @public
+   * Override the type of query to be performed on data store
+   */
+  overrideSearchType?: SearchType;
+}
+
+/**
+ * @public
+ * Search parameters for retrieving from knowledge base.
+ */
+export interface KnowledgeBaseRetrievalConfiguration {
+  /**
+   * @public
+   * Knowledge base vector search configuration
+   */
+  vectorSearchConfiguration: KnowledgeBaseVectorSearchConfiguration | undefined;
+}
+
 /**
  * @public
  * Configurations for retrieval and generation for knowledge base.
@@ -1570,6 +1614,12 @@ export interface KnowledgeBaseRetrieveAndGenerateConfiguration {
    * Arn of a Bedrock model.
    */
   modelArn: string | undefined;
+
+  /**
+   * @public
+   * Search parameters for retrieving from knowledge base.
+   */
+  retrievalConfiguration?: KnowledgeBaseRetrievalConfiguration;
 }
 
 /**
@@ -1679,30 +1729,6 @@ export interface RetrieveAndGenerateResponse {
   citations?: Citation[];
 }
 
-/**
- * @public
- * Knowledge base vector search configuration
- */
-export interface KnowledgeBaseVectorSearchConfiguration {
-  /**
-   * @public
-   * Top-K results to retrieve from knowledge base.
-   */
-  numberOfResults: number | undefined;
-}
-
-/**
- * @public
- * Search parameters for retrieving from knowledge base.
- */
-export interface KnowledgeBaseRetrievalConfiguration {
-  /**
-   * @public
-   * Knowledge base vector search configuration
-   */
-  vectorSearchConfiguration: KnowledgeBaseVectorSearchConfiguration | undefined;
-}
-
 /**
  * @public
  * Knowledge base input query.
@@ -1811,12 +1837,72 @@ export const InvokeAgentRequestFilterSensitiveLog = (obj: InvokeAgentRequest): a
   ...(obj.inputText && { inputText: SENSITIVE_STRING }),
 });
 
+/**
+ * @internal
+ */
+export const TextResponsePartFilterSensitiveLog = (obj: TextResponsePart): any => ({
+  ...obj,
+});
+
+/**
+ * @internal
+ */
+export const GeneratedResponsePartFilterSensitiveLog = (obj: GeneratedResponsePart): any => ({
+  ...obj,
+  ...(obj.textResponsePart && { textResponsePart: SENSITIVE_STRING }),
+});
+
+/**
+ * @internal
+ */
+export const RetrievalResultContentFilterSensitiveLog = (obj: RetrievalResultContent): any => ({
+  ...obj,
+});
+
+/**
+ * @internal
+ */
+export const RetrievalResultLocationFilterSensitiveLog = (obj: RetrievalResultLocation): any => ({
+  ...obj,
+});
+
+/**
+ * @internal
+ */
+export const RetrievedReferenceFilterSensitiveLog = (obj: RetrievedReference): any => ({
+  ...obj,
+  ...(obj.content && { content: SENSITIVE_STRING }),
+  ...(obj.location && { location: SENSITIVE_STRING }),
+});
+
+/**
+ * @internal
+ */
+export const CitationFilterSensitiveLog = (obj: Citation): any => ({
+  ...obj,
+  ...(obj.generatedResponsePart && {
+    generatedResponsePart: GeneratedResponsePartFilterSensitiveLog(obj.generatedResponsePart),
+  }),
+  ...(obj.retrievedReferences && {
+    retrievedReferences: obj.retrievedReferences.map((item) => RetrievedReferenceFilterSensitiveLog(item)),
+  }),
+});
+
+/**
+ * @internal
+ */
+export const AttributionFilterSensitiveLog = (obj: Attribution): any => ({
+  ...obj,
+  ...(obj.citations && { citations: obj.citations.map((item) => CitationFilterSensitiveLog(item)) }),
+});
+
 /**
  * @internal
  */
 export const PayloadPartFilterSensitiveLog = (obj: PayloadPart): any => ({
   ...obj,
   ...(obj.bytes && { bytes: SENSITIVE_STRING }),
+  ...(obj.attribution && { attribution: AttributionFilterSensitiveLog(obj.attribution) }),
 });
 
 /**
@@ -1865,6 +1951,16 @@ export const FinalResponseFilterSensitiveLog = (obj: FinalResponse): any => ({
   ...(obj.text && { text: SENSITIVE_STRING }),
 });
 
+/**
+ * @internal
+ */
+export const KnowledgeBaseLookupOutputFilterSensitiveLog = (obj: KnowledgeBaseLookupOutput): any => ({
+  ...obj,
+  ...(obj.retrievedReferences && {
+    retrievedReferences: obj.retrievedReferences.map((item) => RetrievedReferenceFilterSensitiveLog(item)),
+  }),
+});
+
 /**
  * @internal
  */
@@ -1881,6 +1977,9 @@ export const ObservationFilterSensitiveLog = (obj: Observation): any => ({
   ...(obj.actionGroupInvocationOutput && {
     actionGroupInvocationOutput: ActionGroupInvocationOutputFilterSensitiveLog(obj.actionGroupInvocationOutput),
   }),
+  ...(obj.knowledgeBaseLookupOutput && {
+    knowledgeBaseLookupOutput: KnowledgeBaseLookupOutputFilterSensitiveLog(obj.knowledgeBaseLookupOutput),
+  }),
   ...(obj.finalResponse && { finalResponse: FinalResponseFilterSensitiveLog(obj.finalResponse) }),
   ...(obj.repromptResponse && { repromptResponse: SENSITIVE_STRING }),
 });
@@ -2030,6 +2129,7 @@ export const RetrieveAndGenerateOutputFilterSensitiveLog = (obj: RetrieveAndGene
 export const RetrieveAndGenerateResponseFilterSensitiveLog = (obj: RetrieveAndGenerateResponse): any => ({
   ...obj,
   ...(obj.output && { output: SENSITIVE_STRING }),
+  ...(obj.citations && { citations: obj.citations.map((item) => CitationFilterSensitiveLog(item)) }),
 });
 
 /**
@@ -2047,6 +2147,15 @@ export const RetrieveRequestFilterSensitiveLog = (obj: RetrieveRequest): any =>
   ...(obj.retrievalQuery && { retrievalQuery: SENSITIVE_STRING }),
 });
 
+/**
+ * @internal
+ */
+export const KnowledgeBaseRetrievalResultFilterSensitiveLog = (obj: KnowledgeBaseRetrievalResult): any => ({
+  ...obj,
+  ...(obj.content && { content: SENSITIVE_STRING }),
+  ...(obj.location && { location: SENSITIVE_STRING }),
+});
+
 /**
  * @internal
  */

codegen/sdk-codegen/aws-models/bedrock-agent-runtime.json

@@ -1326,6 +1326,9 @@
           "traits": {
             "smithy.api#required": {}
           }
+        },
+        "retrievalConfiguration": {
+          "target": "com.amazonaws.bedrockagentruntime#KnowledgeBaseRetrievalConfiguration"
         }
       },
       "traits": {
@@ -1338,12 +1341,19 @@
         "numberOfResults": {
           "target": "smithy.api#Integer",
           "traits": {
+            "smithy.api#addedDefault": {},
+            "smithy.api#default": 5,
             "smithy.api#documentation": "Top-K results to retrieve from knowledge base.",
             "smithy.api#range": {
               "min": 1,
-              "max": 10
-            },
-            "smithy.api#required": {}
+              "max": 25
+            }
+          }
+        },
+        "overrideSearchType": {
+          "target": "com.amazonaws.bedrockagentruntime#SearchType",
+          "traits": {
+            "smithy.api#documentation": "Override the type of query to be performed on data store"
           }
         }
       },
@@ -1795,7 +1805,8 @@
         }
       },
       "traits": {
-        "smithy.api#documentation": "Content of a retrieval result."
+        "smithy.api#documentation": "Content of a retrieval result.",
+        "smithy.api#sensitive": {}
       }
     },
     "com.amazonaws.bedrockagentruntime#RetrievalResultLocation": {
@@ -1812,7 +1823,8 @@
         }
       },
       "traits": {
-        "smithy.api#documentation": "The source location of a retrieval result."
+        "smithy.api#documentation": "The source location of a retrieval result.",
+        "smithy.api#sensitive": {}
       }
     },
     "com.amazonaws.bedrockagentruntime#RetrievalResultLocationType": {
@@ -2151,6 +2163,26 @@
         "smithy.api#documentation": "list of retrieved references"
       }
     },
+    "com.amazonaws.bedrockagentruntime#SearchType": {
+      "type": "enum",
+      "members": {
+        "HYBRID": {
+          "target": "smithy.api#Unit",
+          "traits": {
+            "smithy.api#enumValue": "HYBRID"
+          }
+        },
+        "SEMANTIC": {
+          "target": "smithy.api#Unit",
+          "traits": {
+            "smithy.api#enumValue": "SEMANTIC"
+          }
+        }
+      },
+      "traits": {
+        "smithy.api#documentation": "Query type to be performed on data store."
+      }
+    },
     "com.amazonaws.bedrockagentruntime#ServiceQuotaExceededException": {
       "type": "structure",
       "members": {
@@ -2297,7 +2329,8 @@
         }
       },
       "traits": {
-        "smithy.api#documentation": "Text response part"
+        "smithy.api#documentation": "Text response part",
+        "smithy.api#sensitive": {}
       }
     },
     "com.amazonaws.bedrockagentruntime#ThrottlingException": {