Skip to content

Commit 93feebd

Browse files
amychisholm03amychisholm03
committed
@google/genai subscribers
1 parent b1687f6 commit 93feebd

File tree

7 files changed

+440
-1
lines changed

7 files changed

+440
-1
lines changed

‎lib/instrumentations.js‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ const InstrumentationDescriptor = require('./instrumentation-descriptor')
1111
module.exports = function instrumentations() {
1212
return {
1313
'@azure/functions': { type: InstrumentationDescriptor.TYPE_GENERIC },
14-
'@google/genai': { type: InstrumentationDescriptor.TYPE_GENERIC },
14+
// '@google/genai': { type: InstrumentationDescriptor.TYPE_GENERIC },
1515
'@grpc/grpc-js': { module: './instrumentation/grpc-js' },
1616
'@hapi/hapi': { type: InstrumentationDescriptor.TYPE_WEB_FRAMEWORK },
1717
'@hapi/vision': { type: InstrumentationDescriptor.TYPE_WEB_FRAMEWORK },

‎lib/subscriber-configs.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ const subscribers = {
1313
...require('./subscribers/elasticsearch/config'),
1414
...require('./subscribers/express/config'),
1515
...require('./subscribers/fastify/config'),
16+
...require('./subscribers/google-genai/config'),
1617
...require('./subscribers/ioredis/config'),
1718
...require('./subscribers/mcp-sdk/config'),
1819
...require('./subscribers/mysql/config'),

‎lib/subscribers/google-genai/base.js‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
/*
2+
* Copyright 2025 New Relic Corporation. All rights reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
const {
7+
LlmChatCompletionMessage,
8+
LlmChatCompletionSummary,
9+
LlmErrorMessage
10+
} = require('../../../lib/llm-events/google-genai')
11+
const { extractLlmContext } = require('../../util/llm-utils')
12+
const { AI } = require('../../../lib/metrics/names')
13+
const { GEMINI } = AI
14+
const { DESTINATIONS } = require('../../config/attribute-filter')
15+
16+
const Subscriber = require('../base')
17+
18+
class GoogleGenAISubscriber extends Subscriber {
19+
constructor({ agent, logger, channelName, packageName = '@google/genai' }) {
20+
super({ agent, logger, channelName, packageName })
21+
this.events = ['asyncEnd']
22+
}
23+
24+
get enabled() {
25+
return super.enabled && this.config?.ai_monitoring?.enabled
26+
}
27+
28+
/**
29+
* Enqueues a LLM event to the custom event aggregator
30+
*
31+
* @param {object} params input params
32+
* @param {string} params.type LLM event type
33+
* @param {object} params.msg LLM event
34+
*/
35+
recordEvent({ type, msg }) {
36+
const llmContext = extractLlmContext(this.agent)
37+
38+
this.agent.customEventAggregator.add([
39+
{ type, timestamp: Date.now() },
40+
Object.assign({}, msg, llmContext)
41+
])
42+
}
43+
44+
/**
45+
* Generates LlmChatCompletionSummary for a chat completion creation.
46+
* Also iterates over both input messages and the first response message
47+
* and creates LlmChatCompletionMessage.
48+
*
49+
* Also assigns relevant ids by response id for LlmFeedbackEvent creation
50+
*
51+
* @param {object} params input params
52+
* @param {TraceSegment} params.segment active segment from chat completion
53+
* @param {object} params.request chat completion params
54+
* @param {object} [params.response] chat completion response
55+
* @param {boolean} [params.err] err if it exists
56+
* @param {Transaction} params.transaction active transaction
57+
*/
58+
recordChatCompletionMessages({
59+
segment,
60+
request,
61+
response,
62+
err,
63+
transaction
64+
}) {
65+
const agent = this.agent
66+
if (!response) {
67+
// If we get an error, it is possible that `response = null`.
68+
// In that case, we define it to be an empty object.
69+
response = {}
70+
}
71+
72+
// Explicitly end segment to provide consistent duration
73+
// for both LLM events and the segment
74+
segment.end()
75+
const completionSummary = new LlmChatCompletionSummary({
76+
agent,
77+
segment,
78+
transaction,
79+
request,
80+
response,
81+
withError: err != null
82+
})
83+
84+
// Only take the first response message and append to input messages
85+
// request.contents can be a string or an array of strings
86+
// response.candidates is an array of candidates (choices); we only take the first one
87+
const inputMessages = Array.isArray(request.contents) ? request.contents : [request.contents]
88+
const responseMessage = response?.candidates?.[0]?.content
89+
const messages = responseMessage !== undefined ? [...inputMessages, responseMessage] : inputMessages
90+
messages.forEach((message, index) => {
91+
const completionMsg = new LlmChatCompletionMessage({
92+
agent,
93+
segment,
94+
transaction,
95+
request,
96+
response,
97+
index,
98+
completionId: completionSummary.id,
99+
message
100+
})
101+
102+
this.recordEvent({ type: 'LlmChatCompletionMessage', msg: completionMsg })
103+
})
104+
105+
this.recordEvent({ type: 'LlmChatCompletionSummary', msg: completionSummary })
106+
107+
if (err) {
108+
const llmError = new LlmErrorMessage({ cause: err, summary: completionSummary, response })
109+
agent.errors.add(transaction, err, llmError)
110+
}
111+
}
112+
113+
/**
114+
* Increments the tracking metric and sets the llm attribute on transactions
115+
*
116+
* @param {object} params input params
117+
* @param {Transaction} params.transaction active transaction
118+
* @param {string} params.version package version
119+
*/
120+
addLlmMeta({ transaction, version }) {
121+
const TRACKING_METRIC = `${GEMINI.TRACKING_PREFIX}/${version}`
122+
this.agent.metrics.getOrCreateMetric(TRACKING_METRIC).incrementCallCount()
123+
transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
124+
}
125+
}
126+
127+
module.exports = GoogleGenAISubscriber

‎lib/subscribers/google-genai/config.js‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
/*
2+
* Copyright 2025 New Relic Corporation. All rights reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
const generateContentInternal = {
7+
path: './google-genai/generate-content.js',
8+
instrumentations: [
9+
{
10+
channelName: 'nr_generateContentInternal',
11+
module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.cjs' },
12+
functionQuery: {
13+
className: 'Models',
14+
methodName: 'generateContentInternal',
15+
kind: 'Async'
16+
}
17+
},
18+
{
19+
channelName: 'nr_generateContentInternal',
20+
module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.mjs' },
21+
functionQuery: {
22+
className: 'Models',
23+
methodName: 'generateContentInternal',
24+
kind: 'Async'
25+
}
26+
}
27+
]
28+
}
29+
30+
const generateContentStreamInternal = {
31+
path: './google-genai/generate-content-stream.js',
32+
instrumentations: [
33+
{
34+
channelName: 'nr_generateContentStreamInternal',
35+
module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.cjs' },
36+
functionQuery: {
37+
className: 'Models',
38+
methodName: 'generateContentStreamInternal',
39+
kind: 'Async'
40+
}
41+
},
42+
{
43+
channelName: 'nr_generateContentStreamInternal',
44+
module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.mjs' },
45+
functionQuery: {
46+
className: 'Models',
47+
methodName: 'generateContentStreamInternal',
48+
kind: 'Async'
49+
}
50+
}
51+
]
52+
}
53+
54+
const embedContent = {
55+
path: './google-genai/embed-content.js',
56+
instrumentations: [
57+
{
58+
channelName: 'nr_embedContent',
59+
module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.cjs' },
60+
functionQuery: {
61+
className: 'Models',
62+
methodName: 'embedContent',
63+
kind: 'Async'
64+
}
65+
},
66+
{
67+
channelName: 'nr_embedContent',
68+
module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.mjs' },
69+
functionQuery: {
70+
className: 'Models',
71+
methodName: 'embedContent',
72+
kind: 'Async'
73+
}
74+
}
75+
]
76+
}
77+
78+
module.exports = {
79+
'@google/genai': [
80+
generateContentInternal,
81+
generateContentStreamInternal,
82+
embedContent
83+
]
84+
}

‎lib/subscribers/google-genai/embed-content.js‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright 2025 New Relic Corporation. All rights reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
const GoogleGenAISubscriber = require('./base')
7+
const { AI } = require('../../../lib/metrics/names')
8+
const { GEMINI } = AI
9+
const {
10+
LlmErrorMessage,
11+
LlmEmbedding
12+
} = require('../../../lib/llm-events/google-genai')
13+
14+
class GoogleGenAIEmbedContentSubscriber extends GoogleGenAISubscriber {
15+
constructor ({ agent, logger, channelName = 'nr_embedContent' }) {
16+
super({ agent, logger, channelName })
17+
}
18+
19+
handler(data, ctx) {
20+
const segment = this.agent.tracer.createSegment({
21+
name: GEMINI.EMBEDDING,
22+
parent: ctx.segment,
23+
transaction: ctx.transaction
24+
})
25+
return ctx.enterSegment({ segment })
26+
}
27+
28+
asyncEnd(data) {
29+
const agent = this.agent
30+
const ctx = agent.tracer.getContext()
31+
if (!ctx?.segment || !ctx?.transaction) {
32+
return
33+
}
34+
this.addLlmMeta({ transaction: ctx.transaction, version: data.moduleVersion })
35+
let { result: response, arguments: args, error: err } = data
36+
const [request] = args
37+
38+
if (!response) {
39+
// If we get an error, it is possible that `response = null`.
40+
// In that case, we define it to be an empty object.
41+
response = {}
42+
}
43+
44+
// Explicitly end segment to get consistent duration
45+
// for both LLM events and the segment
46+
ctx.segment.end()
47+
48+
const embedding = new LlmEmbedding({
49+
agent,
50+
segment: ctx.segment,
51+
transaction: ctx.transaction,
52+
request,
53+
response,
54+
withError: err != null
55+
})
56+
57+
this.recordEvent({ type: 'LlmEmbedding', msg: embedding })
58+
59+
if (err) {
60+
const llmError = new LlmErrorMessage({ cause: err, embedding, response })
61+
agent.errors.add(ctx.transaction, err, llmError)
62+
}
63+
}
64+
}
65+
66+
module.exports = GoogleGenAIEmbedContentSubscriber

0 commit comments

Comments
 (0)