Skip to content

Commit 5dd389e

Browse files
liustveliustve
authored
feat: Support Gen AI attributes for Amazon Nova foundational model (#133)
*Description of changes:* Added GenAI inference parameters auto instrumentation support for Amazon Nova. <img width="824" alt="image" src="https://github.com/user-attachments/assets/78c172dc-7d3f-48bf-9795-e3369d8849fd" /> Contract test: <img width="855" alt="image" src="https://github.com/user-attachments/assets/8f4db5fc-1ab6-44e6-9d1a-20c2320f48e3" /> Unit test: <img width="750" alt="image" src="https://github.com/user-attachments/assets/f610fb86-85db-4727-971a-e8f929d7cc9a" /> By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.
1 parent a35f89c commit 5dd389e

File tree

4 files changed

+146
-1
lines changed

4 files changed

+146
-1
lines changed

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,16 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
232232
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] =
233233
requestBody.textGenerationConfig.maxTokenCount;
234234
}
235+
} else if (modelId.includes('amazon.nova')) {
236+
if (requestBody.inferenceConfig?.temperature !== undefined) {
237+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.inferenceConfig.temperature;
238+
}
239+
if (requestBody.inferenceConfig?.top_p !== undefined) {
240+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.inferenceConfig.top_p;
241+
}
242+
if (requestBody.inferenceConfig?.max_new_tokens !== undefined) {
243+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.inferenceConfig.max_new_tokens;
244+
}
235245
} else if (modelId.includes('anthropic.claude')) {
236246
if (requestBody.max_tokens !== undefined) {
237247
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
@@ -335,6 +345,18 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
335345
responseBody.results[0].completionReason,
336346
]);
337347
}
348+
} else if (currentModelId.includes('amazon.nova')) {
349+
if (responseBody.usage !== undefined) {
350+
if (responseBody.usage.inputTokens !== undefined) {
351+
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.inputTokens);
352+
}
353+
if (responseBody.usage.outputTokens !== undefined) {
354+
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.outputTokens);
355+
}
356+
}
357+
if (responseBody.stopReason !== undefined) {
358+
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stopReason]);
359+
}
338360
} else if (currentModelId.includes('anthropic.claude')) {
339361
if (responseBody.usage?.input_tokens !== undefined) {
340362
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.input_tokens);

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -400,6 +400,60 @@ describe('BedrockRuntime', () => {
400400
expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT);
401401
});
402402

403+
it('Add Amazon Nova model attributes to span', async () => {
404+
const modelId: string = 'amazon.nova-pro-v1:0';
405+
const prompt: string = 'Campfire story';
406+
const mockRequestBody: string = JSON.stringify({
407+
inputText: prompt,
408+
inferenceConfig: {
409+
max_new_tokens: 500,
410+
temperature: 0.9,
411+
top_p: 0.7,
412+
},
413+
});
414+
const mockResponseBody: any = {
415+
output: { message: { content: [{ text: '' }], role: 'assistant' } },
416+
stopReason: 'max_tokens',
417+
usage: { inputTokens: 432, outputTokens: 681 },
418+
419+
request: {
420+
commandInput: {
421+
modelId: modelId,
422+
},
423+
},
424+
};
425+
426+
nock(`https://bedrock-runtime.${region}.amazonaws.com`)
427+
.post(`/model/${encodeURIComponent(modelId)}/invoke`)
428+
.reply(200, mockResponseBody);
429+
430+
await bedrock
431+
.invokeModel({
432+
modelId: modelId,
433+
body: mockRequestBody,
434+
})
435+
.catch((err: any) => {});
436+
437+
const testSpans: ReadableSpan[] = getTestSpans();
438+
const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => {
439+
return s.name === 'BedrockRuntime.InvokeModel';
440+
});
441+
expect(invokeModelSpans.length).toBe(1);
442+
const invokeModelSpan = invokeModelSpans[0];
443+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined();
444+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined();
445+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined();
446+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws.bedrock');
447+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId);
448+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(500);
449+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0.9);
450+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(0.7);
451+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(432);
452+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(681);
453+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['max_tokens']);
454+
expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT);
455+
});
456+
403457
it('Add Anthropic Claude model attributes to span', async () => {
404458
const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
405459
const prompt: string = 'Complete this text. It was the best of times it was the worst...';

‎contract-tests/images/applications/aws-sdk/server.js

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -631,7 +631,29 @@ async function handleBedrockRequest(req, res, path) {
631631
},
632632
],
633633
}
634-
634+
}
635+
636+
if (path.includes("amazon.nova")) {
637+
638+
modelId = "amazon.nova-pro-v1:0"
639+
640+
request_body = {
641+
messages: [{role: "user", content: [{text: "A camping trip"}]}],
642+
inferenceConfig: {
643+
max_new_tokens: 800,
644+
temperature: 0.9,
645+
top_p: 0.7,
646+
},
647+
}
648+
649+
response_body = {
650+
output: {message: {content: [{text: ""}], role: "assistant"}},
651+
stopReason: "max_tokens",
652+
usage: {
653+
inputTokens: 432,
654+
outputTokens: 681
655+
},
656+
}
635657
}
636658

637659
if (path.includes('anthropic.claude')) {

‎contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -466,6 +466,34 @@ def test_bedrock_runtime_invoke_model_amazon_titan(self):
466466

467467
span_name="BedrockRuntime.InvokeModel"
468468
)
469+
470+
def test_bedrock_runtime_invoke_model_amazon_nova(self):
471+
result = self.do_test_requests(
472+
"bedrock/invokemodel/invoke-model/amazon.nova-pro-v1:0",
473+
"GET",
474+
200,
475+
0,
476+
0,
477+
local_operation="GET /bedrock",
478+
rpc_service="BedrockRuntime",
479+
remote_service="AWS::BedrockRuntime",
480+
remote_operation="InvokeModel",
481+
remote_resource_type="AWS::Bedrock::Model",
482+
remote_resource_identifier='amazon.nova-pro-v1:0',
483+
request_specific_attributes={
484+
_GEN_AI_REQUEST_MODEL: 'amazon.nova-pro-v1:0',
485+
_GEN_AI_REQUEST_MAX_TOKENS: 800,
486+
_GEN_AI_REQUEST_TEMPERATURE: 0.9,
487+
_GEN_AI_REQUEST_TOP_P: 0.7
488+
},
489+
response_specific_attributes={
490+
_GEN_AI_RESPONSE_FINISH_REASONS: ['max_tokens'],
491+
_GEN_AI_USAGE_INPUT_TOKENS: 432,
492+
_GEN_AI_USAGE_OUTPUT_TOKENS: 681
493+
},
494+
495+
span_name="BedrockRuntime.InvokeModel"
496+
)
469497

470498
def test_bedrock_runtime_invoke_model_anthropic_claude(self):
471499
self.do_test_requests(
@@ -1105,6 +1133,25 @@ def _assert_semantic_conventions_attributes(
11051133

11061134
for key, value in response_specific_attributes.items():
11071135
self._assert_attribute(attributes_dict, key, value)
1136+
1137+
def _assert_attribute(self, attributes_dict: Dict[str, AnyValue], key, value) -> None:
1138+
if isinstance(value, str):
1139+
self._assert_str_attribute(attributes_dict, key, value)
1140+
elif isinstance(value, int):
1141+
self._assert_int_attribute(attributes_dict, key, value)
1142+
elif isinstance(value, float):
1143+
self._assert_float_attribute(attributes_dict, key, value)
1144+
else:
1145+
self._assert_array_value_ddb_table_name(attributes_dict, key, value)
1146+
1147+
@override
1148+
def _assert_str_attribute(self, attributes_dict: Dict[str, AnyValue], key: str, expected_value: str):
1149+
self.assertIn(key, attributes_dict)
1150+
actual_value: AnyValue = attributes_dict[key]
1151+
self.assertIsNotNone(actual_value)
1152+
pattern = re.compile(expected_value)
1153+
match = pattern.fullmatch(actual_value.string_value)
1154+
self.assertTrue(match is not None, f"Actual: {actual_value.string_value} does not match Expected: {expected_value}")
11081155

11091156
@override
11101157
def _assert_metric_attributes(

0 commit comments

Comments
 (0)