File size: 4,621 Bytes
95f8c4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """FastAPI application entrypoint."""
from __future__ import annotations
import argparse
import logging
from contextlib import asynccontextmanager
from fastapi import Depends, FastAPI
from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles
from aistudio_api.infrastructure.gateway.client import AIStudioClient
from .dependencies import require_api_key
from .routes_accounts import router as accounts_router
from .routes_gemini import router as gemini_router
from .routes_models import router as models_router
from .routes_system import protected_router as system_protected_router
from .routes_system import public_router as system_public_router
from .state import runtime_state
logging.basicConfig(
level=logging.INFO, format="%(asctime)s [%(name)s] %(message)s", datefmt="%H:%M:%S"
)
logger = logging.getLogger("aistudio.server")
@asynccontextmanager
async def lifespan(app: FastAPI):
import asyncio
from aistudio_api.application.account_rotator import RotationMode, init_rotator
from aistudio_api.application.account_service import AccountService
from aistudio_api.config import settings
from aistudio_api.infrastructure.account.account_store import AccountStore
client = AIStudioClient(
port=runtime_state.browser_port,
)
runtime_state.client = client
from aistudio_api.config import settings as app_settings
runtime_state.busy_lock = asyncio.Semaphore(app_settings.max_concurrency)
# 注入 snapshot 缓存引用，切号时需要清除
from aistudio_api.infrastructure.gateway.client import _snapshot_cache
runtime_state.snapshot_cache = _snapshot_cache
# 初始化账号管理服务
account_store = AccountStore()
account_service = AccountService(account_store)
runtime_state.account_service = account_service
# 初始化账号轮询器
rotation_mode = getattr(settings, "account_rotation_mode", "round_robin")
cooldown = getattr(settings, "account_cooldown_seconds", 60)
rotator = init_rotator(
account_store,
mode=RotationMode(rotation_mode),
cooldown_seconds=cooldown,
)
runtime_state.rotator = rotator
logger.info(
"Client initialized (port=%s, rotation=%s, accounts=%d)",
runtime_state.browser_port,
rotator.mode,
len(account_store.list_accounts()),
)
# 后台预热浏览器，避免首次请求延迟
warmup_task = None
async def _warmup():
try:
await client.warmup()
except Exception as e:
logger.warning("浏览器预热失败: %s", e)
warmup_task = asyncio.create_task(_warmup())
yield
logger.info("Shutting down")
if warmup_task and not warmup_task.done():
warmup_task.cancel()
if client:
try:
await client.close()
except Exception as e:
logger.debug("Error closing client: %s", e)
runtime_state.client = None
runtime_state.busy_lock = None
runtime_state.account_service = None
runtime_state.rotator = None
app = FastAPI(title="AI Studio API", lifespan=lifespan)
app.include_router(system_public_router)
app.include_router(system_protected_router, dependencies=[Depends(require_api_key)])
app.include_router(gemini_router, dependencies=[Depends(require_api_key)])
app.include_router(models_router, dependencies=[Depends(require_api_key)])
app.include_router(accounts_router, dependencies=[Depends(require_api_key)])
# 挂载静态文件
import os
static_dir = os.path.join(os.path.dirname(__file__), "..", "static")
if os.path.isdir(static_dir):
app.mount("/static", StaticFiles(directory=static_dir), name="static")
@app.get("/")
async def root():
return RedirectResponse(url="/static/index.html")
@app.get("/login")
async def login_page():
return RedirectResponse(url="/static/login.html")
@app.get("/auth/check")
async def auth_check():
"""检查认证状态，用于前端判断是否需要登录。"""
from aistudio_api.config import settings
return {"auth_enabled": settings.auth_enabled}
def main():
from aistudio_api.config import settings
parser = argparse.ArgumentParser(description="AI Studio API Server")
parser.add_argument("--port", type=int, default=settings.port)
parser.add_argument("--browser-port", type=int, default=settings.browser_port)
args = parser.parse_args()
runtime_state.browser_port = args.browser_port
import uvicorn
logger.info("Starting server on port %s", args.port)
uvicorn.run(app, host="0.0.0.0", port=args.port, log_level="info")
|