feat(client-bedrock-agent-runtime): Added raw model response and usag…

…e metrics to PreProcessing and PostProcessing Trace
aws · Oct 2, 2024 · 9326e19 · 9326e19
1 parent cf6979c
commit 9326e19
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 17 deletions.
diff --git a/clients/client-bedrock-agent-runtime/src/commands/InvokeAgentCommand.ts b/clients/client-bedrock-agent-runtime/src/commands/InvokeAgentCommand.ts
@@ -393,6 +393,15 @@ export interface InvokeAgentCommandOutput extends InvokeAgentResponse, __Metadat
  * //               rationale: "STRING_VALUE",
  * //               isValid: true || false,
  * //             },
+ * //             rawResponse: { // RawResponse
+ * //               content: "STRING_VALUE",
+ * //             },
+ * //             metadata: { // Metadata
+ * //               usage: { // Usage
+ * //                 inputTokens: Number("int"),
+ * //                 outputTokens: Number("int"),
+ * //               },
+ * //             },
  * //           },
  * //         },
  * //         orchestrationTrace: { // OrchestrationTrace Union: only one key present
@@ -511,11 +520,11 @@ export interface InvokeAgentCommandOutput extends InvokeAgentResponse, __Metadat
  * //           },
  * //           modelInvocationOutput: { // OrchestrationModelInvocationOutput
  * //             traceId: "STRING_VALUE",
- * //             rawResponse: { // RawResponse
+ * //             rawResponse: {
  * //               content: "STRING_VALUE",
  * //             },
- * //             metadata: { // Metadata
- * //               usage: { // Usage
+ * //             metadata: {
+ * //               usage: {
  * //                 inputTokens: Number("int"),
  * //                 outputTokens: Number("int"),
  * //               },
@@ -545,6 +554,15 @@ export interface InvokeAgentCommandOutput extends InvokeAgentResponse, __Metadat
  * //             parsedResponse: { // PostProcessingParsedResponse
  * //               text: "STRING_VALUE",
  * //             },
+ * //             rawResponse: {
+ * //               content: "STRING_VALUE",
+ * //             },
+ * //             metadata: {
+ * //               usage: {
+ * //                 inputTokens: Number("int"),
+ * //                 outputTokens: Number("int"),
+ * //               },
+ * //             },
  * //           },
  * //         },
  * //         failureTrace: { // FailureTrace

diff --git a/clients/client-bedrock-agent-runtime/src/models/models_0.ts b/clients/client-bedrock-agent-runtime/src/models/models_0.ts
@@ -2953,6 +2953,18 @@ export interface PostProcessingModelInvocationOutput {
    * @public
    */
   parsedResponse?: PostProcessingParsedResponse;
+
+  /**
+   * <p>Contains the raw output from the foundation model.</p>
+   * @public
+   */
+  rawResponse?: RawResponse;
+
+  /**
+   * <p>Provides details of the foundation model.</p>
+   * @public
+   */
+  metadata?: Metadata;
 }
 
 /**
@@ -3055,6 +3067,18 @@ export interface PreProcessingModelInvocationOutput {
    * @public
    */
   parsedResponse?: PreProcessingParsedResponse;
+
+  /**
+   * <p>Contains the raw output from the foundation model.</p>
+   * @public
+   */
+  rawResponse?: RawResponse;
+
+  /**
+   * <p>Provides details of the foundation model.</p>
+   * @public
+   */
+  metadata?: Metadata;
 }
 
 /**
@@ -4062,24 +4086,24 @@ export interface ExternalSource {
 }
 
 /**
- * <p>The configurations of the external source wrapper object in the retrieveAndGenerate function.</p>
+ * <p>The configurations of the external source wrapper object in the <code>retrieveAndGenerate</code> function.</p>
  * @public
  */
 export interface ExternalSourcesRetrieveAndGenerateConfiguration {
   /**
-   * <p>The modelArn used with the external source wrapper object in the retrieveAndGenerate function.</p>
+   * <p>The model Amazon Resource Name (ARN) for the external source wrapper object in the <code>retrieveAndGenerate</code> function.</p>
    * @public
    */
   modelArn: string | undefined;
 
   /**
-   * <p>The document used with the external source wrapper object in the retrieveAndGenerate function.</p>
+   * <p>The document for the external source wrapper object in the <code>retrieveAndGenerate</code> function.</p>
    * @public
    */
   sources: ExternalSource[] | undefined;
 
   /**
-   * <p>The prompt used with the external source wrapper object with the retrieveAndGenerate function.</p>
+   * <p>The prompt used with the external source wrapper object with the <code>retrieveAndGenerate</code> function.</p>
    * @public
    */
   generationConfiguration?: ExternalSourcesGenerationConfiguration;
@@ -4912,19 +4936,20 @@ export interface RetrieveRequest {
  */
 export interface RetrieveAndGenerateConfiguration {
   /**
-   * <p>The type of resource that is queried by the request.</p>
+   * <p>The type of resource that contains your data for retrieving information and generating responses.</p>
+   *          <p>If you choose ot use <code>EXTERNAL_SOURCES</code>, then currently only Claude 3 Sonnet models for knowledge bases are supported.</p>
    * @public
    */
   type: RetrieveAndGenerateType | undefined;
 
   /**
-   * <p>Contains details about the resource being queried.</p>
+   * <p>Contains details about the knowledge base for retrieving information and generating responses.</p>
    * @public
    */
   knowledgeBaseConfiguration?: KnowledgeBaseRetrieveAndGenerateConfiguration;
 
   /**
-   * <p>The configuration used with the external source wrapper object in the retrieveAndGenerate function.</p>
+   * <p>The configuration for the external source wrapper object in the <code>retrieveAndGenerate</code> function.</p>
    * @public
    */
   externalSourcesConfiguration?: ExternalSourcesRetrieveAndGenerateConfiguration;
@@ -5546,6 +5571,8 @@ export const PostProcessingModelInvocationOutputFilterSensitiveLog = (
 ): any => ({
   ...obj,
   ...(obj.parsedResponse && { parsedResponse: SENSITIVE_STRING }),
+  ...(obj.rawResponse && { rawResponse: SENSITIVE_STRING }),
+  ...(obj.metadata && { metadata: SENSITIVE_STRING }),
 });
 
 /**
@@ -5571,6 +5598,8 @@ export const PreProcessingParsedResponseFilterSensitiveLog = (obj: PreProcessing
 export const PreProcessingModelInvocationOutputFilterSensitiveLog = (obj: PreProcessingModelInvocationOutput): any => ({
   ...obj,
   ...(obj.parsedResponse && { parsedResponse: SENSITIVE_STRING }),
+  ...(obj.rawResponse && { rawResponse: SENSITIVE_STRING }),
+  ...(obj.metadata && { metadata: SENSITIVE_STRING }),
 });
 
 /**

diff --git a/codegen/sdk-codegen/aws-models/bedrock-agent-runtime.json b/codegen/sdk-codegen/aws-models/bedrock-agent-runtime.json
@@ -1504,26 +1504,26 @@
         "modelArn": {
           "target": "com.amazonaws.bedrockagentruntime#BedrockModelArn",
           "traits": {
-            "smithy.api#documentation": "<p>The modelArn used with the external source wrapper object in the retrieveAndGenerate function.</p>",
+            "smithy.api#documentation": "<p>The model Amazon Resource Name (ARN) for the external source wrapper object in the <code>retrieveAndGenerate</code> function.</p>",
             "smithy.api#required": {}
           }
         },
         "sources": {
           "target": "com.amazonaws.bedrockagentruntime#ExternalSources",
           "traits": {
-            "smithy.api#documentation": "<p>The document used with the external source wrapper object in the retrieveAndGenerate function.</p>",
+            "smithy.api#documentation": "<p>The document for the external source wrapper object in the <code>retrieveAndGenerate</code> function.</p>",
             "smithy.api#required": {}
           }
         },
         "generationConfiguration": {
           "target": "com.amazonaws.bedrockagentruntime#ExternalSourcesGenerationConfiguration",
           "traits": {
-            "smithy.api#documentation": "<p>The prompt used with the external source wrapper object with the retrieveAndGenerate function.</p>"
+            "smithy.api#documentation": "<p>The prompt used with the external source wrapper object with the <code>retrieveAndGenerate</code> function.</p>"
           }
         }
       },
       "traits": {
-        "smithy.api#documentation": "<p>The configurations of the external source wrapper object in the retrieveAndGenerate function.</p>"
+        "smithy.api#documentation": "<p>The configurations of the external source wrapper object in the <code>retrieveAndGenerate</code> function.</p>"
       }
     },
     "com.amazonaws.bedrockagentruntime#FailureReasonString": {
@@ -4217,6 +4217,12 @@
           "traits": {
             "smithy.api#documentation": "<p>Details about the response from the Lambda parsing of the output of the post-processing step.</p>"
           }
+        },
+        "rawResponse": {
+          "target": "com.amazonaws.bedrockagentruntime#RawResponse"
+        },
+        "metadata": {
+          "target": "com.amazonaws.bedrockagentruntime#Metadata"
         }
       },
       "traits": {
@@ -4274,6 +4280,12 @@
           "traits": {
             "smithy.api#documentation": "<p>Details about the response from the Lambda parsing of the output of the pre-processing step.</p>"
           }
+        },
+        "rawResponse": {
+          "target": "com.amazonaws.bedrockagentruntime#RawResponse"
+        },
+        "metadata": {
+          "target": "com.amazonaws.bedrockagentruntime#Metadata"
         }
       },
       "traits": {
@@ -5038,20 +5050,20 @@
         "type": {
           "target": "com.amazonaws.bedrockagentruntime#RetrieveAndGenerateType",
           "traits": {
-            "smithy.api#documentation": "<p>The type of resource that is queried by the request.</p>",
+            "smithy.api#documentation": "<p>The type of resource that contains your data for retrieving information and generating responses.</p>\n         <p>If you choose ot use <code>EXTERNAL_SOURCES</code>, then currently only Claude 3 Sonnet models for knowledge bases are supported.</p>",
             "smithy.api#required": {}
           }
         },
         "knowledgeBaseConfiguration": {
           "target": "com.amazonaws.bedrockagentruntime#KnowledgeBaseRetrieveAndGenerateConfiguration",
           "traits": {
-            "smithy.api#documentation": "<p>Contains details about the resource being queried.</p>"
+            "smithy.api#documentation": "<p>Contains details about the knowledge base for retrieving information and generating responses.</p>"
           }
         },
         "externalSourcesConfiguration": {
           "target": "com.amazonaws.bedrockagentruntime#ExternalSourcesRetrieveAndGenerateConfiguration",
           "traits": {
-            "smithy.api#documentation": "<p>The configuration used with the external source wrapper object in the retrieveAndGenerate function.</p>"
+            "smithy.api#documentation": "<p>The configuration for the external source wrapper object in the <code>retrieveAndGenerate</code> function.</p>"
           }
         }
       },