From 45f960707bb4c0fae46a7e5168f059c27e443727 Mon Sep 17 00:00:00 2001
From: Sam Brenner <106700075+sabrenner@users.noreply.github.com>
Date: Fri, 20 Sep 2024 10:18:01 -0400
Subject: [PATCH] fix(openai): do not add usage tags or emit usage metrics if
 they are not returned (#4706)

fix(openai): do not add usage tags or emit usage metrics if they are not returned
---
 packages/datadog-plugin-openai/src/index.js   | 45 +++++++++-------
 .../datadog-plugin-openai/test/index.spec.js  | 51 +++++++++++++++++++
 2 files changed, 78 insertions(+), 18 deletions(-)

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