diff --git a/packages/datadog-plugin-openai/src/index.js b/packages/datadog-plugin-openai/src/index.js index 5eef54a8de5..f96b44543d2 100644 --- a/packages/datadog-plugin-openai/src/index.js +++ b/packages/datadog-plugin-openai/src/index.js @@ -276,25 +276,34 @@ class OpenApiPlugin extends TracingPlugin { const completionTokens = spanTags['openai.response.usage.completion_tokens'] const completionTokensEstimated = spanTags['openai.response.usage.completion_tokens_estimated'] + const totalTokens = spanTags['openai.response.usage.total_tokens'] + if (!error) { - if (promptTokensEstimated) { - this.metrics.distribution( - 'openai.tokens.prompt', promptTokens, [...tags, 'openai.estimated:true']) - } else { - this.metrics.distribution('openai.tokens.prompt', promptTokens, tags) + if (promptTokens != null) { + if (promptTokensEstimated) { + this.metrics.distribution( + 'openai.tokens.prompt', promptTokens, [...tags, 'openai.estimated:true']) + } else { + this.metrics.distribution('openai.tokens.prompt', promptTokens, tags) + } } - if (completionTokensEstimated) { - this.metrics.distribution( - 'openai.tokens.completion', completionTokens, [...tags, 'openai.estimated:true']) - } else { - this.metrics.distribution('openai.tokens.completion', completionTokens, tags) + + if (completionTokens != null) { + if (completionTokensEstimated) { + this.metrics.distribution( + 'openai.tokens.completion', completionTokens, [...tags, 'openai.estimated:true']) + } else { + this.metrics.distribution('openai.tokens.completion', completionTokens, tags) + } } - if (promptTokensEstimated || completionTokensEstimated) { - this.metrics.distribution( - 'openai.tokens.total', promptTokens + completionTokens, [...tags, 'openai.estimated:true']) - } else { - this.metrics.distribution('openai.tokens.total', promptTokens + completionTokens, tags) + if (totalTokens != null) { + if (promptTokensEstimated || completionTokensEstimated) { + this.metrics.distribution( + 'openai.tokens.total', totalTokens, [...tags, 'openai.estimated:true']) + } else { + this.metrics.distribution('openai.tokens.total', totalTokens, tags) + } } } @@ -777,9 +786,9 @@ function usageExtraction (tags, body, methodName, openaiStore) { if (completionEstimated) tags['openai.response.usage.completion_tokens_estimated'] = true } - if (promptTokens) tags['openai.response.usage.prompt_tokens'] = promptTokens - if (completionTokens) tags['openai.response.usage.completion_tokens'] = completionTokens - if (totalTokens) tags['openai.response.usage.total_tokens'] = totalTokens + if (promptTokens != null) tags['openai.response.usage.prompt_tokens'] = promptTokens + if (completionTokens != null) tags['openai.response.usage.completion_tokens'] = completionTokens + if (totalTokens != null) tags['openai.response.usage.total_tokens'] = totalTokens } function truncateApiKey (apiKey) { diff --git a/packages/datadog-plugin-openai/test/index.spec.js b/packages/datadog-plugin-openai/test/index.spec.js index b10b912dbb4..8df38a11650 100644 --- a/packages/datadog-plugin-openai/test/index.spec.js +++ b/packages/datadog-plugin-openai/test/index.spec.js @@ -560,6 +560,57 @@ describe('Plugin', () => { }) }) + describe('embedding with missing usages', () => { + afterEach(() => { + nock.cleanAll() + }) + + it('makes a successful call', async () => { + nock('https://api.openai.com:443') + .post('/v1/embeddings') + .reply(200, { + object: 'list', + data: [{ + object: 'embedding', + index: 0, + embedding: [-0.0034387498, -0.026400521] + }], + model: 'text-embedding-ada-002-v2', + usage: { + prompt_tokens: 0 + } + }, []) + + const checkTraces = agent + .use(traces => { + expect(traces[0][0].metrics).to.have.property('openai.response.usage.prompt_tokens', 0) + expect(traces[0][0].metrics).to.not.have.property('openai.response.usage.completion_tokens') + expect(traces[0][0].metrics).to.not.have.property('openai.response.usage.total_tokens') + }) + + const params = { + model: 'text-embedding-ada-002', + input: '', + user: 'hunter2' + } + + if (semver.satisfies(realVersion, '>=4.0.0')) { + const result = await openai.embeddings.create(params) + expect(result.model).to.eql('text-embedding-ada-002-v2') + } else { + const result = await openai.createEmbedding(params) + expect(result.data.model).to.eql('text-embedding-ada-002-v2') + } + + await checkTraces + + expect(metricStub).to.have.been.calledWith('openai.request.duration') // timing value not guaranteed + expect(metricStub).to.have.been.calledWith('openai.tokens.prompt') + expect(metricStub).to.not.have.been.calledWith('openai.tokens.completion') + expect(metricStub).to.not.have.been.calledWith('openai.tokens.total') + }) + }) + describe('list models', () => { let scope