♻️ refactor: Remove code related to gemini-pro-vision (lobehub#2860)

LovelyGuYiMeng · Jul 13, 2024 · e2ba302 · e2ba302
1 parent 471721a
commit e2ba302
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 54 deletions.
diff --git a/docs/self-hosting/advanced/model-list.mdx b/docs/self-hosting/advanced/model-list.mdx
@@ -36,7 +36,7 @@ Examples are as follows:
 
 - `chatglm-6b=ChatGLM 6B<4096>`: ChatGLM 6B, maximum context of 4k, no advanced capabilities;
 - `spark-v3.5=讯飞星火 v3.5<8192:fc>`: Xunfei Spark 3.5 model, maximum context of 8k, supports Function Call;
-- `gemini-pro-vision=Gemini Pro Vision<16000:vision>`: Google Vision model, maximum context of 16k, supports image recognition;
+- `gemini-1.5-flash-latest=Gemini 1.5 Flash<16000:vision>`: Google Vision model, maximum context of 16k, supports image recognition;
 - `gpt-4-all=ChatGPT Plus<128000:fc:vision:file>`, hacked version of ChatGPT Plus web, context of 128k, supports image recognition, Function Call, file upload.
 
 Currently supported extension capabilities are:

diff --git a/docs/self-hosting/advanced/model-list.zh-CN.mdx b/docs/self-hosting/advanced/model-list.zh-CN.mdx
@@ -35,7 +35,7 @@ id=displayName<maxToken:vision:fc:file>
 
 - `chatglm-6b=ChatGLM 6B<4096>`：ChatGLM 6B，最大上下文 4k，没有高阶能力；
 - `spark-v3.5=讯飞星火 v3.5<8192:fc>`：讯飞星火 3.5 模型，最大上下文 8k，支持 Function Call；
-- `gemini-pro-vision=Gemini Pro Vision<16000:vision>`：Google 视觉模型，最大上下文 16k，支持图像识别；
+- `gemini-1.5-flash-latest=Gemini 1.5 Flash<16000:vision>`：Google 视觉模型，最大上下文 16k，支持图像识别；
 - `gpt-4-all=ChatGPT Plus<128000:fc:vision:file>`，hack 的 ChatGPT Plus 网页版，上下 128k ，支持图像识别、Function Call、文件上传
 
 目前支持的扩展能力有：

diff --git a/src/libs/agent-runtime/google/index.test.ts b/src/libs/agent-runtime/google/index.test.ts
@@ -2,9 +2,9 @@
 import { FunctionDeclarationSchemaType, FunctionDeclarationsTool } from '@google/generative-ai';
 import { JSONSchema7 } from 'json-schema';
 import OpenAI from 'openai';
-import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { ChatStreamCallbacks, OpenAIChatMessage } from '@/libs/agent-runtime';
+import { OpenAIChatMessage } from '@/libs/agent-runtime';
 
 import * as debugStreamModule from '../utils/debugStream';
 import { LobeGoogleAI } from './index';
@@ -383,7 +383,7 @@ describe('LobeGoogleAI', () => {
             role: 'user',
           },
         ];
-        const model = 'gemini-pro-vision';
+        const model = 'gemini-1.5-flash-latest';
 
         // 调用 buildGoogleMessages 方法
         const contents = instance['buildGoogleMessages'](messages, model);
@@ -398,36 +398,6 @@ describe('LobeGoogleAI', () => {
       });
     });
 
-    describe('convertModel', () => {
-      it('should use default text model when no images are included in messages', () => {
-        const messages: OpenAIChatMessage[] = [
-          { content: 'Hello', role: 'user' },
-          { content: 'Hi', role: 'assistant' },
-        ];
-
-        // 调用 buildGoogleMessages 方法
-        const model = instance['convertModel']('gemini-pro-vision', messages);
-
-        expect(model).toEqual('gemini-pro'); // 假设 'gemini-pro' 是默认文本模型
-      });
-
-      it('should use specified model when images are included in messages', () => {
-        const messages: OpenAIChatMessage[] = [
-          {
-            content: [
-              { type: 'text', text: 'Hello' },
-              { type: 'image_url', image_url: { url: 'data:image/png;base64,...' } },
-            ],
-            role: 'user',
-          },
-        ];
-
-        const model = instance['convertModel']('gemini-pro-vision', messages);
-
-        expect(model).toEqual('gemini-pro-vision');
-      });
-    });
-
     describe('buildGoogleTools', () => {
       it('should return undefined when tools is undefined or empty', () => {
         expect(instance['buildGoogleTools'](undefined)).toBeUndefined();

diff --git a/src/libs/agent-runtime/google/index.ts b/src/libs/agent-runtime/google/index.ts
@@ -50,7 +50,7 @@ export class LobeGoogleAI implements LobeRuntimeAI {
 
   async chat(payload: ChatStreamPayload, options?: ChatCompetitionOptions) {
     try {
-      const model = this.convertModel(payload.model, payload.messages);
+      const model = payload.model;
 
       const contents = this.buildGoogleMessages(payload.messages, model);
 
@@ -197,21 +197,6 @@ export class LobeGoogleAI implements LobeRuntimeAI {
     return contents;
   };
 
-  private convertModel = (model: string, messages: OpenAIChatMessage[]) => {
-    let finalModel: string = model;
-
-    if (model.includes('pro-vision')) {
-      // if message are all text message, use vision will return an error:
-      // "[400 Bad Request] Add an image to use models/gemini-pro-vision, or switch your model to a text model."
-      const noNeedVision = messages.every((m) => typeof m.content === 'string');
-
-      // so we need to downgrade to gemini-pro
-      if (noNeedVision) finalModel = 'gemini-pro';
-    }
-
-    return finalModel;
-  };
-
   private parseErrorMessage(message: string): {
     error: any;
     errorType: ILobeAgentRuntimeErrorType;

diff --git a/src/utils/parseModels.test.ts b/src/utils/parseModels.test.ts
@@ -49,14 +49,14 @@ describe('parseModelString', () => {
 
     it('multi models', () => {
       const result = parseModelString(
-        'gemini-pro-vision=Gemini Pro Vision<16000:vision>,gpt-4-all=ChatGPT Plus<128000:fc:vision:file>',
+        'gemini-1.5-flash-latest=Gemini 1.5 Flash<16000:vision>,gpt-4-all=ChatGPT Plus<128000:fc:vision:file>',
       );
 
       expect(result.add).toEqual([
         {
-          displayName: 'Gemini Pro Vision',
+          displayName: 'Gemini 1.5 Flash',
           vision: true,
-          id: 'gemini-pro-vision',
+          id: 'gemini-1.5-flash-latest',
           tokens: 16000,
         },
         {