final

Author: dlarocque

packages/vertexai/src/requests/stream-reader.test.ts

@@ -17,6 +17,7 @@
 
 import {
   aggregateResponses,
+  deleteEmptyTextParts,
   getResponseStream,
   processStream
 } from './stream-reader';
@@ -33,6 +34,7 @@ import {
   GenerateContentResponse,
   HarmCategory,
   HarmProbability,
+  Part,
   SafetyRating
 } from '../types';
 
@@ -220,6 +222,23 @@ describe('processStream', () => {
     }
     expect(foundCitationMetadata).to.be.true;
   });
+  it('removes empty text parts', async () => {
+    const fakeResponse = getMockResponseStreaming(
+      'streaming-success-empty-text-part.txt'
+    );
+    const result = processStream(fakeResponse as Response);
+    const aggregatedResponse = await result.response;
+    expect(aggregatedResponse.text()).to.equal('1');
+    expect(aggregatedResponse.candidates?.length).to.equal(1);
+    expect(aggregatedResponse.candidates?.[0].content.parts.length).to.equal(1);
+
+    // The chunk with the empty text part will still go through the stream
+    let numChunks = 0;
+    for await (const _ of result.stream) {
+      numChunks++;
+    }
+    expect(numChunks).to.equal(2);
+  });
 });
 
 describe('aggregateResponses', () => {
@@ -404,3 +423,104 @@ describe('aggregateResponses', () => {
     });
   });
 });
+
+describe('deleteEmptyTextParts', () => {
+  it('removes empty text parts from a single candidate', () => {
+    const parts: Part[] = [
+      {
+        text: ''
+      },
+      {
+        text: 'foo'
+      }
+    ];
+    const generateContentResponse: GenerateContentResponse = {
+      candidates: [
+        {
+          index: 0,
+          content: {
+            role: 'model',
+            parts
+          }
+        }
+      ]
+    };
+
+    deleteEmptyTextParts(generateContentResponse);
+    expect(generateContentResponse.candidates?.[0].content.parts).to.deep.equal(
+      [
+        {
+          text: 'foo'
+        }
+      ]
+    );
+  });
+  it('removes empty text parts from all candidates', () => {
+    const parts: Part[] = [
+      {
+        text: ''
+      },
+      {
+        text: 'foo'
+      }
+    ];
+    const generateContentResponse: GenerateContentResponse = {
+      candidates: [
+        {
+          index: 0,
+          content: {
+            role: 'model',
+            parts
+          }
+        },
+        {
+          index: 1,
+          content: {
+            role: 'model',
+            parts
+          }
+        }
+      ]
+    };
+
+    deleteEmptyTextParts(generateContentResponse);
+    expect(generateContentResponse.candidates?.[0].content.parts).to.deep.equal(
+      [
+        {
+          text: 'foo'
+        }
+      ]
+    );
+    expect(generateContentResponse.candidates?.[1].content.parts).to.deep.equal(
+      [
+        {
+          text: 'foo'
+        }
+      ]
+    );
+  });
+  it('does not remove candidate even if all parts are removed', () => {
+    const parts: Part[] = [
+      {
+        text: ''
+      }
+    ];
+    const generateContentResponse: GenerateContentResponse = {
+      candidates: [
+        {
+          index: 0,
+          content: {
+            role: 'model',
+            parts
+          }
+        }
+      ]
+    };
+
+    deleteEmptyTextParts(generateContentResponse);
+    expect(generateContentResponse.candidates?.length).to.equal(1);
+    expect(generateContentResponse.candidates?.[0].content.parts).to.deep.equal(
+      []
+    );
+  });
+});

packages/vertexai/test-utils/mock-response.ts

@@ -49,6 +49,9 @@ export function getMockResponseStreaming(
   filename: string,
   chunkLength: number = 20
 ): Partial<Response> {
+  if (!(filename in mocksLookup)) {
+    throw Error(`Mock response file not found: '${filename}'`);
+  }
   const fullText = mocksLookup[filename];
 
   return {