Skip to content

Commit 1a74c9b

Browse files
sabrennersabrenner
authored
feat(llmobs): add vertexai plugin (#5413)
* add llmobs plugin * wip tests * tests + ci workflow * move helper functions outside of class scope * add integration tag * Update packages/dd-trace/test/llmobs/plugins/google-cloud-vertexai/index.spec.js
1 parent 959c5e6 commit 1a74c9b

File tree

6 files changed

+705
-187
lines changed

6 files changed

+705
-187
lines changed

‎.github/workflows/llmobs.yml‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ jobs:
7373
with:
7474
suffix: llmobs-${{ github.job }}
7575

76-
aws-sdk:
76+
bedrock:
7777
runs-on: ubuntu-latest
7878
env:
7979
PLUGINS: aws-sdk
@@ -92,3 +92,23 @@ jobs:
9292
uses: ./.github/actions/testagent/logs
9393
with:
9494
suffix: llmobs-${{ github.job }}
95+
96+
vertex-ai:
97+
runs-on: ubuntu-latest
98+
env:
99+
PLUGINS: google-cloud-vertexai
100+
steps:
101+
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
102+
- uses: ./.github/actions/testagent/start
103+
- uses: ./.github/actions/node/oldest-maintenance-lts
104+
- uses: ./.github/actions/install
105+
- run: yarn test:llmobs:plugins:ci
106+
shell: bash
107+
- uses: ./.github/actions/node/active-lts
108+
- run: yarn test:llmobs:plugins:ci
109+
shell: bash
110+
- uses: codecov/codecov-action@13ce06bfc6bbe3ecf90edbbf1bc32fe5978ca1d3 # v5.3.1
111+
- if: always()
112+
uses: ./.github/actions/testagent/logs
113+
with:
114+
suffix: llmobs-${{ github.job }}

‎packages/datadog-plugin-google-cloud-vertexai/src/index.js‎

Lines changed: 8 additions & 186 deletions
Original file line numberDiff line numberDiff line change
@@ -1,195 +1,17 @@
11
'use strict'
22

3-
const { MEASURED } = require('../../../ext/tags')
4-
const { storage } = require('../../datadog-core')
5-
const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
6-
const makeUtilities = require('../../dd-trace/src/plugins/util/llm')
3+
const CompositePlugin = require('../../dd-trace/src/plugins/composite')
4+
const GoogleVertexAITracingPlugin = require('./tracing')
5+
const VertexAILLMObsPlugin = require('../../dd-trace/src/llmobs/plugins/vertexai')
76

8-
class GoogleCloudVertexAIPlugin extends TracingPlugin {
7+
class GoogleCloudVertexAIPlugin extends CompositePlugin {
98
static get id () { return 'google-cloud-vertexai' }
10-
static get prefix () {
11-
return 'tracing:apm:vertexai:request'
12-
}
13-
14-
constructor () {
15-
super(...arguments)
16-
17-
Object.assign(this, makeUtilities('vertexai', this._tracerConfig))
18-
}
19-
20-
bindStart (ctx) {
21-
const { instance, request, resource, stream } = ctx
22-
23-
const tags = this.tagRequest(request, instance, stream)
24-
25-
const span = this.startSpan('vertexai.request', {
26-
service: this.config.service,
27-
resource,
28-
kind: 'client',
29-
meta: {
30-
[MEASURED]: 1,
31-
...tags
32-
}
33-
}, false)
34-
35-
const store = storage('legacy').getStore() || {}
36-
ctx.currentStore = { ...store, span }
37-
38-
return ctx.currentStore
39-
}
40-
41-
asyncEnd (ctx) {
42-
const span = ctx.currentStore?.span
43-
if (!span) return
44-
45-
const { result } = ctx
46-
47-
const response = result?.response
48-
if (response) {
49-
const tags = this.tagResponse(response)
50-
span.addTags(tags)
51-
}
52-
53-
span.finish()
54-
}
55-
56-
tagRequest (request, instance, stream) {
57-
const model = extractModel(instance)
58-
const tags = {
59-
'vertexai.request.model': model
60-
}
61-
62-
const history = instance.historyInternal
63-
let contents = typeof request === 'string' || Array.isArray(request) ? request : request.contents
64-
if (history) {
65-
contents = [...history, ...(Array.isArray(contents) ? contents : [contents])]
66-
}
67-
68-
const generationConfig = instance.generationConfig || {}
69-
for (const key of Object.keys(generationConfig)) {
70-
const transformedKey = key.replace(/([a-z0-9])([A-Z])/g, '$1_$2').toLowerCase()
71-
tags[`vertexai.request.generation_config.${transformedKey}`] = JSON.stringify(generationConfig[key])
72-
}
73-
74-
if (stream) {
75-
tags['vertexai.request.stream'] = true
76-
}
77-
78-
if (!this.isPromptCompletionSampled()) return tags
79-
80-
const systemInstructions = extractSystemInstructions(instance)
81-
82-
for (const [idx, systemInstruction] of systemInstructions.entries()) {
83-
tags[`vertexai.request.system_instruction.${idx}.text`] = systemInstruction
84-
}
85-
86-
if (typeof contents === 'string') {
87-
tags['vertexai.request.contents.0.text'] = contents
88-
return tags
89-
}
90-
91-
for (const [contentIdx, content] of contents.entries()) {
92-
this.tagRequestContent(tags, content, contentIdx)
9+
static get plugins () {
10+
return {
11+
llmobs: VertexAILLMObsPlugin,
12+
tracing: GoogleVertexAITracingPlugin
9313
}
94-
95-
return tags
9614
}
97-
98-
tagRequestPart (part, tags, partIdx, contentIdx) {
99-
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.text`] = this.normalize(part.text)
100-
101-
const functionCall = part.functionCall
102-
const functionResponse = part.functionResponse
103-
104-
if (functionCall) {
105-
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.function_call.name`] = functionCall.name
106-
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.function_call.args`] =
107-
this.normalize(JSON.stringify(functionCall.args))
108-
}
109-
if (functionResponse) {
110-
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.function_response.name`] =
111-
functionResponse.name
112-
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.function_response.response`] =
113-
this.normalize(JSON.stringify(functionResponse.response))
114-
}
115-
}
116-
117-
tagRequestContent (tags, content, contentIdx) {
118-
if (typeof content === 'string') {
119-
tags[`vertexai.request.contents.${contentIdx}.text`] = this.normalize(content)
120-
return
121-
}
122-
123-
if (content.text || content.functionCall || content.functionResponse) {
124-
this.tagRequestPart(content, tags, 0, contentIdx)
125-
return
126-
}
127-
128-
const { role, parts } = content
129-
if (role) {
130-
tags[`vertexai.request.contents.${contentIdx}.role`] = role
131-
}
132-
133-
for (const [partIdx, part] of parts.entries()) {
134-
this.tagRequestPart(part, tags, partIdx, contentIdx)
135-
}
136-
}
137-
138-
tagResponse (response) {
139-
const tags = {}
140-
141-
const candidates = response.candidates
142-
for (const [candidateIdx, candidate] of candidates.entries()) {
143-
const finishReason = candidate.finishReason
144-
if (finishReason) {
145-
tags[`vertexai.response.candidates.${candidateIdx}.finish_reason`] = finishReason
146-
}
147-
const candidateContent = candidate.content
148-
const role = candidateContent.role
149-
tags[`vertexai.response.candidates.${candidateIdx}.content.role`] = role
150-
151-
if (!this.isPromptCompletionSampled()) continue
152-
153-
const parts = candidateContent.parts
154-
for (const [partIdx, part] of parts.entries()) {
155-
const text = part.text
156-
tags[`vertexai.response.candidates.${candidateIdx}.content.parts.${partIdx}.text`] =
157-
this.normalize(String(text))
158-
159-
const functionCall = part.functionCall
160-
if (!functionCall) continue
161-
162-
tags[`vertexai.response.candidates.${candidateIdx}.content.parts.${partIdx}.function_call.name`] =
163-
functionCall.name
164-
tags[`vertexai.response.candidates.${candidateIdx}.content.parts.${partIdx}.function_call.args`] =
165-
this.normalize(JSON.stringify(functionCall.args))
166-
}
167-
}
168-
169-
const tokenCounts = response.usageMetadata
170-
if (tokenCounts) {
171-
tags['vertexai.response.usage.prompt_tokens'] = tokenCounts.promptTokenCount
172-
tags['vertexai.response.usage.completion_tokens'] = tokenCounts.candidatesTokenCount
173-
tags['vertexai.response.usage.total_tokens'] = tokenCounts.totalTokenCount
174-
}
175-
176-
return tags
177-
}
178-
}
179-
180-
function extractModel (instance) {
181-
const model = instance.model || instance.resourcePath || instance.publisherModelEndpoint
182-
return model?.split('/').pop()
183-
}
184-
185-
function extractSystemInstructions (instance) {
186-
// systemInstruction is either a string or a Content object
187-
// Content objects have parts (Part[]) and a role
188-
const systemInstruction = instance.systemInstruction
189-
if (!systemInstruction) return []
190-
if (typeof systemInstruction === 'string') return [systemInstruction]
191-
192-
return systemInstruction.parts?.map(part => part.text)
19315
}
19416

19517
module.exports = GoogleCloudVertexAIPlugin

0 commit comments

Comments
 (0)