From 57c1b59a51ecb48082418c99d6d4319ec5b85dd0 Mon Sep 17 00:00:00 2001
From: suyao <sy20010504@gmail.com>
Date: Sun, 30 Mar 2025 23:44:50 +0800
Subject: [PATCH] fix(models): reorganize gemini websearch model lists

---
 src/renderer/src/config/models.ts | 52 ++++++++++++++++++++-----------
 1 file changed, 34 insertions(+), 18 deletions(-)

diff --git a/src/renderer/src/config/models.ts b/src/renderer/src/config/models.ts
index 6f473aa9..0507beba 100644
--- a/src/renderer/src/config/models.ts
+++ b/src/renderer/src/config/models.ts
@@ -165,7 +165,15 @@ const visionAllowedModels = [
   'gemma-3(?:-[\\w-]+)'
 ]
 
-const visionExcludedModels = ['gpt-4-\\d+-preview', 'gpt-4-turbo-preview', 'gpt-4-32k', 'gpt-4-\\d+', 'o1-mini', 'o1-preview', 'AIDC-AI/Marco-o1']
+const visionExcludedModels = [
+  'gpt-4-\\d+-preview',
+  'gpt-4-turbo-preview',
+  'gpt-4-32k',
+  'gpt-4-\\d+',
+  'o1-mini',
+  'o1-preview',
+  'AIDC-AI/Marco-o1'
+]
 export const VISION_REGEX = new RegExp(
   `\\b(?!(?:${visionExcludedModels.join('|')})\\b)(${visionAllowedModels.join('|')})\\b`,
   'i'
@@ -203,7 +211,13 @@ export const FUNCTION_CALLING_MODELS = [
   'gemini(?:-[\\w-]+)?' // 提前排除了gemini的嵌入模型
 ]
 
-const FUNCTION_CALLING_EXCLUDED_MODELS = ['aqa(?:-[\\w-]+)?', 'imagen(?:-[\\w-]+)?', 'o1-mini', 'o1-preview', 'AIDC-AI/Marco-o1']
+const FUNCTION_CALLING_EXCLUDED_MODELS = [
+  'aqa(?:-[\\w-]+)?',
+  'imagen(?:-[\\w-]+)?',
+  'o1-mini',
+  'o1-preview',
+  'AIDC-AI/Marco-o1'
+]
 
 export const FUNCTION_CALLING_REGEX = new RegExp(
   `\\b(?!(?:${FUNCTION_CALLING_EXCLUDED_MODELS.join('|')})\\b)(?:${FUNCTION_CALLING_MODELS.join('|')})\\b`,
@@ -1944,6 +1958,17 @@ export const TEXT_TO_IMAGES_MODELS_SUPPORT_IMAGE_ENHANCEMENT = [
 
 export const GENERATE_IMAGE_MODELS = ['gemini-2.0-flash-exp-image-generation', 'gemini-2.0-flash-exp']
 
+export const GEMINI_SEARCH_MODELS = [
+  'gemini-2.0-flash',
+  'gemini-2.0-flash-lite',
+  'gemini-2.0-flash-exp',
+  'gemini-2.0-flash-001',
+  'gemini-2.0-pro-exp-02-05',
+  'gemini-2.0-pro-exp',
+  'gemini-2.5-pro-exp',
+  'gemini-2.5-pro-exp-03-25'
+]
+
 export function isTextToImageModel(model: Model): boolean {
   return TEXT_TO_IMAGE_REGEX.test(model.id)
 }
@@ -2048,34 +2073,25 @@ export function isWebSearchModel(model: Model): boolean {
     return false
   }
 
+  if (provider.id === 'aihubmix') {
+    const models = ['gemini-2.0-flash-search', 'gemini-2.0-flash-exp-search', 'gemini-2.0-pro-exp-02-05-search']
+    return models.includes(model?.id)
+  }
+
   if (provider?.type === 'openai') {
-    if (model?.id?.includes('gemini-2.0-flash-exp')) {
+    if (GEMINI_SEARCH_MODELS.includes(model?.id)) {
       return true
     }
   }
 
   if (provider.id === 'gemini' || provider?.type === 'gemini') {
-    const models = [
-      'gemini-2.0-flash',
-      'gemini-2.0-flash-exp',
-      'gemini-2.0-flash-001',
-      'gemini-2.0-pro-exp-02-05',
-      'gemini-2.0-pro-exp',
-      'gemini-2.5-pro-exp',
-      'gemini-2.5-pro-exp-03-25'
-    ]
-    return models.includes(model?.id)
+    return GEMINI_SEARCH_MODELS.includes(model?.id)
   }
 
   if (provider.id === 'hunyuan') {
     return model?.id !== 'hunyuan-lite'
   }
 
-  if (provider.id === 'aihubmix') {
-    const models = ['gemini-2.0-flash-search', 'gemini-2.0-flash-exp-search', 'gemini-2.0-pro-exp-02-05-search']
-    return models.includes(model?.id)
-  }
-
   if (provider.id === 'zhipu') {
     return model?.id?.startsWith('glm-4-')
   }