ThunderAI/js/workers/model-worker-openai.js

220 lines
8.6 KiB
JavaScript

/*
* ThunderAI [https://micz.it/thunderbird-addon-thunderai/]
* Copyright (C) 2024 - 2025 Mic (m@micz.it)
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*
*
* This file contains a modified version of the code from the project at https://github.com/boxabirds/chatgpt-frontend-nobuild
* The original code has been released under the Apache License, Version 2.0.
*/
import { OpenAI } from '../api/openai.js';
import { taLogger } from '../mzta-logger.js';
let chatgpt_api_key = null;
let chatgpt_model = '';
let openai = null;
let stopStreaming = false;
let i18nStrings = null;
let do_debug = false;
let taLog = null;
let conversationHistory = [];
let assistantResponseAccumulator = '';
// Token batching configuration - optimized for better performance
const TOKEN_BATCH_SIZE = 25; // Send tokens in batches of 25 characters
const TOKEN_BATCH_DELAY = 25; // Max 25ms between batches
const TOKEN_BATCH_TIMEOUT = 100; // Force flush after 100ms regardless of size
let tokenBatch = '';
let batchTimer = null;
let timeoutTimer = null;
let lastBatchTime = 0;
let batchStartTime = 0;
// Function to send batched tokens
function sendTokenBatch(force = false, reason = 'unknown') {
if (tokenBatch && (force || tokenBatch.length >= TOKEN_BATCH_SIZE || performance.now() - lastBatchTime >= TOKEN_BATCH_DELAY)) {
postMessage({ type: 'tokenBatch', payload: { tokens: tokenBatch } });
// Reset batch state
tokenBatch = '';
lastBatchTime = performance.now();
batchStartTime = 0;
// Clear all timers
if (batchTimer) {
clearTimeout(batchTimer);
batchTimer = null;
}
if (timeoutTimer) {
clearTimeout(timeoutTimer);
timeoutTimer = null;
}
}
}
// Function to add token to batch
function addTokenToBatch(token) {
tokenBatch += token;
// Set batch start time for the first token
if (tokenBatch.length === token.length) {
batchStartTime = performance.now();
}
// Send immediately if batch is full
if (tokenBatch.length >= TOKEN_BATCH_SIZE) {
sendTokenBatch(true, 'size-limit');
} else {
// Set timer to send batch if it's the first token in a new batch
if (tokenBatch.length === token.length && !batchTimer) {
batchTimer = setTimeout(() => sendTokenBatch(true, 'delay-timeout'), TOKEN_BATCH_DELAY);
}
// Set timeout-based flushing if not already set
if (!timeoutTimer) {
timeoutTimer = setTimeout(() => sendTokenBatch(true, 'timeout-flush'), TOKEN_BATCH_TIMEOUT);
}
}
}
self.onmessage = async function(event) {
if (event.data.type === 'init') {
chatgpt_api_key = event.data.chatgpt_api_key;
chatgpt_model = event.data.chatgpt_model;
openai = new OpenAI(chatgpt_api_key, chatgpt_model, event.data.chatgpt_developer_messages, true, event.data.chatgpt_api_store);
do_debug = event.data.do_debug;
i18nStrings = event.data.i18nStrings;
taLog = new taLogger('model-worker-openai', do_debug);
} else if (event.data.type === 'chatMessage') {
conversationHistory.push({ role: 'user', content: event.data.message });
const response = await openai.fetchResponse(conversationHistory); //4096);
postMessage({ type: 'messageSent' });
if (!response.ok) {
let error_message = '';
let errorDetail = '';
if(response.is_exception === true){
error_message = response.error;
}else{
try{
const errorJSON = await response.json();
errorDetail = JSON.stringify(errorJSON);
error_message = errorJSON.error.message;
}catch(e){
error_message = response.statusText;
}
taLog.log("error_message: " + JSON.stringify(error_message));
}
postMessage({ type: 'error', payload: i18nStrings["chatgpt_api_request_failed"] + ": " + response.status + " " + response.statusText + ", Detail: " + error_message + " " + errorDetail });
throw new Error("[ThunderAI] OpenAI ChatGPT API request failed: " + response.status + " " + response.statusText + ", Detail: " + error_message + " " + errorDetail);
}
// Check if this is a streaming or non-streaming response
const isStreaming = response.headers.get('content-type')?.includes('text/event-stream');
if (!isStreaming) {
// Handle non-streaming response
try {
const responseData = await response.json();
const content = responseData.choices[0].message.content;
conversationHistory.push({ role: 'assistant', content: content });
assistantResponseAccumulator = content;
// Send the complete response as a single batch
postMessage({ type: 'tokenBatch', payload: { tokens: content } });
postMessage({ type: 'tokensDone' });
return;
} catch (error) {
console.error(`[ThunderAI Worker Debug] Error processing non-streaming response:`, error);
postMessage({ type: 'error', payload: 'Failed to process non-streaming response: ' + error.message });
return;
}
}
const reader = response.body.getReader();
const decoder = new TextDecoder("utf-8");
let buffer = '';
console.log(`[ThunderAI Worker Debug] Starting to read stream...`);
while (true) {
if (stopStreaming) {
stopStreaming = false;
reader.cancel();
// Send any remaining tokens in the batch
sendTokenBatch(true, 'stream-stop');
conversationHistory.push({ role: 'assistant', content: assistantResponseAccumulator });
assistantResponseAccumulator = '';
postMessage({ type: 'tokensDone' });
break;
}
const { done, value } = await reader.read();
if (done) {
// Send any remaining tokens in the batch
sendTokenBatch(true, 'stream-stop');
conversationHistory.push({ role: 'assistant', content: assistantResponseAccumulator });
assistantResponseAccumulator = '';
postMessage({ type: 'tokensDone' });
break;
}
// lots of low-level OpenAI response parsing stuff
const chunk = decoder.decode(value);
buffer += chunk;
taLog.log("buffer " + buffer);
const lines = buffer.split("\n");
buffer = lines.pop();
let parsedLines = [];
try{
parsedLines = lines
.map((line) => line.replace(/^data: /, "").trim()) // Remove the "data: " prefix
.filter((line) => line !== "" && line !== "[DONE]") // Remove empty lines and "[DONE]"
// .map((line) => JSON.parse(line)); // Parse the JSON string
.map((line) => {
taLog.log("line: " + JSON.stringify(line));
return JSON.parse(line);
});
}catch(e){
taLog.error("Error parsing lines: " + e);
}
for (const parsedLine of parsedLines) {
const { choices } = parsedLine;
const { delta } = choices[0];
const { content } = delta;
// Update the UI with the new content
if (content) {
assistantResponseAccumulator += content;
// Add to batch instead of sending immediately
addTokenToBatch(content);
}
}
}
} else if (event.data.type === 'stop') {
stopStreaming = true;
}
};