Merge pull request wyeeeee#212 from wyeeeee/dev

wyeeeee · web-flow · commit f3cb3da55257 · 2025-05-19T12:50:11.000+08:00
Dev
diff --git a/app/config/persistence.py b/app/config/persistence.py
@@ -93,6 +93,56 @@ def load_settings():
                     else:
                         setattr(settings, name, value)
             
+            # 在加载完设置后，检查是否需要刷新模型配置
+            try:
+                # 如果加载了Google Credentials JSON或Vertex Express API Key，需要刷新模型配置
+                if (hasattr(settings, 'GOOGLE_CREDENTIALS_JSON') and settings.GOOGLE_CREDENTIALS_JSON) or \
+                   (hasattr(settings, 'VERTEX_EXPRESS_API_KEY') and settings.VERTEX_EXPRESS_API_KEY):
+                    log('info', "检测到Google Credentials JSON或Vertex Express API Key，准备刷新模型配置")
+                    
+                    # 导入必要的模块
+                    from app.vertex.model_loader import refresh_models_config_cache
+                    from app.vertex.vertex_ai_init import init_vertex_ai
+                    from app.vertex.credentials_manager import CredentialManager
+                    
+                    # 创建新的CredentialManager实例
+                    credential_manager = CredentialManager()
+                    
+                    # 如果有Google Credentials JSON，加载到CredentialManager
+                    if hasattr(settings, 'GOOGLE_CREDENTIALS_JSON') and settings.GOOGLE_CREDENTIALS_JSON:
+                        from app.vertex.credentials_manager import parse_multiple_json_credentials
+                        parsed_json_objects = parse_multiple_json_credentials(settings.GOOGLE_CREDENTIALS_JSON)
+                        if parsed_json_objects:
+                            loaded_count = credential_manager.load_credentials_from_json_list(parsed_json_objects)
+                            log('info', f"从持久化的Google Credentials JSON中加载了{loaded_count}个凭据")
+                    
+                    # 初始化Vertex AI
+                    import asyncio
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                    try:
+                        success = loop.run_until_complete(init_vertex_ai(credential_manager=credential_manager))
+                        if success:
+                            log('info', "成功初始化Vertex AI服务")
+                        else:
+                            log('warning', "初始化Vertex AI服务失败")
+                    finally:
+                        loop.close()
+                    
+                    # 刷新模型配置缓存
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                    try:
+                        refresh_success = loop.run_until_complete(refresh_models_config_cache())
+                        if refresh_success:
+                            log('info', "成功刷新模型配置缓存")
+                        else:
+                            log('warning', "刷新模型配置缓存失败")
+                    finally:
+                        loop.close()
+            except Exception as e:
+                log('error', f"刷新模型配置时出错: {str(e)}")
+            
             log('info', f"加载设置成功")
             return True
         except Exception as e:
diff --git a/app/vertex/main.py b/app/vertex/main.py
@@ -46,10 +46,35 @@
 
 @app.on_event("startup")
 async def startup_event():
-    if await init_vertex_ai(credential_manager): # Added await
-        vertex_log('info', "Vertex AI credential and model config initialization check completed successfully.")
-    else:
-        vertex_log('error', "Failed to initialize a fallback Vertex AI client. API will likely fail.")
+    try:
+        # 检查是否有Google Credentials JSON
+        if hasattr(settings, 'GOOGLE_CREDENTIALS_JSON') and settings.GOOGLE_CREDENTIALS_JSON:
+            vertex_log('info', "检测到持久化的Google Credentials JSON，准备加载")
+            from app.vertex.credentials_manager import parse_multiple_json_credentials
+            parsed_json_objects = parse_multiple_json_credentials(settings.GOOGLE_CREDENTIALS_JSON)
+            if parsed_json_objects:
+                loaded_count = credential_manager.load_credentials_from_json_list(parsed_json_objects)
+                vertex_log('info', f"从持久化的Google Credentials JSON中加载了{loaded_count}个凭据")
+        
+        # 检查是否有Vertex Express API Key
+        if hasattr(settings, 'VERTEX_EXPRESS_API_KEY') and settings.VERTEX_EXPRESS_API_KEY:
+            vertex_log('info', "检测到持久化的Vertex Express API Key")
+        
+        # 初始化Vertex AI
+        if await init_vertex_ai(credential_manager):
+            vertex_log('info', "Vertex AI credential and model config initialization check completed successfully.")
+            
+            # 刷新模型配置缓存
+            from app.vertex.model_loader import refresh_models_config_cache
+            refresh_success = await refresh_models_config_cache()
+            if refresh_success:
+                vertex_log('info', "成功刷新模型配置缓存")
+            else:
+                vertex_log('warning', "刷新模型配置缓存失败")
+        else:
+            vertex_log('error', "Failed to initialize a fallback Vertex AI client. API will likely fail.")
+    except Exception as e:
+        vertex_log('error', f"启动时初始化Vertex AI服务出错: {str(e)}")
 
 @app.get("/")
 async def root():
diff --git a/app/vertex/routes/models_api.py b/app/vertex/routes/models_api.py
@@ -195,45 +195,70 @@ async def list_models(fastapi_request: Request, api_key: str = Depends(get_api_k
     # Add base models and their variations
     for original_model_id in sorted(list(all_model_ids)):
         current_display_prefix = ""
-        # Only add PAY_PREFIX if the model is not already an EXPRESS model (which has its own prefix)
-        if not original_model_id.startswith("[EXPRESS]") and \
-           has_sa_creds and not has_express_key and EXPERIMENTAL_MARKER not in original_model_id:
+        # 检查是否为[EXPRESS]模型，保留原始前缀
+        is_express_model = original_model_id.startswith("[EXPRESS]")
+        base_model_without_prefix = original_model_id
+        
+        # 如果是EXPRESS模型，移除前缀用于后续处理，但在显示时会保留
+        if is_express_model:
+            # 从原始ID中提取不带[EXPRESS]前缀的模型名
+            base_model_without_prefix = original_model_id[len("[EXPRESS] "):]
+            vertex_log('info', f"处理EXPRESS模型: {original_model_id}, 基础名称: {base_model_without_prefix}")
+        
+        # 只有非EXPRESS模型才考虑添加PAY_PREFIX
+        if not is_express_model and has_sa_creds and EXPERIMENTAL_MARKER not in original_model_id:
+            # 只要有SA凭证就应该显示PAY前缀，不管是否有Express Key
             current_display_prefix = PAY_PREFIX
+            vertex_log('info', f"为非EXPRESS模型添加PAY前缀: {current_display_prefix}{original_model_id}")
         
-        base_display_id = f"{current_display_prefix}{original_model_id}"
+        # 构建显示ID，如果是EXPRESS模型，确保保留[EXPRESS]前缀
+        base_display_id = original_model_id if is_express_model else f"{current_display_prefix}{original_model_id}"
         
+        vertex_log('info', f"添加模型到列表: {base_display_id}")
         dynamic_models_data.append({
             "id": base_display_id, "object": "model", "created": current_time, "owned_by": "google",
-            "permission": [], "root": original_model_id, "parent": None
+            "permission": [], "root": base_model_without_prefix, "parent": None
         })
         
         # Conditionally add common variations (standard suffixes)
-        if not original_model_id.startswith("gemini-2.0"): # Suffix rules based on original_model_id
+        if not base_model_without_prefix.startswith("gemini-2.0"): # Suffix rules based on original_model_id
             standard_suffixes = ["-search", "-encrypt", "-encrypt-full", "-auto"]
             for suffix in standard_suffixes:
-                # Suffix is applied to the original model ID part
-                suffixed_model_part = f"{original_model_id}{suffix}"
-                # Then the whole thing is prefixed
-                final_suffixed_display_id = f"{current_display_prefix}{suffixed_model_part}"
+                # Suffix is applied to the original model ID part (without EXPRESS prefix)
+                suffixed_model_part = f"{base_model_without_prefix}{suffix}"
+                # Then the whole thing is prefixed - retain EXPRESS prefix if it was there
+                final_suffixed_display_id = ""
+                if is_express_model:
+                    final_suffixed_display_id = f"[EXPRESS] {suffixed_model_part}"
+                else:
+                    # 非EXPRESS模型的后缀版本也需要正确显示PAY前缀
+                    final_suffixed_display_id = f"{current_display_prefix}{suffixed_model_part}"
                 
                 # Check if this suffixed ID is already in all_model_ids (unlikely with prefix) or already added
                 if final_suffixed_display_id not in all_model_ids and not any(m['id'] == final_suffixed_display_id for m in dynamic_models_data):
+                    vertex_log('info', f"添加后缀模型到列表: {final_suffixed_display_id}")
                     dynamic_models_data.append({
                         "id": final_suffixed_display_id, "object": "model", "created": current_time, "owned_by": "google",
-                        "permission": [], "root": original_model_id, "parent": None
+                        "permission": [], "root": base_model_without_prefix, "parent": None
                     })
         
         # Apply special suffixes for models starting with "gemini-2.5-flash"
-        if original_model_id.startswith("gemini-2.5-flash"): # Suffix rules based on original_model_id
+        if base_model_without_prefix.startswith("gemini-2.5-flash"): # Suffix rules based on original_model_id
             special_flash_suffixes = ["-nothinking", "-max"]
             for special_suffix in special_flash_suffixes:
-                suffixed_model_part = f"{original_model_id}{special_suffix}"
-                final_special_suffixed_display_id = f"{current_display_prefix}{suffixed_model_part}"
+                suffixed_model_part = f"{base_model_without_prefix}{special_suffix}"
+                # Retain EXPRESS prefix if original model had it
+                if is_express_model:
+                    final_special_suffixed_display_id = f"[EXPRESS] {suffixed_model_part}"
+                else:
+                    # 非EXPRESS模型的特殊后缀版本也需要正确显示PAY前缀
+                    final_special_suffixed_display_id = f"{current_display_prefix}{suffixed_model_part}"
 
                 if final_special_suffixed_display_id not in all_model_ids and not any(m['id'] == final_special_suffixed_display_id for m in dynamic_models_data):
+                    vertex_log('info', f"添加特殊后缀模型到列表: {final_special_suffixed_display_id}")
                     dynamic_models_data.append({
                         "id": final_special_suffixed_display_id, "object": "model", "created": current_time, "owned_by": "google",
-                        "permission": [], "root": original_model_id, "parent": None
+                        "permission": [], "root": base_model_without_prefix, "parent": None
                     })
 
         # Ensure uniqueness again after adding suffixes
@@ -250,6 +275,7 @@ async def list_models(fastapi_request: Request, api_key: str = Depends(get_api_k
                 if EXPERIMENTAL_MARKER in base_model_id_for_openai:
                     display_model_id = f"{base_model_id_for_openai}{OPENAI_DIRECT_SUFFIX}"
                 else:
+                    # OpenAI直接模式下也应该保持PAY前缀
                     display_model_id = f"{PAY_PREFIX}{base_model_id_for_openai}{OPENAI_DIRECT_SUFFIX}"
                 
                 # Check if already added (e.g. if remote config somehow already listed it or added as a base model)
