feat: 添加搜索模型配置并改进Markdown链接处理

在Dockerfile中添加SEARCH_MODELS环境变量，支持gemini-2.0-flash-exp和gemini-2.0-pro-exp模型
改进message_converter中的图片链接正则表达式

.env.example

@@ -1,8 +1,10 @@
 API_KEYS=["AIzaSyxxxxxxxxxxxxxxxxxxx","AIzaSyxxxxxxxxxxxxxxxxxxx"]
 ALLOWED_TOKENS=["sk-123456"]
 # AUTH_TOKEN=sk-123456
-MODEL_SEARCH=["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]
-MODEL_IMAGE=["gemini-2.0-flash-exp"]
+TEST_MODEL=gemini-1.5-flash
+IMAGE_MODELS=["gemini-2.0-flash-exp"]
+SEARCH_MODELS=["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]
+FILTERED_MODELS=["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]
 TOOLS_CODE_EXECUTION_ENABLED=false
 SHOW_SEARCH_LINK=true
 SHOW_THINKING_PROCESS=true

Dockerfile

@@ -11,7 +11,8 @@ ENV API_KEYS='["your_api_key_1"]'
 ENV ALLOWED_TOKENS='["your_token_1"]'
 ENV BASE_URL=https://generativelanguage.googleapis.com/v1beta
 ENV TOOLS_CODE_EXECUTION_ENABLED=false
-ENV MODEL_SEARCH='["gemini-2.0-flash-exp"]'
+ENV IMAGE_MODELS='["gemini-2.0-flash-exp"]'
+ENV SEARCH_MODELS='["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]'
 
 # Expose port
 EXPOSE 8000

README.md

@@ -64,10 +64,13 @@
     AUTH_TOKEN=""  # 超级管理员token，具有所有权限，默认使用 ALLOWED_TOKENS 的第一个
 
     # 模型功能配置
-    MODEL_SEARCH=["gemini-2.0-flash-exp"]  # 支持搜索功能的模型列表
+    TEST_MODEL="gemini-1.5-flash" # 用于测试密钥是否可用的模型名
+    SEARCH_MODELS=["gemini-2.0-flash-exp"]  # 支持搜索功能的模型列表
+    IMAGE_MODELS=["gemini-2.0-flash-exp"] # 支持绘图功能的模型列表
     TOOLS_CODE_EXECUTION_ENABLED=false  # 是否启用代码执行工具，默认false
     SHOW_SEARCH_LINK=true  # 是否在响应中显示搜索结果链接，默认true
     SHOW_THINKING_PROCESS=true  # 是否显示模型思考过程，默认true
+    FILTERED_MODELS=["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"] # 被禁用的模型列表
 
     # 图片生成配置
     PAID_KEY="your-paid-api-key"  # 付费版API Key，用于图片生成等高级功能
@@ -115,9 +118,17 @@
 
    #### 模型功能配置
 
-    - `MODEL_SEARCH`: 搜索功能支持的模型
+    - `TEST_MODEL`: 用于测试密钥可用性的模型
+      - 默认值: `"gemini-1.5-flash"`
+    - `SEARCH_MODELS`: 搜索功能支持的模型
       - 默认值: `["gemini-2.0-flash-exp"]`
       - 说明: 仅列表中的模型可使用搜索功能
+    - `IMAGE_MODELS`: 绘图功能支持的模型
+      - 默认值: `["gemini-2.0-flash-exp"]`
+      - 说明: 仅列表中的模型可使用绘图功能
+    - `FILTERED_MODELS`: 被禁用的模型列表
+      - 默认值: `["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]`
+      - 说明: 列表中的模型将被禁用
     - `TOOLS_CODE_EXECUTION_ENABLED`: 代码执行功能
       - 默认值: `false`
       - 安全提示: 生产环境建议禁用

app/api/gemini_routes.py

@@ -23,7 +23,7 @@ async def get_key_manager():
 async def get_next_working_key_wrapper(key_manager: KeyManager = Depends(get_key_manager)):
     return await key_manager.get_next_working_key()
 
-model_service = ModelService(settings.MODEL_SEARCH,settings.MODEL_IMAGE)
+model_service = ModelService(settings.SEARCH_MODELS,settings.IMAGE_MODELS)
 
 
 @router.get("/models")
@@ -41,8 +41,8 @@ async def list_models(_=Depends(security_service.verify_key),
     model_mapping = {x.get("name", "").split("/", maxsplit=1)[1]: x for x in models_json["models"]}
 
     # 添加搜索模型
-    if settings.MODEL_SEARCH:
-        for name in settings.MODEL_SEARCH:
+    if model_service.search_models:
+        for name in model_service.search_models:
             model = model_mapping.get(name, None)
             if not model:
                 continue
@@ -56,8 +56,8 @@ async def list_models(_=Depends(security_service.verify_key),
             models_json["models"].append(item)
 
     # 添加图像生成模型
-    if settings.MODEL_IMAGE:
-        for name in settings.MODEL_IMAGE:
+    if model_service.image_models:
+        for name in model_service.image_models:
             model = model_mapping.get(name, None)
             if not model:
                 continue

app/api/openai_routes.py

@@ -17,7 +17,7 @@ logger = get_openai_logger()
 
 # 初始化服务
 security_service = SecurityService(settings.ALLOWED_TOKENS, settings.AUTH_TOKEN)
-model_service = ModelService(settings.MODEL_SEARCH,settings.MODEL_IMAGE)
+model_service = ModelService(settings.SEARCH_MODELS,settings.IMAGE_MODELS)
 embedding_service = EmbeddingService(settings.BASE_URL)
 image_create_service = ImageCreateService()
 

app/core/config.py

@@ -6,8 +6,9 @@ class Settings(BaseSettings):
     API_KEYS: List[str]
     ALLOWED_TOKENS: List[str]
     BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
-    MODEL_SEARCH: List[str] = ["gemini-2.0-flash-exp"]
-    MODEL_IMAGE: List[str] = ["gemini-2.0-flash-exp"]
+    SEARCH_MODELS: List[str] = ["gemini-2.0-flash-exp"]
+    IMAGE_MODELS: List[str] = ["gemini-2.0-flash-exp"]
+    FILTERED_MODELS: List[str] = ["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]
     TOOLS_CODE_EXECUTION_ENABLED: bool = False
     SHOW_SEARCH_LINK: bool = True
     SHOW_THINKING_PROCESS: bool = True

app/core/uploader.py

@@ -290,9 +290,9 @@ class CloudFlareImgBedUploader(ImageUploader):
         try:
             # 准备请求URL（添加认证码参数，如果存在）
             if self.auth_code:
-                request_url = f"{self.api_url}?authCode={self.auth_code}"
+                request_url = f"{self.api_url}?authCode={self.auth_code}&uploadNameType=origin"
             else:
-                request_url = self.api_url
+                request_url = f"{self.api_url}?uploadNameType=origin"
             
             # 准备文件数据
             files = {

app/services/chat/message_converter.py

@@ -7,7 +7,7 @@ import requests
 import base64
 
 SUPPORTED_ROLES = ["user", "model", "system"]
-IMAGE_URL_PATTERN = r'\[image\]\((.*?)\)'
+IMAGE_URL_PATTERN = r'\[(.*?)\]\((.*?)\)'
 
 
 class MessageConverter(ABC):
@@ -33,7 +33,7 @@ def _get_mime_type_and_data(base64_string):
         pattern = r'data:([^;]+);base64,(.+)'
         match = re.match(pattern, base64_string)
         if match:
-            mime_type = match.group(1)
+            mime_type = "image/jpeg" if match.group(1) == "image/jpg" else match.group(1)
             encoded_data = match.group(2)
             return mime_type, encoded_data
     
@@ -87,7 +87,7 @@ def _process_text_with_image(text: str) -> List[Dict[str, Any]]:
     img_url_match = re.search(IMAGE_URL_PATTERN, text)
     if img_url_match:
         # 提取URL
-        img_url = img_url_match.group(1)
+        img_url = img_url_match.group(2)
         # 将URL对应的图片转换为base64
         try:
             base64_data = _convert_image_to_base64(img_url)

app/services/chat/response_handler.py

@@ -205,7 +205,7 @@ def _extract_image_data(part: dict) -> str:
     current_date = time.strftime("%Y/%m/%d")
     filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"
     base64_data = part["inlineData"]["data"]
-                #将base64_data转成bytes数组
+    #将base64_data转成bytes数组
     bytes_data = base64.b64decode(base64_data)
     upload_response = image_uploader.upload(bytes_data,filename)
     if upload_response.success:

app/services/gemini_chat_service.py

@@ -71,7 +71,7 @@ def _build_payload(model: str, request: GeminiRequest) -> Dict[str, Any]:
         "systemInstruction": request_dict.get("systemInstruction", "")
     }
     
-    if model.endswith("-image"):
+    if model.endswith("-image") or model.endswith("-image-generation"):
         payload.pop("systemInstruction")
         payload["generationConfig"]["responseModalities"] = ["Text","Image"]
     return payload

app/services/model_service.py

@@ -7,10 +7,11 @@ from app.core.config import settings
 logger = get_model_logger()
 
 class ModelService:
-    def __init__(self, model_search: list, model_image: list):
-        self.model_search = model_search
-        self.model_image = model_image
+    def __init__(self, search_models: list, image_models: list):
+        self.search_models = search_models
+        self.image_models = image_models
         self.base_url = "https://generativelanguage.googleapis.com/v1beta"
+        self.filtered_models = settings.FILTERED_MODELS
 
     def get_gemini_models(self, api_key: str) -> Optional[Dict[str, Any]]:
         url = f"{self.base_url}/models?key={api_key}"
@@ -19,6 +20,16 @@ class ModelService:
             response = requests.get(url)
             if response.status_code == 200:
                 gemini_models = response.json()
+                
+                filtered_models_list = []
+                for model in gemini_models.get("models", []):
+                    model_id = model["name"].split("/")[-1]
+                    if model_id not in self.filtered_models:
+                        filtered_models_list.append(model)
+                    else:
+                        logger.info(f"Filtered out model: {model_id}")
+                
+                gemini_models["models"] = filtered_models_list
                 return gemini_models
             else:
                 logger.error(f"Error: {response.status_code}")
@@ -54,11 +65,11 @@ class ModelService:
             }
             openai_format["data"].append(openai_model)
 
-            if model_id in self.model_search:
+            if model_id in self.search_models:
                 search_model = openai_model.copy()
                 search_model["id"] = f"{model_id}-search"
                 openai_format["data"].append(search_model)
-            if model_id in self.model_image:
+            if model_id in self.image_models:
                 image_model = openai_model.copy()
                 image_model["id"] = f"{model_id}-image"
                 openai_format["data"].append(image_model)
@@ -76,9 +87,9 @@ class ModelService:
         model = model.strip()
         if model.endswith("-search"):
             model = model[:-7]
-            return model in settings.MODEL_SEARCH
+            return model in self.search_models
         if model.endswith("-image"):
             model = model[:-6]
-            return model in settings.MODEL_IMAGE
+            return model in self.image_models
 
-        return True
+        return model not in self.filtered_models

app/services/openai_chat_service.py

@@ -35,7 +35,7 @@ def _build_tools(
 
     if (
             settings.TOOLS_CODE_EXECUTION_ENABLED
-            and not (model.endswith("-search") or "-thinking" in model or model.endswith("-image"))
+            and not (model.endswith("-search") or "-thinking" in model or model.endswith("-image") or model.endswith("-image-generation"))
             and not _has_image_parts(messages)
     ):
         tools.append({"code_execution": {}})
@@ -110,7 +110,7 @@ def _build_payload(
         "tools": _build_tools(request, messages),
         "safetySettings": _get_safety_settings(request.model),
     }
-    if request.model.endswith("-image"):
+    if request.model.endswith("-image") or request.model.endswith("-image-generation"):
         payload["generationConfig"]["responseModalities"] = ["Text","Image"]
         
     if (
@@ -119,6 +119,7 @@ def _build_payload(
         and instruction.get("role") == "system"
         and instruction.get("parts")
         and not request.model.endswith("-image")
+        and not request.model.endswith("-image-generation")
     ):
         payload["systemInstruction"] = instruction