@@ -268,6 +268,8 @@ def extract_attributes(self, attributes: _AttributeMapT):
268
268
269
269
if "amazon.titan" in model_id :
270
270
self ._extract_titan_attributes (attributes , request_body )
271
+ if "amazon.nova" in model_id :
272
+ self ._extract_nova_attributes (attributes , request_body )
271
273
elif "anthropic.claude" in model_id :
272
274
self ._extract_claude_attributes (attributes , request_body )
273
275
elif "meta.llama" in model_id :
@@ -288,6 +290,12 @@ def _extract_titan_attributes(self, attributes, request_body):
288
290
self ._set_if_not_none (attributes , GEN_AI_REQUEST_TOP_P , config .get ("topP" ))
289
291
self ._set_if_not_none (attributes , GEN_AI_REQUEST_MAX_TOKENS , config .get ("maxTokenCount" ))
290
292
293
+ def _extract_nova_attributes (self , attributes , request_body ):
294
+ config = request_body .get ("inferenceConfig" , {})
295
+ self ._set_if_not_none (attributes , GEN_AI_REQUEST_TEMPERATURE , config .get ("temperature" ))
296
+ self ._set_if_not_none (attributes , GEN_AI_REQUEST_TOP_P , config .get ("top_p" ))
297
+ self ._set_if_not_none (attributes , GEN_AI_REQUEST_MAX_TOKENS , config .get ("max_new_tokens" ))
298
+
291
299
def _extract_claude_attributes (self , attributes , request_body ):
292
300
self ._set_if_not_none (attributes , GEN_AI_REQUEST_MAX_TOKENS , request_body .get ("max_tokens" ))
293
301
self ._set_if_not_none (attributes , GEN_AI_REQUEST_TEMPERATURE , request_body .get ("temperature" ))
@@ -324,6 +332,7 @@ def _set_if_not_none(attributes, key, value):
324
332
if value is not None :
325
333
attributes [key ] = value
326
334
335
+ # pylint: disable=too-many-branches
327
336
def on_success (self , span : Span , result : Dict [str , Any ]):
328
337
model_id = self ._call_context .params .get (_MODEL_ID )
329
338
@@ -342,6 +351,8 @@ def on_success(self, span: Span, result: Dict[str, Any]):
342
351
response_body = json .loads (telemetry_content .decode ("utf-8" ))
343
352
if "amazon.titan" in model_id :
344
353
self ._handle_amazon_titan_response (span , response_body )
354
+ if "amazon.nova" in model_id :
355
+ self ._handle_amazon_nova_response (span , response_body )
345
356
elif "anthropic.claude" in model_id :
346
357
self ._handle_anthropic_claude_response (span , response_body )
347
358
elif "meta.llama" in model_id :
@@ -375,6 +386,17 @@ def _handle_amazon_titan_response(self, span: Span, response_body: Dict[str, Any
375
386
if "completionReason" in result :
376
387
span .set_attribute (GEN_AI_RESPONSE_FINISH_REASONS , [result ["completionReason" ]])
377
388
389
+ # pylint: disable=no-self-use
390
+ def _handle_amazon_nova_response (self , span : Span , response_body : Dict [str , Any ]):
391
+ if "usage" in response_body :
392
+ usage = response_body ["usage" ]
393
+ if "inputTokens" in usage :
394
+ span .set_attribute (GEN_AI_USAGE_INPUT_TOKENS , usage ["inputTokens" ])
395
+ if "outputTokens" in usage :
396
+ span .set_attribute (GEN_AI_USAGE_OUTPUT_TOKENS , usage ["outputTokens" ])
397
+ if "stopReason" in response_body :
398
+ span .set_attribute (GEN_AI_RESPONSE_FINISH_REASONS , [response_body ["stopReason" ]])
399
+
378
400
# pylint: disable=no-self-use
379
401
def _handle_anthropic_claude_response (self , span : Span , response_body : Dict [str , Any ]):
380
402
if "usage" in response_body :
0 commit comments