From 708892a79c984738e500b11150d2e18f0705aae3 Mon Sep 17 00:00:00 2001
From: Daniel La Rocque <dlarocque@google.com>
Date: Wed, 29 Jan 2025 12:34:15 -0500
Subject: [PATCH] final

---
 .../src/requests/stream-reader.test.ts        | 120 ++++++++++++++++++
 packages/vertexai/test-utils/mock-response.ts |   3 +
 2 files changed, 123 insertions(+)

diff --git a/packages/vertexai/src/requests/stream-reader.test.ts b/packages/vertexai/src/requests/stream-reader.test.ts
index eea72f9c7a6..b2ab735f62d 100644
--- a/packages/vertexai/src/requests/stream-reader.test.ts
+++ b/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(
+      []
+    );
+  });
+});
diff --git a/packages/vertexai/test-utils/mock-response.ts b/packages/vertexai/test-utils/mock-response.ts
index 8332d9eb36e..747ac4dec7d 100644
--- a/packages/vertexai/test-utils/mock-response.ts
+++ b/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 {