🐛 fix: fix google gemini pro 1.5 lobehub#1801

bentwnghk · Mar 29, 2024 · 4d12bf8 · 4d12bf8
1 parent ecdbc4f
commit 4d12bf8
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 53 deletions.
diff --git a/package.json b/package.json
@@ -85,7 +85,7 @@
     "@azure/openai": "^1.0.0-beta.11",
     "@bentwnghk/ui": "latest",
     "@cfworker/json-schema": "^1",
-    "@google/generative-ai": "^0.2.0",
+    "@google/generative-ai": "^0.3.1",
     "@icons-pack/react-simple-icons": "^9",
     "@lobehub/chat-plugin-sdk": "latest",
     "@lobehub/chat-plugins-gateway": "latest",

diff --git a/src/app/api/chat/google/route.ts b/src/app/api/chat/google/route.ts
@@ -13,19 +13,6 @@ import { POST as UniverseRoute } from '../[provider]/route';
 // so if you want to use with proxy, you need comment the code below
 export const runtime = 'edge';
 
-export const preferredRegion = [
-  'bom1',
-  'cle1',
-  'cpt1',
-  'gru1',
-  'hnd1',
-  'iad1',
-  'icn1',
-  'kix1',
-  'pdx1',
-  'sfo1',
-  'sin1',
-  'syd1',
-];
+export const preferredRegion = ['sfo1'];
 
-export const POST = async (req: Request) => UniverseRoute(req, { params: { provider: 'google' } });
+export const POST = async (req: Request) => UniverseRoute(req, { params: { provider: 'google' } });
diff --git a/src/config/modelProviders/google.ts b/src/config/modelProviders/google.ts
@@ -17,6 +17,7 @@ const Google: ModelProviderCard = {
       displayName: 'Gemini 1.5 Pro',
       id: 'gemini-1.5-pro-latest',
       tokens: 1_048_576,
+      vision: true,
     },
     {
       displayName: 'Gemini Ultra',
@@ -27,4 +28,4 @@ const Google: ModelProviderCard = {
   id: 'google',
 };
 
-export default Google;
+export default Google;
diff --git a/src/libs/agent-runtime/google/index.ts b/src/libs/agent-runtime/google/index.ts
@@ -14,17 +14,6 @@ import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
 import { parseDataUri } from '../utils/uriParser';
 
-type GoogleChatErrors = GoogleChatError[];
-
-interface GoogleChatError {
-  '@type': string;
-  'domain': string;
-  'metadata': {
-    service: string;
-  };
-  'reason': string;
-}
-
 enum HarmCategory {
   HARM_CATEGORY_DANGEROUS_CONTENT = 'HARM_CATEGORY_DANGEROUS_CONTENT',
   HARM_CATEGORY_HARASSMENT = 'HARM_CATEGORY_HARASSMENT',
@@ -49,32 +38,39 @@ export class LobeGoogleAI implements LobeRuntimeAI {
     try {
       const { contents, model } = this.buildGoogleMessages(payload.messages, payload.model);
       const geminiStream = await this.client
-        .getGenerativeModel({
-          generationConfig: {
-            maxOutputTokens: payload.max_tokens,
-            temperature: payload.temperature,
-            topP: payload.top_p,
-          },
-          model,
-          safetySettings: [
-            {
-              category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
-              threshold: HarmBlockThreshold.BLOCK_NONE,
-            },
-            {
-              category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
-              threshold: HarmBlockThreshold.BLOCK_NONE,
-            },
-            {
-              category: HarmCategory.HARM_CATEGORY_HARASSMENT,
-              threshold: HarmBlockThreshold.BLOCK_NONE,
-            },
-            {
-              category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
-              threshold: HarmBlockThreshold.BLOCK_NONE,
+        .getGenerativeModel(
+          {
+            generationConfig: {
+              maxOutputTokens: payload.max_tokens,
+              temperature: payload.temperature,
+              topP: payload.top_p,
             },
-          ],
-        })
+            model,
+            // avoid wide sensitive words
+            // refs: https://github.com/lobehub/lobe-chat/pull/1418
+            safetySettings: [
+              {
+                category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+                threshold: HarmBlockThreshold.BLOCK_NONE,
+              },
+              {
+                category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+                threshold: HarmBlockThreshold.BLOCK_NONE,
+              },
+              {
+                category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+                threshold: HarmBlockThreshold.BLOCK_NONE,
+              },
+              {
+                category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                threshold: HarmBlockThreshold.BLOCK_NONE,
+              },
+            ],
+          },
+          {
+            apiVersion: 'v1beta',
+          },
+        )
         .generateContentStream({ contents });
 
       // Convert the response into a friendly text-stream
@@ -188,6 +184,21 @@ export class LobeGoogleAI implements LobeRuntimeAI {
       return defaultError;
     }
   }
+
+  private shouldUseBetaApi(model: string) {
+    return ['gemini-1.5-pro-latest', 'gemini-ultra-latest'].includes(model);
+  }
 }
 
 export default LobeGoogleAI;
+
+type GoogleChatErrors = GoogleChatError[];
+
+interface GoogleChatError {
+  '@type': string;
+  'domain': string;
+  'metadata': {
+    service: string;
+  };
+  'reason': string;
+}