suppress the exception from unspported InvokeModel streaming APIs

Author: mxiamxia

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts

@@ -738,53 +738,15 @@ describe('BedrockRuntime', () => {
     });
 
     describe('Response Body Type Handling', () => {
-      it('handles string response body correctly', async () => {
+      it('handles normal Anthropic Claude response correctly', async () => {
         const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
         const mockRequestBody: string = JSON.stringify({
           anthropic_version: 'bedrock-2023-05-31',
           max_tokens: 1000,
           messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
         });
 
-        // Mock response body as already converted string
-        const mockResponseBodyString = JSON.stringify({
-          stop_reason: 'end_turn',
-          usage: { input_tokens: 15, output_tokens: 13 },
-        });
-
-        nock(`https://bedrock-runtime.${region}.amazonaws.com`)
-          .post(`/model/${encodeURIComponent(modelId)}/invoke`)
-          .reply(200, mockResponseBodyString);
-
-        await bedrock
-          .invokeModel({
-            modelId: modelId,
-            body: mockRequestBody,
-          })
-          .catch((err: any) => {});
-
-        const testSpans: ReadableSpan[] = getTestSpans();
-        const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => {
-          return s.name === 'BedrockRuntime.InvokeModel';
-        });
-        expect(invokeModelSpans.length).toBe(1);
-        const invokeModelSpan = invokeModelSpans[0];
-
-        // Verify attributes are set correctly despite body being a string
-        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(15);
-        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(13);
-        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['end_turn']);
-      });
-
-      it('handles Anthropic Claude response body correctly', async () => {
-        const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
-        const mockRequestBody: string = JSON.stringify({
-          anthropic_version: 'bedrock-2023-05-31',
-          max_tokens: 1000,
-          messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
-        });
-
-        // Mock response body - use standard object format (AWS SDK will handle type conversion)
+        // Use standard object format - AWS SDK and instrumentation will handle the conversion
         const mockResponseBodyObj = {
           stop_reason: 'end_turn',
           usage: { input_tokens: 20, output_tokens: 15 },
@@ -814,24 +776,18 @@ describe('BedrockRuntime', () => {
         expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['end_turn']);
       });
 
-      it('handles Buffer response body correctly', async () => {
+      it('handles unexpected body type gracefully', async () => {
         const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
         const mockRequestBody: string = JSON.stringify({
           anthropic_version: 'bedrock-2023-05-31',
           max_tokens: 1000,
           messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
         });
 
-        // Mock response body as Buffer
-        const mockResponseBodyObj = {
-          stop_reason: 'max_tokens',
-          usage: { input_tokens: 25, output_tokens: 18 },
-        };
-        const mockResponseBodyBuffer = Buffer.from(JSON.stringify(mockResponseBodyObj), 'utf8');
-
+        // Mock response body as unexpected type - using reply function to return a number
         nock(`https://bedrock-runtime.${region}.amazonaws.com`)
           .post(`/model/${encodeURIComponent(modelId)}/invoke`)
-          .reply(200, mockResponseBodyBuffer);
+          .reply(200, () => 12345 as any);
 
         await bedrock
           .invokeModel({
@@ -847,26 +803,35 @@ describe('BedrockRuntime', () => {
         expect(invokeModelSpans.length).toBe(1);
         const invokeModelSpan = invokeModelSpans[0];
 
-        // Verify attributes are set correctly when body is Buffer
-        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(25);
-        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(18);
-        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual([
-          'max_tokens',
-        ]);
+        // Verify that no AI attributes are set when body type is unexpected
+        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBeUndefined();
+        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBeUndefined();
+        expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toBeUndefined();
+
+        // Note: We can't easily test diag.debug() output in unit tests, but the important part
+        // is that the function returns early and doesn't crash when encountering unexpected types
+        // Debug message will be: "Unexpected body type in Bedrock response: number for commandName InvokeModelCommand"
       });
 
-      it('handles unexpected body type gracefully', async () => {
+      it('handles streaming response (SmithyMessageDecoderStream) gracefully', async () => {
         const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
         const mockRequestBody: string = JSON.stringify({
           anthropic_version: 'bedrock-2023-05-31',
           max_tokens: 1000,
           messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
         });
 
-        // Mock response body as unexpected type - using reply function to return a number
+        // Mock response body as streaming object (constructor name matching)
+        const mockStreamingBody = {
+          constructor: { name: 'SmithyMessageDecoderStream' },
+          [Symbol.asyncIterator]: function* () {
+            yield { chunk: { bytes: new TextEncoder().encode('{"type":"chunk"}') } };
+          },
+        };
+
         nock(`https://bedrock-runtime.${region}.amazonaws.com`)
           .post(`/model/${encodeURIComponent(modelId)}/invoke`)
-          .reply(200, () => 12345 as any);
+          .reply(200, mockStreamingBody);
 
         await bedrock
           .invokeModel({
@@ -882,13 +847,13 @@ describe('BedrockRuntime', () => {
         expect(invokeModelSpans.length).toBe(1);
         const invokeModelSpan = invokeModelSpans[0];
 
-        // Verify that no AI attributes are set when body type is unexpected
+        // Verify that no AI attributes are set when body is streaming (metrics not available in initial response)
         expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBeUndefined();
         expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBeUndefined();
         expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toBeUndefined();
 
-        // Note: We can't easily test diag.debug() output in unit tests, but the important part
-        // is that the function returns early and doesn't crash when encountering unexpected types
+        // Streaming responses should be skipped gracefully without crashing
+        // TODO: support InvokeModel Streaming API and Converse APIs later
       });
     });
   });

sample-applications/simple-express-server/package.json

@@ -11,10 +11,12 @@
   "author": "",
   "license": "Apache-2.0",
   "dependencies": {
+    "@aws-sdk/client-bedrock-runtime": "3.574.0",
     "@aws-sdk/client-dynamodb": "^3.621.0",
     "@aws-sdk/client-s3": "^3.621.0",
     "@aws-sdk/client-sqs": "^3.621.0",
     "@aws-sdk/lib-dynamodb": "^3.616.0",
+    "@aws/aws-distro-opentelemetry-node-autoinstrumentation": "file:../../aws-distro-opentelemetry-node-autoinstrumentation",
     "@types/express": "^4.17.21",
     "@types/node": "^20.14.6",
     "body-parser": "^1.20.2",