WIP, custom API Modes

josStorer · josStorer · commit 8706737a5cff · 2024-08-06T23:10:37.000+08:00
diff --git a/src/components/ConversationCard/index.jsx b/src/components/ConversationCard/index.jsx
@@ -1,11 +1,13 @@
-import { memo, useEffect, useMemo, useRef, useState } from 'react'
+import { memo, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react'
 import PropTypes from 'prop-types'
 import Browser from 'webextension-polyfill'
 import InputBox from '../InputBox'
 import ConversationItem from '../ConversationItem'
 import {
+  apiModeToModelName,
   createElementAtPosition,
-  getApiModesStringArrayFromConfig,
+  getApiModesFromConfig,
+  isApiModeSelected,
   isFirefox,
   isMobile,
   isSafari,
@@ -61,33 +63,33 @@ function ConversationCard(props) {
   const [completeDraggable, setCompleteDraggable] = useState(false)
   // `.some` for multi mode models. e.g. bingFree4-balanced
   const useForegroundFetch = bingWebModelKeys.some((n) => session.modelName.includes(n))
+  const [apiModes, setApiModes] = useState([])
 
   /**
    * @type {[ConversationItemData[], (conversationItemData: ConversationItemData[]) => void]}
    */
-  const [conversationItemData, setConversationItemData] = useState(
-    (() => {
-      if (session.conversationRecords.length === 0)
-        if (props.question && triggered)
-          return [
-            new ConversationItemData(
-              'answer',
-              `<p class="gpt-loading">${t(`Waiting for response...`)}</p>`,
-            ),
-          ]
-        else return []
-      else {
-        const ret = []
-        for (const record of session.conversationRecords) {
-          ret.push(new ConversationItemData('question', record.question, true))
-          ret.push(new ConversationItemData('answer', record.answer, true))
-        }
-        return ret
-      }
-    })(),
-  )
+  const [conversationItemData, setConversationItemData] = useState([])
   const config = useConfig()
 
+  useLayoutEffect(() => {
+    if (session.conversationRecords.length === 0) {
+      if (props.question && triggered)
+        setConversationItemData([
+          new ConversationItemData(
+            'answer',
+            `<p class="gpt-loading">${t(`Waiting for response...`)}</p>`,
+          ),
+        ])
+    } else {
+      const ret = []
+      for (const record of session.conversationRecords) {
+        ret.push(new ConversationItemData('question', record.question, true))
+        ret.push(new ConversationItemData('answer', record.answer, true))
+      }
+      setConversationItemData(ret)
+    }
+  }, [])
+
   useEffect(() => {
     setCompleteDraggable(!isSafari() && !isFirefox() && !isMobile())
   }, [])
@@ -118,6 +120,15 @@ function ConversationCard(props) {
     }
   }, [props.question, triggered]) // usually only triggered once
 
+  useLayoutEffect(() => {
+    setApiModes(getApiModesFromConfig(config, true))
+  }, [
+    config.activeApiModes,
+    config.customApiModes,
+    config.azureDeploymentName,
+    config.ollamaModelName,
+  ])
+
   /**
    * @param {string} value
    * @param {boolean} appended
@@ -369,25 +380,32 @@ function ConversationCard(props) {
             className="normal-button"
             required
             onChange={(e) => {
-              const modelName = e.target.value
-              const newSession = { ...session, modelName, aiName: modelNameToDesc(modelName, t) }
+              const apiMode = apiModes[e.target.value]
+              const modelName = apiModeToModelName(apiMode)
+              const newSession = {
+                ...session,
+                modelName,
+                apiMode,
+                aiName: modelNameToDesc(modelName, t),
+              }
               if (config.autoRegenAfterSwitchModel && conversationItemData.length > 0)
                 getRetryFn(newSession)()
               else setSession(newSession)
             }}
           >
-            {getApiModesStringArrayFromConfig(config, true).map((modelName) => {
+            {apiModes.map((apiMode, index) => {
+              const modelName = apiModeToModelName(apiMode)
               const desc = modelNameToDesc(modelName, t)
-              if (desc)
+              if (desc) {
+                let selected
+                if (isApiModeSelected(apiMode, session)) selected = true
+                else selected = session.modelName === modelName
                 return (
-                  <option
-                    value={modelName}
-                    key={modelName}
-                    selected={modelName === session.modelName}
-                  >
+                  <option value={index} key={index} selected={selected}>
                     {desc}
                   </option>
                 )
+              }
             })}
           </select>
         </span>
diff --git a/src/config/index.mjs b/src/config/index.mjs
@@ -286,6 +286,7 @@ export const defaultConfig = {
   themeMode: 'auto',
   /** @type {keyof Models}*/
   modelName: getNavigatorLanguage() === 'zh' ? 'moonshotWebFree' : 'claude2WebFree',
+  apiMode: null,
 
   preferredLanguage: getNavigatorLanguage(),
   clickIconAction: 'popup',
diff --git a/src/content-script/index.jsx b/src/content-script/index.jsx
@@ -91,7 +91,7 @@ async function mountComponent(siteConfig) {
 
     render(
       <FloatingToolbar
-        session={initSession({ modelName: userConfig.modelName })}
+        session={initSession({ modelName: userConfig.modelName, apiMode: userConfig.apiMode })}
         selection=""
         container={toolbarContainer}
         triggered={triggered}
@@ -107,7 +107,7 @@ async function mountComponent(siteConfig) {
   container.id = 'chatgptbox-container'
   render(
     <DecisionCard
-      session={initSession({ modelName: userConfig.modelName })}
+      session={initSession({ modelName: userConfig.modelName, apiMode: userConfig.apiMode })}
       question={question}
       siteConfig={siteConfig}
       container={container}
@@ -151,9 +151,10 @@ const deleteToolbar = () => {
 
 const createSelectionTools = async (toolbarContainer, selection) => {
   toolbarContainer.className = 'chatgptbox-toolbar-container'
+  const userConfig = await getUserConfig()
   render(
     <FloatingToolbar
-      session={initSession({ modelName: (await getUserConfig()).modelName })}
+      session={initSession({ modelName: userConfig.modelName, apiMode: userConfig.apiMode })}
       selection={selection}
       container={toolbarContainer}
       dockable={true}
@@ -276,9 +277,10 @@ async function prepareForRightClickMenu() {
         : { x: window.innerWidth / 2 - 300, y: window.innerHeight / 2 - 200 }
       const container = createElementAtPosition(position.x, position.y)
       container.className = 'chatgptbox-toolbar-container-not-queryable'
+      const userConfig = await getUserConfig()
       render(
         <FloatingToolbar
-          session={initSession({ modelName: (await getUserConfig()).modelName })}
+          session={initSession({ modelName: userConfig.modelName, apiMode: userConfig.apiMode })}
           selection={data.selectionText}
           container={container}
           triggered={true}
diff --git a/src/popup/sections/ApiModes.jsx b/src/popup/sections/ApiModes.jsx
@@ -1,8 +1,13 @@
 import { useTranslation } from 'react-i18next'
 import PropTypes from 'prop-types'
-import { apiModeToModelName, getApiModesFromConfig, modelNameToDesc } from '../../utils/index.mjs'
+import {
+  apiModeToModelName,
+  getApiModesFromConfig,
+  isApiModeSelected,
+  modelNameToDesc,
+} from '../../utils/index.mjs'
 import { PencilIcon, TrashIcon } from '@primer/octicons-react'
-import { useState } from 'react'
+import { useLayoutEffect, useState } from 'react'
 import {
   AlwaysCustomGroups,
   CustomApiKeyGroups,
@@ -30,7 +35,31 @@ export function ApiModes({ config, updateConfig }) {
   const [editing, setEditing] = useState(false)
   const [editingApiMode, setEditingApiMode] = useState(defaultApiMode)
   const [editingIndex, setEditingIndex] = useState(-1)
-  const apiModes = getApiModesFromConfig(config)
+  const [apiModes, setApiModes] = useState([])
+  const [apiModeStringArray, setApiModeStringArray] = useState([])
+
+  useLayoutEffect(() => {
+    const apiModes = getApiModesFromConfig(config)
+    setApiModes(apiModes)
+    setApiModeStringArray(apiModes.map(apiModeToModelName))
+  }, [
+    config.activeApiModes,
+    config.customApiModes,
+    config.azureDeploymentName,
+    config.ollamaModelName,
+  ])
+
+  const updateWhenApiModeDisabled = (apiMode) => {
+    if (isApiModeSelected(apiMode, config) || config.modelName === apiModeToModelName(apiMode))
+      updateConfig({
+        modelName:
+          apiModeStringArray.includes(config.modelName) &&
+          config.modelName !== apiModeToModelName(apiMode)
+            ? config.modelName
+            : 'customModel',
+        apiMode: null,
+      })
+  }
 
   const editingComponent = (
     <div style={{ display: 'flex', flexDirection: 'column', '--spacing': '4px' }}>
@@ -52,6 +81,12 @@ export function ApiModes({ config, updateConfig }) {
                 customApiModes: [...apiModes, editingApiMode],
               })
             } else {
+              const apiMode = apiModes[editingIndex]
+              if (
+                isApiModeSelected(apiMode, config) ||
+                config.modelName === apiModeToModelName(apiMode) // there is a minor bug here, but it's not a big issue
+              )
+                updateConfig({ apiMode: editingApiMode })
               const customApiModes = [...apiModes]
               customApiModes[editingIndex] = editingApiMode
               updateConfig({ activeApiModes: [], customApiModes })
@@ -142,6 +177,7 @@ export function ApiModes({ config, updateConfig }) {
                 type="checkbox"
                 checked={apiMode.active}
                 onChange={(e) => {
+                  if (!e.target.checked) updateWhenApiModeDisabled(apiMode)
                   const customApiModes = [...apiModes]
                   customApiModes[index] = { ...apiMode, active: e.target.checked }
                   updateConfig({ activeApiModes: [], customApiModes })
@@ -165,6 +201,7 @@ export function ApiModes({ config, updateConfig }) {
                   style={{ cursor: 'pointer' }}
                   onClick={(e) => {
                     e.preventDefault()
+                    updateWhenApiModeDisabled(apiMode)
                     const customApiModes = [...apiModes]
                     customApiModes.splice(index, 1)
                     updateConfig({ activeApiModes: [], customApiModes })
diff --git a/src/popup/sections/GeneralPart.jsx b/src/popup/sections/GeneralPart.jsx
@@ -1,7 +1,13 @@
 import { useTranslation } from 'react-i18next'
-import { useState } from 'react'
+import { useLayoutEffect, useState } from 'react'
 import FileSaver from 'file-saver'
-import { openUrl, modelNameToDesc, getApiModesStringArrayFromConfig } from '../../utils/index.mjs'
+import {
+  openUrl,
+  modelNameToDesc,
+  isApiModeSelected,
+  getApiModesFromConfig,
+  apiModeToModelName,
+} from '../../utils/index.mjs'
 import {
   isUsingOpenAiApiKey,
   isUsingAzureOpenAi,
@@ -83,6 +89,16 @@ async function checkBilling(apiKey, apiUrl) {
 export function GeneralPart({ config, updateConfig }) {
   const { t, i18n } = useTranslation()
   const [balance, setBalance] = useState(null)
+  const [apiModes, setApiModes] = useState([])
+
+  useLayoutEffect(() => {
+    setApiModes(getApiModesFromConfig(config, true))
+  }, [
+    config.activeApiModes,
+    config.customApiModes,
+    config.azureDeploymentName,
+    config.ollamaModelName,
+  ])
 
   const getBalance = async () => {
     const response = await fetch(`${config.customOpenAiApiUrl}/dashboard/billing/credit_grants`, {
@@ -155,22 +171,23 @@ export function GeneralPart({ config, updateConfig }) {
             }
             required
             onChange={(e) => {
-              const modelName = e.target.value
-              updateConfig({ modelName: modelName })
+              const apiMode = apiModes[e.target.value]
+              updateConfig({ apiMode: apiMode })
             }}
           >
-            {getApiModesStringArrayFromConfig(config, true).map((modelName) => {
+            {apiModes.map((apiMode, index) => {
+              const modelName = apiModeToModelName(apiMode)
               const desc = modelNameToDesc(modelName, t)
-              if (desc)
+              if (desc) {
+                let selected
+                if (isApiModeSelected(apiMode, config)) selected = true
+                else selected = config.modelName === modelName
                 return (
-                  <option
-                    value={modelName}
-                    key={modelName}
-                    selected={modelName === config.modelName}
-                  >
+                  <option value={index} key={index} selected={selected}>
                     {desc}
                   </option>
                 )
+              }
             })}
           </select>
           {isUsingMultiModeModel(config) && (
diff --git a/src/services/init-session.mjs b/src/services/init-session.mjs
@@ -36,6 +36,7 @@ import { t } from 'i18next'
  * @param {string|null} sessionName
  * @param {string|null} modelName
  * @param {boolean|null} autoClean
+ * @param {Object|null} apiMode
  * @returns {Session}
  */
 export function initSession({
@@ -44,6 +45,7 @@ export function initSession({
   sessionName = null,
   modelName = null,
   autoClean = false,
+  apiMode = null,
 } = {}) {
   return {
     // common
@@ -57,6 +59,7 @@ export function initSession({
 
     aiName: modelName ? modelNameToDesc(modelName, t) : null,
     modelName,
+    apiMode,
 
     autoClean,
     isRetry: false,
diff --git a/src/services/local-session.mjs b/src/services/local-session.mjs
@@ -4,10 +4,10 @@ import { getUserConfig } from '../config/index.mjs'
 
 export const initDefaultSession = async () => {
   const config = await getUserConfig()
-  const modelName = config.modelName
   return initSession({
     sessionName: new Date().toLocaleString(),
-    modelName: modelName,
+    modelName: config.modelName,
+    apiMode: config.apiMode,
     autoClean: false,
   })
 }
diff --git a/src/utils/model-name-convert.mjs b/src/utils/model-name-convert.mjs
@@ -105,3 +105,9 @@ export function getApiModesFromConfig(config, onlyActive) {
 export function getApiModesStringArrayFromConfig(config, onlyActive) {
   return getApiModesFromConfig(config, onlyActive).map(apiModeToModelName)
 }
+
+export function isApiModeSelected(apiMode, configOrSession) {
+  return (
+    configOrSession.apiMode && JSON.stringify(configOrSession.apiMode) === JSON.stringify(apiMode)
+  )
+}
