feat: Add time_to_first_token for Google Gen AI (#3785)

Author: amychisholm03

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

@@ -20,9 +20,10 @@ module.exports = class GoogleGenAiLlmChatCompletionSummary extends LlmChatComple
    * @param {Transaction} params.transaction Current and active transaction
    * @param {object} params.request Google Gen AI request object
    * @param {object} params.response Google Gen AI response object
+   * @param {number} [params.timeOfFirstToken] Timestamp of when the first token was sent, for streaming only.
    * @param {boolean} [params.error] Set to `true` if an error occurred
    */
-  constructor({ agent, segment, transaction, request, response, error }) {
+  constructor({ agent, segment, transaction, request, response, timeOfFirstToken, error }) {
     super({ agent,
       segment,
       transaction,
@@ -32,6 +33,7 @@ module.exports = class GoogleGenAiLlmChatCompletionSummary extends LlmChatComple
       maxTokens: request.config?.maxOutputTokens,
       temperature: request.config?.temperature,
       vendor: 'gemini',
+      timeOfFirstToken,
       error })
 
     let requestMessagesLength = 0

lib/subscribers/google-genai/generate-content-stream.js

@@ -30,6 +30,7 @@ class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateCont
     let cachedResult = {}
     let err
     let entireMessage = ''
+    let timeOfFirstToken
     response.next = async function wrappedNext(...nextArgs) {
       let result = {}
       try {
@@ -44,6 +45,7 @@ class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateCont
         }
 
         if (result?.value?.text) {
+          if (!timeOfFirstToken) timeOfFirstToken = Date.now()
           entireMessage += result.value.text // readonly variable that equates to result.value.candidates[0].content.parts[0].text
         }
       } catch (streamErr) {
@@ -64,6 +66,7 @@ class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateCont
             ctx,
             request,
             response: cachedResult,
+            timeOfFirstToken,
             err
           })
         }

lib/subscribers/google-genai/generate-content.js

@@ -56,17 +56,19 @@ class GoogleGenAIGenerateContentSubscriber extends AiMonitoringChatSubscriber {
    * @param {Context} params.ctx active context
    * @param {object} params.request request made to method
    * @param {object} params.response response from method
+   * @param {number} [params.timeOfFirstToken] Timestamp of when the first streaming token was sent.
    * @param {object} [params.err] error object if present
    * @returns {object} a llm completion summary instance for Google Gen AI
    */
-  createCompletionSummary({ ctx, request, response = {}, err }) {
+  createCompletionSummary({ ctx, request, response = {}, timeOfFirstToken, err }) {
     const { transaction, segment } = ctx
     return new LlmChatCompletionSummary({
       agent: this.agent,
       segment,
       transaction,
       request,
       response,
+      timeOfFirstToken,
       error: !!err
     })
   }

test/versioned/google-genai/chat-completions.test.js

@@ -285,6 +285,40 @@ test('should call the tokenCountCallback in streaming', (t, end) => {
   })
 })
 
+test('should set time_to_first_token on llm chat completion summary', (t, end) => {
+  const { client, agent } = t.nr
+  helper.runInTransaction(agent, async (tx) => {
+    const content = 'Streamed response'
+    const model = 'gemini-2.0-flash'
+    const stream = await client.models.generateContentStream({
+      config: {
+        maxOutputTokens: 100,
+        temperature: 0.5
+      },
+      model,
+      contents: [content, 'What does 1 plus 1 equal?']
+    })
+
+    let res = ''
+    for await (const chunk of stream) {
+      assert.ok(chunk.text, 'should have text in chunk')
+      res += chunk.text
+    }
+    assert.ok(res)
+
+    const events = agent.customEventAggregator.events.toArray()
+    const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0]
+    assert.equal(chatSummary[0].type, 'LlmChatCompletionSummary')
+    const timeToFirstToken = chatSummary?.[1]?.['time_to_first_token']
+    assert.ok(timeToFirstToken, 'time_to_first_token should exist')
+    assert.equal(typeof timeToFirstToken, 'number', 'time_to_first_token should be a number')
+    assert.ok(timeToFirstToken >= 0, 'time_to_first_token should be >= 0')
+
+    tx.end()
+    end()
+  })
+})
+
 test('handles error in stream', (t, end) => {
   const { client, agent } = t.nr
   helper.runInTransaction(agent, async (tx) => {