rixty
/
SillyTavern
mirror of https://github.com/SillyTavern/SillyTavern.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use SSE streaming for MakerSuite

This commit is contained in:
Cohee 2024-02-10 02:43:50 +02:00
parent 867c42cb6d
commit 18f84979f2
2 changed files with 9 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
public/scripts/openai.js
View File

@ -1692,24 +1692,17 @@ async function sendOpenAIRequest(type, messages, signal) {
throw new Error(`Got response status ${response.status}`);
}
if (stream) {
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();
}
const eventStream = new EventSourceStream();
response.body.pipeThrough(eventStream);
const reader = eventStream.readable.getReader();
return async function* streamData() {
let text = '';
let utf8Decoder = new TextDecoder();
const swipes = [];
while (true) {
const { done, value } = await reader.read();
if (done) return;
const rawData = isSSEStream ? value.data : utf8Decoder.decode(value, { stream: true });
if (isSSEStream && rawData === '[DONE]') return;
const rawData = value.data;
if (rawData === '[DONE]') return;
tryParseStreamingError(response, rawData);
const parsed = JSON.parse(rawData);
@ -1750,7 +1743,7 @@ 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 || '';
return data?.candidates?.[0]?.content?.parts?.[0]?.text || '';
} else {
return data.choices[0]?.delta?.content || data.choices[0]?.message?.content || data.choices[0]?.text || '';
}

34
src/endpoints/backends/chat-completions.js
View File

@ -267,7 +267,7 @@ async function sendMakerSuiteRequest(request, response) {
? (stream ? 'streamGenerateContent' : 'generateContent')
: (isText ? 'generateText' : 'generateMessage');
const generateResponse = await fetch(`https://generativelanguage.googleapis.com/${apiVersion}/models/${model}:${responseType}?key=${apiKey}`, {
const generateResponse = await fetch(`https://generativelanguage.googleapis.com/${apiVersion}/models/${model}:${responseType}?key=${apiKey}${stream ? '&alt=sse' : ''}`, {
body: JSON.stringify(body),
method: 'POST',
headers: {
@ -279,36 +279,8 @@ async function sendMakerSuiteRequest(request, response) {
// have to do this because of their busted ass streaming endpoint
if (stream) {
try {
let partialData = '';
generateResponse.body.on('data', (data) => {
const chunk = data.toString();
if (chunk.startsWith(',') || chunk.endsWith(',') || chunk.startsWith('[') || chunk.endsWith(']')) {
partialData = chunk.slice(1);
} else {
partialData += chunk;
}
while (true) {
let json;
try {
json = JSON.parse(partialData);
} catch (e) {
break;
}
response.write(JSON.stringify(json));
partialData = '';
}
});
request.socket.on('close', function () {
if (generateResponse.body instanceof Readable) generateResponse.body.destroy();
response.end();
});
generateResponse.body.on('end', () => {
console.log('Streaming request finished');
response.end();
});
// Pipe remote SSE stream to Express response
forwardFetchResponse(generateResponse, response);
} catch (error) {
console.log('Error forwarding streaming response:', error);
if (!response.headersSent) {