added streaming for google models

This commit is contained in:
based
2023-12-14 21:03:41 +10:00
parent 3e82a7d439
commit ca87f29771
3 changed files with 86 additions and 40 deletions

View File

@ -1452,7 +1452,7 @@ async function sendOpenAIRequest(type, messages, signal) {
const isQuiet = type === 'quiet';
const isImpersonate = type === 'impersonate';
const isContinue = type === 'continue';
const stream = oai_settings.stream_openai && !isQuiet && !isScale && !isAI21 && !isGoogle;
const stream = oai_settings.stream_openai && !isQuiet && !isScale && !isAI21;
if (isTextCompletion && isOpenRouter) {
messages = convertChatCompletionToInstruct(messages, type);
@ -1571,23 +1571,26 @@ async function sendOpenAIRequest(type, messages, signal) {
tryParseStreamingError(response, await response.text());
throw new Error(`Got response status ${response.status}`);
}
if (stream) {
const eventStream = new EventSourceStream();
response.body.pipeThrough(eventStream);
const reader = eventStream.readable.getReader();
let reader;
let isSSEStream = oai_settings.chat_completion_source !== chat_completion_sources.MAKERSUITE;
if (isSSEStream) {
const eventStream = new EventSourceStream();
response.body.pipeThrough(eventStream);
reader = eventStream.readable.getReader();
} else {
reader = response.body.getReader();
}
return async function* streamData() {
let text = '';
let utf8Decoder = new TextDecoder();
while (true) {
const { done, value } = await reader.read();
if (done) return;
if (value.data === '[DONE]') return;
tryParseStreamingError(response, value.data);
// the first and last messages are undefined, protect against that
text += getStreamingReply(JSON.parse(value.data));
const rawData = isSSEStream ? value.data : utf8Decoder.decode(value, { stream: true });
if (isSSEStream && rawData === '[DONE]') return;
tryParseStreamingError(response, rawData);
text += getStreamingReply(JSON.parse(rawData));
yield { text, swipes: [] };
}
};
@ -1609,6 +1612,8 @@ async function sendOpenAIRequest(type, messages, signal) {
function getStreamingReply(data) {
if (oai_settings.chat_completion_source == chat_completion_sources.CLAUDE) {
return data?.completion || '';
} else if (oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE) {
return data?.candidates[0].content.parts[0].text || '';
} else {
return data.choices[0]?.delta?.content || data.choices[0]?.message?.content || data.choices[0]?.text || '';
}