From 8e1b02663bbabd3d5c4175f852b13eb78aaadf4a Mon Sep 17 00:00:00 2001
From: James Sumners <jsumners@newrelic.com>
Date: Mon, 20 Nov 2023 09:53:36 -0500
Subject: [PATCH 1/2] test: Add streaming responses to mock OpenAI sever

---
 test/lib/openai-mock-server.js | 40 ++++++++++++++++++++++++++++++++--
 test/lib/openai-responses.js   | 19 ++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/test/lib/openai-mock-server.js b/test/lib/openai-mock-server.js
index a9237a504e..f99f31976f 100644
--- a/test/lib/openai-mock-server.js
+++ b/test/lib/openai-mock-server.js
@@ -8,6 +8,7 @@
 module.exports = openaiMockServer
 
 const http = require('node:http')
+const { Readable } = require('node:stream')
 const RESPONSES = require('./openai-responses')
 
 /**
@@ -69,8 +70,43 @@ function handler(req, res) {
       res.setHeader(key, value)
     }
     res.statusCode = code
-    res.write(JSON.stringify(body))
-    res.end()
+
+    if (payload.stream === true) {
+      // OpenAI streamed responses are double newline delimited lines that
+      // are prefixed with the string `data: `. The end of the stream is
+      // terminated with a `done: [DONE]` string.
+      const parts = body.split(' ')
+      let i = 0
+      const outStream = new Readable({
+        read() {
+          if (i < parts.length) {
+            const chunk = JSON.stringify({
+              id: 'chatcmpl-8MzOfSMbLxEy70lYAolSwdCzfguQZ',
+              object: 'chat.completion.chunk',
+              // 2023-11-20T09:00:00-05:00
+              created: 1700488800,
+              model: 'gpt-4',
+              choices: [
+                {
+                  index: 0,
+                  finish_reason: null,
+                  delta: { role: 'assistant', content: parts[i] }
+                }
+              ]
+            })
+            this.push(`data: ${chunk}\n\n`)
+            i += 1
+          } else {
+            this.push('data: [DONE]\n\n')
+            this.push(null)
+          }
+        }
+      })
+      outStream.pipe(res)
+    } else {
+      res.write(JSON.stringify(body))
+      res.end()
+    }
   })
 }
 
diff --git a/test/lib/openai-responses.js b/test/lib/openai-responses.js
index ca022ca5c5..6ab3024722 100644
--- a/test/lib/openai-responses.js
+++ b/test/lib/openai-responses.js
@@ -157,3 +157,22 @@ responses.set('You are a mathematician.', {
     usage: { completion_tokens: 11, prompt_tokens: 53, total_tokens: 64 }
   }
 })
+
+responses.set('Streamed response', {
+  headers: {
+    'Content-Type': 'text/event-stream',
+    'openai-model': 'gpt-3.5-turbo-0613',
+    'openai-organization': 'new-relic-nkmd8b',
+    'openai-processing-ms': '1469',
+    'openai-version': '2020-10-01',
+    'x-ratelimit-limit-requests': '200',
+    'x-ratelimit-limit-tokens': '40000',
+    'x-ratelimit-remaining-requests': '199',
+    'x-ratelimit-remaining-tokens': '39940',
+    'x-ratelimit-reset-requests': '7m12s',
+    'x-ratelimit-reset-tokens': '90ms',
+    'x-request-id': '49dbbffbd3c3f4612aa48def69059aad'
+  },
+  code: 200,
+  body: "A streamed response is a way of transmitting data from a server to a client (e.g. from a website to a user's computer or mobile device) in a continuous flow or stream, rather than all at one time. This means the client can start to process the data before all of it has been received, which can improve performance for large amounts of data or slow connections. Streaming is often used for real-time or near-real-time applications like video or audio playback."
+})

From 5eb8eadd27783db9f4b83d7648acd9b7ad277d26 Mon Sep 17 00:00:00 2001
From: James Sumners <jsumners@newrelic.com>
Date: Mon, 20 Nov 2023 11:51:38 -0500
Subject: [PATCH 2/2] address feedback

---
 test/lib/openai-mock-server.js | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/lib/openai-mock-server.js b/test/lib/openai-mock-server.js
index f99f31976f..a679e4dca2 100644
--- a/test/lib/openai-mock-server.js
+++ b/test/lib/openai-mock-server.js
@@ -80,6 +80,7 @@ function handler(req, res) {
       const outStream = new Readable({
         read() {
           if (i < parts.length) {
+            const content = parts.length - 1 === i ? parts[i] : `${parts[i]} `
             const chunk = JSON.stringify({
               id: 'chatcmpl-8MzOfSMbLxEy70lYAolSwdCzfguQZ',
               object: 'chat.completion.chunk',
@@ -90,7 +91,7 @@ function handler(req, res) {
                 {
                   index: 0,
                   finish_reason: null,
-                  delta: { role: 'assistant', content: parts[i] }
+                  delta: { role: 'assistant', content }
                 }
               ]
             })