Skip to content

dac strip io from vertex #6103

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
+4 −215
Draft

dac strip io from vertex #6103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 0 additions & 126 deletions packages/datadog-plugin-google-cloud-vertexai/src/tracing.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ const makeUtilities = require('../../dd-trace/src/plugins/util/llm')

const {
extractModel,
extractSystemInstructions
} = require('./utils')

class GoogleCloudVertexAITracingPlugin extends TracingPlugin {
Expand Down Expand Up @@ -47,14 +46,6 @@ class GoogleCloudVertexAITracingPlugin extends TracingPlugin {
const span = ctx.currentStore?.span
if (!span) return

const { result } = ctx

const response = result?.response
if (response) {
const tags = this.tagResponse(response, span)
span.addTags(tags)
}

span.finish()
}

Expand All @@ -73,129 +64,12 @@ class GoogleCloudVertexAITracingPlugin extends TracingPlugin {
'vertexai.request.model': model
}

const history = instance.historyInternal

let contents = typeof request === 'string' || Array.isArray(request) ? request : request.contents
if (history) {
contents = [...history, ...(Array.isArray(contents) ? contents : [contents])]
}

const generationConfig = instance.generationConfig || {}
for (const key of Object.keys(generationConfig)) {
const transformedKey = key.replaceAll(/([a-z0-9])([A-Z])/g, '$1_$2').toLowerCase()
tags[`vertexai.request.generation_config.${transformedKey}`] = JSON.stringify(generationConfig[key])
}
Comment on lines 67 to 71
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just double checking that we don't want to remove these as well


if (stream) {
tags['vertexai.request.stream'] = true
}

if (!this.isPromptCompletionSampled(span)) return tags

const systemInstructions = extractSystemInstructions(instance)

for (const [idx, systemInstruction] of systemInstructions.entries()) {
tags[`vertexai.request.system_instruction.${idx}.text`] = systemInstruction
}

if (typeof contents === 'string') {
tags['vertexai.request.contents.0.text'] = contents
return tags
}

for (const [contentIdx, content] of contents.entries()) {
this.tagRequestContent(tags, content, contentIdx)
}

return tags
}

tagRequestPart (part, tags, partIdx, contentIdx) {
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.text`] = this.normalize(part.text)

const functionCall = part.functionCall
const functionResponse = part.functionResponse

if (functionCall) {
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.function_call.name`] = functionCall.name
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.function_call.args`] =
this.normalize(JSON.stringify(functionCall.args))
}
if (functionResponse) {
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.function_response.name`] =
functionResponse.name
tags[`vertexai.request.contents.${contentIdx}.parts.${partIdx}.function_response.response`] =
this.normalize(JSON.stringify(functionResponse.response))
}
}

tagRequestContent (tags, content, contentIdx) {
if (typeof content === 'string') {
tags[`vertexai.request.contents.${contentIdx}.text`] = this.normalize(content)
return
}

if (content.text || content.functionCall || content.functionResponse) {
this.tagRequestPart(content, tags, 0, contentIdx)
return
}

const { role, parts } = content
if (role) {
tags[`vertexai.request.contents.${contentIdx}.role`] = role
}

for (const [partIdx, part] of parts.entries()) {
this.tagRequestPart(part, tags, partIdx, contentIdx)
}
}

/**
* Generate the response tags.
*
* @param {Object} response
* @param {Span} span
* @returns {Object}
*/
tagResponse (response, span) {
const tags = {}
const isSampled = this.isPromptCompletionSampled(span)

const candidates = response.candidates
for (const [candidateIdx, candidate] of candidates.entries()) {
const finishReason = candidate.finishReason
if (finishReason) {
tags[`vertexai.response.candidates.${candidateIdx}.finish_reason`] = finishReason
}
const candidateContent = candidate.content
const role = candidateContent.role
tags[`vertexai.response.candidates.${candidateIdx}.content.role`] = role

if (!isSampled) continue

const parts = candidateContent.parts
for (const [partIdx, part] of parts.entries()) {
const text = part.text
tags[`vertexai.response.candidates.${candidateIdx}.content.parts.${partIdx}.text`] =
this.normalize(String(text))

const functionCall = part.functionCall
if (!functionCall) continue

tags[`vertexai.response.candidates.${candidateIdx}.content.parts.${partIdx}.function_call.name`] =
functionCall.name
tags[`vertexai.response.candidates.${candidateIdx}.content.parts.${partIdx}.function_call.args`] =
this.normalize(JSON.stringify(functionCall.args))
}
}

const tokenCounts = response.usageMetadata
if (tokenCounts) {
tags['vertexai.response.usage.prompt_tokens'] = tokenCounts.promptTokenCount
tags['vertexai.response.usage.completion_tokens'] = tokenCounts.candidatesTokenCount
tags['vertexai.response.usage.total_tokens'] = tokenCounts.totalTokenCount
}

return tags
}
}
Expand Down
93 changes: 4 additions & 89 deletions packages/datadog-plugin-google-cloud-vertexai/test/index.spec.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,38 +100,19 @@ describe('Plugin', () => {
expect(span.meta).to.have.property('span.kind', 'client')

expect(span.meta).to.have.property('vertexai.request.model', 'gemini-1.5-flash-002')
expect(span.meta).to.have.property('vertexai.request.contents.0.role', 'user')
expect(span.meta).to.have.property('vertexai.request.contents.0.parts.0.text', 'Hello, how are you?')
expect(span.meta).to.have.property('vertexai.response.candidates.0.finish_reason', 'STOP')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.parts.0.text',
'Hello! How can I assist you today?')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.role', 'model')

expect(span.metrics).to.have.property('vertexai.response.usage.prompt_tokens', 35)
expect(span.metrics).to.have.property('vertexai.response.usage.completion_tokens', 2)
expect(span.metrics).to.have.property('vertexai.response.usage.total_tokens', 37)

if (model.systemInstruction) {
expect(span.meta).to.have.property('vertexai.request.system_instruction.0.text',
'Please provide an answer')
}
expect(span.meta).to.have.property('vertexai.request.generation_config.max_output_tokens', '50')
expect(span.meta).to.have.property('vertexai.request.generation_config.temperature', '1')
})

const { response } = await model.generateContent({
contents: [{ role: 'user', parts: [{ text: 'Hello, how are you?' }] }]
})

expect(response).to.have.property('candidates')

await checkTraces
})

it('makes a successful call with a string argument', async () => {
const checkTraces = agent.assertSomeTraces(traces => {
expect(traces[0][0].meta).to.have.property('vertexai.request.contents.0.text',
'Hello, how are you?')
expect(traces[0][0].meta).to.have.property('vertexai.request.model', 'gemini-1.5-flash-002')
})

const { response } = await model.generateContent('Hello, how are you?')
Expand All @@ -148,11 +129,7 @@ describe('Plugin', () => {
const checkTraces = agent.assertSomeTraces(traces => {
const span = traces[0][0]

expect(span.meta).to.have.property('vertexai.response.candidates.0.content.parts.0.text', 'undefined')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.parts.0.function_call.name',
'add')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.parts.0.function_call.args',
JSON.stringify({ a: 2, b: 2 }))
expect(span.meta).to.have.property('vertexai.request.model', 'gemini-1.5-flash-002')
})

await model.generateContent('what is 2 + 2?')
Expand All @@ -174,24 +151,6 @@ describe('Plugin', () => {
expect(span.meta).to.have.property('span.kind', 'client')

expect(span.meta).to.have.property('vertexai.request.model', 'gemini-1.5-flash-002')
expect(span.meta).to.have.property('vertexai.request.contents.0.text', 'Hello, how are you?')
expect(span.meta).to.have.property('vertexai.response.candidates.0.finish_reason', 'STOP')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.parts.0.text',
'Hi, how are you doing today my friend?')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.role', 'model')

expect(span.metrics).to.have.property('vertexai.response.usage.prompt_tokens', 5)
expect(span.metrics).to.have.property('vertexai.response.usage.completion_tokens', 10)
expect(span.metrics).to.have.property('vertexai.response.usage.total_tokens', 15)

if (model.systemInstruction) {
expect(span.meta).to.have.property('vertexai.request.system_instruction.0.text',
'Please provide an answer')
}
expect(span.meta).to.have.property('vertexai.request.generation_config.max_output_tokens', '50')
expect(span.meta).to.have.property('vertexai.request.generation_config.temperature', '1')

expect(span.metrics).to.have.property('vertexai.request.stream', 1)
})

const { stream, response } = await model.generateContentStream('Hello, how are you?')
Expand Down Expand Up @@ -226,28 +185,6 @@ describe('Plugin', () => {
expect(span.meta).to.have.property('span.kind', 'client')

expect(span.meta).to.have.property('vertexai.request.model', 'gemini-1.5-flash-002')

expect(span.meta).to.have.property('vertexai.request.contents.0.role', 'user')
expect(span.meta).to.have.property('vertexai.request.contents.0.parts.0.text', 'Foobar?')
expect(span.meta).to.have.property('vertexai.request.contents.1.role', 'model')
expect(span.meta).to.have.property('vertexai.request.contents.1.parts.0.text', 'Foobar!')
expect(span.meta).to.have.property('vertexai.request.contents.2.parts.0.text', 'Hello, how are you?')

expect(span.meta).to.have.property('vertexai.response.candidates.0.finish_reason', 'STOP')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.parts.0.text',
'Hello! How can I assist you today?')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.role', 'model')

expect(span.metrics).to.have.property('vertexai.response.usage.prompt_tokens', 35)
expect(span.metrics).to.have.property('vertexai.response.usage.completion_tokens', 2)
expect(span.metrics).to.have.property('vertexai.response.usage.total_tokens', 37)

if (model.systemInstruction) {
expect(span.meta).to.have.property('vertexai.request.system_instruction.0.text',
'Please provide an answer')
}
expect(span.meta).to.have.property('vertexai.request.generation_config.max_output_tokens', '50')
expect(span.meta).to.have.property('vertexai.request.generation_config.temperature', '1')
})

const chat = model.startChat({
Expand All @@ -265,8 +202,7 @@ describe('Plugin', () => {

it('tags a string input', async () => {
const checkTraces = agent.assertSomeTraces(traces => {
expect(traces[0][0].meta).to.have.property('vertexai.request.contents.0.text',
'Hello, how are you?')
expect(traces[0][0].meta).to.have.property('vertexai.request.model', 'gemini-1.5-flash-002')
})

const chat = model.startChat({})
Expand All @@ -279,10 +215,7 @@ describe('Plugin', () => {

it('tags an array of string inputs', async () => {
const checkTraces = agent.assertSomeTraces(traces => {
expect(traces[0][0].meta).to.have.property('vertexai.request.contents.0.text',
'Hello, how are you?')
expect(traces[0][0].meta).to.have.property('vertexai.request.contents.1.text',
'What should I do today?')
expect(traces[0][0].meta).to.have.property('vertexai.request.model', 'gemini-1.5-flash-002')
})

const chat = model.startChat({})
Expand All @@ -306,24 +239,6 @@ describe('Plugin', () => {
expect(span.meta).to.have.property('span.kind', 'client')

expect(span.meta).to.have.property('vertexai.request.model', 'gemini-1.5-flash-002')
expect(span.meta).to.have.property('vertexai.request.contents.0.text', 'Hello, how are you?')
expect(span.meta).to.have.property('vertexai.response.candidates.0.finish_reason', 'STOP')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.parts.0.text',
'Hi, how are you doing today my friend?')
expect(span.meta).to.have.property('vertexai.response.candidates.0.content.role', 'model')

expect(span.metrics).to.have.property('vertexai.response.usage.prompt_tokens', 5)
expect(span.metrics).to.have.property('vertexai.response.usage.completion_tokens', 10)
expect(span.metrics).to.have.property('vertexai.response.usage.total_tokens', 15)

if (model.systemInstruction) {
expect(span.meta).to.have.property('vertexai.request.system_instruction.0.text',
'Please provide an answer')
}
expect(span.meta).to.have.property('vertexai.request.generation_config.max_output_tokens', '50')
expect(span.meta).to.have.property('vertexai.request.generation_config.temperature', '1')

expect(span.metrics).to.have.property('vertexai.request.stream', 1)
})

const chat = model.startChat({})
Expand Down
Loading