Skip to content

Commit 02770d9

Browse files
liustveliustve
committed
gen_ai attributes support for amazon nova
1 parent 27468e5 commit 02770d9

File tree

4 files changed

+151
-13
lines changed

4 files changed

+151
-13
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
@@ -225,6 +225,16 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
225225
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] =
226226
requestBody.textGenerationConfig.maxTokenCount;
227227
}
228+
} else if (modelId.includes('amazon.nova')) {
229+
if (requestBody.inferenceConfig?.temperature !== undefined) {
230+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.inferenceConfig.temperature;
231+
}
232+
if (requestBody.inferenceConfig?.top_p !== undefined) {
233+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.inferenceConfig.top_p;
234+
}
235+
if (requestBody.inferenceConfig?.max_new_tokens !== undefined) {
236+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.inferenceConfig.max_new_tokens;
237+
}
228238
} else if (modelId.includes('anthropic.claude')) {
229239
if (requestBody.max_tokens !== undefined) {
230240
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
@@ -328,6 +338,18 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
328338
responseBody.results[0].completionReason,
329339
]);
330340
}
341+
} else if (currentModelId.includes('amazon.nova')) {
342+
if (responseBody.usage !== undefined) {
343+
if (responseBody.usage.inputTokens !== undefined) {
344+
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.inputTokens);
345+
}
346+
if (responseBody.usage.outputTokens !== undefined) {
347+
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.outputTokens);
348+
}
349+
}
350+
if (responseBody.stopReason !== undefined) {
351+
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stopReason]);
352+
}
331353
} else if (currentModelId.includes('anthropic.claude')) {
332354
if (responseBody.usage?.input_tokens !== undefined) {
333355
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
@@ -397,6 +397,60 @@ describe('BedrockRuntime', () => {
397397
expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT);
398398
});
399399

400+
it('Add Amazon Nova model attributes to span', async () => {
401+
const modelId: string = 'amazon.nova-pro-v1:0';
402+
const prompt: string = 'Campfire story';
403+
const mockRequestBody: string = JSON.stringify({
404+
inputText: prompt,
405+
inferenceConfig: {
406+
max_new_tokens: 500,
407+
temperature: 0.9,
408+
top_p: 0.7,
409+
},
410+
});
411+
const mockResponseBody: any = {
412+
output: { message: { content: [{ text: '' }], role: 'assistant' } },
413+
stopReason: 'max_tokens',
414+
usage: { inputTokens: 432, outputTokens: 681 },
415+
416+
request: {
417+
commandInput: {
418+
modelId: modelId,
419+
},
420+
},
421+
};
422+
423+
nock(`https://bedrock-runtime.${region}.amazonaws.com`)
424+
.post(`/model/${encodeURIComponent(modelId)}/invoke`)
425+
.reply(200, mockResponseBody);
426+
427+
await bedrock
428+
.invokeModel({
429+
modelId: modelId,
430+
body: mockRequestBody,
431+
})
432+
.catch((err: any) => {});
433+
434+
const testSpans: ReadableSpan[] = getTestSpans();
435+
const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => {
436+
return s.name === 'BedrockRuntime.InvokeModel';
437+
});
438+
expect(invokeModelSpans.length).toBe(1);
439+
const invokeModelSpan = invokeModelSpans[0];
440+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined();
441+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined();
442+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined();
443+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws_bedrock');
444+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId);
445+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(500);
446+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0.9);
447+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(0.7);
448+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(432);
449+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(681);
450+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['max_tokens']);
451+
expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT);
452+
});
453+
400454
it('Add Anthropic Claude model attributes to span', async () => {
401455
const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
402456
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
@@ -582,7 +582,29 @@ async function handleBedrockRequest(req, res, path) {
582582
},
583583
],
584584
}
585-
585+
}
586+
587+
if (path.includes("amazon.nova")) {
588+
589+
modelId = "amazon.nova-pro-v1:0"
590+
591+
request_body = {
592+
messages: [{role: "user", content: [{text: "A camping trip"}]}],
593+
inferenceConfig: {
594+
max_new_tokens: 800,
595+
temperature: 0.9,
596+
top_p: 0.7,
597+
},
598+
}
599+
600+
response_body = {
601+
output: {message: {content: [{text: ""}], role: "assistant"}},
602+
stopReason: "max_tokens",
603+
usage: {
604+
inputTokens: 432,
605+
outputTokens: 681
606+
},
607+
}
586608
}
587609

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

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

Lines changed: 52 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# SPDX-License-Identifier: Apache-2.0
33
from logging import INFO, Logger, getLogger
44
import math
5+
import re
56
from typing import Dict, List
67

78
from docker.types import EndpointConfig
@@ -439,6 +440,34 @@ def test_bedrock_runtime_invoke_model_amazon_titan(self):
439440

440441
span_name="BedrockRuntime.InvokeModel"
441442
)
443+
444+
def test_bedrock_runtime_invoke_model_amazon_nova(self):
445+
result = self.do_test_requests(
446+
"bedrock/invokemodel/invoke-model/amazon.nova-pro-v1:0",
447+
"GET",
448+
200,
449+
0,
450+
0,
451+
local_operation="GET /bedrock",
452+
rpc_service="BedrockRuntime",
453+
remote_service="AWS::BedrockRuntime",
454+
remote_operation="InvokeModel",
455+
remote_resource_type="AWS::Bedrock::Model",
456+
remote_resource_identifier='amazon.nova-pro-v1:0',
457+
request_specific_attributes={
458+
_GEN_AI_REQUEST_MODEL: 'amazon.nova-pro-v1:0',
459+
_GEN_AI_REQUEST_MAX_TOKENS: 800,
460+
_GEN_AI_REQUEST_TEMPERATURE: 0.9,
461+
_GEN_AI_REQUEST_TOP_P: 0.7
462+
},
463+
response_specific_attributes={
464+
_GEN_AI_RESPONSE_FINISH_REASONS: ['max_tokens'],
465+
_GEN_AI_USAGE_INPUT_TOKENS: 432,
466+
_GEN_AI_USAGE_OUTPUT_TOKENS: 681
467+
},
468+
469+
span_name="BedrockRuntime.InvokeModel"
470+
)
442471

443472
def test_bedrock_runtime_invoke_model_anthropic_claude(self):
444473
self.do_test_requests(
@@ -805,19 +834,30 @@ def _assert_semantic_conventions_attributes(
805834
self._assert_int_attribute(attributes_dict, SpanAttributes.HTTP_STATUS_CODE, status_code)
806835
# TODO: aws sdk instrumentation is not respecting PEER_SERVICE
807836
# self._assert_str_attribute(attributes_dict, SpanAttributes.PEER_SERVICE, "backend:8080")
808-
self._assert_specific_attributes(attributes_dict, request_specific_attributes)
809-
self._assert_specific_attributes(attributes_dict, response_specific_attributes)
837+
for key, value in request_specific_attributes.items():
838+
self._assert_attribute(attributes_dict, key, value)
839+
840+
for key, value in response_specific_attributes.items():
841+
self._assert_attribute(attributes_dict, key, value)
810842

811-
def _assert_specific_attributes(self, attributes_dict: Dict[str, AnyValue], specific_attributes: Dict[str, AnyValue]) -> None:
812-
for key, value in specific_attributes.items():
813-
if isinstance(value, str):
814-
self._assert_str_attribute(attributes_dict, key, value)
815-
elif isinstance(value, int):
816-
self._assert_int_attribute(attributes_dict, key, value)
817-
elif isinstance(value, float):
818-
self._assert_float_attribute(attributes_dict, key, value)
819-
else:
820-
self._assert_array_value_ddb_table_name(attributes_dict, key, value)
843+
def _assert_attribute(self, attributes_dict: Dict[str, AnyValue], key, value) -> None:
844+
if isinstance(value, str):
845+
self._assert_str_attribute(attributes_dict, key, value)
846+
elif isinstance(value, int):
847+
self._assert_int_attribute(attributes_dict, key, value)
848+
elif isinstance(value, float):
849+
self._assert_float_attribute(attributes_dict, key, value)
850+
else:
851+
self._assert_array_value_ddb_table_name(attributes_dict, key, value)
852+
853+
@override
854+
def _assert_str_attribute(self, attributes_dict: Dict[str, AnyValue], key: str, expected_value: str):
855+
self.assertIn(key, attributes_dict)
856+
actual_value: AnyValue = attributes_dict[key]
857+
self.assertIsNotNone(actual_value)
858+
pattern = re.compile(expected_value)
859+
match = pattern.fullmatch(actual_value.string_value)
860+
self.assertTrue(match is not None, f"Actual: {actual_value.string_value} does not match Expected: {expected_value}")
821861

822862
@override
823863
def _assert_metric_attributes(

0 commit comments

Comments
 (0)