Spaces:
Running
Running
| """FastAPI路由模块""" | |
| import asyncio | |
| import base64 | |
| import json | |
| import mimetypes | |
| import time | |
| import uuid | |
| import random | |
| import string | |
| from http import HTTPStatus | |
| from typing import Any, AsyncGenerator, Awaitable, cast | |
| from fastapi import FastAPI, Request, UploadFile | |
| from fastapi.responses import StreamingResponse, JSONResponse, Response | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from starlette.datastructures import UploadFile as StarletteUploadFile | |
| from starlette.types import ASGIApp, Message, Receive, Scope, Send | |
| from src.core import MODELS_CONFIG_FILE | |
| from src.core.errors import ( | |
| VcoreError, | |
| InvalidArgumentError, | |
| InternalError, | |
| RateLimitError, | |
| AuthenticationError, | |
| ) | |
| from src.api.vcore_client import VcoreAIClient | |
| from src.api.session_pool import business_session | |
| from src.api.oai_adapter import OAIImageRequestConverter, OAIRequestConverter, OAIResponseConverter | |
| from src.api.transform import ResponseAggregator | |
| from src.core.auth import api_key_manager | |
| from src.utils.logger import get_logger, set_request_id, get_request_id | |
| def _inject_anti429(gemini_payload: dict[str, Any]) -> dict[str, Any]: | |
| """如果配置开启了防429,在 gemini_payload 里插入100位随机数""" | |
| from src.core.config import load_config | |
| cfg = load_config() | |
| if not cfg.get("anti429_enabled", False): | |
| return gemini_payload | |
| rand_str = ''.join(random.choices(string.digits, k=100)) | |
| target = cfg.get("anti429_target", "system") | |
| if target == "system": | |
| si = gemini_payload.get("systemInstruction", {}) | |
| if isinstance(si, dict): | |
| parts = si.get("parts", []) | |
| if not isinstance(parts, list): | |
| parts = [] | |
| elif isinstance(si, str): | |
| parts = [{"text": si}] | |
| else: | |
| parts = [] | |
| parts.insert(0, {"text": rand_str}) | |
| gemini_payload["systemInstruction"] = {"parts": parts} | |
| else: | |
| # 插入到第一条 user 消息的 parts 最前面 | |
| contents = gemini_payload.get("contents", []) | |
| if isinstance(contents, str): | |
| gemini_payload["contents"] = [{"role": "user", "parts": [{"text": rand_str}, {"text": contents}]}] | |
| elif isinstance(contents, list) and contents: | |
| for c in contents: | |
| if isinstance(c, dict) and c.get("role") == "user": | |
| parts = c.get("parts", []) | |
| if not isinstance(parts, list): | |
| parts = [] | |
| parts.insert(0, {"text": rand_str}) | |
| c["parts"] = parts | |
| break | |
| else: | |
| # 没有 user 消息,插入一条新的 | |
| contents.insert(0, {"role": "user", "parts": [{"text": rand_str}]}) | |
| else: | |
| gemini_payload["contents"] = [{"role": "user", "parts": [{"text": rand_str}]}] | |
| return gemini_payload | |
| # 初始化日志 | |
| logger = get_logger(__name__) | |
| def _wire_status_code(status_code: Any, fallback: int = 500) -> int: | |
| """转换为 uvicorn/h11 一定能发送的标准 HTTP 状态码。""" | |
| try: | |
| code = int(status_code) | |
| except (TypeError, ValueError): | |
| return fallback | |
| try: | |
| HTTPStatus(code) | |
| return code | |
| except ValueError: | |
| # 499 常用于日志里的 Client Closed Request,但 uvicorn/h11 不一定支持 | |
| # 非标准状态短语;真实下游状态使用标准码,日志/响应体仍可保留原始 code。 | |
| if code == 499: | |
| return 400 | |
| return fallback | |
| _JSON_THREAD_OFFLOAD_THRESHOLD_BYTES = 256 * 1024 | |
| _CLIENT_DISCONNECT_CONTEXT: dict[str, str] = {} | |
| _REQUEST_LOG_CONTEXT: dict[str, str] = {} | |
| _REQUEST_LOG_OUTCOME: dict[str, dict[str, Any]] = {} | |
| SSE_HEADERS = { | |
| "Cache-Control": "no-cache", | |
| "Connection": "keep-alive", | |
| "X-Accel-Buffering": "no", | |
| } | |
| class DownstreamDisconnectedError(Exception): | |
| """下游客户端已断开,当前请求应立即停止并释放资源。""" | |
| def _json_bytes_compact_sync(data: Any) -> bytes: | |
| return json.dumps(data, ensure_ascii=False, separators=(',', ':')).encode('utf-8') | |
| def _json_bytes_default_sync(data: Any) -> bytes: | |
| return json.dumps(data, ensure_ascii=False).encode('utf-8') | |
| def _rough_json_size(data: Any, limit: int = _JSON_THREAD_OFFLOAD_THRESHOLD_BYTES) -> int: | |
| if limit <= 0: | |
| return limit | |
| if isinstance(data, dict): | |
| total = 0 | |
| for key, value in data.items(): | |
| total += _rough_json_size(key, limit - total) | |
| if total >= limit: | |
| return limit | |
| total += _rough_json_size(value, limit - total) | |
| if total >= limit: | |
| return limit | |
| return total | |
| if isinstance(data, list): | |
| total = 0 | |
| for item in data: | |
| total += _rough_json_size(item, limit - total) | |
| if total >= limit: | |
| return limit | |
| return total | |
| if isinstance(data, str): | |
| return min(len(data) * 4, limit) | |
| return 16 | |
| async def _json_bytes_for_payload(data: Any, compact: bool = True) -> bytes: | |
| dumper = _json_bytes_compact_sync if compact else _json_bytes_default_sync | |
| if _rough_json_size(data) >= _JSON_THREAD_OFFLOAD_THRESHOLD_BYTES: | |
| return await asyncio.to_thread(dumper, data) | |
| return dumper(data) | |
| async def _sse_json_event(data: Any, compact: bool = True) -> bytes: | |
| return b"data: " + await _json_bytes_for_payload(data, compact=compact) + b"\n\n" | |
| def _json_bytes_for_sse_fallback(data: Any, compact: bool = True) -> bytes: | |
| if compact: | |
| return json.dumps(data, ensure_ascii=False, separators=(',', ':')).encode('utf-8') | |
| return json.dumps(data, ensure_ascii=False).encode('utf-8') | |
| def _safe_error_message(value: Any, default: str = "Internal server error", limit: int = 4096) -> str: | |
| try: | |
| text = str(value) | |
| except Exception: | |
| text = default | |
| text = text.strip() or default | |
| if len(text) > limit: | |
| return text[:limit] + "..." | |
| return text | |
| async def _safe_sse_json_event( | |
| data: Any, | |
| compact: bool = True, | |
| fallback: dict[str, Any] | None = None, | |
| context: str = "SSE 错误事件", | |
| ) -> bytes: | |
| """构造 SSE 事件;在错误处理路径中即使序列化失败也返回最小标准错误事件。""" | |
| try: | |
| return await _sse_json_event(data, compact=compact) | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as e: | |
| logger.error(f"{context} 构造失败,使用降级错误事件: {e}") | |
| fallback_payload = fallback or { | |
| "error": { | |
| "code": 500, | |
| "message": "Internal server error while building stream error event", | |
| "status": "INTERNAL", | |
| } | |
| } | |
| try: | |
| return b"data: " + _json_bytes_for_sse_fallback(fallback_payload, compact=compact) + b"\n\n" | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as fallback_error: | |
| logger.error(f"{context} 降级错误事件仍构造失败: {fallback_error}") | |
| return b'data: {"error":{"code":500,"message":"Internal server error","status":"INTERNAL"}}\n\n' | |
| def _gemini_error_fallback(error: VcoreError | Exception) -> dict[str, Any]: | |
| if isinstance(error, VcoreError): | |
| return { | |
| "error": { | |
| "code": _wire_status_code(error.code), | |
| "message": _safe_error_message(error.message), | |
| "status": error.status or "INTERNAL", | |
| } | |
| } | |
| return { | |
| "error": { | |
| "code": 500, | |
| "message": _safe_error_message(error), | |
| "status": "INTERNAL", | |
| } | |
| } | |
| def _oai_error_fallback(message: Any, err_type: str = "server_error") -> dict[str, Any]: | |
| return {"error": {"message": _safe_error_message(message), "type": err_type, "code": None}} | |
| async def _json_response_for_payload(data: Any, status_code: int = 200, compact: bool = True) -> Response: | |
| return Response( | |
| content=await _json_bytes_for_payload(data, compact=compact), | |
| status_code=status_code, | |
| media_type="application/json", | |
| ) | |
| def _cleanup_status_info(message: str | None) -> str: | |
| text = (message or "").strip() or "响应成功" | |
| if text.endswith("清理占用资源"): | |
| return text | |
| return f"{text},清理占用资源" | |
| def _vcore_error_status_info(exc: VcoreError) -> str: | |
| prefix = f"{exc.status}: " if exc.status and exc.status != "UNKNOWN" else "" | |
| return _cleanup_status_info(f"{prefix}{exc.message}") | |
| def _extract_status_info_from_payload(data: Any) -> str | None: | |
| if not isinstance(data, dict): | |
| return None | |
| error_obj = data.get("error") | |
| if isinstance(error_obj, dict): | |
| status = error_obj.get("status") or error_obj.get("type") | |
| message = error_obj.get("message") or error_obj.get("detail") | |
| if status and message: | |
| return f"{status}: {message}" | |
| if message: | |
| return str(message) | |
| if status: | |
| return str(status) | |
| detail = data.get("detail") | |
| if detail: | |
| return str(detail) | |
| message = data.get("message") | |
| if message: | |
| return str(message) | |
| status = data.get("status") | |
| if status: | |
| return str(status) | |
| return None | |
| def _extract_status_info_from_body(body: bytes) -> str | None: | |
| if not body: | |
| return None | |
| text = body.decode("utf-8", errors="replace").strip() | |
| if not text: | |
| return None | |
| try: | |
| payload = json.loads(text) | |
| except json.JSONDecodeError: | |
| return text[:500] | |
| return _extract_status_info_from_payload(payload) | |
| def _set_request_log_outcome( | |
| status_code: int | None = None, | |
| status_info: str | None = None, | |
| request_id: str | None = None, | |
| ) -> None: | |
| rid = request_id or get_request_id() | |
| if not rid: | |
| return | |
| outcome = _REQUEST_LOG_OUTCOME.setdefault(rid, {}) | |
| if status_code is not None: | |
| outcome["status_code"] = int(status_code) | |
| if status_info: | |
| outcome["status_info"] = _cleanup_status_info(status_info) | |
| def _mark_request_client_disconnected(source: str | None = None, request_id: str | None = None) -> None: | |
| rid = request_id or get_request_id() | |
| if not rid: | |
| return | |
| if source: | |
| _CLIENT_DISCONNECT_CONTEXT[rid] = source | |
| _set_request_log_outcome(499, "客户端连接断开", rid) | |
| def _consume_background_task_result(task: asyncio.Task[Any]) -> None: | |
| try: | |
| task.result() | |
| except asyncio.CancelledError: | |
| pass | |
| except Exception as e: | |
| logger.debug(f"后台任务结束时出现异常: {e}") | |
| async def _aclose_source_bounded( | |
| source: AsyncGenerator[Any, None], | |
| timeout: float = 1.0, | |
| source_name: str = "unknown", | |
| ) -> None: | |
| close_task = asyncio.create_task(source.aclose()) | |
| done, pending = await asyncio.wait({close_task}, timeout=timeout) | |
| if done: | |
| await asyncio.gather(*done, return_exceptions=True) | |
| return | |
| close_task.add_done_callback(_consume_background_task_result) | |
| logger.warning(f"关闭下游源生成器超时,已转后台继续关闭: source={source_name}, timeout={timeout:.1f}s") | |
| async def _convert_oai_realtime_chunk_objects( | |
| chunk: dict[str, Any], | |
| model: str, | |
| request_id: str, | |
| is_first: bool, | |
| has_prior_tool_calls: bool, | |
| ) -> list[dict[str, Any]]: | |
| converter = OAIResponseConverter.convert_realtime_chunk_objects | |
| if _rough_json_size(chunk) >= _JSON_THREAD_OFFLOAD_THRESHOLD_BYTES: | |
| return await asyncio.to_thread(converter, chunk, model, request_id, is_first, has_prior_tool_calls) | |
| return converter(chunk, model, request_id, is_first, has_prior_tool_calls) | |
| async def _convert_oai_chat_response(gemini_response: dict[str, Any], model: str) -> dict[str, Any]: | |
| converter = OAIResponseConverter.gemini_json_to_oai_json | |
| if _rough_json_size(gemini_response) >= _JSON_THREAD_OFFLOAD_THRESHOLD_BYTES: | |
| return await asyncio.to_thread(converter, gemini_response, model) | |
| return converter(gemini_response, model) | |
| async def _convert_oai_image_data(gemini_response: dict[str, Any], response_format: str) -> list[dict[str, Any]]: | |
| converter = OAIResponseConverter.gemini_json_to_oai_image_data | |
| if _rough_json_size(gemini_response) >= _JSON_THREAD_OFFLOAD_THRESHOLD_BYTES: | |
| return await asyncio.to_thread(converter, gemini_response, response_format) | |
| return converter(gemini_response, response_format) | |
| async def _convert_oai_image_chunk_data(chunk: dict[str, Any], response_format: str) -> list[dict[str, Any]]: | |
| converter = OAIResponseConverter.gemini_chunk_to_openai_image_partial | |
| if _rough_json_size(chunk) >= _JSON_THREAD_OFFLOAD_THRESHOLD_BYTES: | |
| return await asyncio.to_thread(converter, chunk, response_format) | |
| return converter(chunk, response_format) | |
| def _set_request_log_context(**items: Any) -> str: | |
| parts = [f"{key}={value}" for key, value in items.items() if value is not None and value != ""] | |
| context = ", ".join(parts) | |
| if context: | |
| _REQUEST_LOG_CONTEXT[get_request_id()] = context | |
| return context | |
| def _log_request_start(method: str, path: str, **items: Any) -> None: | |
| context = _set_request_log_context(**items) | |
| suffix = f", {context}" if context else "" | |
| logger.info(f"请求开始: {method} {path}{suffix}") | |
| def _gemini_stream_uses_sse(request: Request) -> bool: | |
| """Gemini streamGenerateContent 默认返回官方 JSON 数组;显式 alt=sse 时返回 SSE。""" | |
| alt = (request.query_params.get("alt") or "").strip().lower() | |
| return alt == "sse" | |
| async def _wait_for_client_disconnect(request: Request, source: str) -> None: | |
| while True: | |
| if await request.is_disconnected(): | |
| _mark_request_client_disconnected(source) | |
| logger.debug(f"下游客户端已断开: source={source}") | |
| return | |
| await asyncio.sleep(0.25) | |
| async def _cancel_and_drain(*tasks: asyncio.Task[Any], timeout: float = 2.0) -> None: | |
| pending = [task for task in tasks if not task.done()] | |
| for task in pending: | |
| task.cancel() | |
| if pending: | |
| done, still_pending = await asyncio.wait(pending, timeout=timeout) | |
| if done: | |
| await asyncio.gather(*done, return_exceptions=True) | |
| if still_pending: | |
| logger.warning(f"取消任务超时,跳过等待: tasks={len(still_pending)}, timeout={timeout:.1f}s") | |
| async def _await_with_client_disconnect(request: Request, awaitable: Awaitable[Any], source: str) -> Any: | |
| operation_task = asyncio.ensure_future(awaitable) | |
| disconnect_task = asyncio.create_task(_wait_for_client_disconnect(request, source)) | |
| try: | |
| done, _ = await asyncio.wait( | |
| {operation_task, disconnect_task}, | |
| return_when=asyncio.FIRST_COMPLETED, | |
| ) | |
| if disconnect_task in done: | |
| await _cancel_and_drain(operation_task) | |
| raise DownstreamDisconnectedError(f"Client disconnected during {source}") | |
| await _cancel_and_drain(disconnect_task) | |
| return await operation_task | |
| except BaseException: | |
| await _cancel_and_drain(operation_task, disconnect_task) | |
| raise | |
| async def _aiter_with_client_disconnect( | |
| request: Request, | |
| source: AsyncGenerator[Any, None], | |
| source_name: str, | |
| ) -> AsyncGenerator[Any, None]: | |
| disconnect_task = asyncio.create_task(_wait_for_client_disconnect(request, source_name)) | |
| try: | |
| while True: | |
| next_task = asyncio.create_task(anext(source)) | |
| try: | |
| done, _ = await asyncio.wait( | |
| {next_task, disconnect_task}, | |
| return_when=asyncio.FIRST_COMPLETED, | |
| ) | |
| if disconnect_task in done: | |
| next_task.cancel() | |
| cancelled, pending = await asyncio.wait({next_task}, timeout=2.0) | |
| if pending: | |
| logger.warning(f"下游断开后上游读取任务取消超时,继续关闭资源: source={source_name}") | |
| if cancelled: | |
| await asyncio.gather(*cancelled, return_exceptions=True) | |
| await _aclose_source_bounded(source, source_name=source_name) | |
| raise DownstreamDisconnectedError(f"Client disconnected during {source_name}") | |
| try: | |
| yield await next_task | |
| except StopAsyncIteration: | |
| break | |
| except BaseException: | |
| await _cancel_and_drain(next_task) | |
| raise | |
| finally: | |
| await _cancel_and_drain(disconnect_task) | |
| await _aclose_source_bounded(source, source_name=source_name) | |
| def extract_api_key_from_request(request: Request) -> str | None: | |
| """ | |
| 从请求中提取API密钥 | |
| 支持三种方式(按优先级): | |
| 1. Authorization: Bearer <key> (OpenAI 标准 Header) | |
| 2. x-goog-api-key: <key> (Google/Gemini 标准 Header) | |
| 3. ?key=<key> (Google/Gemini 标准 Query Param) | |
| """ | |
| # 1. 尝试 OpenAI 标准 Authorization Header | |
| auth_header = request.headers.get("Authorization") | |
| if auth_header and auth_header.lower().startswith("bearer "): | |
| return auth_header[7:].strip() | |
| # 2. 尝试 x-goog-api-key Header | |
| goog_api_key = request.headers.get("x-goog-api-key") | |
| if goog_api_key: | |
| return goog_api_key.strip() | |
| # 3. 尝试 URL Query Parameter | |
| query_key = request.query_params.get("key") | |
| if query_key: | |
| return query_key.strip() | |
| return None | |
| class APIKeyMiddleware: | |
| """API密钥认证中间件""" | |
| def __init__(self, app: ASGIApp, excluded_paths: list[str] | None = None, excluded_prefixes: list[str] | None = None): | |
| self.app = app | |
| self.excluded_paths: list[str] = excluded_paths or ["/", "/health"] | |
| self.excluded_prefixes: list[str] = excluded_prefixes or [] | |
| async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: | |
| if scope["type"] != "http": | |
| await self.app(scope, receive, send) | |
| return | |
| # 为每个请求设置唯一的请求ID | |
| set_request_id() | |
| request = Request(scope, receive) | |
| path = request.url.path | |
| method = request.method | |
| client_ip = request.client.host if request.client else "unknown" | |
| logger.debug(f"收到请求: {method} {path} from {client_ip}") | |
| status_code = 500 | |
| start_time = time.time() | |
| response_body_parts: list[bytes] = [] | |
| response_body_truncated = False | |
| max_captured_body_bytes = 4096 | |
| async def send_wrapper(message: Message) -> None: | |
| nonlocal status_code, response_body_truncated | |
| if message["type"] == "http.response.start": | |
| status_code = int(message.get("status", 500)) | |
| elif message["type"] == "http.response.body" and status_code >= 400 and not response_body_truncated: | |
| body = message.get("body", b"") | |
| if isinstance(body, bytes) and body: | |
| captured_size = sum(len(part) for part in response_body_parts) | |
| remaining = max_captured_body_bytes - captured_size | |
| if remaining > 0: | |
| response_body_parts.append(body[:remaining]) | |
| if len(body) > remaining: | |
| response_body_truncated = True | |
| await send(message) | |
| async def call_app() -> None: | |
| await self.app(scope, receive, send_wrapper) | |
| def log_request_result(process_time: float) -> None: | |
| request_id = get_request_id() | |
| disconnect_source = _CLIENT_DISCONNECT_CONTEXT.pop(request_id, None) | |
| request_context = _REQUEST_LOG_CONTEXT.pop(request_id, "") | |
| outcome = _REQUEST_LOG_OUTCOME.pop(request_id, {}) | |
| final_status_code = int(outcome.get("status_code") or status_code) | |
| if disconnect_source and final_status_code == 200: | |
| final_status_code = 499 | |
| if disconnect_source and not request_context: | |
| request_context = f"source={disconnect_source}" | |
| context_text = f", {request_context}" if request_context else "" | |
| if outcome.get("status_info"): | |
| status_info = str(outcome["status_info"]) | |
| elif final_status_code == 499: | |
| status_info = _cleanup_status_info("客户端连接断开") | |
| elif final_status_code >= 400: | |
| error_info = _extract_status_info_from_body(b"".join(response_body_parts)) | |
| status_info = _cleanup_status_info(error_info or f"响应错误 HTTP {final_status_code}") | |
| else: | |
| status_info = _cleanup_status_info("响应成功") | |
| logger.info(f"响应完成: {method} {path} - {final_status_code}{context_text} ({process_time:.3f}s) {status_info}") | |
| # 检查是否是完全排除的路径 | |
| if self.excluded_paths and path in self.excluded_paths: | |
| logger.debug(f"路径 {path} 在排除列表中,跳过认证") | |
| try: | |
| await call_app() | |
| finally: | |
| process_time = time.time() - start_time | |
| log_request_result(process_time) | |
| return | |
| # 检查前缀排除(管理后台、静态资源) | |
| if any(path.startswith(p) for p in self.excluded_prefixes): | |
| try: | |
| await call_app() | |
| finally: | |
| process_time = time.time() - start_time | |
| log_request_result(process_time) | |
| return | |
| # 获取API密钥 | |
| api_key = extract_api_key_from_request(request) | |
| if not api_key: | |
| logger.warning(f"请求 {path} 缺少 API 密钥") | |
| _set_request_log_outcome(401, "UNAUTHENTICATED: 缺少 API 密钥") | |
| response = JSONResponse( | |
| status_code=401, | |
| content={ | |
| "error": { | |
| "code": 401, | |
| "message": "Method doesn't allow unregistered callers (callers without established identity). Please use API Key or other form of API consumer identity to call this API.", | |
| "status": "UNAUTHENTICATED" | |
| } | |
| } | |
| ) | |
| await response(scope, receive, send_wrapper) | |
| process_time = time.time() - start_time | |
| log_request_result(process_time) | |
| return | |
| # 验证API密钥 | |
| if not api_key_manager.validate_key(api_key): | |
| logger.warning(f"请求 {path} 使用了无效的 API 密钥: {api_key[:8]}...") | |
| _set_request_log_outcome(400, "INVALID_ARGUMENT: API key not valid") | |
| response = JSONResponse( | |
| status_code=400, | |
| content={ | |
| "error": { | |
| "code": 400, | |
| "message": "API key not valid. Please pass a valid API key.", | |
| "status": "INVALID_ARGUMENT" | |
| } | |
| } | |
| ) | |
| await response(scope, receive, send_wrapper) | |
| process_time = time.time() - start_time | |
| log_request_result(process_time) | |
| return | |
| # 将API密钥存储在请求状态中 | |
| scope.setdefault("state", {})["api_key"] = api_key | |
| logger.debug(f"API 密钥验证成功: {api_key[:8]}...") | |
| try: | |
| await call_app() | |
| finally: | |
| process_time = time.time() - start_time | |
| log_request_result(process_time) | |
| def create_app(vcore_client: VcoreAIClient) -> FastAPI: | |
| """创建FastAPI应用""" | |
| logger.info("创建 FastAPI 应用") | |
| app = FastAPI( | |
| title="Vcore AI Proxy (Anonymous)", | |
| description="Vcore AI 代理服务,兼容 Gemini API", | |
| version="1.1.0" | |
| ) | |
| # 添加中间件(顺序很重要) | |
| logger.debug("添加中间件") | |
| app.add_middleware( | |
| APIKeyMiddleware, | |
| excluded_paths=["/", "/health", "/admin"], | |
| excluded_prefixes=["/api/admin/", "/admin/", "/static/"], | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| expose_headers=["*"], | |
| ) | |
| # 挂载管理后台路由 | |
| from src.api.admin import router as admin_router | |
| app.include_router(admin_router) | |
| # ==================== 全局异常处理 ==================== | |
| async def vcore_exception_handler(request: Request, exc: VcoreError): # type: ignore[misc] | |
| """处理所有 VcoreError 及其子类""" | |
| _set_request_log_outcome(exc.code, _vcore_error_status_info(exc)) | |
| logger.error(f"VcoreError: {exc.message} (code={exc.code}, status={exc.status})") | |
| return JSONResponse( | |
| status_code=_wire_status_code(exc.code), | |
| content=exc.to_Dict(), | |
| ) | |
| async def downstream_disconnected_handler(request: Request, exc: DownstreamDisconnectedError): # type: ignore[misc] | |
| _mark_request_client_disconnected() | |
| logger.debug(f"下游连接已关闭,请求资源已清理: {request.method} {request.url.path}") | |
| return JSONResponse( | |
| status_code=_wire_status_code(499, fallback=400), | |
| content={"error": {"code": 499, "message": "Client Closed Request", "status": "CANCELLED"}}, | |
| ) | |
| async def general_exception_handler(request: Request, exc: Exception): # type: ignore[misc] | |
| """处理所有未捕获异常""" | |
| logger.error(f"Unhandled Exception: {exc}", exc_info=True) | |
| error = InternalError(message=str(exc)) | |
| _set_request_log_outcome(error.code, _vcore_error_status_info(error)) | |
| return JSONResponse( | |
| status_code=500, | |
| content=error.to_Dict(), | |
| ) | |
| # ==================== 基础端点 ==================== | |
| async def root() -> dict[str, str]: | |
| """根路径,返回服务信息""" | |
| logger.debug("处理根路径请求") | |
| return { | |
| "message": "Vcore AI Proxy Server (Anonymous Edition)", | |
| "version": "1.1.0", | |
| "auth": "API Key Authentication Required", | |
| "docs": "Only Gemini API Compatible" | |
| } | |
| app.get("/")(root) | |
| async def health_check() -> dict[str, str | int]: | |
| """健康检查端点""" | |
| logger.debug("处理健康检查请求") | |
| api_keys_count = len(api_key_manager.api_keys) | |
| logger.debug(f"当前加载的 API 密钥数量: {api_keys_count}") | |
| return { | |
| "status": "healthy", | |
| "timestamp": int(time.time()), | |
| "api_keys_loaded": api_keys_count | |
| } | |
| app.get("/health")(health_check) | |
| async def list_models_oai() -> dict[str, str | list[dict[str, Any]]]: | |
| """返回可用模型列表 (OpenAI 兼容格式)""" | |
| logger.debug("处理模型列表请求") | |
| current_time = int(time.time()) | |
| models: list[str] = _load_models_config() | |
| logger.debug(f"返回 {len(models)} 个可用模型") | |
| return { | |
| "object": "list", | |
| "data": [ | |
| {"id": m, "object": "model", "created": current_time, "owned_by": "google", "permission": []} | |
| for m in models | |
| ] | |
| } | |
| async def list_models_gemini() -> dict[str, list[dict[str, Any]]]: | |
| """返回可用模型列表 (Gemini 兼容格式)""" | |
| logger.debug("处理 Gemini 模型列表请求") | |
| models: list[str] = _load_models_config() | |
| return { | |
| "models": [ | |
| { | |
| "name": f"models/{m}", | |
| "version": m, | |
| "displayName": m, | |
| "description": "Vcore AI Studio anonymous model", | |
| "inputTokenLimit": 1048576, | |
| "outputTokenLimit": 65536, | |
| "supportedGenerationMethods": _supported_generation_methods(m), | |
| } | |
| for m in models | |
| ] | |
| } | |
| app.get("/v1/models")(list_models_oai) | |
| app.get("/v1beta/models")(list_models_gemini) | |
| # ==================== Gemini 兼容端点 ==================== | |
| async def stream_generate_content(model: str, request: Request) -> StreamingResponse | JSONResponse: | |
| """Gemini 格式的流式生成接口""" | |
| _log_request_start(request.method, request.url.path, source="gemini.streamGenerateContent", model=model, mode="stream") | |
| try: | |
| body_any = await request.json() | |
| except json.JSONDecodeError as e: | |
| raise InvalidArgumentError(f"Invalid JSON in request body: {e}") | |
| # 简单类型检查 | |
| if not isinstance(body_any, dict): | |
| raise InvalidArgumentError("Request body must be a JSON object") | |
| body: dict[str, Any] = cast(dict[str, Any], body_any) | |
| body = _prepare_gemini_downstream_payload(body) | |
| logger.debug(f"请求体大小: {len(str(body))} 字符") | |
| # 完整记录请求内容(用于调试) | |
| logger.debug_json("下游请求体", body) | |
| is_sse = _gemini_stream_uses_sse(request) | |
| async def stream_generator(): | |
| chunk_count = 0 | |
| byte_count = 0 | |
| started_at = time.monotonic() | |
| try: | |
| async with business_session(vcore_client, source="gemini.streamGenerateContent") as session: | |
| progress_context: dict[str, str] = {} | |
| clean_payload = _inject_anti429(body) | |
| upstream = session.stream_gemini( | |
| model=model, gemini_payload=clean_payload, progress_context=progress_context | |
| ) | |
| async for chunk in _aiter_with_client_disconnect(request, upstream, "gemini.streamGenerateContent"): | |
| prefix_bytes = b"" | |
| if not is_sse: | |
| if chunk_count == 0: | |
| prefix_bytes = b"[\n" | |
| else: | |
| prefix_bytes = b",\n" | |
| chunk_count += 1 | |
| if is_sse: | |
| event = await _sse_json_event(chunk, compact=True) | |
| else: | |
| event = await _json_bytes_for_payload(chunk, compact=True) | |
| full_event = prefix_bytes + event + (b"\n" if not is_sse else b"") | |
| event_size = len(full_event) | |
| byte_count += event_size | |
| logger.debug( | |
| f"Gemini下游准备发送: model={model}, chunk={chunk_count}, " | |
| f"bytes={event_size}, keys={list(chunk.keys()) if isinstance(chunk, dict) else type(chunk)}" | |
| ) | |
| yield full_event | |
| logger.debug(f"Gemini下游已交给ASGI: model={model}, chunk={chunk_count}") | |
| prefix = progress_context.get("prefix") | |
| if prefix: | |
| ResponseAggregator.log_stream_forward_complete(prefix, chunk_count, byte_count, time.monotonic() - started_at) | |
| if not is_sse: | |
| if chunk_count == 0: | |
| yield b"[]\n" | |
| else: | |
| yield b"]\n" | |
| except asyncio.CancelledError: | |
| _mark_request_client_disconnected("gemini.streamGenerateContent") | |
| raise | |
| except DownstreamDisconnectedError: | |
| _mark_request_client_disconnected("gemini.streamGenerateContent") | |
| raise | |
| except VcoreError as e: | |
| _set_request_log_outcome(e.code, _vcore_error_status_info(e)) | |
| logger.error(f"流式生成 Vcore 错误: {e.message}") | |
| if is_sse: | |
| yield await _safe_sse_json_event( | |
| e.to_Dict(), | |
| compact=True, | |
| fallback=_gemini_error_fallback(e), | |
| context="Gemini 流式 Vcore 错误事件", | |
| ) | |
| else: | |
| prefix_bytes = b"[\n" if chunk_count == 0 else b",\n" | |
| # 这里如果是数组中发生错误,最好也能抛出对应的 JSON,以使得结构闭合(即使有错误) | |
| event = await _json_bytes_for_payload(e.to_Dict(), compact=True) | |
| yield prefix_bytes + event + b"\n]\n" | |
| logger.debug("Gemini 流式错误事件已发送,随后结束连接并清理资源") | |
| except Exception as e: | |
| logger.error(f"流式生成未知错误: {e}") | |
| # 未知错误包装为 InternalError | |
| error = InternalError(message=str(e)) | |
| _set_request_log_outcome(error.code, _vcore_error_status_info(error)) | |
| if is_sse: | |
| yield await _safe_sse_json_event( | |
| error.to_Dict(), | |
| compact=True, | |
| fallback=_gemini_error_fallback(error), | |
| context="Gemini 流式未知错误事件", | |
| ) | |
| else: | |
| prefix_bytes = b"[\n" if chunk_count == 0 else b",\n" | |
| event = await _json_bytes_for_payload(error.to_Dict(), compact=True) | |
| yield prefix_bytes + event + b"\n]\n" | |
| logger.debug("Gemini 流式未知错误事件已发送,随后结束连接并清理资源") | |
| finally: | |
| logger.debug(f"流式生成完成,共发送 {chunk_count} 个数据块") | |
| media_type = "text/event-stream" if is_sse else "application/json" | |
| headers = SSE_HEADERS if is_sse else None | |
| return StreamingResponse(stream_generator(), media_type=media_type, headers=headers) | |
| app.post("/v1beta/models/{model}:streamGenerateContent", response_model=None)(stream_generate_content) | |
| app.post("/v1/models/{model}:streamGenerateContent", response_model=None)(stream_generate_content) | |
| async def generate_content(model: str, request: Request) -> Response: | |
| """Gemini 格式的非流式生成接口""" | |
| _log_request_start(request.method, request.url.path, source="gemini.generateContent", model=model, mode="non-stream") | |
| try: | |
| body_any = await request.json() | |
| except json.JSONDecodeError as e: | |
| raise InvalidArgumentError(f"Invalid JSON in request body: {e}") | |
| if not isinstance(body_any, dict): | |
| raise InvalidArgumentError("Request body must be a JSON object") | |
| body: dict[str, Any] = cast(dict[str, Any], body_any) | |
| body = _prepare_gemini_downstream_payload(body) | |
| logger.debug(f"请求体大小: {len(str(body))} 字符") | |
| # 完整记录请求内容(用于调试) | |
| logger.debug_json("下游请求体", body) | |
| start_time = time.time() | |
| # 直接获取 Gemini 格式响应。会话必须在断连监听外层创建, | |
| # 这样下游断开时即使 curl/聚合协程没有立刻响应取消,也能立即取消会话托管的上游节点任务。 | |
| gemini_payload = _inject_anti429(body) | |
| async with business_session(vcore_client, source="gemini.generateContent") as session: | |
| response = await _await_with_client_disconnect( | |
| request, | |
| session.complete_gemini( | |
| model=model, | |
| gemini_payload=gemini_payload, | |
| _expected_image_count=_expected_image_count_from_payload(gemini_payload), | |
| ), | |
| "gemini.generateContent", | |
| ) | |
| return await _json_response_for_payload(response, compact=True) | |
| app.post("/v1beta/models/{model}:generateContent", response_model=None)(generate_content) | |
| app.post("/v1/models/{model}:generateContent", response_model=None)(generate_content) | |
| # ==================== OpenAI 兼容端点 ==================== | |
| async def oai_chat_completions(request: Request) -> StreamingResponse | JSONResponse | Response: | |
| """OpenAI 格式的 Chat Completion 接口""" | |
| try: | |
| body_any = await request.json() | |
| except json.JSONDecodeError as e: | |
| return JSONResponse(status_code=400, content={"error": {"message": f"Invalid JSON: {e}", "type": "invalid_request_error", "code": None}}) | |
| if not isinstance(body_any, dict): | |
| return JSONResponse(status_code=400, content={"error": {"message": "Request body must be a JSON object", "type": "invalid_request_error", "code": None}}) | |
| body: dict[str, Any] = cast(dict[str, Any], body_any) | |
| stream = body.get("stream", False) | |
| try: | |
| model, gemini_payload = OAIRequestConverter.convert(body) | |
| except (KeyError, ValueError) as e: | |
| return JSONResponse(status_code=400, content={"error": {"message": str(e), "type": "invalid_request_error", "code": None}}) | |
| gemini_payload = _inject_anti429(gemini_payload) | |
| _log_request_start(request.method, request.url.path, source="openai.chat", model=model, mode="stream" if stream else "non-stream") | |
| if stream: | |
| request_id = uuid.uuid4().hex[:24] | |
| async def oai_stream_generator(): | |
| is_first = True | |
| has_finish = False | |
| has_tool_calls = False | |
| chunk_count = 0 | |
| byte_count = 0 | |
| started_at = time.monotonic() | |
| try: | |
| async with business_session(vcore_client, source="openai.chat.stream") as session: | |
| progress_context: dict[str, str] = {} | |
| upstream = session.stream_gemini(model=model, gemini_payload=gemini_payload, progress_context=progress_context) | |
| async for gemini_chunk in _aiter_with_client_disconnect(request, upstream, "openai.chat.stream"): | |
| chunk_count += 1 | |
| event_payloads = await _convert_oai_realtime_chunk_objects( | |
| gemini_chunk, | |
| model, | |
| request_id, | |
| is_first, | |
| has_tool_calls, | |
| ) | |
| is_first = False | |
| for payload in event_payloads: | |
| event = await _sse_json_event(payload, compact=False) | |
| choices = payload.get("choices") if isinstance(payload, dict) else None | |
| if isinstance(choices, list) and any( | |
| isinstance(choice, dict) | |
| and isinstance(choice.get("delta"), dict) | |
| and choice["delta"].get("tool_calls") | |
| for choice in choices | |
| ): | |
| has_tool_calls = True | |
| if isinstance(choices, list) and any( | |
| isinstance(choice, dict) and choice.get("finish_reason") is not None | |
| for choice in choices | |
| ): | |
| has_finish = True | |
| event_size = len(event) | |
| logger.debug( | |
| f"OAI下游SSE准备发送: model={model}, request_id={request_id}, " | |
| f"gemini_chunk_keys={list(gemini_chunk.keys()) if isinstance(gemini_chunk, dict) else type(gemini_chunk)}, " | |
| f"bytes={event_size}" | |
| ) | |
| byte_count += event_size | |
| yield event | |
| logger.debug(f"OAI下游SSE已交给ASGI: model={model}, request_id={request_id}") | |
| if not has_finish: | |
| base = {"id": f"chatcmpl-{request_id}", "object": "chat.completion.chunk", "created": int(time.time()), "model": model} | |
| final_event = await _sse_json_event({**base, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]}, compact=False) | |
| logger.debug(f"OAI下游补充finish事件准备发送: model={model}, request_id={request_id}, bytes={len(final_event)}") | |
| byte_count += len(final_event) | |
| yield final_event | |
| logger.debug(f"OAI下游补充finish事件已交给ASGI: model={model}, request_id={request_id}") | |
| done_event = b"data: [DONE]\n\n" | |
| logger.debug(f"OAI下游DONE准备发送: model={model}, request_id={request_id}, bytes={len(done_event)}") | |
| byte_count += len(done_event) | |
| yield done_event | |
| logger.debug(f"OAI下游DONE已交给ASGI: model={model}, request_id={request_id}") | |
| prefix = progress_context.get("prefix") | |
| if prefix: | |
| ResponseAggregator.log_stream_forward_complete(prefix, chunk_count, byte_count, time.monotonic() - started_at) | |
| except asyncio.CancelledError: | |
| _mark_request_client_disconnected("openai.chat.stream") | |
| raise | |
| except DownstreamDisconnectedError: | |
| _mark_request_client_disconnected("openai.chat.stream") | |
| raise | |
| except VcoreError as e: | |
| _set_request_log_outcome(e.code, _vcore_error_status_info(e)) | |
| err = _vcore_error_to_oai(e) | |
| yield await _safe_sse_json_event( | |
| err, | |
| compact=False, | |
| fallback=_oai_error_fallback(e.message, err.get("error", {}).get("type", "server_error")), | |
| context="OAI 流式 Vcore 错误事件", | |
| ) | |
| yield b"data: [DONE]\n\n" | |
| logger.debug("OAI 流式错误事件与 DONE 已发送,随后结束连接并清理资源") | |
| except Exception as e: | |
| _set_request_log_outcome(500, str(e)) | |
| logger.error(f"OAI 流式错误: {e}") | |
| err = {"error": {"message": str(e), "type": "server_error", "code": None}} | |
| yield await _safe_sse_json_event( | |
| err, | |
| compact=False, | |
| fallback=_oai_error_fallback(e, "server_error"), | |
| context="OAI 流式未知错误事件", | |
| ) | |
| yield b"data: [DONE]\n\n" | |
| logger.debug("OAI 流式未知错误事件与 DONE 已发送,随后结束连接并清理资源") | |
| return StreamingResponse(oai_stream_generator(), media_type="text/event-stream", headers=SSE_HEADERS) | |
| else: | |
| try: | |
| async with business_session(vcore_client, source="openai.chat") as session: | |
| gemini_response = await _await_with_client_disconnect( | |
| request, | |
| session.complete_gemini( | |
| model=model, | |
| gemini_payload=gemini_payload, | |
| _expected_image_count=_expected_image_count_from_payload(gemini_payload), | |
| ), | |
| "openai.chat", | |
| ) | |
| oai_response = await _convert_oai_chat_response(gemini_response, model) | |
| return await _json_response_for_payload(oai_response, compact=True) | |
| except DownstreamDisconnectedError: | |
| raise | |
| except VcoreError as e: | |
| err = _vcore_error_to_oai(e) | |
| return JSONResponse(status_code=_wire_status_code(e.code), content=err) | |
| except Exception as e: | |
| logger.error(f"OAI 非流式错误: {e}") | |
| return JSONResponse(status_code=500, content={"error": {"message": str(e), "type": "server_error", "code": None}}) | |
| app.post("/v1/chat/completions", response_model=None)(oai_chat_completions) | |
| # ==================== OpenAI 图片兼容端点 ==================== | |
| async def oai_images_generations(request: Request) -> JSONResponse | Response: | |
| """OpenAI Images: text-to-image""" | |
| try: | |
| body = await _read_json_or_form_fields(request) | |
| model, gemini_payload, n, response_format = OAIImageRequestConverter.convert_generation(body) | |
| except ValueError as e: | |
| return _oai_bad_request(str(e)) | |
| except json.JSONDecodeError as e: | |
| return _oai_bad_request(f"Invalid JSON: {e}") | |
| except Exception as e: | |
| return _oai_bad_request(str(e)) | |
| stream = _coerce_bool(body.get("stream")) | |
| _log_request_start(request.method, request.url.path, source="openai.images.generations", model=model, mode="stream" if stream else "non-stream", n=n) | |
| if stream: | |
| return _run_oai_image_stream_request(request, vcore_client, model, gemini_payload, n, response_format, source="openai.images.generations") | |
| return await _run_oai_image_request(request, vcore_client, model, gemini_payload, n, response_format, source="openai.images.generations") | |
| async def oai_images_edits(request: Request) -> JSONResponse | Response: | |
| """OpenAI Images: image edit / image-to-image""" | |
| try: | |
| form = await request.form() | |
| fields = _form_to_plain_dict(form) | |
| image_uploads = _form_uploads(form, "image") | |
| if not image_uploads: | |
| return _oai_bad_request("image is required") | |
| images = [await _upload_to_inline_image(file) for file in image_uploads] | |
| mask_uploads = _form_uploads(form, "mask") | |
| mask = await _upload_to_inline_image(mask_uploads[0]) if mask_uploads else None | |
| model = OAIImageRequestConverter.resolve_model(fields.get("model")) | |
| prompt = str(fields.get("prompt") or "Edit the provided image.") | |
| if fields.get("negative_prompt"): | |
| prompt = f"{prompt.strip()}\nAvoid: {fields.get('negative_prompt')}" | |
| n = _coerce_oai_n(fields.get("n")) | |
| response_format = str(fields.get("response_format") or "b64_json") | |
| gemini_payload = OAIImageRequestConverter.build_payload( | |
| model=model, | |
| prompt=prompt, | |
| n=n, | |
| images=images, | |
| mask=mask, | |
| size=fields.get("size"), | |
| quality=fields.get("quality"), | |
| style=fields.get("style"), | |
| background=fields.get("background"), | |
| output_format=fields.get("output_format"), | |
| mode="edit", | |
| ) | |
| except Exception as e: | |
| return _oai_bad_request(str(e)) | |
| stream = _coerce_bool(fields.get("stream")) | |
| _log_request_start(request.method, request.url.path, source="openai.images.edits", model=model, mode="stream" if stream else "non-stream", n=n, images=len(images)) | |
| if stream: | |
| return _run_oai_image_stream_request(request, vcore_client, model, gemini_payload, n, response_format, source="openai.images.edits") | |
| return await _run_oai_image_request(request, vcore_client, model, gemini_payload, n, response_format, source="openai.images.edits") | |
| async def oai_images_variations(request: Request) -> JSONResponse | Response: | |
| """OpenAI Images: image variation""" | |
| try: | |
| form = await request.form() | |
| fields = _form_to_plain_dict(form) | |
| image_uploads = _form_uploads(form, "image") | |
| if not image_uploads: | |
| return _oai_bad_request("image is required") | |
| images = [await _upload_to_inline_image(file) for file in image_uploads] | |
| model = OAIImageRequestConverter.resolve_model(fields.get("model")) | |
| prompt = str(fields.get("prompt") or "Create a variation of the provided image.") | |
| if fields.get("negative_prompt"): | |
| prompt = f"{prompt.strip()}\nAvoid: {fields.get('negative_prompt')}" | |
| n = _coerce_oai_n(fields.get("n")) | |
| response_format = str(fields.get("response_format") or "b64_json") | |
| gemini_payload = OAIImageRequestConverter.build_payload( | |
| model=model, | |
| prompt=prompt, | |
| n=n, | |
| images=images, | |
| size=fields.get("size"), | |
| quality=fields.get("quality"), | |
| style=fields.get("style"), | |
| output_format=fields.get("output_format"), | |
| mode="variation", | |
| ) | |
| except Exception as e: | |
| return _oai_bad_request(str(e)) | |
| stream = _coerce_bool(fields.get("stream")) | |
| _log_request_start(request.method, request.url.path, source="openai.images.variations", model=model, mode="stream" if stream else "non-stream", n=n) | |
| if stream: | |
| return _run_oai_image_stream_request(request, vcore_client, model, gemini_payload, n, response_format, source="openai.images.variations") | |
| return await _run_oai_image_request(request, vcore_client, model, gemini_payload, n, response_format, source="openai.images.variations") | |
| app.post("/v1/images/generations", response_model=None)(oai_images_generations) | |
| app.post("/v1/images/edits", response_model=None)(oai_images_edits) | |
| app.post("/v1/images/variations", response_model=None)(oai_images_variations) | |
| # ==================== Gemini 辅助端点 ==================== | |
| async def count_tokens(model: str, request: Request) -> JSONResponse: | |
| """Gemini countTokens 兼容端点""" | |
| try: | |
| body_any = await request.json() | |
| except json.JSONDecodeError as e: | |
| raise InvalidArgumentError(f"Invalid JSON in request body: {e}") | |
| if not isinstance(body_any, dict): | |
| raise InvalidArgumentError("Request body must be a JSON object") | |
| body = cast(dict[str, Any], body_any) | |
| request_obj = body.get("generateContentRequest") | |
| _log_request_start(request.method, request.url.path, source="gemini.countTokens", model=model, mode="countTokens") | |
| if isinstance(request_obj, dict): | |
| contents = request_obj.get("contents", []) | |
| else: | |
| contents = body.get("contents", []) | |
| normalized = vcore_client.transformer._normalize_gemini_payload({"contents": contents}) | |
| async with business_session(vcore_client, source="gemini.countTokens") as session: | |
| total_tokens = await _await_with_client_disconnect( | |
| request, | |
| session.count_tokens( | |
| model=model, | |
| contents=cast(list[dict[str, Any]], normalized.get("contents", [])), | |
| ), | |
| "gemini.countTokens", | |
| ) | |
| return JSONResponse(content={"totalTokens": total_tokens}) | |
| app.post("/v1beta/models/{model}:countTokens", response_model=None)(count_tokens) | |
| app.post("/v1/models/{model}:countTokens", response_model=None)(count_tokens) | |
| logger.info("FastAPI 应用创建完成") | |
| return app | |
| # ==================== 辅助函数 ==================== | |
| def _load_models_config() -> list[str]: | |
| """加载模型配置""" | |
| try: | |
| with open(MODELS_CONFIG_FILE, 'r', encoding='utf-8') as f: | |
| config = json.load(f) | |
| return cast(list[str], config.get('models', [])) | |
| except Exception: | |
| return ["gemini-1.5-pro", "gemini-1.5-flash", "gemini-2.0-flash-exp", "gemini-2.0-pro-exp-02-05", "gemini-2.5-flash"] | |
| def _supported_generation_methods(model: str) -> list[str]: | |
| methods = ["generateContent", "streamGenerateContent", "countTokens"] | |
| if "image" in model.lower(): | |
| methods.append("generateImages") | |
| return methods | |
| def _prepare_gemini_downstream_payload(body: dict[str, Any]) -> dict[str, Any]: | |
| """把 Gemini 下游请求整理为标准 generateContent payload,避免修改原请求对象。""" | |
| prepared = dict(body) | |
| gen_config = prepared.get("generationConfig") or prepared.get("generation_config") | |
| if isinstance(gen_config, dict): | |
| normalized_gen_config = dict(gen_config) | |
| image_config = normalized_gen_config.get("imageConfig") or normalized_gen_config.get("image_config") | |
| if isinstance(image_config, dict): | |
| normalized_image_config = dict(image_config) | |
| number_of_images = normalized_image_config.get("numberOfImages") or normalized_image_config.get("number_of_images") | |
| if number_of_images is not None: | |
| try: | |
| normalized_image_config["numberOfImages"] = max(1, min(int(number_of_images), 8)) | |
| except (TypeError, ValueError): | |
| normalized_image_config.pop("numberOfImages", None) | |
| normalized_image_config.pop("number_of_images", None) | |
| normalized_gen_config["imageConfig"] = normalized_image_config | |
| normalized_gen_config.pop("image_config", None) | |
| prepared["generationConfig"] = normalized_gen_config | |
| prepared.pop("generation_config", None) | |
| return prepared | |
| def _vcore_error_to_oai(e: VcoreError) -> dict[str, Any]: | |
| """将 VcoreError 转为 OAI 错误格式""" | |
| if isinstance(e, InvalidArgumentError): | |
| err_type = "invalid_request_error" | |
| elif isinstance(e, RateLimitError): | |
| err_type = "rate_limit_error" | |
| elif isinstance(e, AuthenticationError): | |
| err_type = "authentication_error" | |
| else: | |
| err_type = "server_error" | |
| return {"error": {"message": e.message, "type": err_type, "code": None}} | |
| def _expected_image_count_from_payload(gemini_payload: dict[str, Any]) -> int: | |
| """从 Gemini generationConfig 中读取期望图片数量,默认 1。""" | |
| gen_config = gemini_payload.get("generationConfig") or gemini_payload.get("generation_config") or {} | |
| if not isinstance(gen_config, dict): | |
| return 1 | |
| image_config = gen_config.get("imageConfig") or gen_config.get("image_config") or {} | |
| value: Any = None | |
| if isinstance(image_config, dict): | |
| value = image_config.get("numberOfImages") or image_config.get("number_of_images") | |
| if value is None: | |
| value = gen_config.get("candidateCount") or gen_config.get("candidate_count") or 1 | |
| try: | |
| return max(1, int(value)) | |
| except (TypeError, ValueError): | |
| return 1 | |
| async def _read_json_or_form_fields(request: Request) -> dict[str, Any]: | |
| content_type = request.headers.get("content-type", "").lower() | |
| if "multipart/form-data" in content_type or "application/x-www-form-urlencoded" in content_type: | |
| form = await request.form() | |
| return _form_to_plain_dict(form) | |
| body_any = await request.json() | |
| if not isinstance(body_any, dict): | |
| raise ValueError("Request body must be a JSON object") | |
| return cast(dict[str, Any], body_any) | |
| def _form_to_plain_dict(form: Any) -> dict[str, Any]: | |
| data: dict[str, Any] = {} | |
| for key, value in form.multi_items(): | |
| if _is_upload_file(value): | |
| continue | |
| if key in data: | |
| existing = data[key] | |
| if isinstance(existing, list): | |
| existing.append(value) | |
| else: | |
| data[key] = [existing, value] | |
| else: | |
| data[key] = value | |
| return data | |
| def _form_uploads(form: Any, key: str) -> list[UploadFile]: | |
| uploads: list[UploadFile] = [] | |
| for item_key, value in form.multi_items(): | |
| if item_key != key and item_key != f"{key}[]" and not item_key.startswith(f"{key}["): | |
| continue | |
| if _is_upload_file(value): | |
| uploads.append(cast(UploadFile, value)) | |
| return uploads | |
| def _is_upload_file(value: Any) -> bool: | |
| return isinstance(value, (UploadFile, StarletteUploadFile)) | |
| async def _upload_to_inline_image(file: UploadFile) -> dict[str, str]: | |
| data = await file.read() | |
| if not data: | |
| raise ValueError(f"{file.filename or 'image'} is empty") | |
| mime_type = file.content_type or mimetypes.guess_type(file.filename or "")[0] or "image/png" | |
| return { | |
| "mimeType": mime_type, | |
| "data": base64.b64encode(data).decode("ascii"), | |
| } | |
| def _coerce_oai_n(value: Any) -> int: | |
| try: | |
| n = int(value) | |
| except (TypeError, ValueError): | |
| n = 1 | |
| return max(1, min(n, 8)) | |
| def _coerce_bool(value: Any) -> bool: | |
| if isinstance(value, bool): | |
| return value | |
| if isinstance(value, str): | |
| return value.strip().lower() in {"1", "true", "yes", "on"} | |
| if isinstance(value, (int, float)): | |
| return value != 0 | |
| return False | |
| async def _run_oai_image_request( | |
| request: Request, | |
| vcore_client: VcoreAIClient, | |
| model: str, | |
| gemini_payload: dict[str, Any], | |
| n: int, | |
| response_format: str, | |
| source: str = "openai.images", | |
| ) -> Response: | |
| image_items: list[dict[str, Any]] = [] | |
| normalized_response_format = "url" if response_format == "url" else "b64_json" | |
| try: | |
| async with business_session(vcore_client, source=source) as session: | |
| async def operation() -> list[dict[str, Any]]: | |
| gemini_response = await session.complete_gemini( | |
| model=model, | |
| gemini_payload=gemini_payload, | |
| _raw_image_response=True, | |
| _expected_image_count=n, | |
| ) | |
| image_data = await _convert_oai_image_data(gemini_response, normalized_response_format) | |
| return image_data[:n] | |
| image_items = await _await_with_client_disconnect(request, operation(), source) | |
| except DownstreamDisconnectedError: | |
| raise | |
| except VcoreError as e: | |
| return JSONResponse(status_code=_wire_status_code(e.code), content=_vcore_error_to_oai(e)) | |
| except Exception as e: | |
| logger.error(f"OAI 图片请求错误: {e}") | |
| return JSONResponse(status_code=500, content={"error": {"message": str(e), "type": "server_error", "code": None}}) | |
| if not image_items: | |
| return JSONResponse( | |
| status_code=502, | |
| content={"error": {"message": "Upstream response did not contain image data", "type": "server_error", "code": None}}, | |
| ) | |
| return await _json_response_for_payload({ | |
| "created": int(time.time()), | |
| "data": image_items[:n], | |
| }, compact=True) | |
| def _run_oai_image_stream_request( | |
| request: Request, | |
| vcore_client: VcoreAIClient, | |
| model: str, | |
| gemini_payload: dict[str, Any], | |
| n: int, | |
| response_format: str, | |
| source: str = "openai.images", | |
| ) -> StreamingResponse: | |
| normalized_response_format = "url" if response_format == "url" else "b64_json" | |
| request_id = uuid.uuid4().hex[:24] | |
| async def stream_generator() -> AsyncGenerator[bytes, None]: | |
| sent = 0 | |
| chunk_count = 0 | |
| byte_count = 0 | |
| started_at = time.monotonic() | |
| try: | |
| async with business_session(vcore_client, source=source) as session: | |
| progress_context: dict[str, str] = {} | |
| upstream = session.stream_gemini(model=model, gemini_payload=gemini_payload, progress_context=progress_context) | |
| async for chunk in _aiter_with_client_disconnect(request, upstream, source): | |
| chunk_count += 1 | |
| items = await _convert_oai_image_chunk_data(chunk, normalized_response_format) | |
| for item in items: | |
| if sent >= n: | |
| break | |
| payload = { | |
| "id": f"imggen-{request_id}", | |
| "object": "image.generation.chunk", | |
| "created": int(time.time()), | |
| "model": model, | |
| "data": [{"index": sent, **item}], | |
| } | |
| event = await _sse_json_event(payload, compact=True) | |
| byte_count += len(event) | |
| sent += 1 | |
| yield event | |
| if sent >= n: | |
| break | |
| prefix = progress_context.get("prefix") | |
| if prefix: | |
| ResponseAggregator.log_stream_forward_complete(prefix, chunk_count, byte_count, time.monotonic() - started_at) | |
| done_event = b"data: [DONE]\n\n" | |
| byte_count += len(done_event) | |
| yield done_event | |
| except asyncio.CancelledError: | |
| _mark_request_client_disconnected(source) | |
| raise | |
| except DownstreamDisconnectedError: | |
| _mark_request_client_disconnected(source) | |
| raise | |
| except VcoreError as e: | |
| _set_request_log_outcome(e.code, _vcore_error_status_info(e)) | |
| err = _vcore_error_to_oai(e) | |
| yield await _safe_sse_json_event( | |
| err, | |
| compact=True, | |
| fallback=_oai_error_fallback(e.message, err.get("error", {}).get("type", "server_error")), | |
| context="OAI 图片流式 Vcore 错误事件", | |
| ) | |
| yield b"data: [DONE]\n\n" | |
| except Exception as e: | |
| _set_request_log_outcome(500, str(e)) | |
| logger.error(f"OAI 图片流式请求错误: {e}") | |
| yield await _safe_sse_json_event( | |
| {"error": {"message": str(e), "type": "server_error", "code": None}}, | |
| compact=True, | |
| fallback=_oai_error_fallback(e, "server_error"), | |
| context="OAI 图片流式未知错误事件", | |
| ) | |
| yield b"data: [DONE]\n\n" | |
| return StreamingResponse(stream_generator(), media_type="text/event-stream", headers=SSE_HEADERS) | |
| def _oai_bad_request(message: str) -> JSONResponse: | |
| return JSONResponse( | |
| status_code=400, | |
| content={"error": {"message": message, "type": "invalid_request_error", "code": None}}, | |
| ) | |