@@ -207,86 +207,70 @@ describe("Test Meta with AWS Bedrock Instrumentation", () => {
207207
208208 const command = new bedrock . InvokeModelWithResponseStreamCommand ( input ) ;
209209 const response = await bedrockRuntimeClient . send ( command ) ;
210+
211+ // Collect all chunks and find the final one with metrics
212+ // eslint-disable-next-line @typescript-eslint/no-explicit-any
213+ let finalParsedResponse : any = null ;
210214 if ( response . body ) {
211215 for await ( const value of response . body ! ) {
212216 const jsonString = new TextDecoder ( ) . decode ( value . chunk ?. bytes ) ;
213217 const parsedResponse = JSON . parse ( jsonString ) ;
214-
215- const spans = memoryExporter . getFinishedSpans ( ) ;
216-
217- const attributes = spans [ 0 ] . attributes ;
218-
219- assert . strictEqual ( attributes [ ATTR_GEN_AI_SYSTEM ] , "AWS" ) ;
220- assert . strictEqual (
221- attributes [ SpanAttributes . LLM_REQUEST_TYPE ] ,
222- "completion" ,
223- ) ;
224- assert . strictEqual ( attributes [ ATTR_GEN_AI_REQUEST_MODEL ] , model ) ;
225- assert . strictEqual ( attributes [ ATTR_GEN_AI_REQUEST_TOP_P ] , params . top_p ) ;
226- assert . strictEqual (
227- attributes [ ATTR_GEN_AI_REQUEST_TEMPERATURE ] ,
228- params . temperature ,
229- ) ;
230- assert . strictEqual (
231- attributes [ ATTR_GEN_AI_REQUEST_MAX_TOKENS ] ,
232- params . max_gen_len ,
233- ) ;
234- assert . strictEqual ( attributes [ `${ ATTR_GEN_AI_PROMPT } .0.role` ] , "user" ) ;
235- assert . strictEqual (
236- attributes [ `${ ATTR_GEN_AI_PROMPT } .0.content` ] ,
237- prompt ,
238- ) ;
239- assert . strictEqual ( attributes [ ATTR_GEN_AI_REQUEST_MODEL ] , model ) ;
240- assert . strictEqual (
241- attributes [ `${ ATTR_GEN_AI_COMPLETION } .0.role` ] ,
242- "assistant" ,
243- ) ;
244- assert . strictEqual (
245- attributes [ ATTR_GEN_AI_USAGE_PROMPT_TOKENS ] ,
246- parsedResponse [ "prompt_token_count" ] ,
247- ) ;
248- assert . strictEqual (
249- attributes [ ATTR_GEN_AI_USAGE_COMPLETION_TOKENS ] ,
250- parsedResponse [ "generation_token_count" ] ,
251- ) ;
252- assert . strictEqual (
253- attributes [ SpanAttributes . LLM_USAGE_TOTAL_TOKENS ] ,
254- parsedResponse [ "prompt_token_count" ] +
255- parsedResponse [ "generation_token_count" ] ,
256- ) ;
257- assert . strictEqual (
258- attributes [ `${ ATTR_GEN_AI_COMPLETION } .0.finish_reason` ] ,
259- parsedResponse [ "stop_reason" ] ,
260- ) ;
261- assert . strictEqual (
262- attributes [ `${ ATTR_GEN_AI_COMPLETION } .0.content` ] ,
263- parsedResponse [ "generation" ] ,
264- ) ;
265-
218+ // The final chunk contains amazon-bedrock-invocationMetrics
266219 if ( "amazon-bedrock-invocationMetrics" in parsedResponse ) {
267- assert . strictEqual (
268- attributes [ ATTR_GEN_AI_USAGE_PROMPT_TOKENS ] ,
269- parsedResponse [ "amazon-bedrock-invocationMetrics" ] [
270- "inputTokenCount"
271- ] ,
272- ) ;
273- assert . strictEqual (
274- attributes [ ATTR_GEN_AI_USAGE_COMPLETION_TOKENS ] ,
275- parsedResponse [ "amazon-bedrock-invocationMetrics" ] [
276- "outputTokenCount"
277- ] ,
278- ) ;
279- assert . strictEqual (
280- attributes [ SpanAttributes . LLM_USAGE_TOTAL_TOKENS ] ,
281- parsedResponse [ "amazon-bedrock-invocationMetrics" ] [
282- "inputTokenCount"
283- ] +
284- parsedResponse [ "amazon-bedrock-invocationMetrics" ] [
285- "outputTokenCount"
286- ] ,
287- ) ;
220+ finalParsedResponse = parsedResponse ;
288221 }
289222 }
290223 }
224+
225+ // Run assertions only after all chunks have been processed
226+ const spans = memoryExporter . getFinishedSpans ( ) ;
227+ const attributes = spans [ 0 ] . attributes ;
228+
229+ assert . strictEqual ( attributes [ ATTR_GEN_AI_SYSTEM ] , "AWS" ) ;
230+ assert . strictEqual (
231+ attributes [ SpanAttributes . LLM_REQUEST_TYPE ] ,
232+ "completion" ,
233+ ) ;
234+ assert . strictEqual ( attributes [ ATTR_GEN_AI_REQUEST_MODEL ] , model ) ;
235+ assert . strictEqual ( attributes [ ATTR_GEN_AI_REQUEST_TOP_P ] , params . top_p ) ;
236+ assert . strictEqual (
237+ attributes [ ATTR_GEN_AI_REQUEST_TEMPERATURE ] ,
238+ params . temperature ,
239+ ) ;
240+ assert . strictEqual (
241+ attributes [ ATTR_GEN_AI_REQUEST_MAX_TOKENS ] ,
242+ params . max_gen_len ,
243+ ) ;
244+ assert . strictEqual ( attributes [ `${ ATTR_GEN_AI_PROMPT } .0.role` ] , "user" ) ;
245+ assert . strictEqual ( attributes [ `${ ATTR_GEN_AI_PROMPT } .0.content` ] , prompt ) ;
246+ assert . strictEqual (
247+ attributes [ `${ ATTR_GEN_AI_COMPLETION } .0.role` ] ,
248+ "assistant" ,
249+ ) ;
250+
251+ // Token counts should match the final invocation metrics
252+ if ( finalParsedResponse ) {
253+ assert . strictEqual (
254+ attributes [ ATTR_GEN_AI_USAGE_PROMPT_TOKENS ] ,
255+ finalParsedResponse [ "amazon-bedrock-invocationMetrics" ] [
256+ "inputTokenCount"
257+ ] ,
258+ ) ;
259+ assert . strictEqual (
260+ attributes [ ATTR_GEN_AI_USAGE_COMPLETION_TOKENS ] ,
261+ finalParsedResponse [ "amazon-bedrock-invocationMetrics" ] [
262+ "outputTokenCount"
263+ ] ,
264+ ) ;
265+ assert . strictEqual (
266+ attributes [ SpanAttributes . LLM_USAGE_TOTAL_TOKENS ] ,
267+ finalParsedResponse [ "amazon-bedrock-invocationMetrics" ] [
268+ "inputTokenCount"
269+ ] +
270+ finalParsedResponse [ "amazon-bedrock-invocationMetrics" ] [
271+ "outputTokenCount"
272+ ] ,
273+ ) ;
274+ }
291275 } ) ;
292276} ) ;
0 commit comments