Skip to content

Commit 543d5ee

Browse files
mxiamxiamxiamxia
committed
suppress the exception from unspported InvokeModel streaming APIs
1 parent 47b7770 commit 543d5ee

File tree

2 files changed

+51
-82
lines changed

2 files changed

+51
-82
lines changed

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

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -331,20 +331,28 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
331331
responseHook(response: NormalizedResponse, span: Span, tracer: Tracer, config: AwsSdkInstrumentationConfig): void {
332332
const currentModelId = response.request.commandInput?.modelId;
333333
if (response.data?.body) {
334-
let decodedResponseBody: string;
334+
// Check if this is a streaming response (SmithyMessageDecoderStream)
335+
// Intend to not using instanceOf to avoid import smithy as new dep for this file
336+
// https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/client/bedrock-runtime/command/InvokeModelWithResponseStreamCommand
337+
if (response.data.body.constructor?.name === 'SmithyMessageDecoderStream') {
338+
// TODO: support InvokeModel Streaming API and Converse APIs later
339+
diag.debug('Streaming API for invoking model is not supported', response.request.commandName);
340+
return;
341+
}
335342

336-
if (typeof response.data.body === 'string') {
337-
// Already converted by AWS SDK middleware by Uint8ArrayBlobAdapter
338-
decodedResponseBody = response.data.body;
339-
} else if (response.data.body instanceof Uint8Array) {
343+
let decodedResponseBody: string;
344+
// For InvokeModel API which should always have reponse body with Uint8Array type
345+
// https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/client/bedrock-runtime/command/InvokeModelCommand/
346+
if (response.data.body instanceof Uint8Array) {
340347
// Raw Uint8Array from AWS
341348
decodedResponseBody = new TextDecoder().decode(response.data.body);
342-
} else if (Buffer.isBuffer(response.data.body)) {
343-
// Node.js Buffer convert with toString('utf8')
344-
decodedResponseBody = response.data.body.toString('utf8');
345349
} else {
346350
// Handle unexpected types - log and skip processing
347-
diag.debug(`Unexpected body type in Bedrock response: ${typeof response.data.body}`, response.data.body);
351+
diag.debug(
352+
`Unexpected body type in Bedrock response: ${typeof response.data.body} for commandName ${
353+
response.request.commandName
354+
}`
355+
);
348356
return;
349357
}
350358

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

Lines changed: 34 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -738,53 +738,15 @@ describe('BedrockRuntime', () => {
738738
});
739739

740740
describe('Response Body Type Handling', () => {
741-
it('handles string response body correctly', async () => {
741+
it('handles normal Anthropic Claude response correctly', async () => {
742742
const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
743743
const mockRequestBody: string = JSON.stringify({
744744
anthropic_version: 'bedrock-2023-05-31',
745745
max_tokens: 1000,
746746
messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
747747
});
748748

749-
// Mock response body as already converted string
750-
const mockResponseBodyString = JSON.stringify({
751-
stop_reason: 'end_turn',
752-
usage: { input_tokens: 15, output_tokens: 13 },
753-
});
754-
755-
nock(`https://bedrock-runtime.${region}.amazonaws.com`)
756-
.post(`/model/${encodeURIComponent(modelId)}/invoke`)
757-
.reply(200, mockResponseBodyString);
758-
759-
await bedrock
760-
.invokeModel({
761-
modelId: modelId,
762-
body: mockRequestBody,
763-
})
764-
.catch((err: any) => {});
765-
766-
const testSpans: ReadableSpan[] = getTestSpans();
767-
const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => {
768-
return s.name === 'BedrockRuntime.InvokeModel';
769-
});
770-
expect(invokeModelSpans.length).toBe(1);
771-
const invokeModelSpan = invokeModelSpans[0];
772-
773-
// Verify attributes are set correctly despite body being a string
774-
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(15);
775-
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(13);
776-
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['end_turn']);
777-
});
778-
779-
it('handles Anthropic Claude response body correctly', async () => {
780-
const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
781-
const mockRequestBody: string = JSON.stringify({
782-
anthropic_version: 'bedrock-2023-05-31',
783-
max_tokens: 1000,
784-
messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
785-
});
786-
787-
// Mock response body - use standard object format (AWS SDK will handle type conversion)
749+
// Use standard object format - AWS SDK and instrumentation will handle the conversion
788750
const mockResponseBodyObj = {
789751
stop_reason: 'end_turn',
790752
usage: { input_tokens: 20, output_tokens: 15 },
@@ -814,31 +776,23 @@ describe('BedrockRuntime', () => {
814776
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['end_turn']);
815777
});
816778

817-
it('handles Buffer response body correctly', async () => {
779+
it('handles unexpected body type gracefully', async () => {
818780
const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
819781
const mockRequestBody: string = JSON.stringify({
820782
anthropic_version: 'bedrock-2023-05-31',
821783
max_tokens: 1000,
822784
messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
823785
});
824786

825-
// Mock response body as Buffer
826-
const mockResponseBodyObj = {
827-
stop_reason: 'max_tokens',
828-
usage: { input_tokens: 25, output_tokens: 18 },
829-
};
830-
const mockResponseBodyBuffer = Buffer.from(JSON.stringify(mockResponseBodyObj), 'utf8');
831-
787+
// Mock response body as unexpected type - using reply function to return a number
832788
nock(`https://bedrock-runtime.${region}.amazonaws.com`)
833789
.post(`/model/${encodeURIComponent(modelId)}/invoke`)
834-
.reply(200, mockResponseBodyBuffer);
790+
.reply(200, () => 12345 as any);
835791

836-
await bedrock
837-
.invokeModel({
838-
modelId: modelId,
839-
body: mockRequestBody,
840-
})
841-
.catch((err: any) => {});
792+
await bedrock.invokeModel({
793+
modelId: modelId,
794+
body: mockRequestBody,
795+
});
842796

843797
const testSpans: ReadableSpan[] = getTestSpans();
844798
const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => {
@@ -847,33 +801,40 @@ describe('BedrockRuntime', () => {
847801
expect(invokeModelSpans.length).toBe(1);
848802
const invokeModelSpan = invokeModelSpans[0];
849803

850-
// Verify attributes are set correctly when body is Buffer
851-
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(25);
852-
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(18);
853-
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual([
854-
'max_tokens',
855-
]);
804+
// Verify that no AI attributes are set when body type is unexpected
805+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBeUndefined();
806+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBeUndefined();
807+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toBeUndefined();
808+
809+
// Note: We can't easily test diag.debug() output in unit tests, but the important part
810+
// is that the function returns early and doesn't crash when encountering unexpected types
811+
// Debug message will be: "Unexpected body type in Bedrock response: number for commandName InvokeModelCommand"
856812
});
857813

858-
it('handles unexpected body type gracefully', async () => {
814+
it('handles streaming response (SmithyMessageDecoderStream) gracefully', async () => {
859815
const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
860816
const mockRequestBody: string = JSON.stringify({
861817
anthropic_version: 'bedrock-2023-05-31',
862818
max_tokens: 1000,
863819
messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
864820
});
865821

866-
// Mock response body as unexpected type - using reply function to return a number
822+
// Mock response body as streaming object (constructor name matching)
823+
const mockStreamingBody = {
824+
constructor: { name: 'SmithyMessageDecoderStream' },
825+
[Symbol.asyncIterator]: function* () {
826+
yield { chunk: { bytes: new TextEncoder().encode('{"type":"chunk"}') } };
827+
},
828+
};
829+
867830
nock(`https://bedrock-runtime.${region}.amazonaws.com`)
868831
.post(`/model/${encodeURIComponent(modelId)}/invoke`)
869-
.reply(200, () => 12345 as any);
832+
.reply(200, mockStreamingBody);
870833

871-
await bedrock
872-
.invokeModel({
873-
modelId: modelId,
874-
body: mockRequestBody,
875-
})
876-
.catch((err: any) => {});
834+
await bedrock.invokeModel({
835+
modelId: modelId,
836+
body: mockRequestBody,
837+
});
877838

878839
const testSpans: ReadableSpan[] = getTestSpans();
879840
const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => {
@@ -882,13 +843,13 @@ describe('BedrockRuntime', () => {
882843
expect(invokeModelSpans.length).toBe(1);
883844
const invokeModelSpan = invokeModelSpans[0];
884845

885-
// Verify that no AI attributes are set when body type is unexpected
846+
// Verify that no AI attributes are set when body is streaming (metrics not available in initial response)
886847
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBeUndefined();
887848
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBeUndefined();
888849
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toBeUndefined();
889850

890-
// Note: We can't easily test diag.debug() output in unit tests, but the important part
891-
// is that the function returns early and doesn't crash when encountering unexpected types
851+
// Streaming responses should be skipped gracefully without crashing
852+
// TODO: support InvokeModel Streaming API and Converse APIs later
892853
});
893854
});
894855
});

0 commit comments

Comments
 (0)