Skip to content

Commit 1439ae4

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

File tree

2 files changed

+43
-70
lines changed

2 files changed

+43
-70
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 simithy 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: 26 additions & 61 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,24 +776,18 @@ 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

836792
await bedrock
837793
.invokeModel({
@@ -847,26 +803,35 @@ describe('BedrockRuntime', () => {
847803
expect(invokeModelSpans.length).toBe(1);
848804
const invokeModelSpan = invokeModelSpans[0];
849805

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-
]);
806+
// Verify that no AI attributes are set when body type is unexpected
807+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBeUndefined();
808+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBeUndefined();
809+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toBeUndefined();
810+
811+
// Note: We can't easily test diag.debug() output in unit tests, but the important part
812+
// is that the function returns early and doesn't crash when encountering unexpected types
813+
// Debug message will be: "Unexpected body type in Bedrock response: number for commandName InvokeModelCommand"
856814
});
857815

858-
it('handles unexpected body type gracefully', async () => {
816+
it('handles streaming response (SmithyMessageDecoderStream) gracefully', async () => {
859817
const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
860818
const mockRequestBody: string = JSON.stringify({
861819
anthropic_version: 'bedrock-2023-05-31',
862820
max_tokens: 1000,
863821
messages: [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
864822
});
865823

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

871836
await bedrock
872837
.invokeModel({
@@ -882,13 +847,13 @@ describe('BedrockRuntime', () => {
882847
expect(invokeModelSpans.length).toBe(1);
883848
const invokeModelSpan = invokeModelSpans[0];
884849

885-
// Verify that no AI attributes are set when body type is unexpected
850+
// Verify that no AI attributes are set when body is streaming (metrics not available in initial response)
886851
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBeUndefined();
887852
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBeUndefined();
888853
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toBeUndefined();
889854

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
855+
// Streaming responses should be skipped gracefully without crashing
856+
// TODO: support InvokeModel Streaming API and Converse APIs later
892857
});
893858
});
894859
});

0 commit comments

Comments
 (0)