diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts index 82da47a..68fa3fe 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts @@ -232,6 +232,16 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.textGenerationConfig.maxTokenCount; } + } else if (modelId.includes('amazon.nova')) { + if (requestBody.inferenceConfig?.temperature !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.inferenceConfig.temperature; + } + if (requestBody.inferenceConfig?.top_p !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.inferenceConfig.top_p; + } + if (requestBody.inferenceConfig?.max_new_tokens !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.inferenceConfig.max_new_tokens; + } } else if (modelId.includes('anthropic.claude')) { if (requestBody.max_tokens !== undefined) { spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; @@ -335,6 +345,18 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { responseBody.results[0].completionReason, ]); } + } else if (currentModelId.includes('amazon.nova')) { + if (responseBody.usage !== undefined) { + if (responseBody.usage.inputTokens !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.inputTokens); + } + if (responseBody.usage.outputTokens !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.outputTokens); + } + } + if (responseBody.stopReason !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stopReason]); + } } else if (currentModelId.includes('anthropic.claude')) { if (responseBody.usage?.input_tokens !== undefined) { span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.input_tokens); diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts index f642e09..1759774 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts @@ -400,6 +400,60 @@ describe('BedrockRuntime', () => { expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT); }); + it('Add Amazon Nova model attributes to span', async () => { + const modelId: string = 'amazon.nova-pro-v1:0'; + const prompt: string = 'Campfire story'; + const mockRequestBody: string = JSON.stringify({ + inputText: prompt, + inferenceConfig: { + max_new_tokens: 500, + temperature: 0.9, + top_p: 0.7, + }, + }); + const mockResponseBody: any = { + output: { message: { content: [{ text: '' }], role: 'assistant' } }, + stopReason: 'max_tokens', + usage: { inputTokens: 432, outputTokens: 681 }, + + request: { + commandInput: { + modelId: modelId, + }, + }, + }; + + nock(`https://bedrock-runtime.${region}.amazonaws.com`) + .post(`/model/${encodeURIComponent(modelId)}/invoke`) + .reply(200, mockResponseBody); + + 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]; + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws.bedrock'); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(500); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0.9); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(0.7); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(432); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(681); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['max_tokens']); + expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT); + }); + it('Add Anthropic Claude model attributes to span', async () => { const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0'; const prompt: string = 'Complete this text. It was the best of times it was the worst...'; diff --git a/contract-tests/images/applications/aws-sdk/server.js b/contract-tests/images/applications/aws-sdk/server.js index 7352bfa..288bb96 100644 --- a/contract-tests/images/applications/aws-sdk/server.js +++ b/contract-tests/images/applications/aws-sdk/server.js @@ -631,7 +631,29 @@ async function handleBedrockRequest(req, res, path) { }, ], } - + } + + if (path.includes("amazon.nova")) { + + modelId = "amazon.nova-pro-v1:0" + + request_body = { + messages: [{role: "user", content: [{text: "A camping trip"}]}], + inferenceConfig: { + max_new_tokens: 800, + temperature: 0.9, + top_p: 0.7, + }, + } + + response_body = { + output: {message: {content: [{text: ""}], role: "assistant"}}, + stopReason: "max_tokens", + usage: { + inputTokens: 432, + outputTokens: 681 + }, + } } if (path.includes('anthropic.claude')) { diff --git a/contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py b/contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py index c7cac2a..ba20402 100644 --- a/contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py +++ b/contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py @@ -466,6 +466,34 @@ def test_bedrock_runtime_invoke_model_amazon_titan(self): span_name="BedrockRuntime.InvokeModel" ) + + def test_bedrock_runtime_invoke_model_amazon_nova(self): + result = self.do_test_requests( + "bedrock/invokemodel/invoke-model/amazon.nova-pro-v1:0", + "GET", + 200, + 0, + 0, + local_operation="GET /bedrock", + rpc_service="BedrockRuntime", + remote_service="AWS::BedrockRuntime", + remote_operation="InvokeModel", + remote_resource_type="AWS::Bedrock::Model", + remote_resource_identifier='amazon.nova-pro-v1:0', + request_specific_attributes={ + _GEN_AI_REQUEST_MODEL: 'amazon.nova-pro-v1:0', + _GEN_AI_REQUEST_MAX_TOKENS: 800, + _GEN_AI_REQUEST_TEMPERATURE: 0.9, + _GEN_AI_REQUEST_TOP_P: 0.7 + }, + response_specific_attributes={ + _GEN_AI_RESPONSE_FINISH_REASONS: ['max_tokens'], + _GEN_AI_USAGE_INPUT_TOKENS: 432, + _GEN_AI_USAGE_OUTPUT_TOKENS: 681 + }, + + span_name="BedrockRuntime.InvokeModel" + ) def test_bedrock_runtime_invoke_model_anthropic_claude(self): self.do_test_requests( @@ -1105,6 +1133,25 @@ def _assert_semantic_conventions_attributes( for key, value in response_specific_attributes.items(): self._assert_attribute(attributes_dict, key, value) + + def _assert_attribute(self, attributes_dict: Dict[str, AnyValue], key, value) -> None: + if isinstance(value, str): + self._assert_str_attribute(attributes_dict, key, value) + elif isinstance(value, int): + self._assert_int_attribute(attributes_dict, key, value) + elif isinstance(value, float): + self._assert_float_attribute(attributes_dict, key, value) + else: + self._assert_array_value_ddb_table_name(attributes_dict, key, value) + + @override + def _assert_str_attribute(self, attributes_dict: Dict[str, AnyValue], key: str, expected_value: str): + self.assertIn(key, attributes_dict) + actual_value: AnyValue = attributes_dict[key] + self.assertIsNotNone(actual_value) + pattern = re.compile(expected_value) + match = pattern.fullmatch(actual_value.string_value) + self.assertTrue(match is not None, f"Actual: {actual_value.string_value} does not match Expected: {expected_value}") @override def _assert_metric_attributes(