108 lines
2.1 KiB
TypeScript
108 lines
2.1 KiB
TypeScript
import { GoogleGenerativeAIStream, StreamingTextResponse, StreamData } from '.';
|
|
import { readAllChunks } from '../tests/utils/mock-client';
|
|
|
|
function simulateGenerativeAIResponse(chunks: any[]) {
|
|
chunks = chunks.slice(); // make a copy
|
|
return {
|
|
stream: {
|
|
[Symbol.asyncIterator]() {
|
|
return {
|
|
next() {
|
|
const chunk = chunks.shift();
|
|
if (chunk) {
|
|
return Promise.resolve({
|
|
value: chunk,
|
|
done: false,
|
|
});
|
|
} else {
|
|
return Promise.resolve({ done: true });
|
|
}
|
|
},
|
|
};
|
|
},
|
|
} as AsyncIterable<any>,
|
|
};
|
|
}
|
|
|
|
export const googleGenerativeAIChunks = [
|
|
{
|
|
candidates: [
|
|
{
|
|
content: {
|
|
parts: [{ text: 'Hello' }],
|
|
},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
candidates: [
|
|
{
|
|
content: {
|
|
parts: [{ text: ',' }],
|
|
},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
candidates: [
|
|
{
|
|
content: {
|
|
parts: [{ text: ' world' }],
|
|
},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
candidates: [
|
|
{
|
|
content: {
|
|
parts: [{ text: '.' }],
|
|
},
|
|
},
|
|
],
|
|
},
|
|
];
|
|
|
|
it('should send text', async () => {
|
|
const data = new StreamData();
|
|
|
|
const aiResponse = simulateGenerativeAIResponse(googleGenerativeAIChunks);
|
|
const stream = GoogleGenerativeAIStream(aiResponse, {
|
|
onFinal() {
|
|
data.close();
|
|
},
|
|
});
|
|
|
|
const response = new StreamingTextResponse(stream, {}, data);
|
|
|
|
expect(await readAllChunks(response)).toEqual([
|
|
'0:"Hello"\n',
|
|
'0:","\n',
|
|
'0:" world"\n',
|
|
'0:"."\n',
|
|
]);
|
|
});
|
|
|
|
it('should send text and data', async () => {
|
|
const data = new StreamData();
|
|
|
|
data.append({ t1: 'v1' });
|
|
|
|
const aiResponse = simulateGenerativeAIResponse(googleGenerativeAIChunks);
|
|
const stream = GoogleGenerativeAIStream(aiResponse, {
|
|
onFinal() {
|
|
data.close();
|
|
},
|
|
});
|
|
|
|
const response = new StreamingTextResponse(stream, {}, data);
|
|
|
|
expect(await readAllChunks(response)).toEqual([
|
|
'2:[{"t1":"v1"}]\n',
|
|
'0:"Hello"\n',
|
|
'0:","\n',
|
|
'0:" world"\n',
|
|
'0:"."\n',
|
|
]);
|
|
});
|