From 04eeb59d47429f10e187ffab62d9b4f6e1191958 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 27 Feb 2024 10:21:06 -0500 Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=A0=EF=B8=8F=20chore:=20Abort=20AI=20R?= =?UTF-8?q?equests=20on=20Close=20&=20Remove=20Verbose=20Logs=20for=20Plug?= =?UTF-8?q?ins=20(#1914)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore: remove verbose logging of ChatOpenAI * feat: abort AI requests on request close --- api/app/clients/BaseClient.js | 2 + api/app/clients/llm/createLLM.js | 1 - api/server/controllers/AskController.js | 14 +++ api/server/controllers/EditController.js | 14 +++ api/server/routes/assistants/chat.js | 124 +++++++++++++---------- 5 files changed, 98 insertions(+), 57 deletions(-) diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 6009515f2c..a359ed7193 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -447,6 +447,8 @@ class BaseClient { } const completion = await this.sendCompletion(payload, opts); + this.abortController.requestCompleted = true; + const responseMessage = { messageId: responseMessageId, conversationId, diff --git a/api/app/clients/llm/createLLM.js b/api/app/clients/llm/createLLM.js index c8810b667b..a944d0c32d 100644 --- a/api/app/clients/llm/createLLM.js +++ b/api/app/clients/llm/createLLM.js @@ -67,7 +67,6 @@ function createLLM({ return new ChatOpenAI( { streaming, - verbose: true, credentials, configuration, ...azureOptions, diff --git a/api/server/controllers/AskController.js b/api/server/controllers/AskController.js index 52ed0d30fd..272cb9ce85 100644 --- a/api/server/controllers/AskController.js +++ b/api/server/controllers/AskController.js @@ -92,6 +92,20 @@ const AskController = async (req, res, next, initializeClient, addTitle) => { const { abortController, onStart } = createAbortController(req, res, getAbortData); + res.on('close', () => { + logger.debug('[AskController] Request closed'); + if (!abortController) { + return; + } else if (abortController.signal.aborted) { + return; + } else if (abortController.requestCompleted) { + return; + } + + abortController.abort(); + logger.debug('[AskController] Request aborted on close'); + }); + const messageOptions = { user, parentMessageId, diff --git a/api/server/controllers/EditController.js b/api/server/controllers/EditController.js index 43b82e7193..3cc880fd1c 100644 --- a/api/server/controllers/EditController.js +++ b/api/server/controllers/EditController.js @@ -90,6 +90,20 @@ const EditController = async (req, res, next, initializeClient) => { const { abortController, onStart } = createAbortController(req, res, getAbortData); + res.on('close', () => { + logger.debug('[EditController] Request closed'); + if (!abortController) { + return; + } else if (abortController.signal.aborted) { + return; + } else if (abortController.requestCompleted) { + return; + } + + abortController.abort(); + logger.debug('[EditController] Request aborted on close'); + }); + try { const { client } = await initializeClient({ req, res, endpointOption }); diff --git a/api/server/routes/assistants/chat.js b/api/server/routes/assistants/chat.js index 2df34a9ce3..095446988c 100644 --- a/api/server/routes/assistants/chat.js +++ b/api/server/routes/assistants/chat.js @@ -39,6 +39,7 @@ router.post('/abort', handleAbort()); */ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res) => { logger.debug('[/assistants/chat/] req.body', req.body); + const { text, model, @@ -96,7 +97,73 @@ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res const cache = getLogStores(CacheKeys.ABORT_KEYS); const cacheKey = `${req.user.id}:${conversationId}`; + const handleError = async (error) => { + if (error.message === 'Run cancelled') { + return res.end(); + } + if (error.message === 'Request closed') { + logger.debug('[/assistants/chat/] Request aborted on close'); + } + + logger.error('[/assistants/chat/]', error); + + if (!openai || !thread_id || !run_id) { + return res.status(500).json({ error: 'The Assistant run failed to initialize' }); + } + + try { + await cache.delete(cacheKey); + const cancelledRun = await openai.beta.threads.runs.cancel(thread_id, run_id); + logger.debug('Cancelled run:', cancelledRun); + } catch (error) { + logger.error('[abortRun] Error cancelling run', error); + } + + await sleep(2000); + try { + const run = await openai.beta.threads.runs.retrieve(thread_id, run_id); + await recordUsage({ + ...run.usage, + model: run.model, + user: req.user.id, + conversationId, + }); + } catch (error) { + logger.error('[/assistants/chat/] Error fetching or processing run', error); + } + + try { + const runMessages = await checkMessageGaps({ + openai, + run_id, + thread_id, + conversationId, + latestMessageId: responseMessageId, + }); + + const finalEvent = { + title: 'New Chat', + final: true, + conversation: await getConvo(req.user.id, conversationId), + runMessages, + }; + + if (res.headersSent && finalEvent) { + return sendMessage(res, finalEvent); + } + + res.json(finalEvent); + } catch (error) { + logger.error('[/assistants/chat/] Error finalizing error process', error); + return res.status(500).json({ error: 'The Assistant run failed' }); + } + }; + try { + res.on('close', async () => { + await handleError(new Error('Request closed')); + }); + if (convoId && !_thread_id) { throw new Error('Missing thread_id for existing conversation'); } @@ -318,62 +385,7 @@ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res }); } } catch (error) { - if (error.message === 'Run cancelled') { - return res.end(); - } - - logger.error('[/assistants/chat/]', error); - - if (!openai || !thread_id || !run_id) { - return res.status(500).json({ error: 'The Assistant run failed to initialize' }); - } - - try { - await cache.delete(cacheKey); - const cancelledRun = await openai.beta.threads.runs.cancel(thread_id, run_id); - logger.debug('Cancelled run:', cancelledRun); - } catch (error) { - logger.error('[abortRun] Error cancelling run', error); - } - - await sleep(2000); - try { - const run = await openai.beta.threads.runs.retrieve(thread_id, run_id); - await recordUsage({ - ...run.usage, - model: run.model, - user: req.user.id, - conversationId, - }); - } catch (error) { - logger.error('[/assistants/chat/] Error fetching or processing run', error); - } - - try { - const runMessages = await checkMessageGaps({ - openai, - run_id, - thread_id, - conversationId, - latestMessageId: responseMessageId, - }); - - const finalEvent = { - title: 'New Chat', - final: true, - conversation: await getConvo(req.user.id, conversationId), - runMessages, - }; - - if (res.headersSent && finalEvent) { - return sendMessage(res, finalEvent); - } - - res.json(finalEvent); - } catch (error) { - logger.error('[/assistants/chat/] Error finalizing error process', error); - return res.status(500).json({ error: 'The Assistant run failed' }); - } + await handleError(error); } });