support for amazon nova model

Author: liustve

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py

@@ -295,7 +295,7 @@ def _extract_nova_attributes(self, attributes, request_body):
         self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, config.get("temperature"))
         self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, config.get("top_p"))
         self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get("max_new_tokens"))
-        
+
     def _extract_claude_attributes(self, attributes, request_body):
         self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens"))
         self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature"))
@@ -384,12 +384,15 @@ def _handle_amazon_titan_response(self, span: Span, response_body: Dict[str, Any
                     span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, result["tokenCount"])
                 if "completionReason" in result:
                     span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [result["completionReason"]])
-    
+
+    # pylint: disable=no-self-use
     def _handle_amazon_nova_response(self, span: Span, response_body: Dict[str, Any]):
-        if "inputTokenCount" in response_body:
-            span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body["inputTokenCount"])
-        if "outputTokenCount" in response_body:
-            span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, response_body["outputTokenCount"])
+        if "usage" in response_body:
+            usage = response_body["usage"]
+            if "inputTokens" in usage:
+                span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, usage["inputTokens"])
+            if "outputTokens" in usage:
+                span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"])
         if "stopReason" in response_body:
             span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stopReason"]])