"""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 (OpenAI 标准 Header) 2. x-goog-api-key: (Google/Gemini 标准 Header) 3. ?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) # ==================== 全局异常处理 ==================== @app.exception_handler(VcoreError) 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(), ) @app.exception_handler(DownstreamDisconnectedError) 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"}}, ) @app.exception_handler(Exception) 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}}, )