@@ -211,11 +211,196 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
211211 spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_MODEL ] = modelId ;
212212 }
213213
214+ if ( request . commandInput ?. body ) {
215+ const requestBody = JSON . parse ( request . commandInput . body ) ;
216+ if ( modelId . includes ( 'amazon.titan' ) ) {
217+ if ( requestBody . textGenerationConfig ?. temperature !== undefined ) {
218+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TEMPERATURE ] =
219+ requestBody . textGenerationConfig . temperature ;
220+ }
221+ if ( requestBody . textGenerationConfig ?. topP !== undefined ) {
222+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TOP_P ] = requestBody . textGenerationConfig . topP ;
223+ }
224+ if ( requestBody . textGenerationConfig ?. maxTokenCount !== undefined ) {
225+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_MAX_TOKENS ] =
226+ requestBody . textGenerationConfig . maxTokenCount ;
227+ }
228+ } else if ( modelId . includes ( 'anthropic.claude' ) ) {
229+ if ( requestBody . max_tokens !== undefined ) {
230+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_MAX_TOKENS ] = requestBody . max_tokens ;
231+ }
232+ if ( requestBody . temperature !== undefined ) {
233+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TEMPERATURE ] = requestBody . temperature ;
234+ }
235+ if ( requestBody . top_p !== undefined ) {
236+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TOP_P ] = requestBody . top_p ;
237+ }
238+ } else if ( modelId . includes ( 'meta.llama' ) ) {
239+ if ( requestBody . max_gen_len !== undefined ) {
240+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_MAX_TOKENS ] = requestBody . max_gen_len ;
241+ }
242+ if ( requestBody . temperature !== undefined ) {
243+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TEMPERATURE ] = requestBody . temperature ;
244+ }
245+ if ( requestBody . top_p !== undefined ) {
246+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TOP_P ] = requestBody . top_p ;
247+ }
248+ } else if ( modelId . includes ( 'cohere.command-r' ) ) {
249+ if ( requestBody . max_tokens !== undefined ) {
250+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_MAX_TOKENS ] = requestBody . max_tokens ;
251+ }
252+ if ( requestBody . temperature !== undefined ) {
253+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TEMPERATURE ] = requestBody . temperature ;
254+ }
255+ if ( requestBody . p !== undefined ) {
256+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TOP_P ] = requestBody . p ;
257+ }
258+ if ( requestBody . message !== undefined ) {
259+ // NOTE: We approximate the token count since this value is not directly available in the body
260+ // According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
261+ // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
262+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_USAGE_INPUT_TOKENS ] = Math . ceil ( requestBody . message . length / 6 ) ;
263+ }
264+ } else if ( modelId . includes ( 'cohere.command' ) ) {
265+ if ( requestBody . max_tokens !== undefined ) {
266+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_MAX_TOKENS ] = requestBody . max_tokens ;
267+ }
268+ if ( requestBody . temperature !== undefined ) {
269+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TEMPERATURE ] = requestBody . temperature ;
270+ }
271+ if ( requestBody . p !== undefined ) {
272+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TOP_P ] = requestBody . p ;
273+ }
274+ if ( requestBody . prompt !== undefined ) {
275+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_USAGE_INPUT_TOKENS ] = Math . ceil ( requestBody . prompt . length / 6 ) ;
276+ }
277+ } else if ( modelId . includes ( 'ai21.jamba' ) ) {
278+ if ( requestBody . max_tokens !== undefined ) {
279+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_MAX_TOKENS ] = requestBody . max_tokens ;
280+ }
281+ if ( requestBody . temperature !== undefined ) {
282+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TEMPERATURE ] = requestBody . temperature ;
283+ }
284+ if ( requestBody . top_p !== undefined ) {
285+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TOP_P ] = requestBody . top_p ;
286+ }
287+ } else if ( modelId . includes ( 'mistral' ) ) {
288+ if ( requestBody . prompt !== undefined ) {
289+ // NOTE: We approximate the token count since this value is not directly available in the body
290+ // According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
291+ // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
292+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_USAGE_INPUT_TOKENS ] = Math . ceil ( requestBody . prompt . length / 6 ) ;
293+ }
294+ if ( requestBody . max_tokens !== undefined ) {
295+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_MAX_TOKENS ] = requestBody . max_tokens ;
296+ }
297+ if ( requestBody . temperature !== undefined ) {
298+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TEMPERATURE ] = requestBody . temperature ;
299+ }
300+ if ( requestBody . top_p !== undefined ) {
301+ spanAttributes [ AwsSpanProcessingUtil . GEN_AI_REQUEST_TOP_P ] = requestBody . top_p ;
302+ }
303+ }
304+ }
305+
214306 return {
215307 isIncoming,
216308 spanAttributes,
217309 spanKind,
218310 spanName,
219311 } ;
220312 }
313+
314+ responseHook ( response : NormalizedResponse , span : Span , tracer : Tracer , config : AwsSdkInstrumentationConfig ) : void {
315+ const currentModelId = response . request . commandInput ?. modelId ;
316+ if ( response . data ?. body ) {
317+ const decodedResponseBody = new TextDecoder ( ) . decode ( response . data . body ) ;
318+ const responseBody = JSON . parse ( decodedResponseBody ) ;
319+ if ( currentModelId . includes ( 'amazon.titan' ) ) {
320+ if ( responseBody . inputTextTokenCount !== undefined ) {
321+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_INPUT_TOKENS , responseBody . inputTextTokenCount ) ;
322+ }
323+ if ( responseBody . results ?. [ 0 ] ?. tokenCount !== undefined ) {
324+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_OUTPUT_TOKENS , responseBody . results [ 0 ] . tokenCount ) ;
325+ }
326+ if ( responseBody . results ?. [ 0 ] ?. completionReason !== undefined ) {
327+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_RESPONSE_FINISH_REASONS , [
328+ responseBody . results [ 0 ] . completionReason ,
329+ ] ) ;
330+ }
331+ } else if ( currentModelId . includes ( 'anthropic.claude' ) ) {
332+ if ( responseBody . usage ?. input_tokens !== undefined ) {
333+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_INPUT_TOKENS , responseBody . usage . input_tokens ) ;
334+ }
335+ if ( responseBody . usage ?. output_tokens !== undefined ) {
336+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_OUTPUT_TOKENS , responseBody . usage . output_tokens ) ;
337+ }
338+ if ( responseBody . stop_reason !== undefined ) {
339+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_RESPONSE_FINISH_REASONS , [ responseBody . stop_reason ] ) ;
340+ }
341+ } else if ( currentModelId . includes ( 'meta.llama' ) ) {
342+ if ( responseBody . prompt_token_count !== undefined ) {
343+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_INPUT_TOKENS , responseBody . prompt_token_count ) ;
344+ }
345+ if ( responseBody . generation_token_count !== undefined ) {
346+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_OUTPUT_TOKENS , responseBody . generation_token_count ) ;
347+ }
348+ if ( responseBody . stop_reason !== undefined ) {
349+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_RESPONSE_FINISH_REASONS , [ responseBody . stop_reason ] ) ;
350+ }
351+ } else if ( currentModelId . includes ( 'cohere.command-r' ) ) {
352+ if ( responseBody . text !== undefined ) {
353+ // NOTE: We approximate the token count since this value is not directly available in the body
354+ // According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
355+ // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
356+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_OUTPUT_TOKENS , Math . ceil ( responseBody . text . length / 6 ) ) ;
357+ }
358+ if ( responseBody . finish_reason !== undefined ) {
359+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_RESPONSE_FINISH_REASONS , [ responseBody . finish_reason ] ) ;
360+ }
361+ } else if ( currentModelId . includes ( 'cohere.command' ) ) {
362+ if ( responseBody . generations ?. [ 0 ] ?. text !== undefined ) {
363+ span . setAttribute (
364+ AwsSpanProcessingUtil . GEN_AI_USAGE_OUTPUT_TOKENS ,
365+ // NOTE: We approximate the token count since this value is not directly available in the body
366+ // According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
367+ // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
368+ Math . ceil ( responseBody . generations [ 0 ] . text . length / 6 )
369+ ) ;
370+ }
371+ if ( responseBody . generations ?. [ 0 ] ?. finish_reason !== undefined ) {
372+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_RESPONSE_FINISH_REASONS , [
373+ responseBody . generations [ 0 ] . finish_reason ,
374+ ] ) ;
375+ }
376+ } else if ( currentModelId . includes ( 'ai21.jamba' ) ) {
377+ if ( responseBody . usage ?. prompt_tokens !== undefined ) {
378+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_INPUT_TOKENS , responseBody . usage . prompt_tokens ) ;
379+ }
380+ if ( responseBody . usage ?. completion_tokens !== undefined ) {
381+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_USAGE_OUTPUT_TOKENS , responseBody . usage . completion_tokens ) ;
382+ }
383+ if ( responseBody . choices ?. [ 0 ] ?. finish_reason !== undefined ) {
384+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_RESPONSE_FINISH_REASONS , [
385+ responseBody . choices [ 0 ] . finish_reason ,
386+ ] ) ;
387+ }
388+ } else if ( currentModelId . includes ( 'mistral' ) ) {
389+ if ( responseBody . outputs ?. [ 0 ] ?. text !== undefined ) {
390+ span . setAttribute (
391+ AwsSpanProcessingUtil . GEN_AI_USAGE_OUTPUT_TOKENS ,
392+ // NOTE: We approximate the token count since this value is not directly available in the body
393+ // According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
394+ // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
395+ Math . ceil ( responseBody . outputs [ 0 ] . text . length / 6 )
396+ ) ;
397+ }
398+ if ( responseBody . outputs ?. [ 0 ] ?. stop_reason !== undefined ) {
399+ span . setAttribute ( AwsSpanProcessingUtil . GEN_AI_RESPONSE_FINISH_REASONS , [
400+ responseBody . outputs [ 0 ] . stop_reason ,
401+ ] ) ;
402+ }
403+ }
404+ }
405+ }
221406}
0 commit comments