feat(client-bedrock-agent-runtime): Knowledge Bases for Amazon Bedroc…

…k now supports custom prompts and model parameters in the orchestrationConfiguration of the RetrieveAndGenerate API. The modelArn field accepts Custom Models and Imported Models ARNs.
aws · Oct 21, 2024 · f3c2a2c · f3c2a2c
1 parent 9588015
commit f3c2a2c
Show file tree

Hide file tree

Showing 4 changed files with 108 additions and 16 deletions.
diff --git a/clients/client-bedrock-agent-runtime/src/commands/RetrieveAndGenerateCommand.ts b/clients/client-bedrock-agent-runtime/src/commands/RetrieveAndGenerateCommand.ts
@@ -135,6 +135,22 @@ export interface RetrieveAndGenerateCommandOutput extends RetrieveAndGenerateRes
  *         },
  *       },
  *       orchestrationConfiguration: { // OrchestrationConfiguration
+ *         promptTemplate: {
+ *           textPromptTemplate: "STRING_VALUE",
+ *         },
+ *         inferenceConfig: {
+ *           textInferenceConfig: {
+ *             temperature: Number("float"),
+ *             topP: Number("float"),
+ *             maxTokens: Number("int"),
+ *             stopSequences: [
+ *               "STRING_VALUE",
+ *             ],
+ *           },
+ *         },
+ *         additionalModelRequestFields: {
+ *           "<keys>": "DOCUMENT_VALUE",
+ *         },
  *         queryTransformationConfiguration: { // QueryTransformationConfiguration
  *           type: "QUERY_DECOMPOSITION", // required
  *         },

diff --git a/clients/client-bedrock-agent-runtime/src/models/models_0.ts b/clients/client-bedrock-agent-runtime/src/models/models_0.ts
@@ -2604,7 +2604,7 @@ export interface OrchestrationModelInvocationOutput {
   rawResponse?: RawResponse;
 
   /**
-   * <p>Contains information about the foundation model output.</p>
+   * <p>Contains information about the foundation model output from the orchestration step.</p>
    * @public
    */
   metadata?: Metadata;
@@ -2955,13 +2955,17 @@ export interface PostProcessingModelInvocationOutput {
   parsedResponse?: PostProcessingParsedResponse;
 
   /**
-   * <p>Contains the raw output from the foundation model.</p>
+   * <p>
+   *             Details of the raw response from the foundation model output.
+   *         </p>
    * @public
    */
   rawResponse?: RawResponse;
 
   /**
-   * <p>Provides details of the foundation model.</p>
+   * <p>
+   *             Contains information about the foundation model output from the post-processing step.
+   *         </p>
    * @public
    */
   metadata?: Metadata;
@@ -3069,13 +3073,17 @@ export interface PreProcessingModelInvocationOutput {
   parsedResponse?: PreProcessingParsedResponse;
 
   /**
-   * <p>Contains the raw output from the foundation model.</p>
+   * <p>
+   *             Details of the raw response from the foundation model output.
+   *         </p>
    * @public
    */
   rawResponse?: RawResponse;
 
   /**
-   * <p>Provides details of the foundation model.</p>
+   * <p>
+   *             Contains information about the foundation model output from the pre-processing step.
+   *         </p>
    * @public
    */
   metadata?: Metadata;
@@ -4178,12 +4186,30 @@ export interface QueryTransformationConfiguration {
  * @public
  */
 export interface OrchestrationConfiguration {
+  /**
+   * <p>Contains the template for the prompt that's sent to the model for response generation.</p>
+   * @public
+   */
+  promptTemplate?: PromptTemplate;
+
+  /**
+   * <p> Configuration settings for inference when using RetrieveAndGenerate to generate responses while using a knowledge base as a source. </p>
+   * @public
+   */
+  inferenceConfig?: InferenceConfig;
+
+  /**
+   * <p> Additional model parameters and corresponding values not included in the textInferenceConfig structure for a knowledge base. This allows users to provide custom model parameters specific to the language model being used. </p>
+   * @public
+   */
+  additionalModelRequestFields?: Record<string, __DocumentType>;
+
   /**
    * <p>To split up the prompt and retrieve multiple sources, set the transformation type to
    *     <code>QUERY_DECOMPOSITION</code>.</p>
    * @public
    */
-  queryTransformationConfiguration: QueryTransformationConfiguration | undefined;
+  queryTransformationConfiguration?: QueryTransformationConfiguration;
 }
 
 /**
@@ -5723,6 +5749,14 @@ export const GenerationConfigurationFilterSensitiveLog = (obj: GenerationConfigu
   ...(obj.promptTemplate && { promptTemplate: PromptTemplateFilterSensitiveLog(obj.promptTemplate) }),
 });
 
+/**
+ * @internal
+ */
+export const OrchestrationConfigurationFilterSensitiveLog = (obj: OrchestrationConfiguration): any => ({
+  ...obj,
+  ...(obj.promptTemplate && { promptTemplate: PromptTemplateFilterSensitiveLog(obj.promptTemplate) }),
+});
+
 /**
  * @internal
  */
@@ -5829,6 +5863,9 @@ export const KnowledgeBaseRetrieveAndGenerateConfigurationFilterSensitiveLog = (
   ...(obj.generationConfiguration && {
     generationConfiguration: GenerationConfigurationFilterSensitiveLog(obj.generationConfiguration),
   }),
+  ...(obj.orchestrationConfiguration && {
+    orchestrationConfiguration: OrchestrationConfigurationFilterSensitiveLog(obj.orchestrationConfiguration),
+  }),
 });
 
 /**

diff --git a/clients/client-bedrock-agent-runtime/src/protocols/Aws_restJson1.ts b/clients/client-bedrock-agent-runtime/src/protocols/Aws_restJson1.ts
@@ -1122,7 +1122,7 @@ const se_KnowledgeBaseRetrieveAndGenerateConfiguration = (
     generationConfiguration: (_) => se_GenerationConfiguration(_, context),
     knowledgeBaseId: [],
     modelArn: [],
-    orchestrationConfiguration: _json,
+    orchestrationConfiguration: (_) => se_OrchestrationConfiguration(_, context),
     retrievalConfiguration: (_) => se_KnowledgeBaseRetrievalConfiguration(_, context),
   });
 };
@@ -1141,7 +1141,17 @@ const se_KnowledgeBaseVectorSearchConfiguration = (
   });
 };
 
-// se_OrchestrationConfiguration omitted.
+/**
+ * serializeAws_restJson1OrchestrationConfiguration
+ */
+const se_OrchestrationConfiguration = (input: OrchestrationConfiguration, context: __SerdeContext): any => {
+  return take(input, {
+    additionalModelRequestFields: (_) => se_AdditionalModelRequestFields(_, context),
+    inferenceConfig: (_) => se_InferenceConfig(_, context),
+    promptTemplate: _json,
+    queryTransformationConfiguration: _json,
+  });
+};
 
 // se_PromptSessionAttributesMap omitted.
 

diff --git a/codegen/sdk-codegen/aws-models/bedrock-agent-runtime.json b/codegen/sdk-codegen/aws-models/bedrock-agent-runtime.json
@@ -1064,7 +1064,7 @@
           "min": 1,
           "max": 2048
         },
-        "smithy.api#pattern": "^(arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:(([0-9]{12}:custom-model/[a-z0-9-]{1,63}[.]{1}[a-z0-9-]{1,63}/[a-z0-9]{12})|(:foundation-model/[a-z0-9-]{1,63}[.]{1}[a-z0-9-]{1,63}([.:]?[a-z0-9-]{1,63}))))|(arn:aws(|-us-gov|-cn|-iso|-iso-b):bedrock:(|[0-9a-z-]{1,20}):(|[0-9]{12}):inference-profile/[a-zA-Z0-9-:.]+)|([a-z0-9-]{1,63}[.]{1}[a-z0-9-]{1,63}([.:]?[a-z0-9-]{1,63}))|(([0-9a-zA-Z][_-]?)+)$"
+        "smithy.api#pattern": "^(arn:aws(-[^:]{1,10})?:bedrock:[a-z0-9-]{1,20}:([0-9]{12})?:([a-z-]{1,20}/)?)?([a-z0-9.-]{1,63}){0,2}(([:][a-z0-9-]{1,63}){0,2})?(/[a-z0-9]{1,12})?$"
       }
     },
     "com.amazonaws.bedrockagentruntime#ByteContentBlob": {
@@ -4012,11 +4012,28 @@
     "com.amazonaws.bedrockagentruntime#OrchestrationConfiguration": {
       "type": "structure",
       "members": {
+        "promptTemplate": {
+          "target": "com.amazonaws.bedrockagentruntime#PromptTemplate",
+          "traits": {
+            "smithy.api#documentation": "<p>Contains the template for the prompt that's sent to the model for response generation.</p>"
+          }
+        },
+        "inferenceConfig": {
+          "target": "com.amazonaws.bedrockagentruntime#InferenceConfig",
+          "traits": {
+            "smithy.api#documentation": "<p> Configuration settings for inference when using RetrieveAndGenerate to generate responses while using a knowledge base as a source. </p>"
+          }
+        },
+        "additionalModelRequestFields": {
+          "target": "com.amazonaws.bedrockagentruntime#AdditionalModelRequestFields",
+          "traits": {
+            "smithy.api#documentation": "<p> Additional model parameters and corresponding values not included in the textInferenceConfig structure for a knowledge base. This allows users to provide custom model parameters specific to the language model being used. </p>"
+          }
+        },
         "queryTransformationConfiguration": {
           "target": "com.amazonaws.bedrockagentruntime#QueryTransformationConfiguration",
           "traits": {
-            "smithy.api#documentation": "<p>To split up the prompt and retrieve multiple sources, set the transformation type to\n    <code>QUERY_DECOMPOSITION</code>.</p>",
-            "smithy.api#required": {}
+            "smithy.api#documentation": "<p>To split up the prompt and retrieve multiple sources, set the transformation type to\n    <code>QUERY_DECOMPOSITION</code>.</p>"
           }
         }
       },
@@ -4042,7 +4059,7 @@
         "metadata": {
           "target": "com.amazonaws.bedrockagentruntime#Metadata",
           "traits": {
-            "smithy.api#documentation": "<p>Contains information about the foundation model output.</p>"
+            "smithy.api#documentation": "<p>Contains information about the foundation model output from the orchestration step.</p>"
           }
         }
       },
@@ -4219,10 +4236,16 @@
           }
         },
         "rawResponse": {
-          "target": "com.amazonaws.bedrockagentruntime#RawResponse"
+          "target": "com.amazonaws.bedrockagentruntime#RawResponse",
+          "traits": {
+            "smithy.api#documentation": "<p>\n            Details of the raw response from the foundation model output.\n        </p>"
+          }
         },
         "metadata": {
-          "target": "com.amazonaws.bedrockagentruntime#Metadata"
+          "target": "com.amazonaws.bedrockagentruntime#Metadata",
+          "traits": {
+            "smithy.api#documentation": "<p>\n            Contains information about the foundation model output from the post-processing step.\n        </p>"
+          }
         }
       },
       "traits": {
@@ -4282,10 +4305,16 @@
           }
         },
         "rawResponse": {
-          "target": "com.amazonaws.bedrockagentruntime#RawResponse"
+          "target": "com.amazonaws.bedrockagentruntime#RawResponse",
+          "traits": {
+            "smithy.api#documentation": "<p>\n            Details of the raw response from the foundation model output.\n        </p>"
+          }
         },
         "metadata": {
-          "target": "com.amazonaws.bedrockagentruntime#Metadata"
+          "target": "com.amazonaws.bedrockagentruntime#Metadata",
+          "traits": {
+            "smithy.api#documentation": "<p>\n            Contains information about the foundation model output from the pre-processing step.\n        </p>"
+          }
         }
       },
       "traits": {