From a84988363779c26453d40e68826764071e042ef8 Mon Sep 17 00:00:00 2001
From: Dmitri Nasonov <dmitri.nasonov.b@gmail.com>
Date: Fri, 13 Sep 2024 16:14:37 +0100
Subject: [PATCH] Added support for non-stream models (o1-)

---
 package.json                                  |   2 +-
 public/locales/en/api.json                    |  13 +-
 public/locales/en/main.json                   |   1 +
 .../StopGeneratingButton.tsx                  |  21 +-
 src/constants/modelLoader.ts                  |   4 +-
 src/hooks/useSubmit.ts                        | 195 +++++++++++-------
 src/utils/modelReader.ts                      |  32 ++-
 7 files changed, 187 insertions(+), 81 deletions(-)

diff --git a/package.json b/package.json
index 9921b4b98..014bb0c6a 100644
--- a/package.json
+++ b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "better-chatgpt",
   "private": true,
-  "version": "1.8.4",
+  "version": "1.9.0",
   "type": "module",
   "homepage": "./",
   "main": "electron/index.cjs",
diff --git a/public/locales/en/api.json b/public/locales/en/api.json
index 2485db13f..92a247c0a 100644
--- a/public/locales/en/api.json
+++ b/public/locales/en/api.json
@@ -15,5 +15,14 @@
   },
   "customEndpoint": "Use custom API endpoint",
   "advancedConfig": "View advanced API configuration <0>here</0>",
-  "noApiKeyWarning": "No API key supplied! Please check your API settings."
-}
+  "noApiKeyWarning": "No API key supplied! Please check your API settings.",
+  "errors": {
+    "errorGeneratingTitle": "Error generating title!",
+    "noMessagesSubmitted": "No messages submitted!",
+    "messageExceedMaxToken": "Message exceed max token!",
+    "failedToRetrieveData": "Failed to retrieve valid chat completion data.",
+    "streamLocked": "Oops, the stream is locked right now. Please try again",
+    "cancelledByUser": "Cancelled by user",
+    "generationCompleted": "Generation completed"
+  }
+}
\ No newline at end of file
diff --git a/public/locales/en/main.json b/public/locales/en/main.json
index 1c9ad0a42..97d3da69f 100644
--- a/public/locales/en/main.json
+++ b/public/locales/en/main.json
@@ -48,6 +48,7 @@
   "cloned": "Cloned",
   "enterToSubmit": "Enter to submit",
   "submitPlaceholder": "Type a message or click [/] for prompts...",
+  "stopNonStreamGenerationWarning": "Stopping generation is not recommended for non-stream models. Only use it if UI is stuck. Do you want to proceed?",
   "reduceMessagesWarning": "Reducing messages may result in data loss. It is recommended to download the chat in JSON format if you care about the data. Do you want to proceed?",
   "reduceMessagesFailedImportWarning": "Full import failed as the data hit the maximum storage limit. Import as much as possible?",
   "partialImportWarning": "Full import failed as not all of the expected messages were imported: {{message}}. Would you like to import anyway?",
diff --git a/src/components/StopGeneratingButton/StopGeneratingButton.tsx b/src/components/StopGeneratingButton/StopGeneratingButton.tsx
index 295c0fd02..1ffbd9690 100644
--- a/src/components/StopGeneratingButton/StopGeneratingButton.tsx
+++ b/src/components/StopGeneratingButton/StopGeneratingButton.tsx
@@ -1,14 +1,31 @@
-import React from 'react';
+import React, { useEffect } from 'react';
 import useStore from '@store/store';
+import { useTranslation } from 'react-i18next';
+import { modelStreamSupport } from '@constants/modelLoader';
 
 const StopGeneratingButton = () => {
+  const { t } = useTranslation();
   const setGenerating = useStore((state) => state.setGenerating);
   const generating = useStore((state) => state.generating);
 
+  const currentModel = useStore((state) =>
+    state.chats ? state.chats[state.currentChatIndex].config.model : ''
+  );
+  const handleGeneratingStop = () => {
+    if (modelStreamSupport[currentModel]) {
+      setGenerating(false);
+    } else {
+      const confirmMessage = t('stopNonStreamGenerationWarning');
+      if (window.confirm(confirmMessage)) {
+        setGenerating(false);
+      }
+    }
+  };
+
   return generating ? (
     <div
       className='absolute bottom-6 left-0 right-0 m-auto flex md:w-full md:m-auto gap-0 md:gap-2 justify-center'
-      onClick={() => setGenerating(false)}
+      onClick={() => handleGeneratingStop()}
     >
       <button
         className='btn relative btn-neutral border-0 md:border'
diff --git a/src/constants/modelLoader.ts b/src/constants/modelLoader.ts
index 444278864..a386c0dbb 100644
--- a/src/constants/modelLoader.ts
+++ b/src/constants/modelLoader.ts
@@ -5,6 +5,7 @@ let modelOptions: string[] = [];
 let modelMaxToken: { [key: string]: number } = {};
 let modelCost: ModelCost = {};
 let modelTypes: { [key: string]: string } = {};
+let modelStreamSupport: { [key: string]: boolean } = {};
 
 const initializeModels = async () => {
   const models = await loadModels();
@@ -12,8 +13,9 @@ const initializeModels = async () => {
   modelMaxToken = models.modelMaxToken;
   modelCost = models.modelCost;
   modelTypes = models.modelTypes;
+  modelStreamSupport = models.modelStreamSupport;
 };
 
 initializeModels();
 
-export { modelOptions, modelMaxToken, modelCost, modelTypes };
\ No newline at end of file
+export { modelOptions, modelMaxToken, modelCost, modelTypes, modelStreamSupport };
\ No newline at end of file
diff --git a/src/hooks/useSubmit.ts b/src/hooks/useSubmit.ts
index 2bbc236de..2ab055fd6 100644
--- a/src/hooks/useSubmit.ts
+++ b/src/hooks/useSubmit.ts
@@ -11,6 +11,7 @@ import { parseEventSource } from '@api/helper';
 import { limitMessageTokens, updateTotalTokenUsed } from '@utils/messageUtils';
 import { _defaultChatConfig } from '@constants/chat';
 import { officialAPIEndpoint } from '@constants/auth';
+import { modelStreamSupport } from '@constants/modelLoader';
 
 const useSubmit = () => {
   const { t, i18n } = useTranslation('api');
@@ -22,7 +23,7 @@ const useSubmit = () => {
   const generating = useStore((state) => state.generating);
   const currentChatIndex = useStore((state) => state.currentChatIndex);
   const setChats = useStore((state) => state.setChats);
-  
+
   const generateTitle = async (
     message: MessageInterface[],
     modelConfig: ConfigInterface
@@ -35,8 +36,8 @@ const useSubmit = () => {
           throw new Error(t('noApiKeyWarning') as string);
         }
         const titleChatConfig = {
-          ..._defaultChatConfig,  // Spread the original config
-          model:  useStore.getState().titleModel ?? _defaultChatConfig.model,           // Override the model property
+          ..._defaultChatConfig, // Spread the original config
+          model: useStore.getState().titleModel ?? _defaultChatConfig.model, // Override the model property
         };
         // other endpoints
         data = await getChatCompletion(
@@ -49,8 +50,8 @@ const useSubmit = () => {
         );
       } else if (apiKey) {
         const titleChatConfig = {
-          ...modelConfig,  // Spread the original config
-          model: useStore.getState().titleModel ?? modelConfig.model,           // Override the model property
+          ...modelConfig, // Spread the original config
+          model: useStore.getState().titleModel ?? modelConfig.model, // Override the model property
         };
         // own apikey
         data = await getChatCompletion(
@@ -63,7 +64,9 @@ const useSubmit = () => {
         );
       }
     } catch (error: unknown) {
-      throw new Error(`Error generating title!\n${(error as Error).message}`);
+      throw new Error(
+        `${t('errors.errorGeneratingTitle')}\n${(error as Error).message}`
+      );
     }
     return data.choices[0].message.content;
   };
@@ -88,91 +91,139 @@ const useSubmit = () => {
     setGenerating(true);
 
     try {
+      const isStreamSupported =
+        modelStreamSupport[chats[currentChatIndex].config.model];
+      let data;
       let stream;
       if (chats[currentChatIndex].messages.length === 0)
-        throw new Error('No messages submitted!');
+        throw new Error(t('errors.noMessagesSubmitted') as string);
 
       const messages = limitMessageTokens(
         chats[currentChatIndex].messages,
         chats[currentChatIndex].config.max_tokens,
         chats[currentChatIndex].config.model
       );
-      if (messages.length === 0) throw new Error('Message exceed max token!');
+      if (messages.length === 0)
+        throw new Error(t('errors.messageExceedMaxToken') as string);
+      if (!isStreamSupported) {
+        if (!apiKey || apiKey.length === 0) {
+          // official endpoint
+          if (apiEndpoint === officialAPIEndpoint) {
+            throw new Error(t('noApiKeyWarning') as string);
+          }
+          // other endpoints
+          data = await getChatCompletion(
+            useStore.getState().apiEndpoint,
+            messages,
+            chats[currentChatIndex].config,
+            undefined,
+            undefined,
+            useStore.getState().apiVersion
+          );
+        } else if (apiKey) {
+          data = await getChatCompletion(
+            useStore.getState().apiEndpoint,
+            messages,
+            chats[currentChatIndex].config,
+            apiKey,
+            undefined,
+            useStore.getState().apiVersion
+          );
+        }
 
-      // no api key (free)
-      if (!apiKey || apiKey.length === 0) {
-        // official endpoint
-        if (apiEndpoint === officialAPIEndpoint) {
-          throw new Error(t('noApiKeyWarning') as string);
+        if (
+          !data ||
+          !data.choices ||
+          !data.choices[0] ||
+          !data.choices[0].message ||
+          !data.choices[0].message.content
+        ) {
+          throw new Error(t('errors.failedToRetrieveData') as string);
         }
 
-        // other endpoints
-        stream = await getChatCompletionStream(
-          useStore.getState().apiEndpoint,
-          messages,
-          chats[currentChatIndex].config,
-          undefined,
-          undefined,
-          useStore.getState().apiVersion
-        );
-      } else if (apiKey) {
-        // own apikey
-        stream = await getChatCompletionStream(
-          useStore.getState().apiEndpoint,
-          messages,
-          chats[currentChatIndex].config,
-          apiKey,
-          undefined,
-          useStore.getState().apiVersion
+        const updatedChats: ChatInterface[] = JSON.parse(
+          JSON.stringify(useStore.getState().chats)
         );
-      }
+        const updatedMessages = updatedChats[currentChatIndex].messages;
+        (
+          updatedMessages[updatedMessages.length - 1]
+            .content[0] as TextContentInterface
+        ).text += data.choices[0].message.content;
+        setChats(updatedChats);
+      } else {
+        // no api key (free)
+        if (!apiKey || apiKey.length === 0) {
+          // official endpoint
+          if (apiEndpoint === officialAPIEndpoint) {
+            throw new Error(t('noApiKeyWarning') as string);
+          }
 
-      if (stream) {
-        if (stream.locked)
-          throw new Error(
-            'Oops, the stream is locked right now. Please try again'
+          // other endpoints
+          stream = await getChatCompletionStream(
+            useStore.getState().apiEndpoint,
+            messages,
+            chats[currentChatIndex].config,
+            undefined,
+            undefined,
+            useStore.getState().apiVersion
           );
-        const reader = stream.getReader();
-        let reading = true;
-        let partial = '';
-        while (reading && useStore.getState().generating) {
-          const { done, value } = await reader.read();
-          const result = parseEventSource(
-            partial + new TextDecoder().decode(value)
+        } else if (apiKey) {
+          // own apikey
+          stream = await getChatCompletionStream(
+            useStore.getState().apiEndpoint,
+            messages,
+            chats[currentChatIndex].config,
+            apiKey,
+            undefined,
+            useStore.getState().apiVersion
           );
-          partial = '';
+        }
 
-          if (result === '[DONE]' || done) {
-            reading = false;
-          } else {
-            const resultString = result.reduce((output: string, curr) => {
-              if (typeof curr === 'string') {
-                partial += curr;
-              } else {
-                const content = curr.choices[0]?.delta?.content ?? null;
-                if (content) output += content;
-              }
-              return output;
-            }, '');
-
-            const updatedChats: ChatInterface[] = JSON.parse(
-              JSON.stringify(useStore.getState().chats)
+        if (stream) {
+          if (stream.locked)
+            throw new Error(t('errors.streamLocked') as string);
+          const reader = stream.getReader();
+          let reading = true;
+          let partial = '';
+          while (reading && useStore.getState().generating) {
+            const { done, value } = await reader.read();
+            const result = parseEventSource(
+              partial + new TextDecoder().decode(value)
             );
-            const updatedMessages = updatedChats[currentChatIndex].messages;
-            (
-              updatedMessages[updatedMessages.length - 1]
-                .content[0] as TextContentInterface
-            ).text += resultString;
-            setChats(updatedChats);
+            partial = '';
+
+            if (result === '[DONE]' || done) {
+              reading = false;
+            } else {
+              const resultString = result.reduce((output: string, curr) => {
+                if (typeof curr === 'string') {
+                  partial += curr;
+                } else {
+                  const content = curr.choices[0]?.delta?.content ?? null;
+                  if (content) output += content;
+                }
+                return output;
+              }, '');
+
+              const updatedChats: ChatInterface[] = JSON.parse(
+                JSON.stringify(useStore.getState().chats)
+              );
+              const updatedMessages = updatedChats[currentChatIndex].messages;
+              (
+                updatedMessages[updatedMessages.length - 1]
+                  .content[0] as TextContentInterface
+              ).text += resultString;
+              setChats(updatedChats);
+            }
           }
+          if (useStore.getState().generating) {
+            reader.cancel(t('errors.cancelledByUser') as string);
+          } else {
+            reader.cancel(t('errors.generationCompleted') as string);
+          }
+          reader.releaseLock();
+          stream.cancel();
         }
-        if (useStore.getState().generating) {
-          reader.cancel('Cancelled by user');
-        } else {
-          reader.cancel('Generation completed');
-        }
-        reader.releaseLock();
-        stream.cancel();
       }
 
       // update tokens used in chatting
diff --git a/src/utils/modelReader.ts b/src/utils/modelReader.ts
index 90d6eeb62..44c2635d4 100644
--- a/src/utils/modelReader.ts
+++ b/src/utils/modelReader.ts
@@ -22,6 +22,8 @@ interface ModelData {
     is_moderated: boolean;
   };
   per_request_limits: any;
+  // TODO: Remove workaround once openrouter supports it;
+  is_stream_supported: boolean; // custom field until better workaround or openrouter proper support
 }
 
 interface ModelsJson {
@@ -35,6 +37,7 @@ export const loadModels = async (): Promise<{
   modelMaxToken: { [key: string]: number };
   modelCost: ModelCost;
   modelTypes: { [key: string]: string };
+  modelStreamSupport: { [key: string]: boolean };
 }> => {
   const response = await fetch(modelsJsonUrl);
   const modelsJson: ModelsJson = await response.json();
@@ -43,6 +46,7 @@ export const loadModels = async (): Promise<{
   const modelMaxToken: { [key: string]: number } = {};
   const modelCost: ModelCost = {};
   const modelTypes: { [key: string]: string } = {};
+  const modelStreamSupport: { [key: string]: boolean } = {};
 
   // Prepend specific models
   const specificModels = [
@@ -56,6 +60,7 @@ export const loadModels = async (): Promise<{
         request: '0',
       },
       type: 'text',
+      is_stream_supported: true,
     },
     {
       id: 'gpt-4-turbo-2024-04-09',
@@ -67,6 +72,7 @@ export const loadModels = async (): Promise<{
         request: '0',
       },
       type: 'text',
+      is_stream_supported: false,
     },
   ];
 
@@ -79,6 +85,7 @@ export const loadModels = async (): Promise<{
       image: { price: parseFloat(model.pricing.image), unit: 1 },
     };
     modelTypes[model.id] = model.type;
+    modelStreamSupport[model.id] = model.is_stream_supported;
   });
 
   modelsJson.data.forEach((model) => {
@@ -91,6 +98,13 @@ export const loadModels = async (): Promise<{
       image: { price: 0, unit: 1 }, // default for no image models
     };
 
+    // TODO: Remove workaround once openrouter supports it
+    if (modelId.includes('o1-')) {
+      model.is_stream_supported = false;
+    } else {
+      model.is_stream_supported = true;
+    }
+
     // Detect image capabilities
     if (parseFloat(model.pricing.image) > 0) {
       modelTypes[modelId] = 'image';
@@ -101,23 +115,35 @@ export const loadModels = async (): Promise<{
     } else {
       modelTypes[modelId] = 'text';
     }
+    modelStreamSupport[modelId] = model.is_stream_supported;
   });
 
-  // Sort modelOptions to prioritize gpt-4o models at the top, followed by other OpenAI models
+  // Sort modelOptions to prioritize gpt-4o models at the top, followed by o1 models, and then other OpenAI models
   modelOptions.sort((a, b) => {
     const isGpt4oA = a.startsWith('gpt-4o');
     const isGpt4oB = b.startsWith('gpt-4o');
+    const isO1A = a.startsWith('o1-');
+    const isO1B = b.startsWith('o1-');
     const isOpenAIA = a.startsWith('gpt-');
     const isOpenAIB = b.startsWith('gpt-');
 
+    // Prioritize gpt-4o models
     if (isGpt4oA && !isGpt4oB) return -1;
     if (!isGpt4oA && isGpt4oB) return 1;
+
+    // If both are gpt-4o or neither, prioritize o1 models
+    if (isO1A && !isO1B) return -1;
+    if (!isO1A && isO1B) return 1;
+
+    // If both are gpt-4o or o1 or neither, prioritize other OpenAI models
     if (isOpenAIA && !isOpenAIB) return -1;
     if (!isOpenAIA && isOpenAIB) return 1;
+
+    // If both are the same type or neither, maintain original order
     return 0;
   });
 
-  return { modelOptions, modelMaxToken, modelCost, modelTypes };
+  return { modelOptions, modelMaxToken, modelCost, modelTypes, modelStreamSupport };
 };
 
-export type ModelOptions = string;
\ No newline at end of file
+export type ModelOptions = string;