support moonshot

josStorer · Mar 3, 2024 · 0d7d735 · 0d7d735
1 parent 6b494c0
commit 0d7d735
Show file tree

Hide file tree

Showing 5 changed files with 101 additions and 14 deletions.
diff --git a/src/background/index.mjs b/src/background/index.mjs
@@ -30,6 +30,7 @@ import {
   gptApiModelKeys,
   poeWebModelKeys,
   setUserConfig,
+  moonshotApiModelKeys,
 } from '../config/index.mjs'
 import '../_locales/i18n'
 import { openUrl } from '../utils/open-url'
@@ -44,6 +45,7 @@ import { refreshMenu } from './menus.mjs'
 import { registerCommands } from './commands.mjs'
 import { generateAnswersWithBardWebApi } from '../services/apis/bard-web.mjs'
 import { generateAnswersWithClaudeWebApi } from '../services/apis/claude-web.mjs'
+import { generateAnswersWithMoonshotCompletionApi } from '../services/apis/moonshot-api.mjs'
 
 function setPortProxy(port, proxyTabId) {
   port.proxy = Browser.tabs.connect(proxyTabId)
@@ -151,6 +153,14 @@ async function executeApi(session, port, config) {
       sessionKey,
       session.modelName,
     )
+  } else if (moonshotApiModelKeys.includes(session.modelName)) {
+    await generateAnswersWithMoonshotCompletionApi(
+      port,
+      session.question,
+      session,
+      config.apiKey,
+      session.modelName,
+    )
   }
 }
 

diff --git a/src/config/index.mjs b/src/config/index.mjs
@@ -65,11 +65,14 @@ export const poeWebModelKeys = [
   'poeAiWeb_Llama_2_13b',
   'poeAiWeb_Llama_2_70b',
 ]
+export const moonshotApiModelKeys = ['moonshot_v1_8k', 'moonshot_v1_32k', 'moonshot_v1_128k']
+const moonshotApiKeyGenerateUrl = 'https://platform.moonshot.cn/console/api-keys'
 
 /**
  * @typedef {object} Model
  * @property {string} value
  * @property {string} desc
+ * @property {string} [keyGenerateUrl]
  */
 /**
  * @type {Object.<string,Model>}
@@ -136,6 +139,22 @@ export const Models = {
   poeAiWebChatGpt: { value: 'chatgpt', desc: 'Poe AI (Web, ChatGPT)' },
   poeAiWebChatGpt_16k: { value: 'chatgpt-16k', desc: 'Poe AI (Web, ChatGPT-16k)' },
   poeAiWebCustom: { value: '', desc: 'Poe AI (Web, Custom)' },
+
+  moonshot_v1_8k: {
+    value: 'moonshot-v1-8k',
+    desc: 'Moonshot (8k)',
+    keyGenerateUrl: moonshotApiKeyGenerateUrl,
+  },
+  moonshot_v1_32k: {
+    value: 'moonshot-v1-32k',
+    desc: 'Moonshot (32k)',
+    keyGenerateUrl: moonshotApiKeyGenerateUrl,
+  },
+  moonshot_v1_128k: {
+    value: 'moonshot-v1-128k',
+    desc: 'Moonshot (128k)',
+    keyGenerateUrl: moonshotApiKeyGenerateUrl,
+  },
 }
 
 for (const modelName in Models) {
@@ -293,7 +312,8 @@ export function getNavigatorLanguage() {
 export function isUsingApiKey(configOrSession) {
   return (
     gptApiModelKeys.includes(configOrSession.modelName) ||
-    chatgptApiModelKeys.includes(configOrSession.modelName)
+    chatgptApiModelKeys.includes(configOrSession.modelName) ||
+    moonshotApiModelKeys.includes(configOrSession.modelName)
   )
 }
 
@@ -324,6 +344,14 @@ export function isUsingGithubThirdPartyApi(configOrSession) {
   return githubThirdPartyApiModelKeys.includes(configOrSession.modelName)
 }
 
+export function isSupportBalance(configOrSession) {
+  return (
+    isUsingAzureOpenAi(configOrSession) ||
+    gptApiModelKeys.includes(configOrSession.modelName) ||
+    chatgptApiModelKeys.includes(configOrSession.modelName)
+  )
+}
+
 export async function getPreferredLanguageKey() {
   const config = await getUserConfig()
   if (config.preferredLanguage === 'auto') return config.userLanguage

diff --git a/src/popup/sections/GeneralPart.jsx b/src/popup/sections/GeneralPart.jsx
@@ -1,7 +1,8 @@
 import { useTranslation } from 'react-i18next'
-import { useState } from 'react'
+import { useMemo, useState } from 'react'
 import { openUrl } from '../../utils/index.mjs'
 import {
+  isSupportBalance,
   isUsingApiKey,
   isUsingAzureOpenAi,
   isUsingChatGLMApi,
@@ -81,6 +82,10 @@ async function checkBilling(apiKey, apiUrl) {
 export function GeneralPart({ config, updateConfig }) {
   const { t, i18n } = useTranslation()
   const [balance, setBalance] = useState(null)
+  const [currentModel, setCurrentModel] = useState(null)
+  const showBalance = useMemo(() => {
+    return isSupportBalance(config)
+  }, [config])
 
   const getBalance = async () => {
     const response = await fetch(`${config.customOpenAiApiUrl}/dashboard/billing/credit_grants`, {
@@ -153,6 +158,7 @@ export function GeneralPart({ config, updateConfig }) {
             onChange={(e) => {
               const modelName = e.target.value
               updateConfig({ modelName: modelName })
+              setCurrentModel(Models[modelName])
             }}
           >
             {config.activeApiModes.map((modelName) => {
@@ -207,23 +213,29 @@ export function GeneralPart({ config, updateConfig }) {
               />
               {config.apiKey.length === 0 ? (
                 <a
-                  href="https://platform.openai.com/account/api-keys"
+                  href={
+                    'keyGenerateUrl' in currentModel
+                      ? currentModel.keyGenerateUrl
+                      : 'https://platform.openai.com/account/api-keys'
+                  }
                   target="_blank"
                   rel="nofollow noopener noreferrer"
                 >
                   <button style="white-space: nowrap;" type="button">
                     {t('Get')}
                   </button>
                 </a>
-              ) : balance ? (
-                <button type="button" onClick={getBalance}>
-                  {balance}
-                </button>
-              ) : (
-                <button type="button" onClick={getBalance}>
-                  {t('Balance')}
-                </button>
-              )}
+              ) : showBalance ? (
+                balance ? (
+                  <button type="button" onClick={getBalance}>
+                    {balance}
+                  </button>
+                ) : (
+                  <button type="button" onClick={getBalance}>
+                    {t('Balance')}
+                  </button>
+                )
+              ) : null}
             </span>
           )}
           {isUsingCustomModel(config) && (

diff --git a/src/services/apis/moonshot-api.mjs b/src/services/apis/moonshot-api.mjs
@@ -0,0 +1,19 @@
+import { generateAnswersWithChatgptApiCompat } from './openai-api.mjs'
+
+/**
+ * @param {Browser.Runtime.Port} port
+ * @param {string} question
+ * @param {Session} session
+ * @param {string} apiKey
+ * @param {string} modelName
+ */
+export async function generateAnswersWithMoonshotCompletionApi(
+  port,
+  question,
+  session,
+  apiKey,
+  modelName,
+) {
+  const baseUrl = 'https://api.moonshot.cn'
+  return generateAnswersWithChatgptApiCompat(baseUrl, port, question, session, apiKey, modelName)
+}
diff --git a/src/services/apis/openai-api.mjs b/src/services/apis/openai-api.mjs
@@ -95,6 +95,25 @@ export async function generateAnswersWithGptCompletionApi(
  * @param {string} modelName
  */
 export async function generateAnswersWithChatgptApi(port, question, session, apiKey, modelName) {
+  const config = await getUserConfig()
+  return generateAnswersWithChatgptApiCompat(
+    config.customOpenAiApiUrl,
+    port,
+    question,
+    session,
+    apiKey,
+    modelName,
+  )
+}
+
+export async function generateAnswersWithChatgptApiCompat(
+  baseUrl,
+  port,
+  question,
+  session,
+  apiKey,
+  modelName,
+) {
   const { controller, messageListener, disconnectListener } = setAbortController(port)
 
   const config = await getUserConfig()
@@ -104,10 +123,9 @@ export async function generateAnswersWithChatgptApi(port, question, session, api
   )
   prompt.unshift({ role: 'system', content: await getChatSystemPromptBase() })
   prompt.push({ role: 'user', content: question })
-  const apiUrl = config.customOpenAiApiUrl
 
   let answer = ''
-  await fetchSSE(`${apiUrl}/v1/chat/completions`, {
+  await fetchSSE(`${baseUrl}/v1/chat/completions`, {
     method: 'POST',
     signal: controller.signal,
     headers: {