diff --git a/.changeset/seven-oranges-care.md b/.changeset/seven-oranges-care.md new file mode 100644 index 00000000000..41b59355bc7 --- /dev/null +++ b/.changeset/seven-oranges-care.md @@ -0,0 +1,5 @@ +--- +'@firebase/vertexai': patch +--- + +Filter out empty text parts from streaming responses. diff --git a/packages/vertexai/src/requests/stream-reader.test.ts b/packages/vertexai/src/requests/stream-reader.test.ts index eea72f9c7a6..b68c2423066 100644 --- a/packages/vertexai/src/requests/stream-reader.test.ts +++ b/packages/vertexai/src/requests/stream-reader.test.ts @@ -33,8 +33,10 @@ import { GenerateContentResponse, HarmCategory, HarmProbability, - SafetyRating + SafetyRating, + VertexAIErrorCode } from '../types'; +import { VertexAIError } from '../errors'; use(sinonChai); @@ -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', () => { @@ -403,4 +422,49 @@ describe('aggregateResponses', () => { ).to.equal(150); }); }); + + it('throws if a part has no properties', () => { + const responsesToAggregate: GenerateContentResponse[] = [ + { + candidates: [ + { + index: 0, + content: { + role: 'user', + parts: [{} as any] // Empty + }, + finishReason: FinishReason.STOP, + finishMessage: 'something', + safetyRatings: [ + { + category: HarmCategory.HARM_CATEGORY_HARASSMENT, + probability: HarmProbability.NEGLIGIBLE + } as SafetyRating + ] + } + ], + promptFeedback: { + blockReason: BlockReason.SAFETY, + safetyRatings: [ + { + category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, + probability: HarmProbability.LOW + } as SafetyRating + ] + } + } + ]; + + try { + aggregateResponses(responsesToAggregate); + } catch (e) { + expect((e as VertexAIError).code).includes( + VertexAIErrorCode.INVALID_CONTENT + ); + expect((e as VertexAIError).message).to.include( + 'Part should have at least one property, but there are none. This is likely caused ' + + 'by a malformed response from the backend.' + ); + } + }); }); diff --git a/packages/vertexai/src/requests/stream-reader.ts b/packages/vertexai/src/requests/stream-reader.ts index 8162407d90b..5c419d114e0 100644 --- a/packages/vertexai/src/requests/stream-reader.ts +++ b/packages/vertexai/src/requests/stream-reader.ts @@ -62,6 +62,7 @@ async function getResponsePromise( ); return enhancedResponse; } + allResponses.push(value); } } @@ -184,14 +185,24 @@ export function aggregateResponses( } const newPart: Partial = {}; for (const part of candidate.content.parts) { - if (part.text) { + if (part.text !== undefined) { + // The backend can send empty text parts. If these are sent back + // (e.g. in chat history), the backend will respond with an error. + // To prevent this, ignore empty text parts. + if (part.text === '') { + continue; + } newPart.text = part.text; } if (part.functionCall) { newPart.functionCall = part.functionCall; } if (Object.keys(newPart).length === 0) { - newPart.text = ''; + throw new VertexAIError( + VertexAIErrorCode.INVALID_CONTENT, + 'Part should have at least one property, but there are none. This is likely caused ' + + 'by a malformed response from the backend.' + ); } aggregatedResponse.candidates[i].content.parts.push( newPart as Part diff --git a/scripts/update_vertexai_responses.sh b/scripts/update_vertexai_responses.sh index 101eac90d9f..20b9082861e 100755 --- a/scripts/update_vertexai_responses.sh +++ b/scripts/update_vertexai_responses.sh @@ -17,7 +17,7 @@ # This script replaces mock response files for Vertex AI unit tests with a fresh # clone of the shared repository of Vertex AI test data. -RESPONSES_VERSION='v5.*' # The major version of mock responses to use +RESPONSES_VERSION='v6.*' # The major version of mock responses to use REPO_NAME="vertexai-sdk-test-data" REPO_LINK="https://github.com/FirebaseExtended/$REPO_NAME.git"