Merge pull request #528 from sugarforever/feat/gemini-proxy

gemini 支持使用代理
sugarforever · Jun 19, 2024 · 9401d25 · 9401d25
2 parents 3348a5b + 94f912c
commit 9401d25
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 7 deletions.
diff --git a/components/settings/SettingsServers.vue b/components/settings/SettingsServers.vue
@@ -1,6 +1,6 @@
 <script setup lang="ts">
 import type { ContextKeys } from '~/server/middleware/keys'
-import { keysStore } from '~/utils/settings'
+import { keysStore, DEFAULT_KEYS_STORE } from '~/utils/settings'
 import type { PickupPathKey, TransformTypes } from '~/types/helper'
 
 const { t } = useI18n()
@@ -70,6 +70,8 @@ const LLMList = computed<LLMListItem[]>(() => {
       title: t('settings.gemini'),
       fields: [
         { label: t('settings.apiKey'), value: 'gemini.key', type: 'password', placeholder: t('settings.apiKey') },
+        { label: t('settings.endpoint'), value: 'gemini.endpoint', type: 'input', placeholder: t('global.optional'), rule: 'url' },
+        { label: t('settings.proxy'), value: 'gemini.proxy', type: 'checkbox', placeholder: t('settings.proxyTips') },
       ]
     },
     {
@@ -102,7 +104,7 @@ const validate = (data: typeof state) => {
 }
 
 const onSubmit = async () => {
-  keysStore.value = recursiveObject(keysStore.value, (keyPaths, value) => {
+  keysStore.value = recursiveObject(DEFAULT_KEYS_STORE, (keyPaths, value) => {
     const key = keyPaths.join('.') as keyof typeof state
     return key in state ? state[key] : value
   })

diff --git a/server/middleware/keys.ts b/server/middleware/keys.ts
@@ -28,6 +28,8 @@ export interface ContextKeys {
   },
   gemini: {
     key: string
+    endpoint: string
+    proxy: boolean
   },
   groq: {
     key: string

diff --git a/server/models/genai/generative-ai.ts b/server/models/genai/generative-ai.ts
@@ -33,6 +33,7 @@ export type BaseMessageExamplePair = {
  * An interface defining the input to the ChatGoogleGenerativeAI class.
  */
 export interface GoogleGenerativeAIChatInput extends BaseChatModelParams {
+  baseUrl?: string
   /**
      * Version of API endpoint to call (e.g. "v1" or "v1beta"). If not specified,
      * defaults to latest stable version.
@@ -268,7 +269,7 @@ export class ChatGoogleGenerativeAI
         topP: this.topP,
         topK: this.topK,
       },
-    }, { apiVersion: this.apiVersion })
+    }, { apiVersion: this.apiVersion, baseUrl: fields?.baseUrl })
   }
 
   _combineLLMOutput() {

diff --git a/server/utils/models.ts b/server/utils/models.ts
@@ -88,11 +88,12 @@ export const createChatModel = (modelName: string, family: string, event: H3Even
       modelName: modelName
     })
   } else if (family === MODEL_FAMILIES.gemini && GEMINI_MODELS.includes(modelName)) {
-    console.log(`Chat with Gemini ${modelName}`)
+    console.log(`Chat with Gemini ${modelName}, host:`, keys.gemini.endpoint)
     chat = new ChatGoogleGenerativeAI({
       apiVersion: "v1beta",
       apiKey: keys.gemini.key,
-      modelName: modelName
+      modelName: modelName,
+      baseUrl: getProxyEndpoint(keys.gemini.endpoint, keys.gemini.proxy),
     })
   } else if (family === MODEL_FAMILIES.groq && GROQ_MODELS.includes(modelName)) {
     // @langchain/grop does not support configuring groq's baseURL, but groq sdk supports receiving environment variables.

diff --git a/utils/settings.ts b/utils/settings.ts
@@ -7,7 +7,7 @@ function getLocalValue(key: string) {
   return process.server ? '' : (localStorage.getItem(key) || '')
 }
 
-export const keysStore = useStorage<ContextKeys>('keys', {
+export const DEFAULT_KEYS_STORE: ContextKeys = {
   ollama: {
     endpoint: getLocalValue('ollama.host'),
     username: getLocalValue('ollama.username'),
@@ -35,13 +35,17 @@ export const keysStore = useStorage<ContextKeys>('keys', {
   },
   gemini: {
     key: getLocalValue('keys.gemini_api_key'),
+    proxy: false,
+    endpoint: '',
   },
   groq: {
     key: getLocalValue('keys.groq_api_key'),
     endpoint: getLocalValue('keys.groq_api_host'),
     proxy: false,
   },
-} as ContextKeys)
+}
+
+export const keysStore = useStorage<ContextKeys>('keys', DEFAULT_KEYS_STORE)
 
 export const getKeysHeader = () => ({ 'x-chat-ollama-keys': JSON.stringify(keysStore.value) })