feat: option to regenerate the answer after switching model (#98)

josStorer · Mar 31, 2023 · fc982e4 · fc982e4
1 parent bf71a13
commit fc982e4
Show file tree

Hide file tree

Showing 9 changed files with 75 additions and 82 deletions.
diff --git a/src/_locales/en/main.json b/src/_locales/en/main.json
@@ -86,5 +86,6 @@
   "Confirm": "Confirm",
   "Clear Conversation": "Clear Conversation",
   "Retry": "Retry",
-  "Exceeded maximum context length": "Exceeded maximum context length, please clear the conversation and try again"
+  "Exceeded maximum context length": "Exceeded maximum context length, please clear the conversation and try again",
+  "Regenerate the answer after switching model": "Regenerate the answer after switching model"
 }
diff --git a/src/_locales/zh-hans/main.json b/src/_locales/zh-hans/main.json
@@ -86,5 +86,6 @@
   "Confirm": "确认",
   "Clear Conversation": "清理对话",
   "Retry": "重试",
-  "Exceeded maximum context length": "超出最大上下文长度, 请清理对话并重试"
+  "Exceeded maximum context length": "超出最大上下文长度, 请清理对话并重试",
+  "Regenerate the answer after switching model": "快捷切换模型时自动重新生成回答"
 }
diff --git a/src/_locales/zh-hant/main.json b/src/_locales/zh-hant/main.json
@@ -86,5 +86,6 @@
   "Confirm": "確認",
   "Clear Conversation": "清理對話",
   "Retry": "重試",
-  "Exceeded maximum context length": "超出最大上下文長度, 請清理對話並重試"
+  "Exceeded maximum context length": "超出最大上下文長度, 請清理對話並重試",
+  "Regenerate the answer after switching model": "快捷切換模型時自動重新生成回答"
 }
diff --git a/src/components/ConversationCard/index.jsx b/src/components/ConversationCard/index.jsx
@@ -10,9 +10,10 @@ import FileSaver from 'file-saver'
 import { render } from 'preact'
 import FloatingToolbar from '../FloatingToolbar'
 import { useClampWindowSize } from '../../hooks/use-clamp-window-size'
-import { defaultConfig, getUserConfig, Models } from '../../config/index.mjs'
+import { Models } from '../../config/index.mjs'
 import { useTranslation } from 'react-i18next'
 import DeleteButton from '../DeleteButton'
+import { useConfig } from '../../hooks/use-config.mjs'
 
 const logo = Browser.runtime.getURL('logo.png')
 
@@ -61,11 +62,7 @@ function ConversationCard(props) {
       }
     })(),
   )
-  const [config, setConfig] = useState(defaultConfig)
-
-  useEffect(() => {
-    getUserConfig().then(setConfig)
-  }, [])
+  const config = useConfig()
 
   useEffect(() => {
     if (props.onUpdate) props.onUpdate()
@@ -172,6 +169,20 @@ function ConversationCard(props) {
     }
   }, [conversationItemData])
 
+  const getRetryFn = (session) => () => {
+    updateAnswer(`<p class="gpt-loading">${t('Waiting for response...')}</p>`, false, 'answer')
+    setIsReady(false)
+
+    const newSession = { ...session, isRetry: true }
+    setSession(newSession)
+    try {
+      port.postMessage({ stop: true })
+      port.postMessage({ session: newSession })
+    } catch (e) {
+      updateAnswer(e, false, 'error')
+    }
+  }
+
   return (
     <div className="gpt-inner">
       <div className="gpt-header" style="margin: 15px;">
@@ -203,7 +214,8 @@ function ConversationCard(props) {
             required
             onChange={(e) => {
               const modelName = e.target.value
-              setSession({ ...session, modelName, aiName: t(Models[modelName].desc) })
+              if (config.autoRegenAfterSwitchModel)
+                getRetryFn({ ...session, modelName, aiName: t(Models[modelName].desc) })()
             }}
           >
             {Object.entries(Models).map(([key, model]) => {
@@ -298,27 +310,7 @@ function ConversationCard(props) {
             session={session}
             done={data.done}
             port={port}
-            onRetry={
-              idx === conversationItemData.length - 1
-                ? () => {
-                    updateAnswer(
-                      `<p class="gpt-loading">${t('Waiting for response...')}</p>`,
-                      false,
-                      'answer',
-                    )
-                    setIsReady(false)
-
-                    const newSession = { ...session, isRetry: true }
-                    setSession(newSession)
-                    try {
-                      port.postMessage({ stop: true })
-                      port.postMessage({ session: newSession })
-                    } catch (e) {
-                      updateAnswer(e, false, 'error')
-                    }
-                  }
-                : null
-            }
+            onRetry={idx === conversationItemData.length - 1 ? getRetryFn(session) : null}
           />
         ))}
       </div>

diff --git a/src/components/DecisionCard/index.jsx b/src/components/DecisionCard/index.jsx
@@ -2,41 +2,20 @@ import { LightBulbIcon, SearchIcon } from '@primer/octicons-react'
 import { useState, useEffect } from 'react'
 import PropTypes from 'prop-types'
 import ConversationCard from '../ConversationCard'
-import { defaultConfig, getUserConfig } from '../../config/index.mjs'
-import Browser from 'webextension-polyfill'
 import { getPossibleElementByQuerySelector, endsWithQuestionMark } from '../../utils'
 import { useTranslation } from 'react-i18next'
+import { useConfig } from '../../hooks/use-config.mjs'
 
 function DecisionCard(props) {
   const { t } = useTranslation()
   const [triggered, setTriggered] = useState(false)
-  const [config, setConfig] = useState(defaultConfig)
   const [render, setRender] = useState(false)
+  const config = useConfig(() => {
+    setRender(true)
+  })
 
   const question = props.question
 
-  useEffect(() => {
-    getUserConfig().then((config) => {
-      setConfig(config)
-      setRender(true)
-    })
-  }, [])
-
-  useEffect(() => {
-    const listener = (changes) => {
-      const changedItems = Object.keys(changes)
-      let newConfig = {}
-      for (const key of changedItems) {
-        newConfig[key] = changes[key].newValue
-      }
-      setConfig({ ...config, ...newConfig })
-    }
-    Browser.storage.local.onChanged.addListener(listener)
-    return () => {
-      Browser.storage.local.onChanged.removeListener(listener)
-    }
-  }, [config])
-
   const updatePosition = () => {
     if (!render) return
 

diff --git a/src/components/FloatingToolbar/index.jsx b/src/components/FloatingToolbar/index.jsx
@@ -2,12 +2,12 @@ import Browser from 'webextension-polyfill'
 import { cloneElement, useEffect, useState } from 'react'
 import ConversationCard from '../ConversationCard'
 import PropTypes from 'prop-types'
-import { defaultConfig, getUserConfig } from '../../config/index.mjs'
 import { config as toolsConfig } from '../../content-script/selection-tools'
 import { getClientPosition, isMobile, setElementPositionInViewport } from '../../utils'
 import Draggable from 'react-draggable'
 import { useClampWindowSize } from '../../hooks/use-clamp-window-size'
 import { useTranslation } from 'react-i18next'
+import { useConfig } from '../../hooks/use-config.mjs'
 
 const logo = Browser.runtime.getURL('logo.png')
 
@@ -16,36 +16,15 @@ function FloatingToolbar(props) {
   const [selection, setSelection] = useState(props.selection)
   const [prompt, setPrompt] = useState(props.prompt)
   const [triggered, setTriggered] = useState(props.triggered)
-  const [config, setConfig] = useState(defaultConfig)
   const [render, setRender] = useState(false)
   const [closeable, setCloseable] = useState(props.closeable)
   const [position, setPosition] = useState(getClientPosition(props.container))
   const [virtualPosition, setVirtualPosition] = useState({ x: 0, y: 0 })
   const windowSize = useClampWindowSize([750, 1500], [0, Infinity])
-
-  useEffect(() => {
-    getUserConfig().then((config) => {
-      setConfig(config)
-      setRender(true)
-
-      if (!triggered) props.container.style.position = 'absolute'
-    })
-  }, [])
-
-  useEffect(() => {
-    const listener = (changes) => {
-      const changedItems = Object.keys(changes)
-      let newConfig = {}
-      for (const key of changedItems) {
-        newConfig[key] = changes[key].newValue
-      }
-      setConfig({ ...config, ...newConfig })
-    }
-    Browser.storage.local.onChanged.addListener(listener)
-    return () => {
-      Browser.storage.local.onChanged.removeListener(listener)
-    }
-  }, [config])
+  const config = useConfig(() => {
+    setRender(true)
+    if (!triggered) props.container.style.position = 'absolute'
+  })
 
   useEffect(() => {
     if (isMobile()) {

diff --git a/src/config/index.mjs b/src/config/index.mjs
@@ -66,6 +66,7 @@ export const defaultConfig = {
   preferredLanguage: getNavigatorLanguage(),
   insertAtTop: isMobile(),
   lockWhenAnswer: false,
+  autoRegenAfterSwitchModel: false,
   customModelApiUrl: 'http://localhost:8000/chat/completions',
   customModelName: 'chatglm-6b-int4',
 

diff --git a/src/hooks/use-config.mjs b/src/hooks/use-config.mjs
@@ -0,0 +1,28 @@
+import { useEffect, useState } from 'react'
+import { defaultConfig, getUserConfig } from '../config/index.mjs'
+import Browser from 'webextension-polyfill'
+
+export function useConfig(initFn) {
+  const [config, setConfig] = useState(defaultConfig)
+  useEffect(() => {
+    getUserConfig().then((config) => {
+      setConfig(config)
+      if (initFn) initFn()
+    })
+  }, [])
+  useEffect(() => {
+    const listener = (changes) => {
+      const changedItems = Object.keys(changes)
+      let newConfig = {}
+      for (const key of changedItems) {
+        newConfig[key] = changes[key].newValue
+      }
+      setConfig({ ...config, ...newConfig })
+    }
+    Browser.storage.local.onChanged.addListener(listener)
+    return () => {
+      Browser.storage.local.onChanged.removeListener(listener)
+    }
+  }, [config])
+  return config
+}
diff --git a/src/popup/Popup.jsx b/src/popup/Popup.jsx
@@ -235,6 +235,17 @@ function GeneralPart({ config, updateConfig }) {
         />
         {t('Lock scrollbar while answering')}
       </label>
+      <label>
+        <input
+          type="checkbox"
+          checked={config.autoRegenAfterSwitchModel}
+          onChange={(e) => {
+            const checked = e.target.checked
+            updateConfig({ autoRegenAfterSwitchModel: checked })
+          }}
+        />
+        {t('Regenerate the answer after switching model')}
+      </label>
       <br />
     </>
   )