Skip to content

Commit

Permalink
Filter empty text parts when streaming (#8736)
Browse files Browse the repository at this point in the history
* Filter empty text parts when streaming

* Add changeset

* Remove unused filterEmptyTextParts()

* Move logic into a function

* final

* dont throw if mock response file isn't found

* update responses version to 6

* Only ignore empty text parts in the aggregated response.

* review fixes

* Throw in `aggregateResponses` if `newPart` has no properties
  • Loading branch information
dlarocque authored Feb 12, 2025
1 parent c8e5b3e commit 554c7bd
Show file tree
Hide file tree
Showing 4 changed files with 84 additions and 4 deletions.
5 changes: 5 additions & 0 deletions .changeset/seven-oranges-care.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@firebase/vertexai': patch
---

Filter out empty text parts from streaming responses.
66 changes: 65 additions & 1 deletion packages/vertexai/src/requests/stream-reader.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,10 @@ import {
GenerateContentResponse,
HarmCategory,
HarmProbability,
SafetyRating
SafetyRating,
VertexAIErrorCode
} from '../types';
import { VertexAIError } from '../errors';

use(sinonChai);

Expand Down Expand Up @@ -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', () => {
Expand Down Expand Up @@ -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.'
);
}
});
});
15 changes: 13 additions & 2 deletions packages/vertexai/src/requests/stream-reader.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ async function getResponsePromise(
);
return enhancedResponse;
}

allResponses.push(value);
}
}
Expand Down Expand Up @@ -184,14 +185,24 @@ export function aggregateResponses(
}
const newPart: Partial<Part> = {};
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
Expand Down
2 changes: 1 addition & 1 deletion scripts/update_vertexai_responses.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down

0 comments on commit 554c7bd

Please sign in to comment.