openai: be defensive about response object shape

Author: tlhunter

packages/datadog-plugin-openai/src/index.js

@@ -363,6 +363,8 @@ function retrieveModelResponseExtraction (tags, body) {
   tags['openai.response.parent'] = body.parent
   tags['openai.response.root'] = body.root
 
+  if (!body.permission) return
+
   tags['openai.response.permission.id'] = body.permission[0].id
   tags['openai.response.permission.created'] = body.permission[0].created
   tags['openai.response.permission.allow_create_engine'] = body.permission[0].allow_create_engine
@@ -382,10 +384,14 @@ function commonLookupFineTuneRequestExtraction (tags, body) {
 }
 
 function listModelsResponseExtraction (tags, body) {
+  if (!body.data) return
+
   tags['openai.response.count'] = body.data.length
 }
 
 function commonImageResponseExtraction (tags, body) {
+  if (!body.data) return
+
   tags['openai.response.images_count'] = body.data.length
 
   for (let i = 0; i < body.data.length; i++) {
@@ -400,7 +406,7 @@ function createAudioResponseExtraction (tags, body) {
   tags['openai.response.text'] = body.text
   tags['openai.response.language'] = body.language
   tags['openai.response.duration'] = body.duration
-  tags['openai.response.segments_count'] = body.segments.length
+  tags['openai.response.segments_count'] = defensiveArrayLength(body.segments)
 }
 
 function createFineTuneRequestExtraction (tags, body) {
@@ -417,21 +423,24 @@ function createFineTuneRequestExtraction (tags, body) {
 }
 
 function commonFineTuneResponseExtraction (tags, body) {
-  tags['openai.response.events_count'] = body.events.length
+  tags['openai.response.events_count'] = defensiveArrayLength(body.events)
   tags['openai.response.fine_tuned_model'] = body.fine_tuned_model
-  tags['openai.response.hyperparams.n_epochs'] = body.hyperparams.n_epochs
-  tags['openai.response.hyperparams.batch_size'] = body.hyperparams.batch_size
-  tags['openai.response.hyperparams.prompt_loss_weight'] = body.hyperparams.prompt_loss_weight
-  tags['openai.response.hyperparams.learning_rate_multiplier'] = body.hyperparams.learning_rate_multiplier
-  tags['openai.response.training_files_count'] = body.training_files.length
-  tags['openai.response.result_files_count'] = body.result_files.length
-  tags['openai.response.validation_files_count'] = body.validation_files.length
+  if (body.hyperparams) {
+    tags['openai.response.hyperparams.n_epochs'] = body.hyperparams.n_epochs
+    tags['openai.response.hyperparams.batch_size'] = body.hyperparams.batch_size
+    tags['openai.response.hyperparams.prompt_loss_weight'] = body.hyperparams.prompt_loss_weight
+    tags['openai.response.hyperparams.learning_rate_multiplier'] = body.hyperparams.learning_rate_multiplier
+  }
+  tags['openai.response.training_files_count'] = defensiveArrayLength(body.training_files)
+  tags['openai.response.result_files_count'] = defensiveArrayLength(body.result_files)
+  tags['openai.response.validation_files_count'] = defensiveArrayLength(body.validation_files)
   tags['openai.response.updated_at'] = body.updated_at
   tags['openai.response.status'] = body.status
 }
 
 // the OpenAI package appears to stream the content download then provide it all as a singular string
 function downloadFileResponseExtraction (tags, body) {
+  if (!body.file) return
   tags['openai.response.total_bytes'] = body.file.length
 }
 
@@ -472,20 +481,26 @@ function createRetrieveFileResponseExtraction (tags, body) {
 function createEmbeddingResponseExtraction (tags, body) {
   usageExtraction(tags, body)
 
+  if (!body.data) return
+
   tags['openai.response.embeddings_count'] = body.data.length
   for (let i = 0; i < body.data.length; i++) {
     tags[`openai.response.embedding.${i}.embedding_length`] = body.data[i].embedding.length
   }
 }
 
 function commonListCountResponseExtraction (tags, body) {
+  if (!body.data) return
   tags['openai.response.count'] = body.data.length
 }
 
 // TODO: Is there ever more than one entry in body.results?
 function createModerationResponseExtraction (tags, body) {
   tags['openai.response.id'] = body.id
   // tags[`openai.response.model`] = body.model // redundant, already extracted globally
+
+  if (!body.results) return
+
   tags['openai.response.flagged'] = body.results[0].flagged
 
   for (const [category, match] of Object.entries(body.results[0].categories)) {
@@ -501,6 +516,8 @@ function createModerationResponseExtraction (tags, body) {
 function commonCreateResponseExtraction (tags, body, store) {
   usageExtraction(tags, body)
 
+  if (!body.choices) return
+
   tags['openai.response.choices_count'] = body.choices.length
 
   store.choices = body.choices

packages/datadog-plugin-openai/test/index.spec.js

@@ -55,7 +55,12 @@ describe('Plugin', () => {
       describe('createCompletion()', () => {
         let scope
 
-        before(() => {
+        after(() => {
+          nock.removeInterceptor(scope)
+          scope.done()
+        })
+
+        it('makes a successful call', async () => {
           scope = nock('https://api.openai.com:443')
             .post('/v1/completions')
             .reply(200, {
@@ -87,14 +92,7 @@ describe('Plugin', () => {
               'x-ratelimit-reset-tokens', '3ms',
               'x-request-id', '7df89d8afe7bf24dc04e2c4dd4962d7f'
             ])
-        })
-
-        after(() => {
-          nock.removeInterceptor(scope)
-          scope.done()
-        })
 
-        it('makes a successful call', async () => {
           const checkTraces = agent
             .use(traces => {
               expect(traces[0][0]).to.have.property('name', 'openai.request')
@@ -188,6 +186,44 @@ describe('Plugin', () => {
             ]
           })
         })
+
+        it('should not throw with empty response body', async () => {
+          scope = nock('https://api.openai.com:443')
+            .post('/v1/completions')
+            .reply(200, {}, [
+              'Date', 'Mon, 15 May 2023 17:24:22 GMT',
+              'Content-Type', 'application/json',
+              'Content-Length', '349',
+              'Connection', 'close',
+              'openai-model', 'text-davinci-002',
+              'openai-organization', 'kill-9',
+              'openai-processing-ms', '442',
+              'openai-version', '2020-10-01',
+              'x-ratelimit-limit-requests', '3000',
+              'x-ratelimit-limit-tokens', '250000',
+              'x-ratelimit-remaining-requests', '2999',
+              'x-ratelimit-remaining-tokens', '249984',
+              'x-ratelimit-reset-requests', '20ms',
+              'x-ratelimit-reset-tokens', '3ms',
+              'x-request-id', '7df89d8afe7bf24dc04e2c4dd4962d7f'
+            ])
+
+          const checkTraces = agent
+            .use(traces => {
+              expect(traces[0][0]).to.have.property('name', 'openai.request')
+            })
+
+          const result = await openai.createCompletion({
+            model: 'text-davinci-002',
+            prompt: 'Hello, ',
+            suffix: 'foo',
+            stream: true,
+          })
+
+          await checkTraces
+
+          clock.tick(10 * 1000)
+        })
       })
 
       describe('createEmbedding()', () => {