From 6aab140ec21729da501dc497722def16df6984aa Mon Sep 17 00:00:00 2001 From: snaily Date: Sat, 17 May 2025 00:13:49 +0800 Subject: [PATCH] =?UTF-8?q?feat(vertex):=20=E9=9B=86=E6=88=90=20Vertex=20A?= =?UTF-8?q?I=20Express=20API=20=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 本次更新引入了对 Google Vertex AI Express API 的支持，允许用户配置和使用 Vertex AI 模型。 主要变更包括： 后端： - 新增 `VERTEX_API_KEYS` 和 `VERTEX_EXPRESS_BASE_URL` 至系统配置 ([`.env.example`](.env.example:13), [`app/config/config.py:62`](app/config/config.py:62), [`app/database/models.py`](app/database/models.py), [`app/database/services.py`](app/database/services.py))。 - 实现 `VertexExpressChatService` ([`app/service/chat/vertex_express_chat_service.py`](app/service/chat/vertex_express_chat_service.py)) 用于处理与 Vertex AI Express API 的交互。 - 添加 `vertex_express_routes` ([`app/router/vertex_express_routes.py`](app/router/vertex_express_routes.py)) 来暴露 Vertex AI 相关的 API 端点，并集成到主应用 ([`app/core/application.py:36`](app/core/application.py:36), [`app/router/routes.py:15`](app/router/routes.py:15))。 - 更新密钥管理器 ([`app/service/key/key_manager.py`](app/service/key/key_manager.py)) 以支持 Vertex API 密钥的获取、检查和轮换。 前端 (配置编辑器): - 在配置页面 ([`app/templates/config_editor.html:463`](app/templates/config_editor.html:463)) 添加了 Vertex API 密钥列表和 Vertex Express API 基础 URL 的表单字段。 - 实现了批量添加和删除 Vertex API 密钥的功能，包括相应的模态框和操作逻辑 ([`app/static/js/config_editor.js:550`](app/static/js/config_editor.js:550), [`app/static/js/config_editor.js:1097`](app/static/js/config_editor.js:1097), [`app/templates/config_editor.html:1657`](app/templates/config_editor.html:1657))。 - 确保新的配置项在初始化 ([`app/static/js/config_editor.js:598`](app/static/js/config_editor.js:598)) 和表单填充 ([`app/static/js/config_editor.js:671`](app/static/js/config_editor.js:671)) 时得到正确处理。 - 更新了数组项添加逻辑以识别 `VERTEX_API_KEYS` 为敏感字段 ([`app/static/js/config_editor.js:1235`](app/static/js/config_editor.js:1235))。 此功能扩展了应用支持的 AI 服务范围，为用户提供了更多模型选择。 --- .env.example | 4 + app/config/config.py | 31 +- app/core/application.py | 4 +- app/database/connection.py | 5 +- app/database/models.py | 4 +- app/database/services.py | 10 +- app/log/logger.py | 5 + app/middleware/middleware.py | 11 +- app/router/routes.py | 3 +- app/router/vertex_express_routes.py | 146 +++++++++ app/service/chat/gemini_chat_service.py | 4 +- app/service/chat/openai_chat_service.py | 2 +- .../chat/vertex_express_chat_service.py | 277 ++++++++++++++++++ app/service/config/config_service.py | 12 +- app/service/error_log/error_log_service.py | 7 +- app/service/key/key_manager.py | 218 +++++++++++++- .../openai_compatiable_service.py | 2 +- app/static/js/config_editor.js | 226 +++++++++++++- app/templates/config_editor.html | 146 ++++++++- 19 files changed, 1040 insertions(+), 77 deletions(-) create mode 100644 app/router/vertex_express_routes.py create mode 100644 app/service/chat/vertex_express_chat_service.py diff --git a/.env.example b/.env.example index cded0e2..86138a9 100644 --- a/.env.example +++ b/.env.example @@ -10,6 +10,10 @@ MYSQL_DATABASE=default_db API_KEYS=["AIzaSyxxxxxxxxxxxxxxxxxxx","AIzaSyxxxxxxxxxxxxxxxxxxx"] ALLOWED_TOKENS=["sk-123456"] AUTH_TOKEN=sk-123456 +# For Vertex AI Platform API Keys +VERTEX_API_KEYS=["AQ.Abxxxxxxxxxxxxxxxxxxx"] +# For Vertex AI Platform Express API Base URL +VERTEX_EXPRESS_BASE_URL=https://aiplatform.googleapis.com/v1beta1/publishers/google TEST_MODEL=gemini-1.5-flash THINKING_MODELS=["gemini-2.5-flash-preview-04-17"] THINKING_BUDGET_MAP={"gemini-2.5-flash-preview-04-17": 4000} diff --git a/app/config/config.py b/app/config/config.py index 72b4a5c..82a938b 100644 --- a/app/config/config.py +++ b/app/config/config.py @@ -59,8 +59,10 @@ class Settings(BaseSettings): TEST_MODEL: str = DEFAULT_MODEL TIME_OUT: int = DEFAULT_TIMEOUT MAX_RETRIES: int = MAX_RETRIES - PROXIES: List[str] = [] # 新增：代理服务器列表 - + PROXIES: List[str] = [] + VERTEX_API_KEYS: List[str] = [] + VERTEX_EXPRESS_BASE_URL: str = "https://aiplatform.googleapis.com/v1beta1/publishers/google" + # 模型相关配置 SEARCH_MODELS: List[str] = ["gemini-2.0-flash-exp"] IMAGE_MODELS: List[str] = ["gemini-2.0-flash-exp"] @@ -68,8 +70,8 @@ class Settings(BaseSettings): TOOLS_CODE_EXECUTION_ENABLED: bool = False SHOW_SEARCH_LINK: bool = True SHOW_THINKING_PROCESS: bool = True - THINKING_MODELS: List[str] = [] # 新增：用于思考过程的模型列表 - THINKING_BUDGET_MAP: Dict[str, float] = {} # 新增：模型对应的预算映射 + THINKING_MODELS: List[str] = [] + THINKING_BUDGET_MAP: Dict[str, float] = {} # 图像生成相关配置 PAID_KEY: str = "" @@ -101,12 +103,12 @@ class Settings(BaseSettings): GITHUB_REPO_NAME: str = "gemini-balance" # 日志配置 - LOG_LEVEL: str = "INFO" # 默认日志级别 - AUTO_DELETE_ERROR_LOGS_ENABLED: bool = True # 是否开启自动删除错误日志 - AUTO_DELETE_ERROR_LOGS_DAYS: int = 7 # 自动删除多少天前的错误日志 (1, 7, 30) - AUTO_DELETE_REQUEST_LOGS_ENABLED: bool = False # 是否开启自动删除请求日志 - AUTO_DELETE_REQUEST_LOGS_DAYS: int = 30 # 自动删除多少天前的请求日志 (1, 7, 30) - SAFETY_SETTINGS: List[Dict[str, str]] = DEFAULT_SAFETY_SETTINGS # 新增：安全设置 + LOG_LEVEL: str = "INFO" + AUTO_DELETE_ERROR_LOGS_ENABLED: bool = True + AUTO_DELETE_ERROR_LOGS_DAYS: int = 7 + AUTO_DELETE_REQUEST_LOGS_ENABLED: bool = False + AUTO_DELETE_REQUEST_LOGS_DAYS: int = 30 + SAFETY_SETTINGS: List[Dict[str, str]] = DEFAULT_SAFETY_SETTINGS def __init__(self, **kwargs): super().__init__(**kwargs) @@ -141,7 +143,6 @@ def _parse_db_value(key: str, db_value: str, target_type: Type) -> Any: elif target_type == Dict[str, float]: parsed_dict = {} try: - # First attempt: standard JSON parsing parsed = json.loads(db_value) if isinstance(parsed, dict): parsed_dict = {str(k): float(v) for k, v in parsed.items()} @@ -150,7 +151,6 @@ def _parse_db_value(key: str, db_value: str, target_type: Type) -> Any: f"Parsed DB value for key '{key}' is not a dictionary type. Value: {db_value}" ) except (json.JSONDecodeError, ValueError, TypeError) as e1: - # Second attempt: try replacing single quotes if JSONDecodeError occurred if isinstance(e1, json.JSONDecodeError) and "'" in db_value: logger.warning( f"Failed initial JSON parse for key '{key}'. Attempting to replace single quotes. Error: {e1}" @@ -169,11 +169,10 @@ def _parse_db_value(key: str, db_value: str, target_type: Type) -> Any: f"Could not parse '{db_value}' as Dict[str, float] for key '{key}' even after replacing quotes: {e2}. Returning empty dict." ) else: - # Log other errors (ValueError, TypeError) or JSON errors without single quotes logger.error( f"Could not parse '{db_value}' as Dict[str, float] for key '{key}': {e1}. Returning empty dict." ) - return parsed_dict # Return the parsed dict or an empty one if all attempts fail + return parsed_dict # 处理 List[Dict[str, str]] elif target_type == List[Dict[str, str]]: try: @@ -192,7 +191,7 @@ def _parse_db_value(key: str, db_value: str, target_type: Type) -> Any: logger.warning( f"Invalid structure in List[Dict[str, str]] for key '{key}'. Value: {db_value}" ) - return [] # 或者返回默认值？这里返回空列表 + return [] else: logger.warning( f"Parsed DB value for key '{key}' is not a list type. Value: {db_value}" @@ -374,7 +373,7 @@ async def sync_initial_settings(): data = { "key": key, "value": db_value, - "description": f"{key} configuration setting", # 默认描述 + "description": f"{key} configuration setting", "updated_at": now, } diff --git a/app/core/application.py b/app/core/application.py index 396d9bc..16b074f 100644 --- a/app/core/application.py +++ b/app/core/application.py @@ -15,7 +15,7 @@ from app.router.routes import setup_routers from app.scheduler.scheduled_tasks import start_scheduler, stop_scheduler from app.service.key.key_manager import get_key_manager_instance from app.service.update.update_service import check_for_updates -from app.utils.helpers import get_current_version # Import from helpers +from app.utils.helpers import get_current_version logger = get_application_logger() @@ -43,7 +43,7 @@ async def _setup_database_and_config(app_settings): logger.info("Database initialized successfully") await connect_to_db() await sync_initial_settings() - await get_key_manager_instance(app_settings.API_KEYS) + await get_key_manager_instance(app_settings.API_KEYS, app_settings.VERTEX_API_KEYS) logger.info("Database, config sync, and KeyManager initialized successfully") diff --git a/app/database/connection.py b/app/database/connection.py index 7cd113c..a962503 100644 --- a/app/database/connection.py +++ b/app/database/connection.py @@ -45,11 +45,8 @@ Base = declarative_base(metadata=metadata) if settings.DATABASE_TYPE == "sqlite": database = Database(DATABASE_URL) else: - database = Database(DATABASE_URL, min_size=5, max_size=20, pool_recycle=1800) # Reduced recycle time to 30 mins + database = Database(DATABASE_URL, min_size=5, max_size=20, pool_recycle=1800) -# 移除了 SessionLocal 和 get_db 函数 - -# --- Async connection functions for lifespan/async routes --- async def connect_to_db(): """ 连接到数据库 diff --git a/app/database/models.py b/app/database/models.py index 04ad604..c33ae6a 100644 --- a/app/database/models.py +++ b/app/database/models.py @@ -2,7 +2,7 @@ 数据库模型模块 """ import datetime -from sqlalchemy import Column, Integer, String, Text, DateTime, JSON, Boolean # 添加 Boolean +from sqlalchemy import Column, Integer, String, Text, DateTime, JSON, Boolean from app.database.connection import Base @@ -53,7 +53,7 @@ class RequestLog(Base): id = Column(Integer, primary_key=True, autoincrement=True) request_time = Column(DateTime, default=datetime.datetime.now, comment="请求时间") model_name = Column(String(100), nullable=True, comment="模型名称") - api_key = Column(String(100), nullable=True, comment="使用的API密钥") # 考虑安全性，后续可优化 + api_key = Column(String(100), nullable=True, comment="使用的API密钥") is_success = Column(Boolean, nullable=False, comment="请求是否成功") status_code = Column(Integer, nullable=True, comment="API响应状态码") latency_ms = Column(Integer, nullable=True, comment="请求耗时(毫秒)") diff --git a/app/database/services.py b/app/database/services.py index 723cdc5..893b608 100644 --- a/app/database/services.py +++ b/app/database/services.py @@ -189,7 +189,6 @@ async def get_error_logs( ErrorLog.request_time ) - # Apply filters if key_search: query = query.where(ErrorLog.gemini_key.ilike(f"%{key_search}%")) if error_search: @@ -219,14 +218,14 @@ async def get_error_logs( result = await database.fetch_all(query) return [dict(row) for row in result] except Exception as e: - logger.exception(f"Failed to get error logs with filters: {str(e)}") # Use exception for stack trace + logger.exception(f"Failed to get error logs with filters: {str(e)}") raise async def get_error_logs_count( key_search: Optional[str] = None, error_search: Optional[str] = None, - error_code_search: Optional[str] = None, # Added error code search + error_code_search: Optional[str] = None, start_date: Optional[datetime] = None, end_date: Optional[datetime] = None ) -> int: @@ -294,7 +293,7 @@ async def get_error_log_details(log_id: int) -> Optional[Dict[str, Any]]: try: log_dict['request_msg'] = json.dumps(log_dict['request_msg'], ensure_ascii=False, indent=2) except TypeError: - log_dict['request_msg'] = str(log_dict['request_msg']) # Fallback to string + log_dict['request_msg'] = str(log_dict['request_msg']) return log_dict else: return None @@ -372,7 +371,6 @@ async def delete_all_error_logs() -> int: try: # 1. 获取删除前的总数 count_query = select(func.count()).select_from(ErrorLog) - # fetch_val() is suitable here as we expect a single scalar value total_to_delete = await database.fetch_val(count_query) if total_to_delete == 0: @@ -380,7 +378,6 @@ async def delete_all_error_logs() -> int: return 0 # 2. 执行删除操作 - # This creates a query like "DELETE FROM error_log" delete_query = delete(ErrorLog) await database.execute(delete_query) @@ -388,7 +385,6 @@ async def delete_all_error_logs() -> int: return total_to_delete except Exception as e: logger.error(f"Failed to delete all error logs: {str(e)}", exc_info=True) - # Re-raise the exception so it can be caught by the service layer or route handler raise diff --git a/app/log/logger.py b/app/log/logger.py index 37e74fc..1614a46 100644 --- a/app/log/logger.py +++ b/app/log/logger.py @@ -226,3 +226,8 @@ def get_error_log_logger(): def get_request_log_logger(): return Logger.setup_logger("request_log") + + +def get_vertex_express_logger(): + return Logger.setup_logger("vertex_express") + diff --git a/app/middleware/middleware.py b/app/middleware/middleware.py index 495d823..05dded3 100644 --- a/app/middleware/middleware.py +++ b/app/middleware/middleware.py @@ -32,6 +32,7 @@ class AuthMiddleware(BaseHTTPMiddleware): and not request.url.path.startswith("/hf") and not request.url.path.startswith("/openai") and not request.url.path.startswith("/api/version/check") + and not request.url.path.startswith("/vertex-express") ): auth_token = request.cookies.get("auth_token") @@ -60,7 +61,7 @@ def setup_middlewares(app: FastAPI) -> None: # 配置CORS中间件 app.add_middleware( CORSMiddleware, - allow_origins=["*"], # 生产环境建议配置具体的域名 + allow_origins=["*"], allow_credentials=True, allow_methods=[ "GET", @@ -68,8 +69,8 @@ def setup_middlewares(app: FastAPI) -> None: "PUT", "DELETE", "OPTIONS", - ], # 明确指定允许的HTTP方法 - allow_headers=["*"], # 生产环境建议配置具体的请求头 - expose_headers=["*"], # 允许前端访问的响应头 - max_age=600, # 预检请求缓存时间(秒) + ], + allow_headers=["*"], + expose_headers=["*"], + max_age=600, ) diff --git a/app/router/routes.py b/app/router/routes.py index 4723139..46004d7 100644 --- a/app/router/routes.py +++ b/app/router/routes.py @@ -8,7 +8,7 @@ from fastapi.templating import Jinja2Templates from app.core.security import verify_auth_token from app.log.logger import get_routes_logger -from app.router import error_log_routes, gemini_routes, openai_routes, config_routes, scheduler_routes, stats_routes, version_routes, openai_compatiable_routes +from app.router import error_log_routes, gemini_routes, openai_routes, config_routes, scheduler_routes, stats_routes, version_routes, openai_compatiable_routes, vertex_express_routes from app.service.key.key_manager import get_key_manager_instance from app.service.stats.stats_service import StatsService @@ -33,6 +33,7 @@ def setup_routers(app: FastAPI) -> None: app.include_router(stats_routes.router) app.include_router(version_routes.router) app.include_router(openai_compatiable_routes.router) + app.include_router(vertex_express_routes.router) setup_page_routes(app) diff --git a/app/router/vertex_express_routes.py b/app/router/vertex_express_routes.py new file mode 100644 index 0000000..3b7f0f8 --- /dev/null +++ b/app/router/vertex_express_routes.py @@ -0,0 +1,146 @@ +from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import StreamingResponse +from copy import deepcopy +from app.config.config import settings +from app.log.logger import get_vertex_express_logger +from app.core.security import SecurityService +from app.domain.gemini_models import GeminiRequest +from app.service.chat.vertex_express_chat_service import GeminiChatService +from app.service.key.key_manager import KeyManager, get_key_manager_instance +from app.service.model.model_service import ModelService +from app.handler.retry_handler import RetryHandler +from app.handler.error_handler import handle_route_errors +from app.core.constants import API_VERSION + +router = APIRouter(prefix=f"/vertex-express/{API_VERSION}") +logger = get_vertex_express_logger() + +security_service = SecurityService() +model_service = ModelService() + + +async def get_key_manager(): + """获取密钥管理器实例""" + return await get_key_manager_instance() + + +async def get_next_working_key(key_manager: KeyManager = Depends(get_key_manager)): + """获取下一个可用的API密钥""" + return await key_manager.get_next_working_vertex_key() + + +async def get_chat_service(key_manager: KeyManager = Depends(get_key_manager)): + """获取Gemini聊天服务实例""" + return GeminiChatService(settings.VERTEX_EXPRESS_BASE_URL, key_manager) + + +@router.get("/models") +async def list_models( + _=Depends(security_service.verify_key_or_goog_api_key), + key_manager: KeyManager = Depends(get_key_manager) +): + """获取可用的 Gemini 模型列表，并根据配置添加衍生模型（搜索、图像、非思考）。""" + operation_name = "list_gemini_models" + logger.info("-" * 50 + operation_name + "-" * 50) + logger.info("Handling Gemini models list request") + + try: + api_key = await key_manager.get_first_valid_key() + if not api_key: + raise HTTPException(status_code=503, detail="No valid API keys available to fetch models.") + logger.info(f"Using API key: {api_key}") + + models_data = await model_service.get_gemini_models(api_key) + if not models_data or "models" not in models_data: + raise HTTPException(status_code=500, detail="Failed to fetch base models list.") + + models_json = deepcopy(models_data) + model_mapping = {x.get("name", "").split("/", maxsplit=1)[-1]: x for x in models_json.get("models", [])} + + def add_derived_model(base_name, suffix, display_suffix): + model = model_mapping.get(base_name) + if not model: + logger.warning(f"Base model '{base_name}' not found for derived model '{suffix}'.") + return + item = deepcopy(model) + item["name"] = f"models/{base_name}{suffix}" + display_name = f'{item.get("displayName", base_name)}{display_suffix}' + item["displayName"] = display_name + item["description"] = display_name + models_json["models"].append(item) + + if settings.SEARCH_MODELS: + for name in settings.SEARCH_MODELS: + add_derived_model(name, "-search", " For Search") + if settings.IMAGE_MODELS: + for name in settings.IMAGE_MODELS: + add_derived_model(name, "-image", " For Image") + if settings.THINKING_MODELS: + for name in settings.THINKING_MODELS: + add_derived_model(name, "-non-thinking", " Non Thinking") + + logger.info("Gemini models list request successful") + return models_json + except HTTPException as http_exc: + raise http_exc + except Exception as e: + logger.error(f"Error getting Gemini models list: {str(e)}") + raise HTTPException( + status_code=500, detail="Internal server error while fetching Gemini models list" + ) from e + + +@router.post("/models/{model_name}:generateContent") +@RetryHandler(key_arg="api_key") +async def generate_content( + model_name: str, + request: GeminiRequest, + _=Depends(security_service.verify_key_or_goog_api_key), + api_key: str = Depends(get_next_working_key), + key_manager: KeyManager = Depends(get_key_manager), + chat_service: GeminiChatService = Depends(get_chat_service) +): + """处理 Gemini 非流式内容生成请求。""" + operation_name = "gemini_generate_content" + async with handle_route_errors(logger, operation_name, failure_message="Content generation failed"): + logger.info(f"Handling Gemini content generation request for model: {model_name}") + logger.debug(f"Request: \n{request.model_dump_json(indent=2)}") + logger.info(f"Using API key: {api_key}") + + if not await model_service.check_model_support(model_name): + raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported") + + response = await chat_service.generate_content( + model=model_name, + request=request, + api_key=api_key + ) + return response + + +@router.post("/models/{model_name}:streamGenerateContent") +@RetryHandler(key_arg="api_key") +async def stream_generate_content( + model_name: str, + request: GeminiRequest, + _=Depends(security_service.verify_key_or_goog_api_key), + api_key: str = Depends(get_next_working_key), + key_manager: KeyManager = Depends(get_key_manager), + chat_service: GeminiChatService = Depends(get_chat_service) +): + """处理 Gemini 流式内容生成请求。""" + operation_name = "gemini_stream_generate_content" + async with handle_route_errors(logger, operation_name, failure_message="Streaming request initiation failed"): + logger.info(f"Handling Gemini streaming content generation for model: {model_name}") + logger.debug(f"Request: \n{request.model_dump_json(indent=2)}") + logger.info(f"Using API key: {api_key}") + + if not await model_service.check_model_support(model_name): + raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported") + + response_stream = chat_service.stream_generate_content( + model=model_name, + request=request, + api_key=api_key + ) + return StreamingResponse(response_stream, media_type="text/event-stream") \ No newline at end of file diff --git a/app/service/chat/gemini_chat_service.py b/app/service/chat/gemini_chat_service.py index 313cb89..8fa7425 100644 --- a/app/service/chat/gemini_chat_service.py +++ b/app/service/chat/gemini_chat_service.py @@ -131,7 +131,7 @@ class GeminiChatService: self, original_response: Dict[str, Any], text: str ) -> Dict[str, Any]: """创建包含指定文本的响应""" - response_copy = json.loads(json.dumps(original_response)) # 深拷贝 + response_copy = json.loads(json.dumps(original_response)) if response_copy.get("candidates") and response_copy["candidates"][0].get( "content", {} ).get("parts"): @@ -200,7 +200,7 @@ class GeminiChatService: request_datetime = datetime.datetime.now() start_time = time.perf_counter() current_attempt_key = api_key - final_api_key = current_attempt_key # Update final key used + final_api_key = current_attempt_key try: async for line in self.api_client.stream_generate_content( payload, model, current_attempt_key diff --git a/app/service/chat/openai_chat_service.py b/app/service/chat/openai_chat_service.py index 75ccc3b..2b2da4d 100644 --- a/app/service/chat/openai_chat_service.py +++ b/app/service/chat/openai_chat_service.py @@ -51,7 +51,7 @@ def _build_tools( or model.endswith("-image") or model.endswith("-image-generation") ) - and not _has_media_parts(messages) # Use the updated check + and not _has_media_parts(messages) ): tool["codeExecution"] = {} logger.debug("Code execution tool enabled.") diff --git a/app/service/chat/vertex_express_chat_service.py b/app/service/chat/vertex_express_chat_service.py new file mode 100644 index 0000000..313cb89 --- /dev/null +++ b/app/service/chat/vertex_express_chat_service.py @@ -0,0 +1,277 @@ +# app/services/chat_service.py + +import json +import re +import datetime +import time +from typing import Any, AsyncGenerator, Dict, List +from app.config.config import settings +from app.core.constants import GEMINI_2_FLASH_EXP_SAFETY_SETTINGS +from app.domain.gemini_models import GeminiRequest +from app.handler.response_handler import GeminiResponseHandler +from app.handler.stream_optimizer import gemini_optimizer +from app.log.logger import get_gemini_logger +from app.service.client.api_client import GeminiApiClient +from app.service.key.key_manager import KeyManager +from app.database.services import add_error_log, add_request_log + +logger = get_gemini_logger() + + +def _has_image_parts(contents: List[Dict[str, Any]]) -> bool: + """判断消息是否包含图片部分""" + for content in contents: + if "parts" in content: + for part in content["parts"]: + if "image_url" in part or "inline_data" in part: + return True + return False + + +def _build_tools(model: str, payload: Dict[str, Any]) -> List[Dict[str, Any]]: + """构建工具""" + + def _merge_tools(tools: List[Dict[str, Any]]) -> Dict[str, Any]: + record = dict() + for item in tools: + if not item or not isinstance(item, dict): + continue + + for k, v in item.items(): + if k == "functionDeclarations" and v and isinstance(v, list): + functions = record.get("functionDeclarations", []) + functions.extend(v) + record["functionDeclarations"] = functions + else: + record[k] = v + return record + + tool = dict() + if payload and isinstance(payload, dict) and "tools" in payload: + if payload.get("tools") and isinstance(payload.get("tools"), dict): + payload["tools"] = [payload.get("tools")] + items = payload.get("tools", []) + if items and isinstance(items, list): + tool.update(_merge_tools(items)) + + if ( + settings.TOOLS_CODE_EXECUTION_ENABLED + and not (model.endswith("-search") or "-thinking" in model) + and not _has_image_parts(payload.get("contents", [])) + ): + tool["codeExecution"] = {} + if model.endswith("-search"): + tool["googleSearch"] = {} + + # 解决 "Tool use with function calling is unsupported" 问题 + if tool.get("functionDeclarations"): + tool.pop("googleSearch", None) + tool.pop("codeExecution", None) + + return [tool] if tool else [] + + +def _get_safety_settings(model: str) -> List[Dict[str, str]]: + """获取安全设置""" + if model == "gemini-2.0-flash-exp": + return GEMINI_2_FLASH_EXP_SAFETY_SETTINGS + return settings.SAFETY_SETTINGS + + +def _build_payload(model: str, request: GeminiRequest) -> Dict[str, Any]: + """构建请求payload""" + request_dict = request.model_dump() + if request.generationConfig: + if request.generationConfig.maxOutputTokens is None: + # 如果未指定最大输出长度，则不传递该字段，解决截断的问题 + request_dict["generationConfig"].pop("maxOutputTokens") + + payload = { + "contents": request_dict.get("contents", []), + "tools": _build_tools(model, request_dict), + "safetySettings": _get_safety_settings(model), + "generationConfig": request_dict.get("generationConfig"), + "systemInstruction": request_dict.get("systemInstruction"), + } + + if model.endswith("-image") or model.endswith("-image-generation"): + payload.pop("systemInstruction") + payload["generationConfig"]["responseModalities"] = ["Text", "Image"] + + if model.endswith("-non-thinking"): + payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 0} + if model in settings.THINKING_BUDGET_MAP: + payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": settings.THINKING_BUDGET_MAP.get(model,1000)} + + return payload + + +class GeminiChatService: + """聊天服务""" + + def __init__(self, base_url: str, key_manager: KeyManager): + self.api_client = GeminiApiClient(base_url, settings.TIME_OUT) + self.key_manager = key_manager + self.response_handler = GeminiResponseHandler() + + def _extract_text_from_response(self, response: Dict[str, Any]) -> str: + """从响应中提取文本内容""" + if not response.get("candidates"): + return "" + + candidate = response["candidates"][0] + content = candidate.get("content", {}) + parts = content.get("parts", []) + + if parts and "text" in parts[0]: + return parts[0].get("text", "") + return "" + + def _create_char_response( + self, original_response: Dict[str, Any], text: str + ) -> Dict[str, Any]: + """创建包含指定文本的响应""" + response_copy = json.loads(json.dumps(original_response)) # 深拷贝 + if response_copy.get("candidates") and response_copy["candidates"][0].get( + "content", {} + ).get("parts"): + response_copy["candidates"][0]["content"]["parts"][0]["text"] = text + return response_copy + + async def generate_content( + self, model: str, request: GeminiRequest, api_key: str + ) -> Dict[str, Any]: + """生成内容""" + payload = _build_payload(model, request) + start_time = time.perf_counter() + request_datetime = datetime.datetime.now() + is_success = False + status_code = None + response = None + + try: + response = await self.api_client.generate_content(payload, model, api_key) + is_success = True + status_code = 200 + return self.response_handler.handle_response(response, model, stream=False) + except Exception as e: + is_success = False + error_log_msg = str(e) + logger.error(f"Normal API call failed with error: {error_log_msg}") + match = re.search(r"status code (\d+)", error_log_msg) + if match: + status_code = int(match.group(1)) + else: + status_code = 500 + + await add_error_log( + gemini_key=api_key, + model_name=model, + error_type="gemini-chat-non-stream", + error_log=error_log_msg, + error_code=status_code, + request_msg=payload + ) + raise e + finally: + end_time = time.perf_counter() + latency_ms = int((end_time - start_time) * 1000) + await add_request_log( + model_name=model, + api_key=api_key, + is_success=is_success, + status_code=status_code, + latency_ms=latency_ms, + request_time=request_datetime + ) + + async def stream_generate_content( + self, model: str, request: GeminiRequest, api_key: str + ) -> AsyncGenerator[str, None]: + """流式生成内容""" + retries = 0 + max_retries = settings.MAX_RETRIES + payload = _build_payload(model, request) + is_success = False + status_code = None + final_api_key = api_key + + while retries < max_retries: + request_datetime = datetime.datetime.now() + start_time = time.perf_counter() + current_attempt_key = api_key + final_api_key = current_attempt_key # Update final key used + try: + async for line in self.api_client.stream_generate_content( + payload, model, current_attempt_key + ): + # print(line) + if line.startswith("data:"): + line = line[6:] + response_data = self.response_handler.handle_response( + json.loads(line), model, stream=True + ) + text = self._extract_text_from_response(response_data) + # 如果有文本内容，且开启了流式输出优化器，则使用流式输出优化器处理 + if text and settings.STREAM_OPTIMIZER_ENABLED: + # 使用流式输出优化器处理文本输出 + async for ( + optimized_chunk + ) in gemini_optimizer.optimize_stream_output( + text, + lambda t: self._create_char_response(response_data, t), + lambda c: "data: " + json.dumps(c) + "\n\n", + ): + yield optimized_chunk + else: + # 如果没有文本内容（如工具调用等），整块输出 + yield "data: " + json.dumps(response_data) + "\n\n" + logger.info("Streaming completed successfully") + is_success = True + status_code = 200 + break + except Exception as e: + retries += 1 + is_success = False + error_log_msg = str(e) + logger.warning( + f"Streaming API call failed with error: {error_log_msg}. Attempt {retries} of {max_retries}" + ) + match = re.search(r"status code (\d+)", error_log_msg) + if match: + status_code = int(match.group(1)) + else: + status_code = 500 + + await add_error_log( + gemini_key=current_attempt_key, + model_name=model, + error_type="gemini-chat-stream", + error_log=error_log_msg, + error_code=status_code, + request_msg=payload + ) + + api_key = await self.key_manager.handle_api_failure(current_attempt_key, retries) + if api_key: + logger.info(f"Switched to new API key: {api_key}") + else: + logger.error(f"No valid API key available after {retries} retries.") + break + + if retries >= max_retries: + logger.error( + f"Max retries ({max_retries}) reached for streaming." + ) + break + finally: + end_time = time.perf_counter() + latency_ms = int((end_time - start_time) * 1000) + await add_request_log( + model_name=model, + api_key=final_api_key, + is_success=is_success, + status_code=status_code, + latency_ms=latency_ms, + request_time=request_datetime + ) diff --git a/app/service/config/config_service.py b/app/service/config/config_service.py index f47690b..10dfa21 100644 --- a/app/service/config/config_service.py +++ b/app/service/config/config_service.py @@ -55,7 +55,7 @@ class ConfigService: # 处理不同类型的值 if isinstance(value, list): db_value = json.dumps(value) - elif isinstance(value, dict): # 新增对 dict 类型的处理 + elif isinstance(value, dict): db_value = json.dumps(value) elif isinstance(value, bool): db_value = str(value).lower() @@ -115,7 +115,7 @@ class ConfigService: # 重置并重新初始化 KeyManager try: await reset_key_manager_instance() - await get_key_manager_instance(settings.API_KEYS) + await get_key_manager_instance(settings.API_KEYS, settings.VERTEX_API_KEYS) logger.info("KeyManager instance re-initialized with updated settings.") except Exception as e: logger.error(f"Failed to re-initialize KeyManager: {str(e)}") @@ -154,7 +154,7 @@ class ConfigService: deleted_count = 0 not_found_keys: List[str] = [] - current_api_keys = list(settings.API_KEYS) # 创建副本以进行修改 + current_api_keys = list(settings.API_KEYS) keys_actually_removed: List[str] = [] for key_to_del in keys_to_delete: @@ -166,7 +166,7 @@ class ConfigService: not_found_keys.append(key_to_del) if deleted_count > 0: - settings.API_KEYS = current_api_keys # 更新内存中的 settings + settings.API_KEYS = current_api_keys await ConfigService.update_config({"API_KEYS": settings.API_KEYS}) logger.info( f"成功删除 {deleted_count} 个密钥。密钥: {keys_actually_removed}" @@ -182,9 +182,9 @@ class ConfigService: } else: message = "没有密钥被删除。" - if not_found_keys: # 如果提供了密钥但都未找到 + if not_found_keys: message = f"所有 {len(not_found_keys)} 个指定的密钥均未找到: {not_found_keys}。" - elif not keys_to_delete: # 如果 keys_to_delete 列表为空 + elif not keys_to_delete: message = "未指定要删除的密钥。" logger.warning(message) return { diff --git a/app/service/error_log/error_log_service.py b/app/service/error_log/error_log_service.py index cc43995..3bd84c7 100644 --- a/app/service/error_log/error_log_service.py +++ b/app/service/error_log/error_log_service.py @@ -161,10 +161,9 @@ async def process_delete_all_error_logs() -> int: 返回删除的日志数量。 """ try: - # 确保数据库已连接 (如果适用，类似于 delete_old_error_logs) - # if not database.is_connected: - # await database.connect() - # logger.info("Database connection established for deleting all error logs.") + if not database.is_connected: + await database.connect() + logger.info("Database connection established for deleting all error logs.") deleted_count = await db_services.delete_all_error_logs() logger.info( diff --git a/app/service/key/key_manager.py b/app/service/key/key_manager.py index 1d00f98..827ead3 100644 --- a/app/service/key/key_manager.py +++ b/app/service/key/key_manager.py @@ -9,12 +9,19 @@ logger = get_key_manager_logger() class KeyManager: - def __init__(self, api_keys: list): + def __init__(self, api_keys: list, vertex_api_keys: list): self.api_keys = api_keys + self.vertex_api_keys = vertex_api_keys self.key_cycle = cycle(api_keys) + self.vertex_key_cycle = cycle(vertex_api_keys) self.key_cycle_lock = asyncio.Lock() + self.vertex_key_cycle_lock = asyncio.Lock() self.failure_count_lock = asyncio.Lock() + self.vertex_failure_count_lock = asyncio.Lock() self.key_failure_counts: Dict[str, int] = {key: 0 for key in api_keys} + self.vertex_key_failure_counts: Dict[str, int] = { + key: 0 for key in vertex_api_keys + } self.MAX_FAILURES = settings.MAX_FAILURES self.paid_key = settings.PAID_KEY @@ -26,17 +33,33 @@ class KeyManager: async with self.key_cycle_lock: return next(self.key_cycle) + async def get_next_vertex_key(self) -> str: + """获取下一个 Vertex API key""" + async with self.vertex_key_cycle_lock: + return next(self.vertex_key_cycle) + async def is_key_valid(self, key: str) -> bool: """检查key是否有效""" async with self.failure_count_lock: return self.key_failure_counts[key] < self.MAX_FAILURES + async def is_vertex_key_valid(self, key: str) -> bool: + """检查 Vertex key 是否有效""" + async with self.vertex_failure_count_lock: + return self.vertex_key_failure_counts[key] < self.MAX_FAILURES + async def reset_failure_counts(self): """重置所有key的失败计数""" async with self.failure_count_lock: for key in self.key_failure_counts: self.key_failure_counts[key] = 0 + async def reset_vertex_failure_counts(self): + """重置所有 Vertex key 的失败计数""" + async with self.vertex_failure_count_lock: + for key in self.vertex_key_failure_counts: + self.vertex_key_failure_counts[key] = 0 + async def reset_key_failure_count(self, key: str) -> bool: """重置指定key的失败计数""" async with self.failure_count_lock: @@ -49,6 +72,18 @@ class KeyManager: ) return False + async def reset_vertex_key_failure_count(self, key: str) -> bool: + """重置指定 Vertex key 的失败计数""" + async with self.vertex_failure_count_lock: + if key in self.vertex_key_failure_counts: + self.vertex_key_failure_counts[key] = 0 + logger.info(f"Reset failure count for Vertex key: {key}") + return True + logger.warning( + f"Attempt to reset failure count for non-existent Vertex key: {key}" + ) + return False + async def get_next_working_key(self) -> str: """获取下一可用的API key""" initial_key = await self.get_next_key() @@ -62,6 +97,19 @@ class KeyManager: if current_key == initial_key: return current_key + async def get_next_working_vertex_key(self) -> str: + """获取下一可用的 Vertex API key""" + initial_key = await self.get_next_vertex_key() + current_key = initial_key + + while True: + if await self.is_vertex_key_valid(current_key): + return current_key + + current_key = await self.get_next_vertex_key() + if current_key == initial_key: + return current_key + async def handle_api_failure(self, api_key: str, retries: int) -> str: """处理API调用失败""" async with self.failure_count_lock: @@ -75,10 +123,23 @@ class KeyManager: else: return "" + async def handle_vertex_api_failure(self, api_key: str, retries: int) -> str: + """处理 Vertex API 调用失败""" + async with self.vertex_failure_count_lock: + self.vertex_key_failure_counts[api_key] += 1 + if self.vertex_key_failure_counts[api_key] >= self.MAX_FAILURES: + logger.warning( + f"Vertex API key {api_key} has failed {self.MAX_FAILURES} times" + ) + def get_fail_count(self, key: str) -> int: """获取指定密钥的失败次数""" return self.key_failure_counts.get(key, 0) + def get_vertex_fail_count(self, key: str) -> int: + """获取指定 Vertex 密钥的失败次数""" + return self.vertex_key_failure_counts.get(key, 0) + async def get_keys_by_status(self) -> dict: """获取分类后的API key列表，包括失败次数""" valid_keys = {} @@ -94,6 +155,20 @@ class KeyManager: return {"valid_keys": valid_keys, "invalid_keys": invalid_keys} + async def get_vertex_keys_by_status(self) -> dict: + """获取分类后的 Vertex API key 列表，包括失败次数""" + valid_keys = {} + invalid_keys = {} + + async with self.vertex_failure_count_lock: + for key in self.vertex_api_keys: + fail_count = self.vertex_key_failure_counts[key] + if fail_count < self.MAX_FAILURES: + valid_keys[key] = fail_count + else: + invalid_keys[key] = fail_count + return {"valid_keys": valid_keys, "invalid_keys": invalid_keys} + async def get_first_valid_key(self) -> str: """获取第一个有效的API key""" async with self.failure_count_lock: @@ -111,19 +186,24 @@ class KeyManager: _singleton_instance = None _singleton_lock = asyncio.Lock() _preserved_failure_counts: Dict[str, int] | None = None +_preserved_vertex_failure_counts: Dict[str, int] | None = None _preserved_old_api_keys_for_reset: list | None = None +_preserved_vertex_old_api_keys_for_reset: list | None = None _preserved_next_key_in_cycle: str | None = None +_preserved_vertex_next_key_in_cycle: str | None = None -async def get_key_manager_instance(api_keys: list = None) -> KeyManager: +async def get_key_manager_instance( + api_keys: list = None, vertex_api_keys: list = None +) -> KeyManager: """ 获取 KeyManager 单例实例。 - 如果尚未创建实例，将使用提供的 api_keys 初始化 KeyManager。 + 如果尚未创建实例，将使用提供的 api_keys,vertex_api_keys 初始化 KeyManager。 如果已创建实例，则忽略 api_keys 参数，返回现有单例。 如果在重置后调用，会尝试恢复之前的状态（失败计数、循环位置）。 """ - global _singleton_instance, _preserved_failure_counts, _preserved_old_api_keys_for_reset, _preserved_next_key_in_cycle + global _singleton_instance, _preserved_failure_counts, _preserved_vertex_failure_counts, _preserved_old_api_keys_for_reset, _preserved_vertex_old_api_keys_for_reset, _preserved_next_key_in_cycle, _preserved_vertex_next_key_in_cycle async with _singleton_lock: if _singleton_instance is None: @@ -131,14 +211,23 @@ async def get_key_manager_instance(api_keys: list = None) -> KeyManager: raise ValueError( "API keys are required to initialize or re-initialize the KeyManager instance." ) + if vertex_api_keys is None: + raise ValueError( + "Vertex API keys are required to initialize or re-initialize the KeyManager instance." + ) + if not api_keys: logger.warning( "Initializing KeyManager with an empty list of API keys." ) + if not vertex_api_keys: + logger.warning( + "Initializing KeyManager with an empty list of Vertex API keys." + ) - _singleton_instance = KeyManager(api_keys) + _singleton_instance = KeyManager(api_keys, vertex_api_keys) logger.info( - f"KeyManager instance created/re-created with {len(api_keys)} API keys." + f"KeyManager instance created/re-created with {len(api_keys)} API keys and {len(vertex_api_keys)} Vertex API keys." ) # 1. 恢复失败计数 @@ -153,6 +242,19 @@ async def get_key_manager_instance(api_keys: list = None) -> KeyManager: logger.info("Inherited failure counts for applicable keys.") _preserved_failure_counts = None + if _preserved_vertex_failure_counts: + current_vertex_failure_counts = { + key: 0 for key in _singleton_instance.vertex_api_keys + } + for key, count in _preserved_vertex_failure_counts.items(): + if key in current_vertex_failure_counts: + current_vertex_failure_counts[key] = count + _singleton_instance.vertex_key_failure_counts = ( + current_vertex_failure_counts + ) + logger.info("Inherited failure counts for applicable Vertex keys.") + _preserved_vertex_failure_counts = None + # 2. 调整 key_cycle 的起始点 start_key_for_new_cycle = None if ( @@ -201,9 +303,7 @@ async def get_key_manager_instance(api_keys: list = None) -> KeyManager: f"Determined start key '{start_key_for_new_cycle}' not found in new API keys during cycle advancement. " "New cycle will start from the beginning." ) - except ( - StopIteration - ): + except StopIteration: logger.error( "StopIteration while advancing key cycle, implies empty new API key list previously missed." ) @@ -225,6 +325,76 @@ async def get_key_manager_instance(api_keys: list = None) -> KeyManager: _preserved_old_api_keys_for_reset = None _preserved_next_key_in_cycle = None + # 3. 调整 vertex_key_cycle 的起始点 + start_key_for_new_vertex_cycle = None + if ( + _preserved_vertex_old_api_keys_for_reset + and _preserved_vertex_next_key_in_cycle + and _singleton_instance.vertex_api_keys + ): + try: + start_idx_in_old = _preserved_vertex_old_api_keys_for_reset.index( + _preserved_vertex_next_key_in_cycle + ) + + for i in range(len(_preserved_vertex_old_api_keys_for_reset)): + current_old_key_idx = (start_idx_in_old + i) % len( + _preserved_vertex_old_api_keys_for_reset + ) + key_candidate = _preserved_vertex_old_api_keys_for_reset[ + current_old_key_idx + ] + if key_candidate in _singleton_instance.vertex_api_keys: + start_key_for_new_vertex_cycle = key_candidate + break + except ValueError: + logger.warning( + f"Preserved next key '{_preserved_vertex_next_key_in_cycle}' not found in preserved old Vertex API keys. " + "New cycle will start from the beginning of the new list." + ) + except Exception as e: + logger.error( + f"Error determining start key for new Vertex key cycle from preserved state: {e}. " + "New cycle will start from the beginning." + ) + + if start_key_for_new_vertex_cycle and _singleton_instance.vertex_api_keys: + try: + target_idx = _singleton_instance.vertex_api_keys.index( + start_key_for_new_vertex_cycle + ) + for _ in range(target_idx): + next(_singleton_instance.vertex_key_cycle) + logger.info( + f"Vertex key cycle in new instance advanced. Next call to get_next_vertex_key() will yield: {start_key_for_new_vertex_cycle}" + ) + except ValueError: + logger.warning( + f"Determined start key '{start_key_for_new_vertex_cycle}' not found in new Vertex API keys during cycle advancement. " + "New cycle will start from the beginning." + ) + except StopIteration: + logger.error( + "StopIteration while advancing Vertex key cycle, implies empty new Vertex API key list previously missed." + ) + except Exception as e: + logger.error( + f"Error advancing new Vertex key cycle: {e}. Cycle will start from beginning." + ) + else: + if _singleton_instance.vertex_api_keys: + logger.info( + "New Vertex key cycle will start from the beginning of the new Vertex API key list (no specific start key determined or needed)." + ) + else: + logger.info( + "New Vertex key cycle not applicable as the new Vertex API key list is empty." + ) + + # 清理所有保存的状态 + _preserved_vertex_old_api_keys_for_reset = None + _preserved_vertex_next_key_in_cycle = None + return _singleton_instance @@ -234,34 +404,52 @@ async def reset_key_manager_instance(): 将保存当前实例的状态（失败计数、旧 API keys、下一个 key 提示） 以供下一次 get_key_manager_instance 调用时恢复。 """ - global _singleton_instance, _preserved_failure_counts, _preserved_old_api_keys_for_reset, _preserved_next_key_in_cycle + global _singleton_instance, _preserved_failure_counts, _preserved_vertex_failure_counts, _preserved_old_api_keys_for_reset, _preserved_vertex_old_api_keys_for_reset, _preserved_next_key_in_cycle, _preserved_vertex_next_key_in_cycle async with _singleton_lock: if _singleton_instance: # 1. 保存失败计数 _preserved_failure_counts = _singleton_instance.key_failure_counts.copy() + _preserved_vertex_failure_counts = _singleton_instance.vertex_key_failure_counts.copy() # 2. 保存旧的 API keys 列表 _preserved_old_api_keys_for_reset = _singleton_instance.api_keys.copy() + _preserved_vertex_old_api_keys_for_reset = _singleton_instance.vertex_api_keys.copy() # 3. 保存 key_cycle 的下一个 key 提示 try: if _singleton_instance.api_keys: - _preserved_next_key_in_cycle = ( + _preserved_next_key_in_cycle = ( await _singleton_instance.get_next_key() ) else: _preserved_next_key_in_cycle = None - except ( - StopIteration - ): + except StopIteration: logger.warning( "Could not preserve next key hint: key cycle was empty or exhausted in old instance." ) - _preserved_next_key_in_cycle = None + _preserved_next_key_in_cycle = None except Exception as e: logger.error(f"Error preserving next key hint during reset: {e}") _preserved_next_key_in_cycle = None + # 4. 保存 vertex_key_cycle 的下一个 key 提示 + try: + if _singleton_instance.vertex_api_keys: + _preserved_vertex_next_key_in_cycle = ( + await _singleton_instance.get_next_vertex_key() + ) + else: + _preserved_vertex_next_key_in_cycle = None + except StopIteration: + logger.warning( + "Could not preserve next key hint: Vertex key cycle was empty or exhausted in old instance." + ) + _preserved_vertex_next_key_in_cycle = None + except Exception as e: + logger.error(f"Error preserving next key hint during reset: {e}") + _preserved_vertex_next_key_in_cycle = None + + _singleton_instance = None logger.info( "KeyManager instance has been reset. State (failure counts, old keys, next key hint) preserved for next instantiation." diff --git a/app/service/openai_compatiable/openai_compatiable_service.py b/app/service/openai_compatiable/openai_compatiable_service.py index 104c61f..51e062b 100644 --- a/app/service/openai_compatiable/openai_compatiable_service.py +++ b/app/service/openai_compatiable/openai_compatiable_service.py @@ -131,7 +131,7 @@ class OpenAICompatiableService: logger.info("Streaming completed successfully") is_success = True status_code = 200 - break # 成功后退出循环 + break except Exception as e: retries += 1 is_success = False diff --git a/app/static/js/config_editor.js b/app/static/js/config_editor.js index ac1471d..34ac593 100644 --- a/app/static/js/config_editor.js +++ b/app/static/js/config_editor.js @@ -10,8 +10,9 @@ const SHOW_CLASS = "show"; // For modals const API_KEY_REGEX = /AIzaSy\S{33}/g; const PROXY_REGEX = /(?:https?|socks5):\/\/(?:[^:@\/]+(?::[^@\/]+)?@)?(?:[^:\/\s]+)(?::\d+)?/g; +const VERTEX_API_KEY_REGEX = /AQ\.[a-zA-Z0-9_]{50}/g; // 新增 Vertex API Key 正则 const MASKED_VALUE = "••••••••"; - + // DOM Elements - Global Scope for frequently accessed elements const safetySettingsContainer = document.getElementById( "SAFETY_SETTINGS_container" @@ -30,7 +31,17 @@ const bulkDeleteProxyModal = document.getElementById("bulkDeleteProxyModal"); const bulkDeleteProxyInput = document.getElementById("bulkDeleteProxyInput"); const resetConfirmModal = document.getElementById("resetConfirmModal"); const configForm = document.getElementById("configForm"); // Added for frequent use - + +// Vertex API Key Modal Elements +const vertexApiKeyModal = document.getElementById("vertexApiKeyModal"); +const vertexApiKeyBulkInput = document.getElementById("vertexApiKeyBulkInput"); +const bulkDeleteVertexApiKeyModal = document.getElementById( + "bulkDeleteVertexApiKeyModal" +); +const bulkDeleteVertexApiKeyInput = document.getElementById( + "bulkDeleteVertexApiKeyInput" +); + // Model Helper Modal Elements const modelHelperModal = document.getElementById("modelHelperModal"); const modelHelperTitleElement = document.getElementById("modelHelperTitle"); @@ -244,6 +255,8 @@ document.addEventListener("DOMContentLoaded", function () { bulkDeleteApiKeyModal, proxyModal, bulkDeleteProxyModal, + vertexApiKeyModal, // 新增 + bulkDeleteVertexApiKeyModal, // 新增 modelHelperModal, ]; modals.forEach((modal) => { @@ -371,7 +384,71 @@ document.addEventListener("DOMContentLoaded", function () { } initializeSensitiveFields(); // Initialize sensitive field handling - + + // Vertex API Key Modal Elements and Events + const addVertexApiKeyBtn = document.getElementById("addVertexApiKeyBtn"); + const closeVertexApiKeyModalBtn = document.getElementById( + "closeVertexApiKeyModalBtn" + ); + const cancelAddVertexApiKeyBtn = document.getElementById( + "cancelAddVertexApiKeyBtn" + ); + const confirmAddVertexApiKeyBtn = document.getElementById( + "confirmAddVertexApiKeyBtn" + ); + const bulkDeleteVertexApiKeyBtn = document.getElementById( + "bulkDeleteVertexApiKeyBtn" + ); + const closeBulkDeleteVertexModalBtn = document.getElementById( + "closeBulkDeleteVertexModalBtn" + ); + const cancelBulkDeleteVertexApiKeyBtn = document.getElementById( + "cancelBulkDeleteVertexApiKeyBtn" + ); + const confirmBulkDeleteVertexApiKeyBtn = document.getElementById( + "confirmBulkDeleteVertexApiKeyBtn" + ); + + if (addVertexApiKeyBtn) { + addVertexApiKeyBtn.addEventListener("click", () => { + openModal(vertexApiKeyModal); + if (vertexApiKeyBulkInput) vertexApiKeyBulkInput.value = ""; + }); + } + if (closeVertexApiKeyModalBtn) + closeVertexApiKeyModalBtn.addEventListener("click", () => + closeModal(vertexApiKeyModal) + ); + if (cancelAddVertexApiKeyBtn) + cancelAddVertexApiKeyBtn.addEventListener("click", () => + closeModal(vertexApiKeyModal) + ); + if (confirmAddVertexApiKeyBtn) + confirmAddVertexApiKeyBtn.addEventListener( + "click", + handleBulkAddVertexApiKeys + ); + + if (bulkDeleteVertexApiKeyBtn) { + bulkDeleteVertexApiKeyBtn.addEventListener("click", () => { + openModal(bulkDeleteVertexApiKeyModal); + if (bulkDeleteVertexApiKeyInput) bulkDeleteVertexApiKeyInput.value = ""; + }); + } + if (closeBulkDeleteVertexModalBtn) + closeBulkDeleteVertexModalBtn.addEventListener("click", () => + closeModal(bulkDeleteVertexApiKeyModal) + ); + if (cancelBulkDeleteVertexApiKeyBtn) + cancelBulkDeleteVertexApiKeyBtn.addEventListener("click", () => + closeModal(bulkDeleteVertexApiKeyModal) + ); + if (confirmBulkDeleteVertexApiKeyBtn) + confirmBulkDeleteVertexApiKeyBtn.addEventListener( + "click", + handleBulkDeleteVertexApiKeys + ); + // Model Helper Modal Event Listeners if (closeModelHelperModalBtn) { closeModelHelperModalBtn.addEventListener("click", () => @@ -591,6 +668,14 @@ async function initConfig() { ) { config.FILTERED_MODELS = ["gemini-1.0-pro-latest"]; } + // --- 新增：处理 VERTEX_API_KEYS 默认值 --- + if (!config.VERTEX_API_KEYS || !Array.isArray(config.VERTEX_API_KEYS)) { + config.VERTEX_API_KEYS = []; + } + // --- 新增：处理 VERTEX_EXPRESS_BASE_URL 默认值 --- + if (typeof config.VERTEX_EXPRESS_BASE_URL === "undefined") { + config.VERTEX_EXPRESS_BASE_URL = ""; + } // --- 新增：处理 PROXIES 默认值 --- if (!config.PROXIES || !Array.isArray(config.PROXIES)) { config.PROXIES = []; // 默认为空数组 @@ -666,10 +751,12 @@ async function initConfig() { SEARCH_MODELS: ["gemini-1.5-flash-latest"], FILTERED_MODELS: ["gemini-1.0-pro-latest"], UPLOAD_PROVIDER: "smms", - PROXIES: [], // 添加默认值 + PROXIES: [], + VERTEX_API_KEYS: [], // 确保默认值存在 + VERTEX_EXPRESS_BASE_URL: "", // 确保默认值存在 THINKING_MODELS: [], THINKING_BUDGET_MAP: {}, - AUTO_DELETE_ERROR_LOGS_ENABLED: false, // 新增默认值 + AUTO_DELETE_ERROR_LOGS_ENABLED: false, AUTO_DELETE_ERROR_LOGS_DAYS: 7, // 新增默认值 AUTO_DELETE_REQUEST_LOGS_ENABLED: false, // 新增默认值 AUTO_DELETE_REQUEST_LOGS_DAYS: 30, // 新增默认值 @@ -678,7 +765,7 @@ async function initConfig() { FAKE_STREAM_EMPTY_DATA_INTERVAL_SECONDS: 5, // --- 结束：处理假流式配置的默认值 --- }; - + populateForm(defaultConfig); if (configForm) { // Ensure form exists @@ -1090,7 +1177,126 @@ function handleBulkDeleteProxies() { } bulkDeleteProxyInput.value = ""; } - + +/** + * Handles the bulk addition of Vertex API keys from the modal input. + */ +function handleBulkAddVertexApiKeys() { + const vertexApiKeyContainer = document.getElementById( + "VERTEX_API_KEYS_container" + ); + if ( + !vertexApiKeyBulkInput || + !vertexApiKeyContainer || + !vertexApiKeyModal + ) { + return; + } + + const bulkText = vertexApiKeyBulkInput.value; + const extractedKeys = bulkText.match(VERTEX_API_KEY_REGEX) || []; + + const currentKeyInputs = vertexApiKeyContainer.querySelectorAll( + `.${ARRAY_INPUT_CLASS}.${SENSITIVE_INPUT_CLASS}` + ); + let currentKeys = Array.from(currentKeyInputs) + .map((input) => { + return input.hasAttribute("data-real-value") + ? input.getAttribute("data-real-value") + : input.value; + }) + .filter((key) => key && key.trim() !== "" && key !== MASKED_VALUE); + + const combinedKeys = new Set([...currentKeys, ...extractedKeys]); + const uniqueKeys = Array.from(combinedKeys); + + vertexApiKeyContainer.innerHTML = ""; // Clear existing items + + uniqueKeys.forEach((key) => { + addArrayItemWithValue("VERTEX_API_KEYS", key); // VERTEX_API_KEYS are sensitive + }); + + // Ensure new sensitive inputs are masked + const newKeyInputs = vertexApiKeyContainer.querySelectorAll( + `.${ARRAY_INPUT_CLASS}.${SENSITIVE_INPUT_CLASS}` + ); + newKeyInputs.forEach((input) => { + if (configForm && typeof initializeSensitiveFields === "function") { + const focusoutEvent = new Event("focusout", { + bubbles: true, + cancelable: true, + }); + input.dispatchEvent(focusoutEvent); + } + }); + + closeModal(vertexApiKeyModal); + showNotification( + `添加/更新了 ${uniqueKeys.length} 个唯一 Vertex 密钥`, + "success" + ); + vertexApiKeyBulkInput.value = ""; +} + +/** + * Handles the bulk deletion of Vertex API keys based on input from the modal. + */ +function handleBulkDeleteVertexApiKeys() { + const vertexApiKeyContainer = document.getElementById( + "VERTEX_API_KEYS_container" + ); + if ( + !bulkDeleteVertexApiKeyInput || + !vertexApiKeyContainer || + !bulkDeleteVertexApiKeyModal + ) { + return; + } + + const bulkText = bulkDeleteVertexApiKeyInput.value; + if (!bulkText.trim()) { + showNotification("请粘贴需要删除的 Vertex API 密钥", "warning"); + return; + } + + const keysToDelete = new Set(bulkText.match(VERTEX_API_KEY_REGEX) || []); + + if (keysToDelete.size === 0) { + showNotification( + "未在输入内容中提取到有效的 Vertex API 密钥格式", + "warning" + ); + return; + } + + const keyItems = vertexApiKeyContainer.querySelectorAll(`.${ARRAY_ITEM_CLASS}`); + let deleteCount = 0; + + keyItems.forEach((item) => { + const input = item.querySelector( + `.${ARRAY_INPUT_CLASS}.${SENSITIVE_INPUT_CLASS}` + ); + const realValue = + input && + (input.hasAttribute("data-real-value") + ? input.getAttribute("data-real-value") + : input.value); + if (realValue && keysToDelete.has(realValue)) { + item.remove(); + deleteCount++; + } + }); + + closeModal(bulkDeleteVertexApiKeyModal); + + if (deleteCount > 0) { + showNotification(`成功删除了 ${deleteCount} 个匹配的 Vertex 密钥`, "success"); + } else { + showNotification("列表中未找到您输入的任何 Vertex 密钥进行删除", "info"); + } + bulkDeleteVertexApiKeyInput.value = ""; +} + /** * Switches the active configuration tab. * @param {string} tabId - The ID of the tab to switch to. @@ -1231,9 +1437,11 @@ function addArrayItemWithValue(key, value) { const isThinkingModel = key === "THINKING_MODELS"; const isAllowedToken = key === "ALLOWED_TOKENS"; - const isSensitive = key === "API_KEYS" || isAllowedToken; + const isVertexApiKey = key === "VERTEX_API_KEYS"; // 新增判断 + const isSensitive = + key === "API_KEYS" || isAllowedToken || isVertexApiKey; // 更新敏感判断 const modelId = isThinkingModel ? generateUUID() : null; - + const arrayItem = document.createElement("div"); arrayItem.className = `${ARRAY_ITEM_CLASS} flex items-center mb-2 gap-2`; if (isThinkingModel) { diff --git a/app/templates/config_editor.html b/app/templates/config_editor.html index 671049d..71b820c 100644 --- a/app/templates/config_editor.html +++ b/app/templates/config_editor.html @@ -456,7 +456,51 @@ endblock %} {% block head_extra_styles %} /> Gemini API的基础URL - + + +
+ +
+ +
+
+ + +
+ Vertex AI Platform API密钥列表。点击按钮可批量添加或删除。 +
+ + +
+ + + Vertex Express API的基础URL +
+
- + + + + + + +