message test passing

Author: amychisholm03

lib/instrumentation/@google/genai.js

@@ -94,7 +94,9 @@ function recordChatCompletionMessages({
   })
 
   // Only take the first response message and append to input messages
-  const messages = [request.contents, response?.candidates?.[0]?.content?.parts?.[0]]
+  const inputMessages = Array.isArray(request.contents) ? request.contents : [request.contents]
+  const responseMessage = response?.candidates?.[0]?.content
+  const messages = responseMessage !== undefined ? [...inputMessages, responseMessage] : inputMessages
   messages.forEach((message, index) => {
     const completionMsg = new LlmChatCompletionMessage({
       agent,

lib/llm-events/google-genai/chat-completion-message.js

@@ -5,6 +5,7 @@
 
 'use strict'
 const LlmEvent = require('./event')
+const { makeId } = require('../../util/hashes')
 
 module.exports = class LlmChatCompletionMessage extends LlmEvent {
   constructor({
@@ -18,20 +19,20 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
     transaction
   }) {
     super({ agent, segment, request, response, transaction })
-    this.id = `${response.responseId}-${index}`
-    this.role = message?.role
+    this.id = makeId(36)
+    this.role = message?.role ?? 'user' // Role is only defined if it is 'model'
     this.sequence = index
     this.completion_id = completionId
-    this.is_response = response?.candidates?.[0]?.content?.parts?.[0]?.text === message?.content
+    this.is_response = response?.candidates?.[0]?.content?.parts?.[0]?.text === message?.parts?.[0]?.text
 
     if (agent.config.ai_monitoring.record_content.enabled === true) {
-      this.content = message?.content
+      this.content = message
     }
 
     if (this.is_response) {
-      this.token_count = response.candidatesTokenCount
+      this.token_count = response?.usageMetadata?.candidatesTokenCount
     } else {
-      this.token_count = response.promptTokenCount
+      this.token_count = response?.usageMetadata?.promptTokenCount
     }
   }
 }

test/unit/llm-events/google-genai/chat-completion-message.test.js

@@ -26,17 +26,18 @@ test('should create a LlmChatCompletionMessage event', (t, end) => {
   helper.runInTransaction(agent, (tx) => {
     api.startSegment('fakeSegment', false, () => {
       const segment = api.shim.getActiveSegment()
+      const summaryId = 'chat-summary-id'
       const chatMessageEvent = new LlmChatCompletionMessage({
         transaction: tx,
         agent,
         segment,
         request: req,
         response: res,
+        completionId: summaryId,
         message: req.contents,
         index: 0
       })
-      // TODO: add the expected result
-      const expected = getExpectedResult(tx, {}, 'message')
+      const expected = getExpectedResult(tx, chatMessageEvent, 'message', summaryId)
       assert.deepEqual(chatMessageEvent, expected)
       end()
     })

test/unit/llm-events/google-genai/common.js

@@ -45,12 +45,13 @@ function getExpectedResult(tx, event, type, completionId) {
     trace_id: tx.traceId,
     span_id: spanId,
     'request.model': 'gemini-2.0-flash',
+    'response.model': 'gemini-2.0-flash',
     vendor: 'gemini',
-    ingest_source: 'Node'
+    ingest_source: 'Node',
+    token_count: 10,
   }
   const resKeys = {
     duration: child.getDurationInMillis(),
-    'response.model': 'gemini-2.0-flash',
     // TODO: response.headers?
   }
 
@@ -78,7 +79,8 @@ function getExpectedResult(tx, event, type, completionId) {
         content: 'Why is the sky blue?',
         sequence: 0,
         completion_id: completionId,
-        is_response: false
+        is_response: false,
+        role: 'user'
       }
   }