[Security solution] Bedrock streaming and token tracking

x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts

@@ -276,11 +276,16 @@ describe('OpenAIConnector', () => {
     describe('invokeStream', () => {
       let stream;
       beforeEach(() => {
+        const chunk1 =
+          'data: {"object":"chat.completion.chunk","choices":[{"delta":{"content":"My"}}]}\ndata: {"object":"chat.completion.chunk","choices":[{"delta":{"content":" new"}}]}';
+
         stream = createStreamMock();
-        stream.write(`data: ${JSON.stringify(mockResponse.data)}`);
+        stream.write(chunk1);
+        // .write(`data: ${JSON.stringify(mockResponse.data)}`);
         mockRequest = jest.fn().mockResolvedValue({ ...mockResponse, data: stream.transform });
         // @ts-ignore
         connector.request = mockRequest;
+        stream.complete();
       });
 
       it('the API call is successful with correct request parameters', async () => {
@@ -314,7 +319,53 @@ describe('OpenAIConnector', () => {
           responseBody += data.toString();
         });
         await waitFor(() => {
-          expect(responseBody).toEqual(mockResponseString);
+          expect(responseBody).toEqual('My new');
+        });
+      });
+      it('correctly buffers stream of json lines', async () => {
+        stream = createStreamMock();
+        const chunk1 =
+          'data: {"object":"chat.completion.chunk","choices":[{"delta":{"content":"My"}}]}\ndata: {"object":"chat.completion.chunk","choices":[{"delta":{"content":" new"}}]}';
+        const chunk2 =
+          '\ndata: {"object":"chat.completion.chunk","choices":[{"delta":{"content":" message"}}]}\ndata: [DONE]';
+
+        stream.write(chunk1);
+        stream.write(chunk2);
+        mockRequest = jest.fn().mockResolvedValue({ ...mockResponse, data: stream.transform });
+        // @ts-ignore
+        connector.request = mockRequest;
+
+        const response = await connector.invokeStream(sampleOpenAiBody);
+
+        let responseBody: string = '';
+        response.on('data', (data: string) => {
+          responseBody += data.toString();
+        });
+        await waitFor(() => {
+          expect(responseBody).toEqual('My new message');
+        });
+      });
+      it('correctly buffers partial lines', async () => {
+        stream = createStreamMock();
+        const chunk1 =
+          'data: {"object":"chat.completion.chunk","choices":[{"delta":{"content":"My"}}]}\ndata: {"object":"chat.completion.chunk","choices":[{"delta":{"content":" new"';
+        const chunk2 =
+          '}}]}\ndata: {"object":"chat.completion.chunk","choices":[{"delta":{"content":" message"}}]}\ndata: [DONE]';
+
+        stream.write(chunk1);
+        stream.write(chunk2);
+        mockRequest = jest.fn().mockResolvedValue({ ...mockResponse, data: stream.transform });
+        // @ts-ignore
+        connector.request = mockRequest;
+
+        const response = await connector.invokeStream(sampleOpenAiBody);
+
+        let responseBody: string = '';
+        response.on('data', (data: string) => {
+          responseBody += data.toString();
+        });
+        await waitFor(() => {
+          expect(responseBody).toEqual('My new message');
         });
       });
     });

x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts

@@ -203,6 +203,7 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
       body: JSON.stringify(body),
       stream: true,
     })) as unknown as IncomingMessage;
+
     return res.pipe(new PassThrough()).pipe(transformToString());
   }
 
@@ -236,20 +237,20 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
  */
 const transformToString = () => {
   let lineBuffer: string = '';
+  const decoder = new TextDecoder();
+  const encoder = new TextEncoder();
 
   return new Transform({
     transform(chunk, encoding, callback) {
-      const decoder = new TextDecoder();
-      const encoder = new TextEncoder();
-      const lines = decoder.decode(chunk).split('\n');
+      const chunks = decoder.decode(chunk);
+      const lines = chunks.split('\n');
+      if (lines[lines.length - 1] === '') {
+        lines.pop();
+      }
       lines[0] = lineBuffer + lines[0];
-
       lineBuffer = lines.pop() || '';
-
       const nextChunk = lines
-        // every line starts with "data: ", we remove it and are left with stringified JSON or the string "[DONE]"
         .map((str) => str.substring(6))
-        // filter out empty lines and the "[DONE]" string
         .filter((str) => !!str && str !== '[DONE]')
         .map((line) => {
           const openaiResponse = JSON.parse(line);
@@ -259,5 +260,24 @@ const transformToString = () => {
       const newChunk = encoder.encode(nextChunk);
       callback(null, newChunk);
     },
+    flush(callback) {
+      // Emit an additional chunk with the content of lineBuffer if it has length
+      if (lineBuffer.length > 0) {
+        const nextChunk = [lineBuffer]
+          .map((str) => str.substring(6))
+          .filter((str) => !!str && str !== '[DONE]')
+          .map((line) => {
+            const openaiResponse = JSON.parse(line);
+            return openaiResponse.choices[0]?.delta.content ?? '';
+          })
+          .join('');
+        const newChunk = encoder.encode(nextChunk);
+        callback(null, newChunk);
+        // const additionalChunk = encoder.encode(lineBuffer);
+        // callback(null, additionalChunk);
+      } else {
+        callback();
+      }
+    },
   });
 };