diff --git a/packages/datadog-plugin-google-cloud-vertexai/src/tracing.js b/packages/datadog-plugin-google-cloud-vertexai/src/tracing.js index 8b0fb04f696..66411a2f276 100644 --- a/packages/datadog-plugin-google-cloud-vertexai/src/tracing.js +++ b/packages/datadog-plugin-google-cloud-vertexai/src/tracing.js @@ -7,7 +7,6 @@ const makeUtilities = require('../../dd-trace/src/plugins/util/llm') const { extractModel, - extractSystemInstructions } = require('./utils') class GoogleCloudVertexAITracingPlugin extends TracingPlugin { @@ -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() } @@ -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]) } - 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 } } diff --git a/packages/datadog-plugin-google-cloud-vertexai/test/index.spec.js b/packages/datadog-plugin-google-cloud-vertexai/test/index.spec.js index 22d0a76affa..b28d32b126d 100644 --- a/packages/datadog-plugin-google-cloud-vertexai/test/index.spec.js +++ b/packages/datadog-plugin-google-cloud-vertexai/test/index.spec.js @@ -100,29 +100,11 @@ 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 @@ -130,8 +112,7 @@ describe('Plugin', () => { 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?') @@ -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?') @@ -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?') @@ -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({ @@ -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({}) @@ -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({}) @@ -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({})