Reviewed By: bhamodi Differential Revision: D33331422 fbshipit-source-id: 016e8dcc0c0c7f1fc353a348b54fda0d5e2ddc01
64 lines
1.9 KiB
TypeScript
64 lines
1.9 KiB
TypeScript
/**
|
|
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
*
|
|
* This source code is licensed under the MIT license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*
|
|
* @format
|
|
*/
|
|
|
|
import type {PartialResponse, ResponseInfo} from './types';
|
|
import {Base64} from 'js-base64';
|
|
|
|
export function assembleChunksIfResponseIsComplete(
|
|
partialResponseEntry: PartialResponse | undefined,
|
|
): ResponseInfo | undefined {
|
|
const numChunks = partialResponseEntry?.initialResponse?.totalChunks;
|
|
if (
|
|
!numChunks ||
|
|
!partialResponseEntry?.initialResponse ||
|
|
Object.keys(partialResponseEntry.followupChunks).length + 1 < numChunks
|
|
) {
|
|
// Partial response not yet complete, do nothing.
|
|
return;
|
|
}
|
|
|
|
// Partial response has all required chunks, convert it to a full Response.
|
|
const response: ResponseInfo = partialResponseEntry.initialResponse;
|
|
const allChunks: string[] =
|
|
response.data != null
|
|
? [
|
|
response.data,
|
|
...Object.entries(partialResponseEntry.followupChunks)
|
|
// It's important to parseInt here or it sorts lexicographically
|
|
.sort((a, b) => parseInt(a[0], 10) - parseInt(b[0], 10))
|
|
.map(([_k, v]: [string, string]) => v),
|
|
]
|
|
: [];
|
|
const data = combineBase64Chunks(allChunks);
|
|
|
|
return {
|
|
...response,
|
|
// Currently data is always decoded at render time, so re-encode it to match the single response format.
|
|
data,
|
|
};
|
|
}
|
|
|
|
export function combineBase64Chunks(chunks: string[]): string {
|
|
const byteArray = chunks.map((b64Chunk) => {
|
|
return Base64.toUint8Array(b64Chunk);
|
|
});
|
|
const size = byteArray
|
|
.map((b) => b.byteLength)
|
|
.reduce((prev, curr) => prev + curr, 0);
|
|
|
|
const buffer = new Uint8Array(size);
|
|
let offset = 0;
|
|
for (let i = 0; i < byteArray.length; i++) {
|
|
buffer.set(byteArray[i], offset);
|
|
offset += byteArray[i].byteLength;
|
|
}
|
|
|
|
return Base64.fromUint8Array(buffer);
|
|
}
|