61 lines
1.2 KiB
TypeScript
61 lines
1.2 KiB
TypeScript
import {
|
|
createCallbacksTransformer,
|
|
readableFromAsyncIterable,
|
|
type AIStreamCallbacksAndOptions,
|
|
} from './ai-stream';
|
|
import { createStreamDataTransformer } from './stream-data';
|
|
|
|
interface GenerateContentResponse {
|
|
candidates?: GenerateContentCandidate[];
|
|
}
|
|
|
|
interface GenerateContentCandidate {
|
|
index: number;
|
|
content: Content;
|
|
}
|
|
|
|
interface Content {
|
|
role: string;
|
|
parts: Part[];
|
|
}
|
|
|
|
type Part = TextPart | InlineDataPart;
|
|
|
|
interface InlineDataPart {
|
|
text?: never;
|
|
}
|
|
|
|
interface TextPart {
|
|
text: string;
|
|
inlineData?: never;
|
|
}
|
|
|
|
async function* streamable(response: {
|
|
stream: AsyncIterable<GenerateContentResponse>;
|
|
}) {
|
|
for await (const chunk of response.stream) {
|
|
const parts = chunk.candidates?.[0]?.content?.parts;
|
|
|
|
if (parts === undefined) {
|
|
continue;
|
|
}
|
|
|
|
const firstPart = parts[0];
|
|
|
|
if (typeof firstPart.text === 'string') {
|
|
yield firstPart.text;
|
|
}
|
|
}
|
|
}
|
|
|
|
export function GoogleGenerativeAIStream(
|
|
response: {
|
|
stream: AsyncIterable<GenerateContentResponse>;
|
|
},
|
|
cb?: AIStreamCallbacksAndOptions,
|
|
): ReadableStream {
|
|
return readableFromAsyncIterable(streamable(response))
|
|
.pipeThrough(createCallbacksTransformer(cb))
|
|
.pipeThrough(createStreamDataTransformer());
|
|
}
|