google-genai Gemini skeleton

Author: amychisholm03

lib/instrumentation/@google/genai.js

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const { geminiApiKey } = require('../../lib/symbols')
+const {
+    LlmChatCompletionMessage,
+    LlmChatCompletionSummary,
+    LlmEmbedding,
+    LlmErrorMessage
+} = require('../../lib/llm-events/google-genai')
+const { RecorderSpec } = require('../../lib/shim/specs')
+const { extractLlmContext } = require('../util/llm-utils')
+
+const { AI } = require('../../lib/metrics/names')
+const { GEMINI } = AI
+let TRACKING_METRIC = GEMINI.TRACKING_PREFIX
+
+
+module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
+    if (agent?.config?.ai_monitoring?.enabled !== true) {
+        shim.logger.debug('config.ai_monitoring.enabled is set to false. Skipping instrumentation.')
+        return
+    }
+
+    // Update the tracking metric name with the version of the library
+    // being instrumented. We do not have access to the version when
+    // initially declaring the variable.
+    TRACKING_METRIC = `${TRACKING_METRIC}/${shim.pkgVersion}`
+}

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const LlmEvent = require('./event')
+
+module.exports = class LlmChatCompletionMessage extends LlmEvent {
+  constructor({
+    agent,
+    segment,
+    request = {},
+    response = {},
+    index = 0,
+    message,
+    completionId,
+    transaction
+  }) {
+    super({ agent, segment, request, response, transaction })
+    this.id = `${response.id}-${index}`
+    this.role = message?.role
+    this.sequence = index
+    this.completion_id = completionId
+    this.is_response = response?.choices?.[0]?.message?.content === message?.content
+
+    if (agent.config.ai_monitoring.record_content.enabled === true) {
+      this.content = message?.content
+    }
+
+    const tokenCB = agent.llm?.tokenCountCallback
+    if (typeof tokenCB !== 'function') {
+      return
+    }
+
+    if (this.is_response) {
+      this.token_count = tokenCB(this['response.model'], message?.content)
+    } else {
+      this.token_count = tokenCB(request.model || request.engine, message?.content)
+    }
+  }
+}

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

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const LlmEvent = require('./event')
+
+module.exports = class LlmChatCompletionSummary extends LlmEvent {
+  constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
+    super({ agent, segment, request, response, responseAttrs: true, transaction })
+    this.error = withError
+    this['request.max_tokens'] = request.max_tokens
+    this['request.temperature'] = request.temperature
+    this['response.number_of_messages'] = request?.messages?.length + response?.choices?.length
+    this['response.choices.finish_reason'] = response?.choices?.[0]?.finish_reason
+  }
+}

lib/llm-events/google-genai/embedding.js

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const LlmEvent = require('./event')
+
+/**
+ * @typedef {object} LlmEmbeddingParams
+ * @augments LlmEventParams
+ * @property {string} input - The input message for the embedding call
+ */
+/**
+ * @type {LlmEmbeddingParams}
+ */
+const defaultParams = {}
+
+class LlmEmbedding extends LlmEvent {
+  constructor(params = defaultParams) {
+    super(params)
+    const { agent, input } = params
+    const tokenCb = agent?.llm?.tokenCountCallback
+
+    this.input = agent.config?.ai_monitoring?.record_content?.enabled
+      ? input
+      : undefined
+    this.error = params.isError
+    this.duration = params.segment.getDurationInMillis()
+
+    // Even if not recording content, we should use the local token counting callback to record token usage
+    if (typeof tokenCb === 'function') {
+      this.token_count = tokenCb(this.bedrockCommand.modelId, input)
+    }
+  }
+}
+
+module.exports = LlmEmbedding

lib/llm-events/google-genai/event.js

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const BaseEvent = require('../event')
+const { makeId } = require('../../util/hashes')
+
+module.exports = class LlmEvent extends BaseEvent {
+  constructor({ agent, segment, request, response, responseAttrs = false, transaction }) {
+    super()
+
+    this.id = makeId(36)
+    this.appName = agent.config.applications()[0]
+    this.request_id = response?.headers?.['x-request-id']
+    this.trace_id = transaction?.traceId
+    this.span_id = segment?.id
+    this['response.model'] = response.model // TODO: request.model?
+    this.vendor = 'gemini'
+    this.ingest_source = 'Node'
+    this.metadata = agent
+
+    /**
+     * Used in embedding, and chat completion summary.
+     * The flag will include response attributes but also
+     * other attributes from request like model, and api key.
+     * Lastly, it includes the active span's duration.
+     */
+    if (responseAttrs) {
+      this['request.model'] = request.model || request.engine
+      this.duration = segment?.getDurationInMillis()
+      this.responseAttrs(response)
+    }
+  }
+
+  responseAttrs(response) {
+    this['response.organization'] = response?.headers?.['google-organization']
+    this['response.headers.llmVersion'] = response?.headers?.['gemini-version']
+    this['response.headers.ratelimitLimitRequests'] =
+      response?.headers?.['x-ratelimit-limit-requests']
+    this['response.headers.ratelimitLimitTokens'] = response?.headers?.['x-ratelimit-limit-tokens']
+    this['response.headers.ratelimitResetTokens'] = response?.headers?.['x-ratelimit-reset-tokens']
+    this['response.headers.ratelimitRemainingTokens'] =
+      response?.headers?.['x-ratelimit-remaining-tokens']
+    this['response.headers.ratelimitRemainingRequests'] =
+      response?.headers?.['x-ratelimit-remaining-requests']
+  }
+}

lib/llm-events/google-genai/index.js

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const LlmChatCompletionSummary = require('./chat-completion-summary')
+const LlmChatCompletionMessage = require('./chat-completion-message')
+const LlmEmbedding = require('./embedding')
+const LlmErrorMessage = require('../error-message')
+
+module.exports = {
+  LlmChatCompletionMessage,
+  LlmChatCompletionSummary,
+  LlmEmbedding,
+  LlmErrorMessage
+}

lib/metrics/names.js

@@ -174,6 +174,12 @@ const AI = {
   VECTORSTORE: 'Llm/vectorstore'
 }
 
+AI.GEMINI = {
+  TRACKING_PREFIX:`${AI.TRACKING_PREFIX}/Gemini`,
+  EMBEDDING: `${AI.EMBEDDING}/Gemini`,
+  COMPLETION: `${AI.COMPLETION}/Gemini`
+}
+
 AI.OPENAI = {
   TRACKING_PREFIX: `${AI.TRACKING_PREFIX}/OpenAI`,
   EMBEDDING: `${AI.EMBEDDING}/OpenAI/create`,

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

@@ -0,0 +1 @@
+// TODO

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

@@ -0,0 +1 @@
+// TODO

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

@@ -0,0 +1 @@
+// TODO